Use flatbuffers for some messages from Redis. (#341)

* Compile the Ray redis module with C++.

* Redo parsing of object table notifications with flatbuffers.

* Update redis module python tests.

* Redo parsing of task table notifications with flatbuffers.

* Fix linting.

* Redo parsing of db client notifications with flatbuffers.

* Redo publishing of local scheduler heartbeats with flatbuffers.

* Fix linting.

* Remove usage of fixed-width formatting of scheduling state in channel name.

* Reply with flatbuffer object to task table queries, also simplify redis string to flatbuffer string conversion.

* Fix linting and tests.

* fix

* cleanup

* simplify logic in ReplyWithTask

src/common/CMakeLists.txt

@@ -24,6 +24,15 @@ add_custom_command(
   COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
   VERBATIM)
 
+# Generate Python bindings for the flatbuffers objects.
+set(PYTHON_OUTPUT_DIR ${CMAKE_BINARY_DIR}/generated/)
+add_custom_command(
+  TARGET gen_common_fbs
+  COMMAND ${FLATBUFFERS_COMPILER} -p -o ${PYTHON_OUTPUT_DIR} ${COMMON_FBS_SRC}
+  DEPENDS ${FBS_DEPENDS}
+  COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
+  VERBATIM)
+
 add_dependencies(gen_common_fbs flatbuffers_ep)
 
 add_custom_target(

src/common/format/common.fbs

@@ -50,3 +50,62 @@ table TaskInfo {
 }
 
 root_type TaskInfo;
+
+table SubscribeToNotificationsReply {
+  // The object ID of the object that the notification is about.
+  object_id: string;
+  // The size of the object.
+  object_size: long;
+  // The IDs of the managers that contain this object.
+  manager_ids: [string];
+}
+
+root_type SubscribeToNotificationsReply;
+
+table TaskReply {
+  // The task ID of the task that the message is about.
+  task_id: string;
+  // The state of the task. This is encoded as a bit mask of scheduling_state
+  // enum values in task.h.
+  state: long;
+  // A local scheduler ID.
+  local_scheduler_id: string;
+  // A string of bytes representing the task specification.
+  task_spec: string;
+}
+
+root_type TaskReply;
+
+table SubscribeToDBClientTableReply {
+  // The db client ID of the client that the message is about.
+  db_client_id: string;
+  // The type of the client.
+  client_type: string;
+  // If the client is a local scheduler, this is the address of the plasma
+  // manager that the local scheduler is connected to. Otherwise, it is empty.
+  aux_address: string;
+  // True if the message is about the addition of a client and false if it is
+  // about the deletion of a client.
+  is_insertion: bool;
+}
+
+root_type SubscribeToDBClientTableReply;
+
+table LocalSchedulerInfoMessage {
+  // The db client ID of the client that the message is about.
+  db_client_id: string;
+  // The total number of workers that are connected to this local scheduler.
+  total_num_workers: long;
+  // The number of tasks queued in this local scheduler.
+  task_queue_length: long;
+  // The number of workers that are available and waiting for tasks.
+  available_workers: long;
+  // The resource vector of resources generally available to this local
+  // scheduler.
+  static_resources: [double];
+  // The resource vector of resources currently available to this local
+  // scheduler.
+  dynamic_resources: [double];
+}
+
+root_type LocalSchedulerInfoMessage;

src/common/redis_module/CMakeLists.txt

@@ -3,15 +3,15 @@ cmake_minimum_required(VERSION 2.8)
 project(ray_redis_module)
 
 if(APPLE)
-  set(REDIS_MODULE_CFLAGS -W -Wall -dynamic -fno-common -g -ggdb -std=c99 -O2)
+  set(REDIS_MODULE_CFLAGS -W -Wall -dynamic -fno-common -g -ggdb -std=c++11 -O2)
   set(REDIS_MODULE_LDFLAGS "-undefined dynamic_lookup")
   set(CMAKE_SHARED_LIBRARY_SUFFIX ".so")
 else()
-  set(REDIS_MODULE_CFLAGS -W -Wall -fno-common -g -ggdb -std=c99 -O2)
+  set(REDIS_MODULE_CFLAGS -W -Wall -fno-common -g -ggdb -std=c++11 -O2)
   set(REDIS_MODULE_LDFLAGS -shared)
 endif()
 
-add_library(ray_redis_module SHARED ray_redis_module.c)
+add_library(ray_redis_module SHARED ray_redis_module.cc)
 
 target_compile_options(ray_redis_module PUBLIC ${REDIS_MODULE_CFLAGS} -fPIC)
 target_link_libraries(ray_redis_module ${REDIS_MODULE_LDFLAGS})

src/common/redis_module/ray_redis_module.cc

@@ -5,6 +5,10 @@
 
 #include "redis_string.h"
 
+#include "format/common_generated.h"
+
+#include "common_protocol.h"
+
 /**
  * Various tables are maintained in redis:
  *
@@ -42,11 +46,29 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
                                 int mode) {
   RedisModuleString *prefixed_keyname =
       RedisString_Format(ctx, "%s%S", prefix, keyname);
-  RedisModuleKey *key = RedisModule_OpenKey(ctx, prefixed_keyname, mode);
+  RedisModuleKey *key =
+      (RedisModuleKey *) RedisModule_OpenKey(ctx, prefixed_keyname, mode);
   RedisModule_FreeString(ctx, prefixed_keyname);
   return key;
 }
 
+/**
+ * This is a helper method to convert a redis module string to a flatbuffer
+ * string.
+ *
+ * @param fbb The flatbuffer builder.
+ * @param redis_string The redis string.
+ * @return The flatbuffer string.
+ */
+flatbuffers::Offset<flatbuffers::String> RedisStringToFlatbuf(
+    flatbuffers::FlatBufferBuilder &fbb,
+    RedisModuleString *redis_string) {
+  size_t redis_string_size;
+  const char *redis_string_str =
+      RedisModule_StringPtrLen(redis_string, &redis_string_size);
+  return fbb.CreateString(redis_string_str, redis_string_size);
+}
+
 /**
  * Publish a notification to a client's notification channel about an insertion
  * or deletion to the db client table.
@@ -69,16 +91,23 @@ bool PublishDBClientNotification(RedisModuleCtx *ctx,
   /* Construct strings to publish on the db client channel. */
   RedisModuleString *channel_name =
       RedisModule_CreateString(ctx, "db_clients", strlen("db_clients"));
-  RedisModuleString *client_info;
-  const char *is_insertion_string = is_insertion ? "1" : "0";
-  if (aux_address) {
-    client_info =
-        RedisString_Format(ctx, "%S:%S %S %s", ray_client_id, client_type,
-                           aux_address, is_insertion_string);
+  /* Construct the flatbuffers object to publish over the channel. */
+  flatbuffers::FlatBufferBuilder fbb;
+  /* Use an empty aux address if one is not passed in. */
+  flatbuffers::Offset<flatbuffers::String> aux_address_str;
+  if (aux_address != NULL) {
+    aux_address_str = RedisStringToFlatbuf(fbb, aux_address);
   } else {
-    client_info = RedisString_Format(ctx, "%S:%S : %s", ray_client_id,
-                                     client_type, is_insertion_string);
+    aux_address_str = fbb.CreateString("", strlen(""));
   }
+  /* Create the flatbuffers message. */
+  auto message = CreateSubscribeToDBClientTableReply(
+      fbb, RedisStringToFlatbuf(fbb, ray_client_id),
+      RedisStringToFlatbuf(fbb, client_type), aux_address_str, is_insertion);
+  fbb.Finish(message);
+  /* Create a Redis string to publish by serializing the flatbuffers object. */
+  RedisModuleString *client_info = RedisModule_CreateString(
+      ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize());
 
   /* Publish the client info on the db client channel. */
   RedisModuleCallReply *reply;
@@ -328,48 +357,42 @@ bool PublishObjectNotification(RedisModuleCtx *ctx,
                                RedisModuleString *object_id,
                                RedisModuleString *data_size,
                                RedisModuleKey *key) {
-  /* Create a string formatted as "<object id> MANAGERS <size> <manager id1>
-   * <manager id2> ..." */
+  flatbuffers::FlatBufferBuilder fbb;
+
   long long data_size_value;
   if (RedisModule_StringToLongLong(data_size, &data_size_value) !=
       REDISMODULE_OK) {
     return RedisModule_ReplyWithError(ctx, "data_size must be integer");
   }
 
-  RedisModuleString *manager_list = RedisString_Format(ctx, "%S ", object_id);
-
-  /* Append binary data size for this object. */
-  /* TODO(pcm): Replace by a formatted fix length version of the size. */
-  RedisModule_StringAppendBuffer(ctx, manager_list,
-                                 (const char *) &data_size_value,
-                                 sizeof(data_size_value));
-
-  RedisModule_StringAppendBuffer(ctx, manager_list, " MANAGERS",
-                                 strlen(" MANAGERS"));
-
+  std::vector<flatbuffers::Offset<flatbuffers::String>> manager_ids;
   CHECK_ERROR(
       RedisModule_ZsetFirstInScoreRange(key, REDISMODULE_NEGATIVE_INFINITE,
                                         REDISMODULE_POSITIVE_INFINITE, 1, 1),
       "Unable to initialize zset iterator");
-
   /* Loop over the managers in the object table for this object ID. */
   do {
     RedisModuleString *curr = RedisModule_ZsetRangeCurrentElement(key, NULL);
-    RedisModule_StringAppendBuffer(ctx, manager_list, " ", 1);
-    size_t size;
-    const char *val = RedisModule_StringPtrLen(curr, &size);
-    RedisModule_StringAppendBuffer(ctx, manager_list, val, size);
+    manager_ids.push_back(RedisStringToFlatbuf(fbb, curr));
   } while (RedisModule_ZsetRangeNext(key));
 
+  auto message = CreateSubscribeToNotificationsReply(
+      fbb, RedisStringToFlatbuf(fbb, object_id), data_size_value,
+      fbb.CreateVector(manager_ids));
+  fbb.Finish(message);
+
   /* Publish the notification to the clients notification channel.
    * TODO(rkn): These notifications could be batched together. */
   RedisModuleString *channel_name =
       RedisString_Format(ctx, "%s%S", OBJECT_CHANNEL_PREFIX, client_id);
 
+  RedisModuleString *payload = RedisModule_CreateString(
+      ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize());
+
   RedisModuleCallReply *reply;
-  reply = RedisModule_Call(ctx, "PUBLISH", "ss", channel_name, manager_list);
+  reply = RedisModule_Call(ctx, "PUBLISH", "ss", channel_name, payload);
   RedisModule_FreeString(ctx, channel_name);
-  RedisModule_FreeString(ctx, manager_list);
+  RedisModule_FreeString(ctx, payload);
   if (reply == NULL) {
     return false;
   }
@@ -668,35 +691,6 @@ int ResultTableAdd_RedisCommand(RedisModuleCtx *ctx,
   return REDISMODULE_OK;
 }
 
-int ParseTaskState(RedisModuleString *state) {
-  size_t state_length;
-  const char *state_string = RedisModule_StringPtrLen(state, &state_length);
-  int state_integer;
-  int scanned = sscanf(state_string, "%2d", &state_integer);
-  if (scanned != 1 || state_length != 2) {
-    return -1;
-  }
-  return state_integer;
-}
-
-RedisModuleString *NormalizeTaskState(RedisModuleCtx *ctx,
-                                      RedisModuleString *state) {
-  /* Pad the state integer to a fixed-width integer, and make sure it has width
-   * less than or equal to 2. */
-  long long state_integer;
-  int status = RedisModule_StringToLongLong(state, &state_integer);
-  if (status != REDISMODULE_OK) {
-    return NULL;
-  }
-  state = RedisModule_CreateStringPrintf(ctx, "%2d", state_integer);
-  size_t length;
-  RedisModule_StringPtrLen(state, &length);
-  if (length != 2) {
-    return NULL;
-  }
-  return state;
-}
-
 /**
  * Reply with information about a task ID. This is used by
  * RAY.RESULT_TABLE_LOOKUP and RAY.TASK_TABLE_GET.
@@ -716,8 +710,9 @@ int ReplyWithTask(RedisModuleCtx *ctx, RedisModuleString *task_id) {
     RedisModuleString *state = NULL;
     RedisModuleString *local_scheduler_id = NULL;
     RedisModuleString *task_spec = NULL;
-    RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state, "node",
-                        &local_scheduler_id, "TaskSpec", &task_spec, NULL);
+    RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state,
+                        "local_scheduler_id", &local_scheduler_id, "TaskSpec",
+                        &task_spec, NULL);
     if (state == NULL || local_scheduler_id == NULL || task_spec == NULL) {
       /* We must have either all fields or no fields. */
       RedisModule_CloseKey(key);
@@ -725,21 +720,26 @@ int ReplyWithTask(RedisModuleCtx *ctx, RedisModuleString *task_id) {
           ctx, "Missing fields in the task table entry");
     }
 
-    int state_integer = ParseTaskState(state);
-    if (state_integer < 0) {
+    long long state_integer;
+    if (RedisModule_StringToLongLong(state, &state_integer) != REDISMODULE_OK ||
+        state_integer < 0) {
       RedisModule_CloseKey(key);
       RedisModule_FreeString(ctx, state);
       RedisModule_FreeString(ctx, local_scheduler_id);
       RedisModule_FreeString(ctx, task_spec);
-      return RedisModule_ReplyWithError(ctx,
-                                        "Found invalid scheduling state (must "
-                                        "be an integer of width 2");
+      return RedisModule_ReplyWithError(ctx, "Found invalid scheduling state.");
     }
 
-    RedisModule_ReplyWithArray(ctx, 3);
-    RedisModule_ReplyWithLongLong(ctx, state_integer);
-    RedisModule_ReplyWithString(ctx, local_scheduler_id);
-    RedisModule_ReplyWithString(ctx, task_spec);
+    flatbuffers::FlatBufferBuilder fbb;
+    auto message =
+        CreateTaskReply(fbb, RedisStringToFlatbuf(fbb, task_id), state_integer,
+                        RedisStringToFlatbuf(fbb, local_scheduler_id),
+                        RedisStringToFlatbuf(fbb, task_spec));
+    fbb.Finish(message);
+
+    RedisModuleString *reply = RedisModule_CreateString(
+        ctx, (char *) fbb.GetBufferPointer(), fbb.GetSize());
+    RedisModule_ReplyWithString(ctx, reply);
 
     RedisModule_FreeString(ctx, state);
     RedisModule_FreeString(ctx, local_scheduler_id);
@@ -801,54 +801,59 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx,
 int TaskTableWrite(RedisModuleCtx *ctx,
                    RedisModuleString *task_id,
                    RedisModuleString *state,
-                   RedisModuleString *node_id,
+                   RedisModuleString *local_scheduler_id,
                    RedisModuleString *task_spec) {
-  /* Pad the state integer to a fixed-width integer, and make sure it has width
-   * less than or equal to 2. */
-  state = NormalizeTaskState(ctx, state);
-  if (state == NULL) {
-    return RedisModule_ReplyWithError(
-        ctx,
-        "Invalid scheduling state (must be an integer of width at most 2)");
+  /* Extract the scheduling state. */
+  long long state_value;
+  if (RedisModule_StringToLongLong(state, &state_value) != REDISMODULE_OK) {
+    return RedisModule_ReplyWithError(ctx, "scheduling state must be integer");
   }
-
   /* Add the task to the task table. If no spec was provided, get the existing
    * spec out of the task table so we can publish it. */
   RedisModuleString *existing_task_spec = NULL;
   RedisModuleKey *key =
       OpenPrefixedKey(ctx, TASK_PREFIX, task_id, REDISMODULE_WRITE);
   if (task_spec == NULL) {
-    RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
-                        node_id, NULL);
+    RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state,
+                        "local_scheduler_id", local_scheduler_id, NULL);
     RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "TaskSpec",
                         &existing_task_spec, NULL);
     if (existing_task_spec == NULL) {
       RedisModule_CloseKey(key);
-      RedisModule_FreeString(ctx, state);
       return RedisModule_ReplyWithError(
           ctx, "Cannot update a task that doesn't exist yet");
     }
   } else {
-    RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
-                        node_id, "TaskSpec", task_spec, NULL);
+    RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state,
+                        "local_scheduler_id", local_scheduler_id, "TaskSpec",
+                        task_spec, NULL);
   }
   RedisModule_CloseKey(key);
 
   /* Build the PUBLISH topic and message for task table subscribers. The topic
-   * is a string in the format "TASK_PREFIX:<node ID>:<state>". The
-   * message is a string in the format: "<task ID> <state> <node ID> <task
-   * specification>". */
-  RedisModuleString *publish_topic =
-      RedisString_Format(ctx, "%s%S:%S", TASK_PREFIX, node_id, state);
-  RedisModuleString *publish_message;
+   * is a string in the format "TASK_PREFIX:<local scheduler ID>:<state>". The
+   * message is a serialized SubscribeToTasksReply flatbuffer object. */
+  RedisModuleString *publish_topic = RedisString_Format(
+      ctx, "%s%S:%S", TASK_PREFIX, local_scheduler_id, state);
+
+  /* Construct the flatbuffers object for the payload. */
+  flatbuffers::FlatBufferBuilder fbb;
+  /* Use the old task spec if the current one is NULL. */
+  RedisModuleString *task_spec_to_use;
   if (task_spec != NULL) {
-    publish_message = RedisString_Format(ctx, "%S %S %S %S", task_id, state,
-                                         node_id, task_spec);
+    task_spec_to_use = task_spec;
   } else {
-    publish_message = RedisString_Format(ctx, "%S %S %S %S", task_id, state,
-                                         node_id, existing_task_spec);
+    task_spec_to_use = existing_task_spec;
   }
-  RedisModule_FreeString(ctx, state);
+  /* Create the flatbuffers message. */
+  auto message =
+      CreateTaskReply(fbb, RedisStringToFlatbuf(fbb, task_id), state_value,
+                      RedisStringToFlatbuf(fbb, local_scheduler_id),
+                      RedisStringToFlatbuf(fbb, task_spec_to_use));
+  fbb.Finish(message);
+
+  RedisModuleString *publish_message = RedisModule_CreateString(
+      ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize());
 
   RedisModuleCallReply *reply =
       RedisModule_Call(ctx, "PUBLISH", "ss", publish_topic, publish_message);
@@ -878,10 +883,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
  *
  * @param task_id A string that is the ID of the task.
  * @param state A string that is the current scheduling state (a
- *        scheduling_state enum instance). The string's value must be a
- *        nonnegative integer less than 100, so that it has width at most 2. If
- *        less than 2, the value will be left-padded with spaces to a width of
- *        2.
+ *        scheduling_state enum instance).
  * @param local_scheduler_id A string that is the ray client ID of the
  *        associated local scheduler, if any.
  * @param task_spec A string that is the specification of the task, which can
