Add redis store client AsyncGetAll/AsyncBatchDelete/AsyncDeleteByIndex API (#8390)

src/ray/common/ray_config_def.h

@@ -248,6 +248,9 @@ RAY_CONFIG(int32_t, num_actor_checkpoints_to_keep, 20)
 /// Maximum number of ids in one batch to send to GCS to delete keys.
 RAY_CONFIG(uint32_t, maximum_gcs_deletion_batch_size, 1000)
 
+/// Maximum number of items in one batch to scan from GCS storage.
+RAY_CONFIG(uint32_t, maximum_gcs_scan_batch_size, 1000)
+
 /// When getting objects from object store, print a warning every this number of attempts.
 RAY_CONFIG(uint32_t, object_store_get_warn_per_num_attempts, 50)
 

src/ray/gcs/callback.h

@@ -16,6 +16,7 @@
 #define RAY_GCS_CALLBACK_H
 
 #include <boost/optional/optional.hpp>
+#include <unordered_map>
 #include <vector>
 #include "ray/common/status.h"
 
@@ -53,15 +54,10 @@ using SubscribeCallback = std::function<void(const ID &id, const Data &result)>;
 template <typename Data>
 using ItemCallback = std::function<void(const Data &result)>;
 
-/// This callback is used to receive a large amount of results.
-/// \param status Status indicates whether the scan was successful.
-/// \param has_more Whether more data will be called back.
-/// If `has_more == true`, there are more data to be received. This callback will
-/// be called again.
-/// \param result The items returned by storage.
-template <typename Data>
-using SegmentedCallback =
-    std::function<void(Status status, bool has_more, const std::vector<Data> &result)>;
+/// This callback is used to receive multiple key-value items from GCS.
+/// \param result The key-value items returned by GCS.
+template <typename Key, typename Value>
+using MapCallback = std::function<void(const std::unordered_map<Key, Value> &result)>;
 
 }  // namespace gcs
 

src/ray/gcs/gcs_server/gcs_table_storage.cc

@@ -44,18 +44,15 @@ Status GcsTable<Key, Data>::Get(const Key &key,
 }
 
 template <typename Key, typename Data>
