Use Redis lists in the GCS instead of zset (#4023)

* Convert zset to list

* Remove object evictions map from the object directory, yay

* comments

* Fix tests

src/ray/gcs/format/gcs.fbs

@@ -130,8 +130,6 @@ table ObjectTableData {
   manager: string;
   // Whether this entry is an addition or a deletion.
   is_eviction: bool;
-  // The number of times this object has been evicted from this node so far.
-  num_evictions: int;
   // In-line object data.
   inline_object_data: [ubyte];
   // In-line object metadata.

src/ray/gcs/redis_module/ray_redis_module.cc

@@ -70,6 +70,34 @@ Status FormatPubsubChannel(RedisModuleString **out, RedisModuleCtx *ctx,
   return Status::OK();
 }
 
+/// Parse a Redis string into a TablePrefix channel.
+Status ParseTablePrefix(const RedisModuleString *table_prefix_str, TablePrefix *out) {
+  long long table_prefix_long;
+  if (RedisModule_StringToLongLong(table_prefix_str, &table_prefix_long) !=
+      REDISMODULE_OK) {
+    return Status::RedisError("Prefix must be a valid TablePrefix integer");
+  }
+  if (table_prefix_long > static_cast<long long>(TablePrefix::MAX) ||
+      table_prefix_long < static_cast<long long>(TablePrefix::MIN)) {
+    return Status::RedisError("Prefix must be in the TablePrefix range");
+  } else {
+    *out = static_cast<TablePrefix>(table_prefix_long);
+    return Status::OK();
+  }
+}
+
+/// Format the string for a table key. `prefix_enum` must be a valid
+/// TablePrefix as a RedisModuleString. `keyname` is usually a UniqueID as a
+/// RedisModuleString.
+RedisModuleString *PrefixedKeyString(RedisModuleCtx *ctx, RedisModuleString *prefix_enum,
+                                     RedisModuleString *keyname) {
+  TablePrefix prefix;
+  if (!ParseTablePrefix(prefix_enum, &prefix).ok()) {
+    return nullptr;
+  }
+  return RedisString_Format(ctx, "%s%S", EnumNameTablePrefix(prefix), keyname);
+}
+
 // TODO(swang): This helper function should be deprecated by the version below,
 // which uses enums for table prefixes.
 RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx, const char *prefix,
@@ -88,15 +116,8 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx, const char *prefix,
 Status OpenPrefixedKey(RedisModuleKey **out, RedisModuleCtx *ctx,
                        RedisModuleString *prefix_enum, RedisModuleString *keyname,
                        int mode, RedisModuleString **mutated_key_str) {
-  long long prefix_long;
-  if (RedisModule_StringToLongLong(prefix_enum, &prefix_long) != REDISMODULE_OK) {
-    return Status::RedisError("Prefix must be a valid TablePrefix integer.");
-  }
-  if (prefix_long > static_cast<long long>(TablePrefix::MAX) ||
-      prefix_long < static_cast<long long>(TablePrefix::MIN)) {
-    return Status::RedisError("Prefix must be in the TablePrefix range.");
-  }
-  auto prefix = static_cast<TablePrefix>(prefix_long);
+  TablePrefix prefix;
+  RAY_RETURN_NOT_OK(ParseTablePrefix(prefix_enum, &prefix));
   *out =
       OpenPrefixedKey(ctx, EnumNameTablePrefix(prefix), keyname, mode, mutated_key_str);
   return Status::OK();
@@ -283,6 +304,11 @@ int TableAppend_DoWrite(RedisModuleCtx *ctx, RedisModuleString **argv, int argc,
   RedisModuleKey *key;
   REPLY_AND_RETURN_IF_NOT_OK(OpenPrefixedKey(
       &key, ctx, prefix_str, id, REDISMODULE_READ | REDISMODULE_WRITE, mutated_key_str));
+  int type = RedisModule_KeyType(key);
+  REPLY_AND_RETURN_IF_FALSE(
+      type == REDISMODULE_KEYTYPE_LIST || type == REDISMODULE_KEYTYPE_EMPTY,
+      "TABLE_APPEND entries must be a list or an empty list");
+
   // Determine the index at which the data should be appended. If no index is
   // requested, then is the current length of the log.
   size_t index = RedisModule_ValueLength(key);
@@ -300,34 +326,13 @@ int TableAppend_DoWrite(RedisModuleCtx *ctx, RedisModuleString **argv, int argc,
   if (index == RedisModule_ValueLength(key)) {
     // The requested index matches the current length of the log or no index
     // was requested. Perform the append.
-    int flags = REDISMODULE_ZADD_NX;
-    RedisModule_ZsetAdd(key, index, data, &flags);
-    // Check that we actually add a new entry during the append. This is only
-    // necessary since we implement the log with a sorted set, so all entries
-    // must be unique, or else we will have gaps in the log.
-    // TODO(rkn): We need to get rid of this uniqueness requirement. We can
-    // easily have multiple log events with the same message.
-    if (flags != REDISMODULE_ZADD_ADDED) {
-      // The following code is a workaround to store the data at a new unique
-      // key. This is so redis doesn't crash (we currently have duplicate keys
-      // for error conditions, which get delivered via pubsub).
-      size_t len;
-      const char *id_str = RedisModule_StringPtrLen(id, &len);
-      RAY_LOG(INFO) << "Duplicate key: " << std::string(id_str, len);
-      // Store the value into a unique new key, just to keep track of it and
-      // make sure the log size grows.
-      std::string postfix = std::to_string(index);
-      RedisModuleString *new_id =
-          RedisString_Format(ctx, "%S:%b", id, postfix.data(), postfix.size());
-      RedisModuleKey *new_key;
-      REPLY_AND_RETURN_IF_NOT_OK(OpenPrefixedKey(&new_key, ctx, prefix_str, new_id,
-                                                 REDISMODULE_READ | REDISMODULE_WRITE,
-                                                 mutated_key_str));
-      RedisModule_ZsetAdd(new_key, index, data, &flags);
-      REPLY_AND_RETURN_IF_FALSE(flags == REDISMODULE_ZADD_ADDED,
-                                "Appended a duplicate entry");
+    if (RedisModule_ListPush(key, REDISMODULE_LIST_TAIL, data) == REDISMODULE_OK) {
+      return REDISMODULE_OK;
+    } else {
+      static const char *reply = "Unexpected error during TABLE_APPEND";
+      RedisModule_ReplyWithError(ctx, reply);
+      return REDISMODULE_ERR;
     }
-    return REDISMODULE_OK;
   } else {
     // The requested index did not match the current length of the log. Return
     // an error message as a string.
@@ -393,7 +398,17 @@ int ChainTableAppend_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
 
 /// A helper function to create and finish a GcsTableEntry, based on the
 /// current value or values at the given key.
-Status TableEntryToFlatbuf(RedisModuleKey *table_key, RedisModuleString *entry_id,
+///
+/// \param ctx The Redis module context.
+/// \param table_key The Redis key whose entry should be read out. The key must
+/// be open when this function is called and may be closed in this function.
+/// The key's name format is <prefix_str><entry_id>.
+/// \param prefix_str The string prefix associated with the open Redis key.
+/// When parsed, this is expected to be a TablePrefix.
+/// \param entry_id The UniqueID associated with the open Redis key.
+/// \param fbb A flatbuffer builder used to build the GcsTableEntry.
+Status TableEntryToFlatbuf(RedisModuleCtx *ctx, RedisModuleKey *table_key,
+                           RedisModuleString *prefix_str, RedisModuleString *entry_id,
                            flatbuffers::FlatBufferBuilder &fbb) {
   auto key_type = RedisModule_KeyType(table_key);
   switch (key_type) {
@@ -406,18 +421,26 @@ Status TableEntryToFlatbuf(RedisModuleKey *table_key, RedisModuleString *entry_i
                                        fbb.CreateVector(&data, 1));
     fbb.Finish(message);
   } break;
-  case REDISMODULE_KEYTYPE_ZSET: {
+  case REDISMODULE_KEYTYPE_LIST: {
+    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);
     // Build the flatbuffer from the set of log entries.
-    if (RedisModule_ZsetFirstInScoreRange(table_key, REDISMODULE_NEGATIVE_INFINITE,
-                                          REDISMODULE_POSITIVE_INFINITE, 1,
-                                          1) != REDISMODULE_OK) {
-      return Status::RedisError("Empty zset or wrong type");
+    if (RedisModule_CallReplyType(reply) != REDISMODULE_REPLY_ARRAY) {
+      return Status::RedisError("Empty list or wrong type");
     }
     std::vector<flatbuffers::Offset<flatbuffers::String>> data;
-    for (; !RedisModule_ZsetRangeEndReached(table_key);
-         RedisModule_ZsetRangeNext(table_key)) {
-      data.push_back(RedisStringToFlatbuf(
-          fbb, RedisModule_ZsetRangeCurrentElement(table_key, NULL)));
+    for (size_t i = 0; i < RedisModule_CallReplyLength(reply); i++) {
+      RedisModuleCallReply *element = RedisModule_CallReplyArrayElement(reply, i);
+      size_t len;
+      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));
@@ -467,7 +490,7 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int
   } else {
     // Serialize the data to a flatbuffer to return to the client.
     flatbuffers::FlatBufferBuilder fbb;
-    REPLY_AND_RETURN_IF_NOT_OK(TableEntryToFlatbuf(table_key, id, fbb));
+    REPLY_AND_RETURN_IF_NOT_OK(TableEntryToFlatbuf(ctx, table_key, prefix_str, id, fbb));
     RedisModule_ReplyWithStringBuffer(
         ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()), fbb.GetSize());
   }
@@ -524,7 +547,7 @@ int TableRequestNotifications_RedisCommand(RedisModuleCtx *ctx, RedisModuleStrin
   // notifications. An empty notification will be published if the key is
   // empty.
   flatbuffers::FlatBufferBuilder fbb;
-  REPLY_AND_RETURN_IF_NOT_OK(TableEntryToFlatbuf(table_key, id, fbb));
+  REPLY_AND_RETURN_IF_NOT_OK(TableEntryToFlatbuf(ctx, table_key, prefix_str, id, fbb));
   RedisModule_Call(ctx, "PUBLISH", "sb", client_channel,
                    reinterpret_cast<const char *>(fbb.GetBufferPointer()), fbb.GetSize());
 

src/ray/object_manager/object_directory.cc

@@ -123,7 +123,6 @@ ray::Status ObjectDirectory::ReportObjectAdded(
   auto data = std::make_shared<ObjectTableDataT>();
   data->manager = client_id.binary();
   data->is_eviction = false;
-  data->num_evictions = object_evictions_[object_id];
   data->object_size = object_info.data_size;
   data->inline_object_flag = inline_object_flag;
   if (inline_object_flag) {
@@ -144,14 +143,8 @@ ray::Status ObjectDirectory::ReportObjectRemoved(const ObjectID &object_id,
   auto data = std::make_shared<ObjectTableDataT>();
   data->manager = client_id.binary();
   data->is_eviction = true;
-  data->num_evictions = object_evictions_[object_id];
   ray::Status status =
       gcs_client_->object_table().Append(JobID::nil(), object_id, data, nullptr);
-  // Increment the number of times we've evicted this object. NOTE(swang): This
-  // is only necessary because the Ray redis module expects unique entries in a
-  // log. We track the number of evictions so that the next eviction, if there
-  // is one, is unique.
-  object_evictions_[object_id]++;
   return status;
 };
 
@@ -305,7 +298,6 @@ std::string ObjectDirectory::DebugString() const {
   std::stringstream result;
   result << "ObjectDirectory:";
   result << "\n- num listeners: " << listeners_.size();
-  result << "\n- num eviction entries: " << object_evictions_.size();
   return result.str();
 }
 

src/ray/object_manager/object_directory.h

@@ -205,9 +205,6 @@ class ObjectDirectory : public ObjectDirectoryInterface {
   std::shared_ptr<gcs::AsyncGcsClient> gcs_client_;
   /// Info about subscribers to object locations.
   std::unordered_map<ObjectID, LocationListenerState> listeners_;
-  /// Map from object ID to the number of times it's been evicted on this
-  /// node before.
-  std::unordered_map<ObjectID, int> object_evictions_;
 };
 
 }  // namespace ray

test/failure_test.py

@@ -648,13 +648,17 @@ def test_redis_module_failure(shutdown_only):
                      "RAY.TABLE_ADD", 1, 10000, 1, 1)
     run_failure_test("Pubsub channel must be a valid integer", "RAY.TABLE_ADD",
                      1, b"a", 1, 1)
-    run_failure_test("Index is less than 0.", "RAY.TABLE_APPEND", 1, 1, 1, 1,
+    # Change the key from 1 to 2, since the previous command should have
+    # succeeded at writing the key, but not publishing it.
+    run_failure_test("Index is less than 0.", "RAY.TABLE_APPEND", 1, 1, 2, 1,
                      -1)
-    run_failure_test("Index is not a number.", "RAY.TABLE_APPEND", 1, 1, 1, 1,
+    run_failure_test("Index is not a number.", "RAY.TABLE_APPEND", 1, 1, 2, 1,
                      b"a")
-    run_one_command("RAY.TABLE_APPEND", 1, 1, 1, 1)
-    run_failure_test("Appended a duplicate entry", "RAY.TABLE_APPEND", 1, 1, 1,
-                     1, 1)
+    run_one_command("RAY.TABLE_APPEND", 1, 1, 2, 1)
+    # It's okay to add duplicate entries.
+    run_one_command("RAY.TABLE_APPEND", 1, 1, 2, 1)
+    run_one_command("RAY.TABLE_APPEND", 1, 1, 2, 1, 0)
+    run_one_command("RAY.TABLE_APPEND", 1, 1, 2, 1, 1)
 
 
 @pytest.fixture