@@ -908,10 +910,7 @@ int TaskTableAddTask_RedisCommand(RedisModuleCtx *ctx,
  *
  * @param task_id A string that is the ID of the task.
  * @param state A string that is the current scheduling state (a
- *        scheduling_state enum instance). The string's value must be a
- *        nonnegative integer less than 100, so that it has width at most 2. If
- *        less than 2, the value will be left-padded with spaces to a width of
- *        2.
+ *        scheduling_state enum instance).
  * @param ray_client_id A string that is the ray client ID of the associated
  *        local scheduler, if any.
  * @return OK if the operation was successful.
@@ -941,18 +940,15 @@ int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx,
  *        scheduling state. The update happens if and only if the current
  *        scheduling state AND-ed with the bitmask is greater than 0.
  * @param state A string that is the scheduling state (a scheduling_state enum
- *        instance) to update the task entry with. The string's value must be a
- *        nonnegative integer less than 100, so that it has width at most 2. If
- *        less than 2, the value will be left-padded with spaces to a width of
- *        2.
+ *        instance) to update the task entry with.
  * @param ray_client_id A string that is the ray client ID of the associated
  *        local scheduler, if any, to update the task entry with.
  * @return If the current scheduling state does not match the test bitmask,
  *         returns nil. Else, returns the same as RAY.TASK_TABLE_GET: an array
  *         of strings representing the updated task fields in the following
- *         order: 1) (integer) scheduling state 2) (string) associated node ID,
- *         if any 3) (string) the task specification, which can be casted to a
- *         task_spec.
+ *         order: 1) (integer) scheduling state 2) (string) associated local
+ *         scheduler ID, if any 3) (string) the task specification, which can be
+ *         cast to a task_spec.
  */
 int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
                                         RedisModuleString **argv,
@@ -961,18 +957,12 @@ int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
     return RedisModule_WrongArity(ctx);
   }
 
-  RedisModuleString *state = NormalizeTaskState(ctx, argv[3]);
-  if (state == NULL) {
-    return RedisModule_ReplyWithError(
-        ctx,
-        "Invalid scheduling state (must be an integer of width at most 2)");
-  }
+  RedisModuleString *state = argv[3];
 
   RedisModuleKey *key = OpenPrefixedKey(ctx, TASK_PREFIX, argv[1],
                                         REDISMODULE_READ | REDISMODULE_WRITE);
   if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) {
     RedisModule_CloseKey(key);
-    RedisModule_FreeString(ctx, state);
     return RedisModule_ReplyWithNull(ctx);
   }
 
@@ -980,19 +970,20 @@ int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
   RedisModuleString *current_state = NULL;
   RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &current_state,
                       NULL);