-Status GcsTable<Key, Data>::GetAll(
-    const SegmentedCallback<std::pair<Key, Data>> &callback) {
-  auto on_done = [callback](
-                     const Status &status, bool has_more,
-                     const std::vector<std::pair<std::string, std::string>> &result) {
-    std::vector<std::pair<Key, Data>> values;
+Status GcsTable<Key, Data>::GetAll(const MapCallback<Key, Data> &callback) {
+  auto on_done = [callback](const std::unordered_map<std::string, std::string> &result) {
+    std::unordered_map<Key, Data> values;
     for (auto &item : result) {
       Data data;
       data.ParseFromString(item.second);
-      values.emplace_back(std::move(std::make_pair(Key::FromBinary(item.first), data)));
+      values[Key::FromBinary(item.first)] = data;
     }
-    callback(status, has_more, values);
+    callback(values);
   };
   return store_client_->AsyncGetAll(table_name_, on_done);
 }
@@ -68,19 +65,13 @@ Status GcsTable<Key, Data>::Delete(const Key &key, const StatusCallback &callbac
 template <typename Key, typename Data>
 Status GcsTable<Key, Data>::BatchDelete(const std::vector<Key> &keys,
                                         const StatusCallback &callback) {
-  // TODO(ffbin): We will use redis store client batch delete interface directly later.
-  auto finished_count = std::make_shared<int>(0);
-  int size = keys.size();
-  for (Key key : keys) {
-    auto done = [finished_count, size, callback](const Status &status) {
-      ++(*finished_count);
-      if (*finished_count == size) {
-        callback(Status::OK());
-      }
-    };
-    RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key.Binary(), done));
+  std::vector<std::string> keys_to_delete;
+  keys_to_delete.reserve(keys.size());
+  for (auto &key : keys) {
+    keys_to_delete.emplace_back(std::move(key.Binary()));
   }
-  return Status::OK();
+  return this->store_client_->AsyncBatchDelete(this->table_name_, keys_to_delete,
+                                               callback);
 }
 
 template <typename Key, typename Data>
@@ -92,8 +83,8 @@ Status GcsTableWithJobId<Key, Data>::Put(const Key &key, const Data &value,
 }
 
 template <typename Key, typename Data>
-Status GcsTableWithJobId<Key, Data>::GetByJobId(
-    const JobID &job_id, const SegmentedCallback<std::pair<Key, Data>> &callback) {
+Status GcsTableWithJobId<Key, Data>::GetByJobId(const JobID &job_id,
+                                                const MapCallback<Key, Data> &callback) {
   // TODO(ffbin): We will add this function after redis store client support
   // AsyncGetByIndex interface.
   return Status::NotImplemented("GetByJobId not implemented");

src/ray/gcs/gcs_server/gcs_table_storage.h

@@ -74,9 +74,8 @@ class GcsTable {
   /// Get all data from the table asynchronously.
   ///
   /// \param callback Callback that will be called after data has been received.
-  /// If the callback return `has_more == true` mean there's more data to be received.
   /// \return Status
-  Status GetAll(const SegmentedCallback<std::pair<Key, Data>> &callback);
+  Status GetAll(const MapCallback<Key, Data> &callback);
 
   /// Delete data from the table asynchronously.
   ///
@@ -122,8 +121,7 @@ class GcsTableWithJobId : public GcsTable<Key, Data> {
   /// \param job_id The key to lookup from the table.
   /// \param callback Callback that will be called after read finishes.
   /// \return Status
-  Status GetByJobId(const JobID &job_id,
-                    const SegmentedCallback<std::pair<Key, Data>> &callback);
+  Status GetByJobId(const JobID &job_id, const MapCallback<Key, Data> &callback);
 
   /// Delete all the data of the specified job id from the table asynchronously.
   ///

src/ray/gcs/store_client/in_memory_store_client.cc

@@ -58,14 +58,12 @@ Status InMemoryStoreClient::AsyncGet(const std::string &table_name,
 
 Status InMemoryStoreClient::AsyncGetAll(
     const std::string &table_name,
-    const SegmentedCallback<std::pair<std::string, std::string>> &callback) {
+    const MapCallback<std::string, std::string> &callback) {
   auto table = GetOrCreateTable(table_name);
   absl::MutexLock lock(&(table->mutex_));
-  std::vector<std::pair<std::string, std::string>> result;
-  for (auto &record : table->records_) {
-    result.emplace_back(std::make_pair(record.first, record.second));
-  }
-  main_io_service_.post([result, callback]() { callback(Status::OK(), false, result); });
+  std::unordered_map<std::string, std::string> result;
+  result.insert(table->records_.begin(), table->records_.end());
+  main_io_service_.post([result, callback]() { callback(result); });
   return Status::OK();
 }
 
@@ -79,6 +77,18 @@ Status InMemoryStoreClient::AsyncDelete(const std::string &table_name,
   return Status::OK();
 }
 
+Status InMemoryStoreClient::AsyncBatchDelete(const std::string &table_name,
+                                             const std::vector<std::string> &keys,
+                                             const StatusCallback &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  for (auto &key : keys) {
+    table->records_.erase(key);
+  }
+  main_io_service_.post([callback]() { callback(Status::OK()); });
+  return Status::OK();
+}
+
 Status InMemoryStoreClient::AsyncDeleteByIndex(const std::string &table_name,
                                                const std::string &index_key,
                                                const StatusCallback &callback) {

src/ray/gcs/store_client/in_memory_store_client.h

@@ -42,13 +42,16 @@ class InMemoryStoreClient : public StoreClient {
   Status AsyncGet(const std::string &table_name, const std::string &key,
                   const OptionalItemCallback<std::string> &callback) override;
 
-  Status AsyncGetAll(
-      const std::string &table_name,
-      const SegmentedCallback<std::pair<std::string, std::string>> &callback) override;
+  Status AsyncGetAll(const std::string &table_name,
+                     const MapCallback<std::string, std::string> &callback) override;
 
   Status AsyncDelete(const std::string &table_name, const std::string &key,
                      const StatusCallback &callback) override;
 
+  Status AsyncBatchDelete(const std::string &table_name,
+                          const std::vector<std::string> &keys,
+                          const StatusCallback &callback) override;
+
   Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key,
                             const StatusCallback &callback) override;
 

src/ray/gcs/store_client/redis_store_client.cc

@@ -23,11 +23,12 @@ namespace ray {
 
 namespace gcs {
 
+std::string RedisStoreClient::separator_ = ":";
+
 Status RedisStoreClient::AsyncPut(const std::string &table_name, const std::string &key,
                                   const std::string &data,
                                   const StatusCallback &callback) {
-  std::string full_key = table_name + key;
-  return DoPut(full_key, data, callback);
+  return DoPut(GenRedisKey(table_name, key), data, callback);
 }
 
 Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name,
@@ -45,8 +46,7 @@ Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name,
     }
 
     // Write data to Redis.
-    std::string full_key = table_name + key;
-    status = DoPut(full_key, data, callback);
+    status = DoPut(GenRedisKey(table_name, key), data, callback);
 
     if (!status.ok()) {
       // Run callback if failed.
@@ -57,25 +57,10 @@ Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name,
   };
 
   // Write index to Redis.
-  std::string index_table_key = index_key + table_name + key;
+  std::string index_table_key = GenRedisKey(table_name, key, index_key);
   return DoPut(index_table_key, key, write_callback);
 }
 
-Status RedisStoreClient::DoPut(const std::string &key, const std::string &data,
-                               const StatusCallback &callback) {
-  std::vector<std::string> args = {"SET", key, data};
-  RedisCallback write_callback = nullptr;
-  if (callback) {
-    write_callback = [callback](const std::shared_ptr<CallbackReply> &reply) {
-      auto status = reply->ReadAsStatus();
-      callback(status);
-    };
-  }
-
-  auto shard_context = redis_client_->GetShardContext(key);
-  return shard_context->RunArgvAsync(args, write_callback);
-}
-
 Status RedisStoreClient::AsyncGet(const std::string &table_name, const std::string &key,
                                   const OptionalItemCallback<std::string> &callback) {
   RAY_CHECK(callback != nullptr);
@@ -91,18 +76,24 @@ Status RedisStoreClient::AsyncGet(const std::string &table_name, const std::stri
     callback(Status::OK(), result);
   };
 
-  std::string full_key = table_name + key;
-  std::vector<std::string> args = {"GET", full_key};
+  std::string redis_key = GenRedisKey(table_name, key);
+  std::vector<std::string> args = {"GET", redis_key};
 
-  auto shard_context = redis_client_->GetShardContext(full_key);
+  auto shard_context = redis_client_->GetShardContext(redis_key);
   return shard_context->RunArgvAsync(args, redis_callback);
 }
 
 Status RedisStoreClient::AsyncGetAll(
     const std::string &table_name,
-    const SegmentedCallback<std::pair<std::string, std::string>> &callback) {
-  RAY_CHECK(0) << "Not implemented! Will implement this function in next PR.";
-  return Status::OK();
+    const MapCallback<std::string, std::string> &callback) {
+  RAY_CHECK(callback);
+  std::string match_pattern = GenRedisMatchPattern(table_name);
+  auto scanner = std::make_shared<RedisScanner>(redis_client_, table_name, match_pattern);
+  auto on_done = [callback,
+                  scanner](const std::unordered_map<std::string, std::string> &result) {
+    callback(result);
+  };
+  return scanner->ScanKeysAndValues(on_done);
 }
 
 Status RedisStoreClient::AsyncDelete(const std::string &table_name,
@@ -110,27 +101,279 @@ Status RedisStoreClient::AsyncDelete(const std::string &table_name,
                                      const StatusCallback &callback) {
   RedisCallback delete_callback = nullptr;
   if (callback) {
-    delete_callback = [callback](std::shared_ptr<CallbackReply> reply) {
-      int64_t deleted_count = reply->ReadAsInteger();
-      RAY_LOG(DEBUG) << "Delete done, total delete count " << deleted_count;
+    delete_callback = [callback](const std::shared_ptr<CallbackReply> &reply) {
       callback(Status::OK());
     };
   }
 
-  std::string full_key = table_name + key;
-  std::vector<std::string> args = {"DEL", full_key};
+  std::string redis_key = GenRedisKey(table_name, key);
+  std::vector<std::string> args = {"DEL", redis_key};
 
-  auto shard_context = redis_client_->GetShardContext(full_key);
+  auto shard_context = redis_client_->GetShardContext(redis_key);
   return shard_context->RunArgvAsync(args, delete_callback);
 }
 
+Status RedisStoreClient::AsyncBatchDelete(const std::string &table_name,
+                                          const std::vector<std::string> &keys,
+                                          const StatusCallback &callback) {
+  std::vector<std::string> redis_keys;
+  redis_keys.reserve(keys.size());
+  for (auto &key : keys) {
+    redis_keys.push_back(GenRedisKey(table_name, key));
+  }
+  return DeleteByKeys(redis_keys, callback);
+}
+
 Status RedisStoreClient::AsyncDeleteByIndex(const std::string &table_name,
                                             const std::string &index_key,
                                             const StatusCallback &callback) {
-  RAY_CHECK(0) << "Not implemented! Will implement this function in next PR.";
+  std::string match_pattern = GenRedisMatchPattern(table_name, index_key);
+  auto scanner = std::make_shared<RedisScanner>(redis_client_, table_name, match_pattern);
+  auto on_done = [this, table_name, index_key, callback, scanner](
+                     const Status &status, const std::vector<std::string> &result) {
+    if (!result.empty()) {
+      std::vector<std::string> keys;
+      keys.reserve(result.size());
+      for (auto &item : result) {
+        keys.push_back(GetKeyFromRedisKey(item, table_name, index_key));
+      }
+      auto batch_delete_callback = [this, result, callback](const Status &status) {
+        RAY_CHECK_OK(status);
+        // Delete index keys.
+        RAY_CHECK_OK(DeleteByKeys(result, callback));
+      };
+      RAY_CHECK_OK(AsyncBatchDelete(table_name, keys, batch_delete_callback));
+    } else {
+      callback(status);
+    }
+  };
+
+  return scanner->ScanKeys(on_done);
+}
+
+Status RedisStoreClient::DoPut(const std::string &key, const std::string &data,
+                               const StatusCallback &callback) {
+  std::vector<std::string> args = {"SET", key, data};
+  RedisCallback write_callback = nullptr;
+  if (callback) {
+    write_callback = [callback](const std::shared_ptr<CallbackReply> &reply) {
+      auto status = reply->ReadAsStatus();
+      callback(status);
+    };
+  }
+
+  auto shard_context = redis_client_->GetShardContext(key);
+  return shard_context->RunArgvAsync(args, write_callback);
+}
+
+Status RedisStoreClient::DeleteByKeys(const std::vector<std::string> &keys,
+                                      const StatusCallback &callback) {
+  // The `DEL` command for each shard.
+  auto del_commands_by_shards = GenCommandsByShards(redis_client_, "DEL", keys);
+
+  auto finished_count = std::make_shared<int>(0);
+  int size = del_commands_by_shards.size();
+  for (auto &item : del_commands_by_shards) {
+    auto delete_callback = [finished_count, size,
+                            callback](const std::shared_ptr<CallbackReply> &reply) {
+      ++(*finished_count);
+      if (*finished_count == size) {
+        callback(Status::OK());
+      }
+    };
+    RAY_CHECK_OK(item.first->RunArgvAsync(item.second, delete_callback));
+  }
   return Status::OK();
 }
 
+std::unordered_map<RedisContext *, std::vector<std::string>>
+RedisStoreClient::GenCommandsByShards(const std::shared_ptr<RedisClient> &redis_client,
+                                      const std::string &command,
+                                      const std::vector<std::string> &keys) {
+  std::unordered_map<RedisContext *, std::vector<std::string>> commands_by_shards;
+  for (auto &key : keys) {
+    auto shard_context = redis_client->GetShardContext(key).get();
+    auto it = commands_by_shards.find(shard_context);
+    if (it == commands_by_shards.end()) {
+      commands_by_shards[shard_context].push_back(command);
+      commands_by_shards[shard_context].push_back(key);
+    } else {
+      it->second.push_back(key);
+    }
+  }
+  return commands_by_shards;
+}
+
+std::string RedisStoreClient::GenRedisKey(const std::string &table_name,
+                                          const std::string &key) {
+  std::stringstream ss;
+  ss << table_name << separator_ << key;
+  return ss.str();
+}
+
+std::string RedisStoreClient::GenRedisKey(const std::string &table_name,
+                                          const std::string &key,
+                                          const std::string &index_key) {
+  std::stringstream ss;
+  ss << table_name << separator_ << index_key << separator_ << key;
+  return ss.str();
+}
+
+std::string RedisStoreClient::GenRedisMatchPattern(const std::string &table_name) {
+  std::stringstream ss;
+  ss << table_name << separator_ << "*";
+  return ss.str();
+}
+
+std::string RedisStoreClient::GenRedisMatchPattern(const std::string &table_name,
+                                                   const std::string &index_key) {
+  std::stringstream ss;
+  ss << table_name << separator_ << index_key << separator_ << "*";
+  return ss.str();
+}
+
+std::string RedisStoreClient::GetKeyFromRedisKey(const std::string &redis_key,
+                                                 const std::string &table_name) {
+  auto pos = table_name.size() + separator_.size();
+  return redis_key.substr(pos, redis_key.size() - pos);
+}
+
+std::string RedisStoreClient::GetKeyFromRedisKey(const std::string &redis_key,
+                                                 const std::string &table_name,
+                                                 const std::string &index_key) {
+  auto pos = table_name.size() + separator_.size() * 2 + index_key.size();
+  return redis_key.substr(pos, redis_key.size() - pos);
+}
+
+RedisStoreClient::RedisScanner::RedisScanner(std::shared_ptr<RedisClient> redis_client,
+                                             std::string table_name,
+                                             std::string match_pattern)
+    : table_name_(std::move(table_name)),
+      match_pattern_(std::move(match_pattern)),
+      redis_client_(std::move(redis_client)) {
+  for (size_t index = 0; index < redis_client_->GetShardContexts().size(); ++index) {
+    shard_to_cursor_[index] = 0;
+  }
+}
+
+Status RedisStoreClient::RedisScanner::ScanKeysAndValues(
+    const ItemCallback<std::unordered_map<std::string, std::string>> &callback) {
+  auto on_done = [this, callback](const Status &status,
+                                  const std::vector<std::string> &result) {
+    if (result.empty()) {
+      callback(std::unordered_map<std::string, std::string>());
+    } else {
+      MGetValues(result, callback);
+    }
+  };
+  return ScanKeys(on_done);
+}
+
+Status RedisStoreClient::RedisScanner::ScanKeys(
+    const MultiItemCallback<std::string> &callback) {
+  auto on_done = [this, callback](const Status &status) {
+    std::vector<std::string> result;
+    result.insert(result.begin(), keys_.begin(), keys_.end());
+    callback(status, result);
+  };
+  Scan(on_done);
+  return Status::OK();
+}
+
+void RedisStoreClient::RedisScanner::Scan(const StatusCallback &callback) {
+  if (shard_to_cursor_.empty()) {
+    callback(Status::OK());
+    return;
+  }
+
+  size_t batch_count = RayConfig::instance().maximum_gcs_scan_batch_size();
+  for (const auto &item : shard_to_cursor_) {
+    ++pending_request_count_;
+
+    size_t shard_index = item.first;
+    size_t cursor = item.second;
+
+    auto scan_callback = [this, shard_index,
+                          callback](const std::shared_ptr<CallbackReply> &reply) {
+      OnScanCallback(shard_index, reply, callback);
+    };
+
+    // Scan by prefix from Redis.
+    std::vector<std::string> args = {"SCAN",  std::to_string(cursor),
+                                     "MATCH", match_pattern_,
+                                     "COUNT", std::to_string(batch_count)};
+    auto shard_context = redis_client_->GetShardContexts()[shard_index];
+    Status status = shard_context->RunArgvAsync(args, scan_callback);
+
+    if (!status.ok()) {
+      RAY_LOG(FATAL) << "Scan failed, status " << status.ToString();
+    }
+  }
+}
+
+void RedisStoreClient::RedisScanner::OnScanCallback(
+    size_t shard_index, const std::shared_ptr<CallbackReply> &reply,
+    const StatusCallback &callback) {
+  RAY_CHECK(reply);
+  std::vector<std::string> scan_result;
+  size_t cursor = reply->ReadAsScanArray(&scan_result);
+
+  // Update shard cursors and keys_.
+  {
+    absl::MutexLock lock(&mutex_);
+    auto shard_it = shard_to_cursor_.find(shard_index);
+    RAY_CHECK(shard_it != shard_to_cursor_.end());
+    // If cursor is equal to 0, it means that the scan of this shard is finished, so we
+    // erase it from shard_to_cursor_.
+    if (cursor == 0) {
+      shard_to_cursor_.erase(shard_it);
+    } else {
+      shard_it->second = cursor;
+    }
+
+    keys_.insert(scan_result.begin(), scan_result.end());
+  }
+
+  // If pending_request_count_ is equal to 0, it means that the scan of this batch is
+  // completed and the next batch is started if any.
+  if (--pending_request_count_ == 0) {
+    Scan(callback);
+  }
+}
+
+void RedisStoreClient::RedisScanner::MGetValues(
+    const std::vector<std::string> &keys,
+    const ItemCallback<std::unordered_map<std::string, std::string>> &callback) {
+  // The `MGET` command for each shard.
+  auto mget_commands_by_shards = GenCommandsByShards(redis_client_, "MGET", keys);
+
+  auto finished_count = std::make_shared<int>(0);
+  int size = mget_commands_by_shards.size();
+  for (auto &item : mget_commands_by_shards) {
+    auto mget_keys = item.second;
+    auto mget_callback = [this, finished_count, size, mget_keys,
+                          callback](const std::shared_ptr<CallbackReply> &reply) {
+      if (!reply->IsNil()) {
+        auto value = reply->ReadAsStringArray();
+        {
+          absl::MutexLock lock(&mutex_);
+          // The 0 th element of mget_keys is "MGET", so we start from the 1 th element.
+          for (int index = 0; index < (int)value.size(); ++index) {
+            key_value_map_[GetKeyFromRedisKey(mget_keys[index + 1], table_name_)] =
+                value[index];
+          }
+        }
+      }
+
+      ++(*finished_count);
+      if (*finished_count == size) {
+        callback(key_value_map_);
+      }
+    };
+    RAY_CHECK_OK(item.first->RunArgvAsync(item.second, mget_callback));
+  }
+}
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/store_client/redis_store_client.h

@@ -15,8 +15,7 @@
 #ifndef RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H
 #define RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H
 
-#include <memory>
-#include <unordered_set>
+#include "absl/container/flat_hash_set.h"
 #include "ray/gcs/redis_client.h"
 #include "ray/gcs/redis_context.h"
 #include "ray/gcs/store_client/store_client.h"
@@ -28,7 +27,7 @@ namespace gcs {
 
 class RedisStoreClient : public StoreClient {
  public:
-  RedisStoreClient(std::shared_ptr<RedisClient> redis_client)
+  explicit RedisStoreClient(std::shared_ptr<RedisClient> redis_client)
       : redis_client_(std::move(redis_client)) {}
 
   Status AsyncPut(const std::string &table_name, const std::string &key,
@@ -41,20 +40,97 @@ class RedisStoreClient : public StoreClient {
   Status AsyncGet(const std::string &table_name, const std::string &key,
                   const OptionalItemCallback<std::string> &callback) override;
 
-  Status AsyncGetAll(
-      const std::string &table_name,
-      const SegmentedCallback<std::pair<std::string, std::string>> &callback) override;
+  Status AsyncGetAll(const std::string &table_name,
+                     const MapCallback<std::string, std::string> &callback) override;
 
   Status AsyncDelete(const std::string &table_name, const std::string &key,
                      const StatusCallback &callback) override;
 
+  Status AsyncBatchDelete(const std::string &table_name,
+                          const std::vector<std::string> &keys,
+                          const StatusCallback &callback) override;
+
   Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key,
                             const StatusCallback &callback) override;
 
  private:
+  /// \class RedisScanner
+  /// This class is used to scan data from Redis.
+  ///
+  /// If you called one method, should never call the other methods.
+  /// Otherwise it will disturb the status of the RedisScanner.
+  class RedisScanner {
+   public:
+    explicit RedisScanner(std::shared_ptr<RedisClient> redis_client,
+                          std::string table_name, std::string match_pattern);
+
+    Status ScanKeysAndValues(const MapCallback<std::string, std::string> &callback);
+
+    Status ScanKeys(const MultiItemCallback<std::string> &callback);
+
+   private:
+    void Scan(const StatusCallback &callback);
+
+    void OnScanCallback(size_t shard_index, const std::shared_ptr<CallbackReply> &reply,
+                        const StatusCallback &callback);
+
+    void MGetValues(const std::vector<std::string> &keys,
+                    const MapCallback<std::string, std::string> &callback);
+
+    std::string table_name_;
+
+    /// The scan match pattern.
+    std::string match_pattern_;
+
+    /// Mutex to protect the shard_to_cursor_ field and the keys_ field and the
+    /// key_value_map_ field.
+    absl::Mutex mutex_;
+
+    /// All keys that scanned from redis.
+    absl::flat_hash_set<std::string> keys_;
+
+    /// Key-Value pairs that scanned from redis.
+    std::unordered_map<std::string, std::string> key_value_map_;
+
+    /// The scan cursor for each shard.
+    std::unordered_map<size_t, size_t> shard_to_cursor_;
+
+    /// The pending shard scan count.
+    std::atomic<size_t> pending_request_count_{0};
+
+    std::shared_ptr<RedisClient> redis_client_;
+  };
+
   Status DoPut(const std::string &key, const std::string &data,
                const StatusCallback &callback);
 
+  Status DeleteByKeys(const std::vector<std::string> &keys,
+                      const StatusCallback &callback);
+
+  static std::unordered_map<RedisContext *, std::vector<std::string>> GenCommandsByShards(
+      const std::shared_ptr<RedisClient> &redis_client, const std::string &command,
+      const std::vector<std::string> &keys);
+
+  /// The separator is used when building redis key.
+  static std::string separator_;
+
+  static std::string GenRedisKey(const std::string &table_name, const std::string &key);
+
+  static std::string GenRedisKey(const std::string &table_name, const std::string &key,
+                                 const std::string &index_key);
+
+  static std::string GenRedisMatchPattern(const std::string &table_name);
+
+  static std::string GenRedisMatchPattern(const std::string &table_name,
+                                          const std::string &index_key);
+
+  static std::string GetKeyFromRedisKey(const std::string &redis_key,
+                                        const std::string &table_name);
+
+  static std::string GetKeyFromRedisKey(const std::string &redis_key,
+                                        const std::string &table_name,
+                                        const std::string &index_key);
+
   std::shared_ptr<RedisClient> redis_client_;
 };
 

src/ray/gcs/store_client/store_client.h

@@ -69,11 +69,9 @@ class StoreClient {
   ///
   /// \param table_name The name of the table to be read.
   /// \param callback Callback that will be called after data has been received.
-  /// If the callback return `has_more == true` mean there's more data to be received.
   /// \return Status
-  virtual Status AsyncGetAll(
-      const std::string &table_name,
-      const SegmentedCallback<std::pair<std::string, std::string>> &callback) = 0;
+  virtual Status AsyncGetAll(const std::string &table_name,
+                             const MapCallback<std::string, std::string> &callback) = 0;
 
   /// Delete data from the given table asynchronously.
   ///
@@ -84,6 +82,16 @@ class StoreClient {
   virtual Status AsyncDelete(const std::string &table_name, const std::string &key,
                              const StatusCallback &callback) = 0;
 
+  /// Batch delete data from the given table asynchronously.
+  ///
+  /// \param table_name The name of the table from which data is to be deleted.
+  /// \param keys The keys that will be deleted from the table.
+  /// \param callback Callback that will be called after delete finishes.
+  /// \return Status
+  virtual Status AsyncBatchDelete(const std::string &table_name,
+                                  const std::vector<std::string> &keys,
+                                  const StatusCallback &callback) = 0;
+
   /// Delete by index from the given table asynchronously.
   ///
   /// \param table_name The name of the table from which data is to be deleted.

src/ray/gcs/store_client/test/in_memory_store_client_test.cc

@@ -30,14 +30,14 @@ class InMemoryStoreClientTest : public StoreClientTestBase {
 
 TEST_F(InMemoryStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); }
 
-TEST_F(InMemoryStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); }
-
-TEST_F(InMemoryStoreClientTest, AsyncGetAllTest) { TestAsyncGetAll(); }
-
 TEST_F(InMemoryStoreClientTest, AsyncPutAndDeleteWithIndexTest) {
   TestAsyncPutAndDeleteWithIndex();
 }
 
+TEST_F(InMemoryStoreClientTest, AsyncGetAllAndBatchDeleteTest) {
+  TestAsyncGetAllAndBatchDelete();
+}
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/store_client/test/redis_store_client_test.cc

@@ -46,9 +46,13 @@ class RedisStoreClientTest : public StoreClientTestBase {
 
 TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); }
 
-TEST_F(RedisStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); }
+TEST_F(RedisStoreClientTest, AsyncPutAndDeleteWithIndexTest) {
+  TestAsyncPutAndDeleteWithIndex();
+}
 
-TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { TestAsyncGetAll(); }
+TEST_F(RedisStoreClientTest, AsyncGetAllAndBatchDeleteTest) {
+  TestAsyncGetAllAndBatchDelete();
+}
 
 }  // namespace gcs
 

src/ray/gcs/store_client/test/store_client_test_base.h

@@ -31,7 +31,7 @@ namespace gcs {
 
 class StoreClientTestBase : public ::testing::Test {
  public:
-  StoreClientTestBase() {}
+  StoreClientTestBase() = default;
 
   virtual ~StoreClientTestBase() {}
 
@@ -57,8 +57,7 @@ class StoreClientTestBase : public ::testing::Test {
   virtual void DisconnectStoreClient() = 0;
 
  protected:
-  void TestAsyncPutAndAsyncGet() {
-    // AsyncPut without index.
+  void Put() {
     auto put_calllback = [this](const Status &status) {
       RAY_CHECK_OK(status);
       --pending_count_;
@@ -70,8 +69,22 @@ class StoreClientTestBase : public ::testing::Test {
                                            put_calllback));
     }
     WaitPendingDone();
+  }
 
-    // AsyncGet
+  void Delete() {
+    auto delete_calllback = [this](const Status &status) {
+      RAY_CHECK_OK(status);
+      --pending_count_;
+    };
+    for (const auto &elem : key_to_value_) {
+      ++pending_count_;
+      RAY_CHECK_OK(
+          store_client_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback));
+    }
+    WaitPendingDone();
+  }
+
+  void Get() {
     auto get_callback = [this](const Status &status,
                                const boost::optional<std::string> &result) {
       RAY_CHECK_OK(status);
@@ -91,101 +104,49 @@ class StoreClientTestBase : public ::testing::Test {
     WaitPendingDone();
   }
 
-  void TestAsyncDelete() {
-    // AsyncPut
-    auto put_calllback = [this](const Status &status) { --pending_count_; };
+  void GetEmpty() {
     for (const auto &elem : key_to_value_) {
-      ++pending_count_;
-      RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(),
-                                           elem.second.SerializeAsString(),
-                                           put_calllback));
-    }
-    WaitPendingDone();
+      auto key = elem.first.Binary();
+      auto get_callback = [this, key](const Status &status,
+                                      const boost::optional<std::string> &result) {
+        RAY_CHECK_OK(status);
+        RAY_CHECK(!result);
+        --pending_count_;
+      };
 
-    // AsyncDelete
-    auto delete_calllback = [this](const Status &status) {
-      RAY_CHECK_OK(status);
-      --pending_count_;
-    };
-    for (const auto &elem : key_to_value_) {
       ++pending_count_;
-      RAY_CHECK_OK(
-          store_client_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback));
-    }
-    WaitPendingDone();
-
-    // AsyncGet
-    auto get_callback = [this](const Status &status,
-                               const boost::optional<std::string> &result) {
-      RAY_CHECK_OK(status);
-      RAY_CHECK(!result);
-      --pending_count_;
-    };
-    for (const auto &elem : key_to_value_) {
-      ++pending_count_;
-      RAY_CHECK_OK(
-          store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback));
+      RAY_CHECK_OK(store_client_->AsyncGet(table_name_, key, get_callback));
     }
     WaitPendingDone();
   }
 
-  void TestAsyncPutAndDeleteWithIndex() {
-    // AsyncPut with index
+  void PutWithIndex() {
     auto put_calllback = [this](const Status &status) { --pending_count_; };
     for (const auto &elem : key_to_value_) {
       ++pending_count_;
       RAY_CHECK_OK(store_client_->AsyncPutWithIndex(
-          table_name_, elem.first.Binary(), key_to_index_[elem.first].Binary(),
+          table_name_, elem.first.Binary(), key_to_index_[elem.first].Hex(),
           elem.second.SerializeAsString(), put_calllback));
     }
     WaitPendingDone();
+  }
 
-    // AsyncDelete by index
+  void DeleteByIndex() {
     auto delete_calllback = [this](const Status &status) {
       RAY_CHECK_OK(status);
       --pending_count_;
     };
     for (const auto &elem : index_to_keys_) {
       ++pending_count_;
-      RAY_CHECK_OK(store_client_->AsyncDeleteByIndex(table_name_, elem.first.Binary(),
+      RAY_CHECK_OK(store_client_->AsyncDeleteByIndex(table_name_, elem.first.Hex(),
                                                      delete_calllback));
     }
     WaitPendingDone();
-
-    // AsyncGet
-    auto get_callback = [this](const Status &status,
-                               const boost::optional<std::string> &result) {
-      RAY_CHECK_OK(status);
-      RAY_CHECK(!result);
-      --pending_count_;
-    };
-    for (const auto &elem : key_to_value_) {
-      ++pending_count_;
-      RAY_CHECK_OK(
-          store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback));
-    }
-    WaitPendingDone();
   }
 
-  void TestAsyncGetAll() {
-    // AsyncPut
-    auto put_calllback = [this](const Status &status) { --pending_count_; };
-    for (const auto &elem : key_to_value_) {
-      ++pending_count_;
-      // Get index
-      auto it = key_to_index_.find(elem.first);
-      const JobID &index = it->second;
-      RAY_CHECK_OK(store_client_->AsyncPutWithIndex(
-          table_name_, elem.first.Binary(), index.Binary(),
-          elem.second.SerializeAsString(), put_calllback));
-    }
-    WaitPendingDone();
-
-    // AsyncGetAll
+  void GetAll() {
     auto get_all_callback =
-        [this](const Status &status, bool has_more,
-               const std::vector<std::pair<std::string, std::string>> &result) {
-          RAY_CHECK_OK(status);
+        [this](const std::unordered_map<std::string, std::string> &result) {
           static std::unordered_set<ActorID> received_keys;
           for (const auto &item : result) {
             const ActorID &actor_id = ActorID::FromBinary(item.first);
@@ -196,9 +157,7 @@ class StoreClientTestBase : public ::testing::Test {
             auto map_it = key_to_value_.find(actor_id);
             RAY_CHECK(map_it != key_to_value_.end());
           }
-          if (!has_more) {
-            RAY_CHECK(received_keys.size() == key_to_value_.size());
-          }
+          RAY_CHECK(received_keys.size() == key_to_value_.size());
           pending_count_ -= result.size();
         };
 
@@ -207,6 +166,58 @@ class StoreClientTestBase : public ::testing::Test {
     WaitPendingDone();
   }
 
+  void BatchDelete() {
+    auto delete_calllback = [this](const Status &status) {
+      RAY_CHECK_OK(status);
+      --pending_count_;
+    };
+    ++pending_count_;
+    std::vector<std::string> keys;
+    for (auto &elem : key_to_value_) {
+      keys.push_back(elem.first.Binary());
+    }
+    RAY_CHECK_OK(store_client_->AsyncBatchDelete(table_name_, keys, delete_calllback));
+    WaitPendingDone();
+  }
+
+  void TestAsyncPutAndAsyncGet() {
+    // AsyncPut without index.
+    Put();
+
+    // AsyncGet
+    Get();
+
+    // AsyncDelete
+    Delete();
+
+    GetEmpty();
+  }
+
+  void TestAsyncPutAndDeleteWithIndex() {
+    // AsyncPut with index
+    PutWithIndex();
+
+    // AsyncDelete by index
+    DeleteByIndex();
+
+    // AsyncGet
+    GetEmpty();
+  }
+
+  void TestAsyncGetAllAndBatchDelete() {
+    // AsyncPut
+    Put();
+
+    // AsyncGetAll
+    GetAll();
+
+    // AsyncBatchDelete
+    BatchDelete();
+
+    // AsyncGet
+    GetEmpty();
+  }
+
   void GenTestData() {
     for (size_t i = 0; i < key_count_; i++) {
       rpc::ActorTableData actor;