Introduce set data structure in GCS

python/ray/experimental/state.py

@@ -243,14 +243,8 @@ def _object_table(self, object_id):
         object_info = {
             "DataSize": entry.ObjectSize(),
             "Manager": entry.Manager(),
-            "IsEviction": [entry.IsEviction()],
         }
 
-        for i in range(1, gcs_entry.EntriesLength()):
-            entry = ray.gcs_utils.ObjectTableData.GetRootAsObjectTableData(
-                gcs_entry.Entries(i), 0)
-            object_info["IsEviction"].append(entry.IsEviction())
-
         return object_info
 
     def object_table(self, object_id=None):

python/ray/tests/test_basic.py

@@ -2467,10 +2467,6 @@ def wait_for_object_table():
     object_table = ray.global_state.object_table()
     assert len(object_table) == 2
 
-    assert object_table[x_id]["IsEviction"][0] is False
-
-    assert object_table[result_id]["IsEviction"][0] is False
-
     assert object_table[x_id] == ray.global_state.object_table(x_id)
     object_table_entry = ray.global_state.object_table(result_id)
     assert object_table[result_id] == object_table_entry

src/ray/gcs/client.cc

@@ -112,7 +112,7 @@ AsyncGcsClient::AsyncGcsClient(const std::string &address, int port,
   driver_table_.reset(new DriverTable({primary_context_}, this));
   heartbeat_batch_table_.reset(new HeartbeatBatchTable({primary_context_}, this));
   // Tables below would be sharded.
-  object_table_.reset(new ObjectTable(shard_contexts_, this, command_type));
+  object_table_.reset(new ObjectTable(shard_contexts_, this));
   raylet_task_table_.reset(new raylet::TaskTable(shard_contexts_, this, command_type));
   task_reconstruction_log_.reset(new TaskReconstructionLog(shard_contexts_, this));
   task_lease_table_.reset(new TaskLeaseTable(shard_contexts_, this));

src/ray/gcs/format/gcs.fbs

@@ -108,7 +108,13 @@ table ResourcePair {
   value: double;
 }
 
+enum GcsTableNotificationMode:int {
+  APPEND_OR_ADD = 0,
+  REMOVE,
+}
+
 table GcsTableEntry {
+  mode: GcsTableNotificationMode;
   id: string;
   entries: [string];
 }
@@ -124,8 +130,6 @@ table ObjectTableData {
   object_size: long;
   // The node manager ID that this object appeared on or was evicted by.
   manager: string;
-  // Whether this entry is an addition or a deletion.
-  is_eviction: bool;
 }
 
 table TaskReconstructionData {

src/ray/gcs/redis_module/ray_redis_module.cc

@@ -181,22 +181,24 @@ flatbuffers::Offset<flatbuffers::String> RedisStringToFlatbuf(
   return fbb.CreateString(redis_string_str, redis_string_size);
 }
 
-/// Publish a notification for a new entry at a key. This publishes a
+/// Publish a notification for an entry update at a key. This publishes a
 /// notification to all subscribers of the table, as well as every client that
 /// has requested notifications for this key.
 ///
 /// \param pubsub_channel_str The pubsub channel name that notifications for
 ///        this key should be published to. When publishing to a specific
 ///        client, the channel name should be <pubsub_channel>:<client_id>.
 /// \param id The ID of the key that the notification is about.
-/// \param data The data to publish.
+/// \param mode the update mode, such as append or remove.
+/// \param data The appended/removed data.
 /// \return OK if there is no error during a publish.
-int PublishTableAdd(RedisModuleCtx *ctx, RedisModuleString *pubsub_channel_str,
-                    RedisModuleString *id, RedisModuleString *data) {
+int PublishTableUpdate(RedisModuleCtx *ctx, RedisModuleString *pubsub_channel_str,
+                       RedisModuleString *id, GcsTableNotificationMode mode,
+                       RedisModuleString *data) {
   // Serialize the notification to send.
   flatbuffers::FlatBufferBuilder fbb;
   auto data_flatbuf = RedisStringToFlatbuf(fbb, data);
-  auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, id),
+  auto message = CreateGcsTableEntry(fbb, mode, RedisStringToFlatbuf(fbb, id),
                                      fbb.CreateVector(&data_flatbuf, 1));
   fbb.Finish(message);
 
@@ -265,7 +267,8 @@ int TableAdd_DoPublish(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
 
   if (pubsub_channel != TablePubsub::NO_PUBLISH) {
     // All other pubsub channels write the data back directly onto the channel.
-    return PublishTableAdd(ctx, pubsub_channel_str, id, data);
+    return PublishTableUpdate(ctx, pubsub_channel_str, id,
+                              GcsTableNotificationMode::APPEND_OR_ADD, data);
   } else {
     return RedisModule_ReplyWithSimpleString(ctx, "OK");
   }
@@ -364,7 +367,8 @@ int TableAppend_DoPublish(RedisModuleCtx *ctx, RedisModuleString **argv, int /*a
   if (pubsub_channel != TablePubsub::NO_PUBLISH) {
     // All other pubsub channels write the data back directly onto the
     // channel.
-    return PublishTableAdd(ctx, pubsub_channel_str, id, data);
+    return PublishTableUpdate(ctx, pubsub_channel_str, id,
+                              GcsTableNotificationMode::APPEND_OR_ADD, data);
   } else {
     return RedisModule_ReplyWithSimpleString(ctx, "OK");
   }
@@ -407,6 +411,110 @@ int ChainTableAppend_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
 }
 #endif
 
+int Set_DoPublish(RedisModuleCtx *ctx, RedisModuleString **argv, bool is_add) {
+  RedisModuleString *pubsub_channel_str = argv[2];
+  RedisModuleString *id = argv[3];
+  RedisModuleString *data = argv[4];
+  // Publish a message on the requested pubsub channel if necessary.
+  TablePubsub pubsub_channel;
+  REPLY_AND_RETURN_IF_NOT_OK(ParseTablePubsub(&pubsub_channel, pubsub_channel_str));
+  if (pubsub_channel != TablePubsub::NO_PUBLISH) {
+    // All other pubsub channels write the data back directly onto the
+    // channel.
+    return PublishTableUpdate(ctx, pubsub_channel_str, id,
+                              is_add ? GcsTableNotificationMode::APPEND_OR_ADD
+                                     : GcsTableNotificationMode::REMOVE,
+                              data);
+  } else {
+    return RedisModule_ReplyWithSimpleString(ctx, "OK");
+  }
+}
+
+int Set_DoWrite(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, bool is_add, bool &changed) {
+  if (argc != 5) {
+    return RedisModule_WrongArity(ctx);
+  }
+
+  RedisModuleString *prefix_str = argv[1];
+  RedisModuleString *id = argv[3];
+  RedisModuleString *data = argv[4];
+
+  RedisModuleString *key_string = PrefixedKeyString(ctx, prefix_str, id);
+  RedisModuleCallReply *reply =
+      RedisModule_Call(ctx, is_add ? "SADD" : "SREM", "ss", key_string, data);
+  if (RedisModule_CallReplyType(reply) != REDISMODULE_REPLY_ERROR) {
+    changed = RedisModule_CallReplyInteger(reply) > 0;
+    if (!is_add) {
+      // try to delete the empty set.
+      RedisModuleKey *key;
+      REPLY_AND_RETURN_IF_NOT_OK(
+          OpenPrefixedKey(&key, ctx, prefix_str, id, REDISMODULE_WRITE));
+      auto size = RedisModule_ValueLength(key);
+      if (size == 0) {
+        REPLY_AND_RETURN_IF_FALSE(RedisModule_DeleteKey(key) == REDISMODULE_OK,
+                                  "Failed to delete empty set.");
+      }
+    }
+    return REDISMODULE_OK;
+  } else {
+    // the SADD/SREM command failed
+    RedisModule_ReplyWithCallReply(ctx, reply);
+    return REDISMODULE_ERR;
+  }
+}
+
+/// Add an entry to the set stored at a key. Publishes a notification about
+/// the update to all subscribers, if a pubsub channel is provided.
+///
+/// This is called from a client with the command:
+//
+///    RAY.SET_ADD <table_prefix> <pubsub_channel> <id> <data>
+///
+/// \param table_prefix The prefix string for keys in this set.
+/// \param pubsub_channel The pubsub channel name that notifications for
+///        this key should be published to. When publishing to a specific
+///        client, the channel name should be <pubsub_channel>:<client_id>.
+/// \param id The ID of the key to add to.
+/// \param data The data to add to the key.
+/// \return OK if the add succeeds, or an error message string if the add
+///         fails.
+int SetAdd_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+  bool changed;
+  if (Set_DoWrite(ctx, argv, argc, /*is_add=*/true, changed) != REDISMODULE_OK) {
+    return REDISMODULE_ERR;
+  }
+  if (changed) {
+    return Set_DoPublish(ctx, argv, /*is_add=*/true);
+  }
+  return REDISMODULE_OK;
+}
+
+/// Remove an entry from the set stored at a key. Publishes a notification about
+/// the update to all subscribers, if a pubsub channel is provided.
+///
+/// This is called from a client with the command:
+//
+///    RAY.SET_REMOVE <table_prefix> <pubsub_channel> <id> <data>
+///
+/// \param table_prefix The prefix string for keys in this table.
+/// \param pubsub_channel The pubsub channel name that notifications for
+///        this key should be published to. When publishing to a specific
+///        client, the channel name should be <pubsub_channel>:<client_id>.
+/// \param id The ID of the key to remove from.
+/// \param data The data to remove from the key.
+/// \return OK if the remove succeeds, or an error message string if the remove
+///         fails.
+int SetRemove_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+  bool changed;
+  if (Set_DoWrite(ctx, argv, argc, /*is_add=*/false, changed) != REDISMODULE_OK) {
+    return REDISMODULE_ERR;
+  }
+  if (changed) {
+    return Set_DoPublish(ctx, argv, /*is_add=*/false);
+  }
+  return REDISMODULE_OK;
+}
+
 /// A helper function to create and finish a GcsTableEntry, based on the
 /// current value or values at the given key.
 ///
@@ -428,22 +536,31 @@ Status TableEntryToFlatbuf(RedisModuleCtx *ctx, RedisModuleKey *table_key,
     size_t data_len = 0;
     char *data_buf = RedisModule_StringDMA(table_key, &data_len, REDISMODULE_READ);
     auto data = fbb.CreateString(data_buf, data_len);
-    auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, entry_id),
+    auto message = CreateGcsTableEntry(fbb, GcsTableNotificationMode::APPEND_OR_ADD,
+                                       RedisStringToFlatbuf(fbb, entry_id),
                                        fbb.CreateVector(&data, 1));
     fbb.Finish(message);
   } break;
