Start integrating new GCS APIs (#1379)

* Start integrating new GCS calls

* fixes

* tests

* cleanup

* cleanup and valgrind fix

* update tests

* fix valgrind

* fix more valgrind

* fixes

* add separate tests for GCS

* fix linting

* update tests

* cleanup

* fix python linting

* more fixes

* fix linting

* add plasma manager callback

* add some documentation

* fix linting

* fix linting

* fixes

* update

* fix linting

* fix

* add spillback count

* fixes

* linting

* fixes

* fix linting

* fix

* fix

* fix

src/common/redis_module/ray_redis_module.cc

@@ -3,10 +3,11 @@
 
 #include "redis_string.h"
 
-#include "format/common_generated.h"
-#include "task.h"
-
 #include "common_protocol.h"
+#include "format/common_generated.h"
+#include "ray/gcs/format/gcs_generated.h"
+#include "ray/id.h"
+#include "task.h"
 
 // Various tables are maintained in redis:
 //
@@ -406,6 +407,46 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
   RedisModule_StringSet(key, data);
   RedisModule_CloseKey(key);
 
+  size_t len = 0;
+  const char *buf = RedisModule_StringPtrLen(data, &len);
+
+  auto message = flatbuffers::GetRoot<TaskTableData>(buf);
+
+  if (message->scheduling_state() == SchedulingState_WAITING ||
+      message->scheduling_state() == SchedulingState_SCHEDULED) {
+    /* Build the PUBLISH topic and message for task table subscribers. The topic
+     * is a string in the format "TASK_PREFIX:<local scheduler ID>:<state>". The
+     * message is a serialized SubscribeToTasksReply flatbuffer object. */
+    std::string state = std::to_string(message->scheduling_state());
+    RedisModuleString *publish_topic = RedisString_Format(
+        ctx, "%s%b:%s", TASK_PREFIX, message->scheduler_id()->str().data(),
+        sizeof(DBClientID), state.c_str());
+
+    /* Construct the flatbuffers object for the payload. */
+    flatbuffers::FlatBufferBuilder fbb;
+    /* Create the flatbuffers message. */
+    auto msg = CreateTaskReply(
+        fbb, RedisStringToFlatbuf(fbb, id), message->scheduling_state(),
+        fbb.CreateString(message->scheduler_id()),
+        fbb.CreateString(message->execution_dependencies()),
+        fbb.CreateString(message->task_info()), message->spillback_count(),
+        true /* not used */);
+    fbb.Finish(msg);
+
+    RedisModuleString *publish_message = RedisModule_CreateString(
+        ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize());
+
+    RedisModuleCallReply *reply =
+        RedisModule_Call(ctx, "PUBLISH", "ss", publish_topic, publish_message);
+
+    /* See how many clients received this publish. */
+    long long num_clients = RedisModule_CallReplyInteger(reply);
+    CHECKM(num_clients <= 1, "Published to %lld clients.", num_clients);
+
+    RedisModule_FreeString(ctx, publish_message);
+    RedisModule_FreeString(ctx, publish_topic);
+  }
+
   return RedisModule_ReplyWithSimpleString(ctx, "OK");
 }
 
@@ -431,6 +472,63 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx,
   return REDISMODULE_OK;
 }
 