-  int current_state_integer = ParseTaskState(current_state);
+  long long current_state_integer;
+  if (RedisModule_StringToLongLong(current_state, &current_state_integer) !=
+      REDISMODULE_OK) {
+    return RedisModule_ReplyWithError(ctx, "current_state must be integer");
+  }
+
   if (current_state_integer < 0) {
     RedisModule_CloseKey(key);
-    RedisModule_FreeString(ctx, state);
-    return RedisModule_ReplyWithError(ctx,
-                                      "Found invalid scheduling state (must "
-                                      "be an integer of width 2");
+    return RedisModule_ReplyWithError(ctx, "Found invalid scheduling state.");
   }
   long long test_state_bitmask;
   int status = RedisModule_StringToLongLong(argv[2], &test_state_bitmask);
   if (status != REDISMODULE_OK) {
     RedisModule_CloseKey(key);
-    RedisModule_FreeString(ctx, state);
     return RedisModule_ReplyWithError(
         ctx, "Invalid test value for scheduling state");
   }
@@ -1000,16 +991,14 @@ int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
     /* The current value does not match the test bitmask, so do not perform the
      * update. */
     RedisModule_CloseKey(key);
-    RedisModule_FreeString(ctx, state);
     return RedisModule_ReplyWithNull(ctx);
   }
 
   /* The test passed, so perform the update. */
-  RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
-                      argv[4], NULL);
+  RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state,
+                      "local_scheduler_id", argv[4], NULL);
   /* Clean up. */
   RedisModule_CloseKey(key);
-  RedisModule_FreeString(ctx, state);
   /* Construct a reply by getting the task from the task ID. */
   return ReplyWithTask(ctx, argv[1]);
 }
@@ -1023,9 +1012,9 @@ int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
  *
  * @param task_id A string of the task ID to look up.
  * @return An array of strings representing the task fields in the following
- *         order: 1) (integer) scheduling state 2) (string) associated node ID,
- *         if any 3) (string) the task specification, which can be casted to a
- *         task_spec. If the task ID is not in the table, returns nil.
+ *         order: 1) (integer) scheduling state 2) (string) associated local
+ *         scheduler ID, if any 3) (string) the task specification, which can be
+ *         cast to a task_spec. If the task ID is not in the table, returns nil.
  */
 int TaskTableGet_RedisCommand(RedisModuleCtx *ctx,
                               RedisModuleString **argv,
@@ -1038,6 +1027,8 @@ int TaskTableGet_RedisCommand(RedisModuleCtx *ctx,
   return ReplyWithTask(ctx, argv[1]);
 }
 
+extern "C" {
+
 /* This function must be present on each Redis module. It is used in order to
  * register the commands into the Redis server. */
 int RedisModule_OnLoad(RedisModuleCtx *ctx,
@@ -1135,3 +1126,5 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
 
   return REDISMODULE_OK;
 }
+
+} /* extern "C" */

src/common/state/redis.cc

@@ -27,6 +27,10 @@ extern "C" {
 #include "redis.h"
 #include "io.h"
 
+#include "format/common_generated.h"
+
+#include "common_protocol.h"
+
 #ifndef _WIN32
 /* This function is actually not declared in standard POSIX, so declare it. */
 extern int usleep(useconds_t usec);
@@ -358,27 +362,14 @@ Task *parse_and_construct_task_from_redis_reply(redisReply *reply) {
   if (reply->type == REDIS_REPLY_NIL) {
     /* There is no task in the reply, so return NULL. */
     task = NULL;
-  } else if (reply->type == REDIS_REPLY_ARRAY) {
-    /* Check that the reply is as expected. The 0th element is the scheduling
-     * state. The 1st element is the db_client_id of the associated local
-     * scheduler, and the 2nd element is the TaskSpec. */
-    CHECK(reply->elements == 3);
-    CHECK(reply->element[0]->type == REDIS_REPLY_INTEGER);
-    CHECK(reply->element[1]->type == REDIS_REPLY_STRING);
-    CHECK(reply->element[2]->type == REDIS_REPLY_STRING);
-    /* Parse the scheduling state. */
-    long long state = reply->element[0]->integer;
-    /* Parse the local scheduler db_client_id. */
-    DBClientID local_scheduler_id;
-    CHECK(sizeof(local_scheduler_id) == reply->element[1]->len);
-    memcpy(local_scheduler_id.id, reply->element[1]->str,
-           reply->element[1]->len);
-    /* Parse the task spec. */
-    TaskSpec *spec = (TaskSpec *) malloc(reply->element[2]->len);
-    memcpy(spec, reply->element[2]->str, reply->element[2]->len);
-    task = Task_alloc(spec, reply->element[2]->len, state, local_scheduler_id);
-    /* Free the task spec. */
-    TaskSpec_free(spec);
+  } else if (reply->type == REDIS_REPLY_STRING) {
+    /* The reply is a flatbuffer TaskReply object. Parse it and construct the
+     * task. */
+    auto message = flatbuffers::GetRoot<TaskReply>(reply->str);
+    TaskSpec *spec = (TaskSpec *) message->task_spec()->data();
+    int64_t task_spec_size = message->task_spec()->size();
+    task = Task_alloc(spec, task_spec_size, message->state(),
+                      from_flatbuf(message->local_scheduler_id()));
   } else {
     LOG_FATAL("Unexpected reply type %d", reply->type);
   }
@@ -502,81 +493,6 @@ void redis_object_table_lookup_callback(redisAsyncContext *c,
   }
 }
 
-/**
- * This will parse a payload string published on the object notification
- * channel. The string must have the format:
- *
- *    <object id> MANAGERS <manager id1> <manager id2> ...
- *
- * where there may be any positive number of manager IDs.
- *
- * @param db The db handle.
- * @param payload The payload string.
- * @param length The length of the string.
- * @param manager_count This method will write the number of managers at this
- *        address.
- * @param manager_vector This method will allocate an array of pointers to
- *        manager addresses and write the address of the array at this address.
- *        The caller is responsible for freeing this array.
- * @return The object ID that the notification is about.
- */
-ObjectID parse_subscribe_to_notifications_payload(
-    DBHandle *db,
-    char *payload,
-    int length,
-    int64_t *data_size,
-    int *manager_count,
-    const char ***manager_vector) {
-  long long data_size_value = 0;
-  int num_managers = (length - sizeof(ObjectID) - 1 - sizeof(data_size_value) -
-                      1 - strlen("MANAGERS")) /
-                     (1 + sizeof(DBClientID));
-
-  int64_t rval = sizeof(ObjectID) + 1 + sizeof(data_size_value) + 1 +
-                 strlen("MANAGERS") + num_managers * (1 + sizeof(DBClientID));
-
-  CHECKM(length == rval,
-         "length mismatch: num_managers = %d, length = %d, rval = %" PRId64,
-         num_managers, length, rval);
-  CHECK(num_managers > 0);
-  ObjectID obj_id;
-  /* Track our current offset in the payload. */
-  int offset = 0;
-  /* Parse the object ID. */
-  memcpy(&obj_id.id, &payload[offset], sizeof(obj_id.id));
-  offset += sizeof(obj_id.id);
-  /* The next part of the payload is a space. */
-  const char *space_str = " ";
-  CHECK(memcmp(&payload[offset], space_str, strlen(space_str)) == 0);
-  offset += strlen(space_str);
-  /* The next part of the payload is binary data_size. */
-  memcpy(&data_size_value, &payload[offset], sizeof(data_size_value));
-  offset += sizeof(data_size_value);
-  /* The next part of the payload is the string " MANAGERS" with leading ' '. */
-  const char *managers_str = " MANAGERS";
-  CHECK(memcmp(&payload[offset], managers_str, strlen(managers_str)) == 0);
-  offset += strlen(managers_str);
-  /* Parse the managers. */
-  const char **managers = (const char **) malloc(num_managers * sizeof(char *));
-  for (int i = 0; i < num_managers; ++i) {
-    /* First there is a space. */
-    CHECK(memcmp(&payload[offset], " ", strlen(" ")) == 0);
-    offset += strlen(" ");
-    /* Get the manager ID. */
-    DBClientID manager_id;
-    memcpy(&manager_id.id, &payload[offset], sizeof(manager_id.id));
-    offset += sizeof(manager_id.id);
-    /* Write the address of the corresponding manager to the returned array. */
-    redis_get_cached_db_client(db, manager_id, &managers[i]);
-  }
-  CHECK(offset == length);
-  /* Return the manager array and the object ID. */
-  *manager_count = num_managers;
-  *manager_vector = managers;
-  *data_size = data_size_value;
-  return obj_id;
-}
-
 void object_table_redis_subscribe_to_notifications_callback(
     redisAsyncContext *c,
     void *r,
@@ -603,13 +519,22 @@ void object_table_redis_subscribe_to_notifications_callback(
             message_type->str);
 
   if (strcmp(message_type->str, "message") == 0) {
-    /* Handle an object notification. */
-    int64_t data_size = 0;
-    int manager_count;
-    const char **manager_vector;
-    ObjectID obj_id = parse_subscribe_to_notifications_payload(
-        db, reply->element[2]->str, reply->element[2]->len, &data_size,
-        &manager_count, &manager_vector);
+    /* We received an object notification. Parse the payload. */
+    auto message = flatbuffers::GetRoot<SubscribeToNotificationsReply>(
+        reply->element[2]->str);
+    /* Extract the object ID. */
+    ObjectID obj_id = from_flatbuf(message->object_id());
+    /* Extract the data size. */
+    int64_t data_size = message->object_size();
+    int manager_count = message->manager_ids()->size();
+    /* Construct the manager vector from the flatbuffers object. */
+    const char **manager_vector =
+        (const char **) malloc(manager_count * sizeof(char *));
+    for (int i = 0; i < manager_count; ++i) {
+      DBClientID manager_id = from_flatbuf(message->manager_ids()->Get(i));
+      redis_get_cached_db_client(db, manager_id, &manager_vector[i]);
+    }
+
     /* Call the subscribe callback. */
     ObjectTableSubscribeData *data =
         (ObjectTableSubscribeData *) callback_data->data;
@@ -641,7 +566,7 @@ void redis_object_table_subscribe_to_notifications(
     TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   /* The object channel prefix must match the value defined in
-   * src/common/redismodule/ray_redis_module.c. */
+   * src/common/redismodule/ray_redis_module.cc. */
   const char *object_channel_prefix = "OC:";
   const char *object_channel_bcast = "BCAST";
   int status = REDIS_OK;
@@ -869,47 +794,6 @@ void redis_task_table_test_and_update(TableCallbackData *callback_data) {
   }
 }
 
-/* The format of the payload is described in ray_redis_module.c and is
- * "<task ID> <state> <local scheduler ID> <task specification>". TODO(rkn):
- * Make this code nicer. */
-void parse_task_table_subscribe_callback(char *payload,
-                                         int length,
-                                         TaskID *task_id,
-                                         int *state,
-                                         DBClientID *local_scheduler_id,
-                                         TaskSpec **spec,
-                                         int64_t *task_spec_size) {
-  /* Note that the state is padded with spaces to consist of precisely two
-   * characters. */
-  int task_spec_payload_size =
-      length - sizeof(*task_id) - 1 - 2 - 1 - sizeof(*local_scheduler_id) - 1;
-  int offset = 0;
-  /* Read in the task ID. */
-  memcpy(task_id, &payload[offset], sizeof(*task_id));
-  offset += sizeof(*task_id);
-  /* Read in a space. */
-  const char *space_str = (const char *) " ";
-  CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0);
-  offset += strlen(space_str);
-  /* Read in the state, which is an integer left-padded with spaces to two
-   * characters. */
-  CHECK(sscanf(&payload[offset], "%2d", state) == 1);
-  offset += 2;
-  /* Read in a space. */
-  CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0);
-  offset += strlen(space_str);
-  /* Read in the local scheduler ID. */
-  memcpy(local_scheduler_id, &payload[offset], sizeof(*local_scheduler_id));
-  offset += sizeof(*local_scheduler_id);
-  /* Read in a space. */
-  CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0);
-  offset += strlen(space_str);
-  /* Read in the task spec. */
-  *spec = (TaskSpec *) malloc(task_spec_payload_size);
-  memcpy(*spec, &payload[offset], task_spec_payload_size);
-  *task_spec_size = task_spec_payload_size;
-}
-
 void redis_task_table_subscribe_callback(redisAsyncContext *c,
                                          void *r,
                                          void *privdata) {
@@ -917,7 +801,7 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c,
   redisReply *reply = (redisReply *) r;
 
   CHECK(reply->type == REDIS_REPLY_ARRAY);
-  /* The number of elements is 3 for a reply to SUBSCRIBE, and 4  for a reply to
+  /* The number of elements is 3 for a reply to SUBSCRIBE, and 4 for a reply to
    * PSUBSCRIBE. */
   CHECKM(reply->elements == 3 || reply->elements == 4, "reply->elements is %zu",
          reply->elements);
@@ -929,20 +813,22 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c,
   if (strcmp(message_type->str, "message") == 0 ||
       strcmp(message_type->str, "pmessage") == 0) {
     /* Handle a task table event. Parse the payload and call the callback. */
+    auto message = flatbuffers::GetRoot<TaskReply>(payload->str);
+    /* Extract the task ID. */
+    TaskID task_id = from_flatbuf(message->task_id());
+    /* Extract the scheduling state. */
+    int64_t state = message->state();
+    /* Extract the local scheduler ID. */
+    DBClientID local_scheduler_id = from_flatbuf(message->local_scheduler_id());
+    /* Extract the task spec. */
+    TaskSpec *spec = (TaskSpec *) message->task_spec()->data();
+    int64_t task_spec_size = message->task_spec()->size();
+    /* Create a task. */
+    Task *task = Task_alloc(spec, task_spec_size, state, local_scheduler_id);
+
+    /* Call the subscribe callback if there is one. */
     TaskTableSubscribeData *data =
         (TaskTableSubscribeData *) callback_data->data;
-    /* Read out the information from the payload. */
-    TaskID task_id;
-    int state;
-    DBClientID local_scheduler_id;
-    TaskSpec *spec;
-    int64_t task_spec_size;
-    parse_task_table_subscribe_callback(payload->str, payload->len, &task_id,
-                                        &state, &local_scheduler_id, &spec,
-                                        &task_spec_size);
-    Task *task = Task_alloc(spec, task_spec_size, state, local_scheduler_id);
-    TaskSpec_free(spec);
-    /* Call the subscribe callback if there is one. */
     if (data->subscribe_callback != NULL) {
       data->subscribe_callback(task, data->subscribe_context);
     }
@@ -969,24 +855,24 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c,
 void redis_task_table_subscribe(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   TaskTableSubscribeData *data = (TaskTableSubscribeData *) callback_data->data;
-  /* TASK_CHANNEL_PREFIX is defined in ray_redis_module.c and must be kept in
+  /* TASK_CHANNEL_PREFIX is defined in ray_redis_module.cc and must be kept in
    * sync with that file. */
   const char *TASK_CHANNEL_PREFIX = "TT:";
   int status;
   if (IS_NIL_ID(data->local_scheduler_id)) {
     /* TODO(swang): Implement the state_filter by translating the bitmask into
      * a Redis key-matching pattern. */
-    status = redisAsyncCommand(
-        db->sub_context, redis_task_table_subscribe_callback,
-        (void *) callback_data->timer_id, "PSUBSCRIBE %s*:%2d",
-        TASK_CHANNEL_PREFIX, data->state_filter);
+    status =
+        redisAsyncCommand(db->sub_context, redis_task_table_subscribe_callback,
+                          (void *) callback_data->timer_id, "PSUBSCRIBE %s*:%d",
+                          TASK_CHANNEL_PREFIX, data->state_filter);
   } else {
     DBClientID local_scheduler_id = data->local_scheduler_id;
-    status = redisAsyncCommand(
-        db->sub_context, redis_task_table_subscribe_callback,
-        (void *) callback_data->timer_id, "SUBSCRIBE %s%b:%2d",
-        TASK_CHANNEL_PREFIX, (char *) local_scheduler_id.id,
-        sizeof(local_scheduler_id.id), data->state_filter);
+    status =
+        redisAsyncCommand(db->sub_context, redis_task_table_subscribe_callback,
+                          (void *) callback_data->timer_id, "SUBSCRIBE %s%b:%d",
+                          TASK_CHANNEL_PREFIX, (char *) local_scheduler_id.id,
+                          sizeof(local_scheduler_id.id), data->state_filter);
   }
   if ((status == REDIS_ERR) || db->sub_context->err) {
     LOG_REDIS_DEBUG(db->sub_context, "error in redis_task_table_subscribe");
@@ -1052,40 +938,23 @@ void redis_db_client_table_subscribe_callback(redisAsyncContext *c,
     return;
   }
   /* Otherwise, parse the payload and call the callback. */
-  DBClientTableSubscribeData *data =
-      (DBClientTableSubscribeData *) callback_data->data;
-  DBClientID client;
-  memcpy(client.id, payload->str, sizeof(client.id));
-  /* We subtract 1 + sizeof(client.id) to compute the length of the
-   * client_type string, and we add 1 to null-terminate the string. */
-  int client_type_length = payload->len - 1 - sizeof(client.id) + 1;
-  CHECK(client_type_length > 0);
+  auto message =
+      flatbuffers::GetRoot<SubscribeToDBClientTableReply>(payload->str);
+  DBClientID client = from_flatbuf(message->db_client_id());
 
   /* Parse the client type and auxiliary address from the response. If there is
    * only client type, then the update was a delete. */
-  char *client_type = (char *) malloc(client_type_length);
-  char *aux_address = (char *) malloc(client_type_length);
-  int is_insertion;
-  memset(aux_address, 0, client_type_length);
-  /* Published message format: <client_id:client_type aux_addr> */
-  int rv = sscanf(&payload->str[1 + sizeof(client.id)], "%s %s %d", client_type,
-                  aux_address, &is_insertion);
-  CHECKM(rv == 3,
-         "redis_db_client_table_subscribe_callback: expected 2 parsed args, "
-         "Got %d instead.",
-         rv);
-  CHECKM(is_insertion == 1 || is_insertion == 0,
-         "redis_db_client_table_subscribe_callback: expected 0 or 1 for "
-         "insertion field, got %d instead.",
-         is_insertion);
+  char *client_type = (char *) message->client_type()->data();
+  char *aux_address = (char *) message->aux_address()->data();
+  bool is_insertion = message->is_insertion();
 
   /* Call the subscription callback. */
+  DBClientTableSubscribeData *data =
+      (DBClientTableSubscribeData *) callback_data->data;
   if (data->subscribe_callback) {
-    data->subscribe_callback(client, client_type, aux_address,
-                             (bool) is_insertion, data->subscribe_context);
+    data->subscribe_callback(client, client_type, aux_address, is_insertion,
+                             data->subscribe_context);
   }
-  free(client_type);
-  free(aux_address);
 }
 
 void redis_db_client_table_subscribe(TableCallbackData *callback_data) {
@@ -1114,15 +983,26 @@ void redis_local_scheduler_table_subscribe_callback(redisAsyncContext *c,
   if (strcmp(message_type->str, "message") == 0) {
     /* Handle a local scheduler heartbeat. Parse the payload and call the
      * subscribe callback. */
-    redisReply *payload = reply->element[2];
+    auto message =
+        flatbuffers::GetRoot<LocalSchedulerInfoMessage>(reply->element[2]->str);
+
+    /* Extract the client ID. */
+    DBClientID client_id = from_flatbuf(message->db_client_id());
+    /* Extract the fields of the local scheduler info struct. */
+    LocalSchedulerInfo info;
+    info.total_num_workers = message->total_num_workers();
+    info.task_queue_length = message->task_queue_length();
+    info.available_workers = message->available_workers();
+    for (int i = 0; i < ResourceIndex_MAX; ++i) {
+      info.static_resources[i] = message->static_resources()->Get(i);
+    }
+    for (int i = 0; i < ResourceIndex_MAX; ++i) {
+      info.dynamic_resources[i] = message->dynamic_resources()->Get(i);
+    }
+
+    /* Call the subscribe callback. */
     LocalSchedulerTableSubscribeData *data =
         (LocalSchedulerTableSubscribeData *) callback_data->data;
-    DBClientID client_id;
-    LocalSchedulerInfo info;
-    /* The payload should be the concatenation of these two structs. */
-    CHECK(sizeof(client_id) + sizeof(info) == payload->len);
-    memcpy(&client_id, payload->str, sizeof(client_id));
-    memcpy(&info, payload->str + sizeof(client_id), sizeof(info));
     if (data->subscribe_callback) {
       data->subscribe_callback(client_id, info, data->subscribe_context);
     }
@@ -1167,10 +1047,22 @@ void redis_local_scheduler_table_send_info(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   LocalSchedulerTableSendInfoData *data =
       (LocalSchedulerTableSendInfoData *) callback_data->data;
+
+  /* Create a flatbuffer object to serialize and publish. */
+  flatbuffers::FlatBufferBuilder fbb;
+  /* Create the flatbuffers message. */
+  LocalSchedulerInfo info = data->info;
+  auto message = CreateLocalSchedulerInfoMessage(
+      fbb, to_flatbuf(fbb, db->client), info.total_num_workers,
+      info.task_queue_length, info.available_workers,
+      fbb.CreateVector(info.static_resources, ResourceIndex_MAX),
+      fbb.CreateVector(info.dynamic_resources, ResourceIndex_MAX));
+  fbb.Finish(message);
+
   int status = redisAsyncCommand(
       db->context, redis_local_scheduler_table_send_info_callback,
-      (void *) callback_data->timer_id, "PUBLISH local_schedulers %b%b",
-      db->client.id, sizeof(db->client.id), &data->info, sizeof(data->info));
+      (void *) callback_data->timer_id, "PUBLISH local_schedulers %b",
+      fbb.GetBufferPointer(), fbb.GetSize());
   if ((status == REDIS_ERR) || db->context->err) {
     LOG_REDIS_DEBUG(db->context,
                     "error in redis_local_scheduler_table_send_info");