-  case REDISMODULE_KEYTYPE_LIST: {
+  case REDISMODULE_KEYTYPE_LIST:
+  case REDISMODULE_KEYTYPE_SET: {
     RedisModule_CloseKey(table_key);
     // Close the key before executing the command. NOTE(swang): According to
     // https://github.com/RedisLabs/RedisModulesSDK/blob/master/API.md, "While
     // a key is open, it should only be accessed via the low level key API."
     RedisModuleString *table_key_str = PrefixedKeyString(ctx, prefix_str, entry_id);
     // TODO(swang): This could potentially be replaced with the native redis
     // server list iterator, once it is implemented for redis modules.
-    RedisModuleCallReply *reply =
-        RedisModule_Call(ctx, "LRANGE", "sll", table_key_str, 0, -1);
+    RedisModuleCallReply *reply = nullptr;
+    switch (key_type) {
+    case REDISMODULE_KEYTYPE_LIST:
+      reply = RedisModule_Call(ctx, "LRANGE", "sll", table_key_str, 0, -1);
+      break;
+    case REDISMODULE_KEYTYPE_SET:
+      reply = RedisModule_Call(ctx, "SMEMBERS", "s", table_key_str);
+      break;
+    }
     // Build the flatbuffer from the set of log entries.
-    if (RedisModule_CallReplyType(reply) != REDISMODULE_REPLY_ARRAY) {
+    if (reply == nullptr || RedisModule_CallReplyType(reply) != REDISMODULE_REPLY_ARRAY) {
       return Status::RedisError("Empty list or wrong type");
     }
     std::vector<flatbuffers::Offset<flatbuffers::String>> data;
@@ -453,13 +570,14 @@ Status TableEntryToFlatbuf(RedisModuleCtx *ctx, RedisModuleKey *table_key,
       const char *element_str = RedisModule_CallReplyStringPtr(element, &len);
       data.push_back(fbb.CreateString(element_str, len));
     }
-    auto message = CreateGcsTableEntry(fbb, RedisStringToFlatbuf(fbb, entry_id),
-                                       fbb.CreateVector(data));
+    auto message =
+        CreateGcsTableEntry(fbb, GcsTableNotificationMode::APPEND_OR_ADD,
+                            RedisStringToFlatbuf(fbb, entry_id), fbb.CreateVector(data));
     fbb.Finish(message);
   } break;
   case REDISMODULE_KEYTYPE_EMPTY: {
     auto message = CreateGcsTableEntry(
-        fbb, RedisStringToFlatbuf(fbb, entry_id),
+        fbb, GcsTableNotificationMode::APPEND_OR_ADD, RedisStringToFlatbuf(fbb, entry_id),
         fbb.CreateVector(std::vector<flatbuffers::Offset<flatbuffers::String>>()));
     fbb.Finish(message);
   } break;
@@ -752,6 +870,8 @@ int DebugString_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int
 // Wrap all Redis commands with Redis' auto memory management.
 AUTO_MEMORY(TableAdd_RedisCommand);
 AUTO_MEMORY(TableAppend_RedisCommand);
+AUTO_MEMORY(SetAdd_RedisCommand);
+AUTO_MEMORY(SetRemove_RedisCommand);
 AUTO_MEMORY(TableLookup_RedisCommand);
 AUTO_MEMORY(TableRequestNotifications_RedisCommand);
 AUTO_MEMORY(TableDelete_RedisCommand);
@@ -785,6 +905,16 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     return REDISMODULE_ERR;
   }
 
+  if (RedisModule_CreateCommand(ctx, "ray.set_add", SetAdd_RedisCommand, "write", 0, 0,
+                                0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
+  if (RedisModule_CreateCommand(ctx, "ray.set_remove", SetRemove_RedisCommand, "write", 0,
+                                0, 0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
   if (RedisModule_CreateCommand(ctx, "ray.table_lookup", TableLookup_RedisCommand,
                                 "readonly", 0, 0, 0) == REDISMODULE_ERR) {
     return REDISMODULE_ERR;

src/ray/gcs/tables.cc

@@ -112,6 +112,20 @@ template <typename ID, typename Data>
 Status Log<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
                                 const Callback &subscribe,
                                 const SubscriptionCallback &done) {
+  auto subscribeWrapper = [subscribe](AsyncGcsClient *client, const ID &id,
+                                      const GcsTableNotificationMode mode,
+                                      const std::vector<DataT> &data) {
+    RAY_CHECK(mode != GcsTableNotificationMode::REMOVE);
+    subscribe(client, id, data);
+  };
+  return Subscribe(job_id, client_id, subscribeWrapper, done);
+}
+
+template <typename ID, typename Data>
+Status Log<ID, Data>::Subscribe(const JobID &job_id,
+                                                    const ClientID &client_id,
+                                                    const NotificationCallback &subscribe,
+                                                    const SubscriptionCallback &done) {
   RAY_CHECK(subscribe_callback_index_ == -1)
       << "Client called Subscribe twice on the same table";
   auto callback = [this, subscribe, done](const std::string &data) {
@@ -137,7 +151,7 @@ Status Log<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
           data_root->UnPackTo(&result);
           results.emplace_back(std::move(result));
         }
-        subscribe(client_, id, results);
+        subscribe(client_, id, root->mode(), results);
       }
     }
     // We do not delete the callback after calling it since there may be
@@ -274,6 +288,50 @@ std::string Table<ID, Data>::DebugString() const {
   return result.str();
 }
 
+template <typename ID, typename Data>
+Status Set<ID, Data>::Add(const JobID &job_id, const ID &id,
+                          std::shared_ptr<DataT> &dataT, const WriteCallback &done) {
+  num_adds_++;
+  auto callback = [this, id, dataT, done](const std::string &data) {
+    if (done != nullptr) {
+      (done)(client_, id, *dataT);
+    }
+    return true;
+  };
+  flatbuffers::FlatBufferBuilder fbb;
+  fbb.ForceDefaults(true);
+  fbb.Finish(Data::Pack(fbb, dataT.get()));
+  return GetRedisContext(id)->RunAsync("RAY.SET_ADD", id, fbb.GetBufferPointer(),
+                                       fbb.GetSize(), prefix_, pubsub_channel_,
+                                       std::move(callback));
+}
+
+template <typename ID, typename Data>
+Status Set<ID, Data>::Remove(const JobID &job_id, const ID &id,
+                             std::shared_ptr<DataT> &dataT, const WriteCallback &done) {
+  num_removes_++;
+  auto callback = [this, id, dataT, done](const std::string &data) {
+    if (done != nullptr) {
+      (done)(client_, id, *dataT);
+    }
+    return true;
+  };
+  flatbuffers::FlatBufferBuilder fbb;
+  fbb.ForceDefaults(true);
+  fbb.Finish(Data::Pack(fbb, dataT.get()));
+  return GetRedisContext(id)->RunAsync("RAY.SET_REMOVE", id, fbb.GetBufferPointer(),
+                                       fbb.GetSize(), prefix_, pubsub_channel_,
+                                       std::move(callback));
+}
+
+template <typename ID, typename Data>
+std::string Set<ID, Data>::DebugString() const {
+  std::stringstream result;
+  result << "num lookups: " << num_lookups_ << ", num adds: " << num_adds_
+         << ", num removes: " << num_removes_;
+  return result.str();
+}
+
 Status ErrorTable::PushErrorToDriver(const JobID &job_id, const std::string &type,
                                      const std::string &error_message, double timestamp) {
   auto data = std::make_shared<ErrorTableDataT>();
@@ -508,6 +566,7 @@ Status ActorCheckpointIdTable::AddCheckpointId(const JobID &job_id,
 }
 
 template class Log<ObjectID, ObjectTableData>;
+template class Set<ObjectID, ObjectTableData>;
 template class Log<TaskID, ray::protocol::Task>;
 template class Table<TaskID, ray::protocol::Task>;
 template class Table<TaskID, TaskTableData>;