+bool is_nil(const std::string &data) {
+  CHECK(data.size() == kUniqueIDSize);
+  const uint8_t *d = reinterpret_cast<const uint8_t *>(data.data());
+  for (int i = 0; i < kUniqueIDSize; ++i) {
+    if (d[i] != 255) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// This is a temporary redis command that will be removed once
+// the GCS uses https://github.com/pcmoritz/credis.
+// Be careful, this only supports Task Table payloads.
+int TableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
+                                    RedisModuleString **argv,
+                                    int argc) {
+  if (argc != 3) {
+    return RedisModule_WrongArity(ctx);
+  }
+  RedisModuleString *id = argv[1];
+  RedisModuleString *update_data = argv[2];
+
+  RedisModuleKey *key =
+      OpenPrefixedKey(ctx, "T:", id, REDISMODULE_READ | REDISMODULE_WRITE);
+
+  size_t value_len = 0;
+  char *value_buf = RedisModule_StringDMA(key, &value_len, REDISMODULE_READ);
+
+  size_t update_len = 0;
+  const char *update_buf = RedisModule_StringPtrLen(update_data, &update_len);
+
+  auto data = flatbuffers::GetMutableRoot<TaskTableData>(
+      reinterpret_cast<void *>(value_buf));
+
+  auto update = flatbuffers::GetRoot<TaskTableTestAndUpdate>(update_buf);
+
+  bool do_update = data->scheduling_state() & update->test_state_bitmask();
+
+  if (!is_nil(update->test_scheduler_id()->str())) {
+    do_update =
+        do_update &&
+        update->test_scheduler_id()->str() == data->scheduler_id()->str();
+  }
+
+  if (do_update) {
+    CHECK(data->mutate_scheduling_state(update->update_state()));
+  }
+  CHECK(data->mutate_updated(do_update));
+
+  int result = RedisModule_ReplyWithStringBuffer(ctx, value_buf, value_len);
+
+  RedisModule_CloseKey(key);
+
+  return result;
+}
+
 /**
  * Add a new entry to the object table or update an existing one.
  *
@@ -1239,6 +1337,12 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
     return REDISMODULE_ERR;
   }
 
+  if (RedisModule_CreateCommand(ctx, "ray.table_test_and_update",
+                                TableTestAndUpdate_RedisCommand, "write", 0, 0,
+                                0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
   if (RedisModule_CreateCommand(ctx, "ray.object_table_lookup",
                                 ObjectTableLookup_RedisCommand, "readonly", 0,
                                 0, 0) == REDISMODULE_ERR) {

src/common/redis_module/redis_string.h

@@ -43,6 +43,12 @@ RedisModuleString *RedisString_Format(RedisModuleCtx *ctx,
         RedisModule_StringAppendBuffer(ctx, result, s, strlen(s));
         i += 1;
         break;
+      case 'b':
+        s = va_arg(ap, const char *);
+        l = va_arg(ap, size_t);
+        RedisModule_StringAppendBuffer(ctx, result, s, l);
+        i += 1;
+        break;
       default: /* Handle %% and generally %<unknown>. */
         RedisModule_StringAppendBuffer(ctx, result, &next, 1);
         i += 1;

src/common/state/redis.cc

@@ -1161,7 +1161,13 @@ void redis_task_table_subscribe(TableCallbackData *callback_data) {
   /* 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:";
+#if !RAY_USE_NEW_GCS
   for (auto subscribe_context : db->subscribe_contexts) {
+#else
+  /* In the new code path, subscriptions currently go through the
+   * primary redis shard. */
+  for (auto subscribe_context : {db->subscribe_context}) {
+#endif
     int status;
     if (data->local_scheduler_id.is_nil()) {
       /* TODO(swang): Implement the state_filter by translating the bitmask into

src/common/task.cc

@@ -366,7 +366,7 @@ void TaskSpec_free(TaskSpec *spec) {
 
 TaskExecutionSpec::TaskExecutionSpec(
     const std::vector<ObjectID> &execution_dependencies,
-    TaskSpec *spec,
+    const TaskSpec *spec,
     int64_t task_spec_size,
     int spillback_count)
     : execution_dependencies_(execution_dependencies),
@@ -380,7 +380,7 @@ TaskExecutionSpec::TaskExecutionSpec(
 
 TaskExecutionSpec::TaskExecutionSpec(
     const std::vector<ObjectID> &execution_dependencies,
-    TaskSpec *spec,
+    const TaskSpec *spec,
     int64_t task_spec_size)
     : TaskExecutionSpec(execution_dependencies, spec, task_spec_size, 0) {}
 
@@ -394,7 +394,7 @@ TaskExecutionSpec::TaskExecutionSpec(TaskExecutionSpec *other)
   spec_ = std::unique_ptr<TaskSpec[]>(spec_copy);
 }
 
-std::vector<ObjectID> TaskExecutionSpec::ExecutionDependencies() {
+std::vector<ObjectID> TaskExecutionSpec::ExecutionDependencies() const {
   return execution_dependencies_;
 }
 
@@ -423,18 +423,18 @@ void TaskExecutionSpec::SetLastTimeStamp(int64_t new_timestamp) {
   last_timestamp_ = new_timestamp;
 }
 
-TaskSpec *TaskExecutionSpec::Spec() {
+TaskSpec *TaskExecutionSpec::Spec() const {
   return spec_.get();
 }
 
-int64_t TaskExecutionSpec::NumDependencies() {
+int64_t TaskExecutionSpec::NumDependencies() const {
   TaskSpec *spec = Spec();
   int64_t num_dependencies = TaskSpec_num_args(spec);
   num_dependencies += execution_dependencies_.size();
   return num_dependencies;
 }
 
-int TaskExecutionSpec::DependencyIdCount(int64_t dependency_index) {
+int TaskExecutionSpec::DependencyIdCount(int64_t dependency_index) const {
   TaskSpec *spec = Spec();
   /* The first dependencies are the arguments of the task itself, followed by
    * the execution dependencies. Find the total number of task arguments so
@@ -453,7 +453,7 @@ int TaskExecutionSpec::DependencyIdCount(int64_t dependency_index) {
 }
 
 ObjectID TaskExecutionSpec::DependencyId(int64_t dependency_index,
-                                         int64_t id_index) {
+                                         int64_t id_index) const {
   TaskSpec *spec = Spec();
   /* The first dependencies are the arguments of the task itself, followed by
    * the execution dependencies. Find the total number of task arguments so
@@ -470,7 +470,7 @@ ObjectID TaskExecutionSpec::DependencyId(int64_t dependency_index,
   }
 }
 
-bool TaskExecutionSpec::DependsOn(ObjectID object_id) {
+bool TaskExecutionSpec::DependsOn(ObjectID object_id) const {
   // Iterate through the task arguments to see if it contains object_id.
   TaskSpec *spec = Spec();
   int64_t num_args = TaskSpec_num_args(spec);
@@ -494,7 +494,7 @@ bool TaskExecutionSpec::DependsOn(ObjectID object_id) {
   return false;
 }
 
-bool TaskExecutionSpec::IsStaticDependency(int64_t dependency_index) {
+bool TaskExecutionSpec::IsStaticDependency(int64_t dependency_index) const {
   TaskSpec *spec = Spec();
   /* The first dependencies are the arguments of the task itself, followed by
    * the execution dependencies. If the requested dependency index is a task
@@ -505,7 +505,7 @@ bool TaskExecutionSpec::IsStaticDependency(int64_t dependency_index) {
 
 /* TASK INSTANCES */
 
-Task *Task_alloc(TaskSpec *spec,
+Task *Task_alloc(const TaskSpec *spec,
                  int64_t task_spec_size,
                  int state,
                  DBClientID local_scheduler_id,

src/common/task.h

@@ -18,10 +18,10 @@ typedef char TaskSpec;
 class TaskExecutionSpec {
  public:
   TaskExecutionSpec(const std::vector<ObjectID> &execution_dependencies,
-                    TaskSpec *spec,
+                    const TaskSpec *spec,
                     int64_t task_spec_size);
   TaskExecutionSpec(const std::vector<ObjectID> &execution_dependencies,
-                    TaskSpec *spec,
+                    const TaskSpec *spec,
                     int64_t task_spec_size,
                     int spillback_count);
   TaskExecutionSpec(TaskExecutionSpec *execution_spec);
@@ -30,7 +30,7 @@ class TaskExecutionSpec {
   ///
   /// @return A vector of object IDs representing this task's execution
   ///         dependencies.
-  std::vector<ObjectID> ExecutionDependencies();
+  std::vector<ObjectID> ExecutionDependencies() const;
 
   /// Set the task's execution dependencies.
   ///
@@ -70,33 +70,33 @@ class TaskExecutionSpec {
   /// Get the task spec.
   ///
   /// @return A pointer to the immutable task spec.
-  TaskSpec *Spec();
+  TaskSpec *Spec() const;
 
   /// Get the number of dependencies. This comprises the immutable task
   /// arguments and the mutable execution dependencies.
   ///
   /// @return The number of dependencies.
-  int64_t NumDependencies();
+  int64_t NumDependencies() const;
 
   /// Get the number of object IDs at the given dependency index.
   ///
   /// @param dependency_index The dependency index whose object IDs to count.
   /// @return The number of object IDs at the given dependency_index.
-  int DependencyIdCount(int64_t dependency_index);
+  int DependencyIdCount(int64_t dependency_index) const;
 
   /// Get the object ID of a given dependency index.
   ///
   /// @param dependency_index The index at which we should look up the object
   ///        ID.
   /// @param id_index The index of the object ID.
-  ObjectID DependencyId(int64_t dependency_index, int64_t id_index);
+  ObjectID DependencyId(int64_t dependency_index, int64_t id_index) const;
 
   /// Compute whether the task is dependent on an object ID.
   ///
   /// @param object_id The object ID that the task may be dependent on.
   /// @return bool This returns true if the task is dependent on the given
   ///         object ID and false otherwise.
-  bool DependsOn(ObjectID object_id);
+  bool DependsOn(ObjectID object_id) const;
 
   /// Returns whether the given dependency index is a static dependency (an
   /// argument of the immutable task).
@@ -104,7 +104,7 @@ class TaskExecutionSpec {
   /// @param dependency_index The requested dependency index.
   /// @return bool This returns true if the requested dependency index is
   ///         immutable (an argument of the task).
-  bool IsStaticDependency(int64_t dependency_index);
+  bool IsStaticDependency(int64_t dependency_index) const;
 
  private:
   /** A list of object IDs representing this task's dependencies at execution
@@ -532,7 +532,7 @@ struct Task {
  * @param local_scheduler_id The ID of the local scheduler that the task is
  *        scheduled on, if any.
  */
-Task *Task_alloc(TaskSpec *spec,
+Task *Task_alloc(const TaskSpec *spec,
                  int64_t task_spec_size,
                  int state,
                  DBClientID local_scheduler_id,

src/common/test/run_tests.sh

@@ -13,11 +13,14 @@ sleep 1s
 ./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
 ./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
 
-./src/common/db_tests
-./src/common/io_tests
-./src/common/task_tests
-./src/common/redis_tests
-./src/common/task_table_tests
-./src/common/object_table_tests
+if [ -z "$RAY_USE_NEW_GCS" ]; then
+  ./src/common/db_tests
+  ./src/common/io_tests
+  ./src/common/task_tests
+  ./src/common/redis_tests
+  ./src/common/task_table_tests
+  ./src/common/object_table_tests
+fi
+
 ./src/common/thirdparty/redis/src/redis-cli -p 6379 shutdown
 ./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown

src/common/test/run_valgrind.sh

@@ -15,12 +15,14 @@ sleep 1s
 ./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
 ./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
 
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/db_tests
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/io_tests
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_tests
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/redis_tests
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_table_tests
-valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/object_table_tests
+if [ -z "$RAY_USE_NEW_GCS" ]; then
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/db_tests
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/io_tests
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_tests
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/redis_tests
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_table_tests
+  valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/object_table_tests
+fi
 
 ./src/common/thirdparty/redis/src/redis-cli shutdown
 ./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown