[GCS]Add in-memory store client (#8144)

BUILD.bazel

@@ -1086,6 +1086,23 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "gcs_in_memory_store_client",
+    srcs = [
+        "src/ray/gcs/store_client/in_memory_store_client.cc",
+    ],
+    hdrs = [
+        "src/ray/gcs/callback.h",
+        "src/ray/gcs/store_client/in_memory_store_client.h",
+        "src/ray/gcs/store_client/store_client.h",
+    ],
+    copts = COPTS,
+    deps = [
+        ":ray_common",
+        ":ray_util",
+    ],
+)
+
 cc_library(
     name = "store_client_test_lib",
     hdrs = [
@@ -1113,6 +1130,17 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "in_memory_store_client_test",
+    srcs = ["src/ray/gcs/store_client/test/in_memory_store_client_test.cc"],
+    copts = COPTS,
+    deps = [
+        ":gcs_in_memory_store_client",
+        ":store_client_test_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "gcs",
     srcs = glob(

src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc

@@ -27,6 +27,10 @@ class GcsTableStorageTest : public gcs::StoreClientTestBase {
 
   virtual ~GcsTableStorageTest() {}
 
+  static void SetUpTestCase() { RedisServiceManagerForTest::SetUpTestCase(); }
+
+  static void TearDownTestCase() { RedisServiceManagerForTest::TearDownTestCase(); }
+
   void InitStoreClient() override {
     gcs::RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true);
     redis_client_ = std::make_shared<gcs::RedisClient>(options);

src/ray/gcs/store_client/in_memory_store_client.cc

@@ -0,0 +1,113 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/gcs/store_client/in_memory_store_client.h"
+
+namespace ray {
+
+namespace gcs {
+
+Status InMemoryStoreClient::AsyncPut(const std::string &table_name,
+                                     const std::string &key, const std::string &data,
+                                     const StatusCallback &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  table->records_[key] = data;
+  main_io_service_.post([callback]() { callback(Status::OK()); });
+  return Status::OK();
+}
+
+Status InMemoryStoreClient::AsyncPutWithIndex(const std::string &table_name,
+                                              const std::string &key,
+                                              const std::string &index_key,
+                                              const std::string &data,
+                                              const StatusCallback &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  table->records_[key] = data;
+  table->index_keys_[index_key].emplace_back(key);
+  main_io_service_.post([callback]() { callback(Status::OK()); });
+  return Status::OK();
+}
+
+Status InMemoryStoreClient::AsyncGet(const std::string &table_name,
+                                     const std::string &key,
+                                     const OptionalItemCallback<std::string> &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  auto iter = table->records_.find(key);
+  if (iter != table->records_.end()) {
+    auto data = iter->second;
+    main_io_service_.post([callback, data]() { callback(Status::OK(), data); });
+  } else {
+    main_io_service_.post([callback]() { callback(Status::OK(), boost::none); });
+  }
+  return Status::OK();
+}
+
+Status InMemoryStoreClient::AsyncGetAll(
+    const std::string &table_name,
+    const SegmentedCallback<std::pair<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); });
+  return Status::OK();
+}
+
+Status InMemoryStoreClient::AsyncDelete(const std::string &table_name,
+                                        const std::string &key,
+                                        const StatusCallback &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  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) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  auto iter = table->index_keys_.find(index_key);
+  if (iter != table->index_keys_.end()) {
+    for (auto &key : iter->second) {
+      table->records_.erase(key);
+    }
+    table->index_keys_.erase(iter);
+  }
+  main_io_service_.post([callback]() { callback(Status::OK()); });
+  return Status::OK();
+}
+
+std::shared_ptr<InMemoryStoreClient::InMemoryTable> InMemoryStoreClient::GetOrCreateTable(
+    const std::string &table_name) {
+  absl::MutexLock lock(&mutex_);
+  auto iter = tables_.find(table_name);
+  if (iter != tables_.end()) {
+    return iter->second;
+  } else {
+    auto table = std::make_shared<InMemoryTable>();
+    tables_[table_name] = table;
+    return table;
+  }
+}
+
+}  // namespace gcs
+
+}  // namespace ray

src/ray/gcs/store_client/in_memory_store_client.h

@@ -0,0 +1,83 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef RAY_GCS_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H
+#define RAY_GCS_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/synchronization/mutex.h"
+#include "ray/gcs/store_client/store_client.h"
+#include "ray/protobuf/gcs.pb.h"
+
+namespace ray {
+
+namespace gcs {
+
+/// \class InMemoryStoreClient
+///
+/// This class is thread safe.
+class InMemoryStoreClient : public StoreClient {
+ public:
+  explicit InMemoryStoreClient(boost::asio::io_service &main_io_service)
+      : main_io_service_(main_io_service) {}
+
+  Status AsyncPut(const std::string &table_name, const std::string &key,
+                  const std::string &data, const StatusCallback &callback) override;
+
+  Status AsyncPutWithIndex(const std::string &table_name, const std::string &key,
+                           const std::string &index_key, const std::string &data,
+                           const StatusCallback &callback) override;
+
+  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 AsyncDelete(const std::string &table_name, const std::string &key,
+                     const StatusCallback &callback) override;
+
+  Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key,
+                            const StatusCallback &callback) override;
+
+ private:
+  struct InMemoryTable {
+    /// Mutex to protect the records_ field and the index_keys_ field.
+    absl::Mutex mutex_;
+    // Mapping from key to data.
+    absl::flat_hash_map<std::string, std::string> records_ GUARDED_BY(mutex_);
+    // Mapping from index key to keys.
+    absl::flat_hash_map<std::string, std::vector<std::string>> index_keys_
+        GUARDED_BY(mutex_);
+  };
+
+  std::shared_ptr<InMemoryStoreClient::InMemoryTable> GetOrCreateTable(
+      const std::string &table_name);
+
+  /// Mutex to protect the tables_ field.
+  absl::Mutex mutex_;
+  absl::flat_hash_map<std::string, std::shared_ptr<InMemoryTable>> tables_
+      GUARDED_BY(mutex_);
+
+  /// Async API Callback needs to post to main_io_service_ to ensure the orderly execution
+  /// of the callback.
+  boost::asio::io_service &main_io_service_;
+};
+
+}  // namespace gcs
+
+}  // namespace ray
+
+#endif  // RAY_GCS_STORE_CLIENT_IN_MEMORY_STORE_CLIENT_H

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

@@ -0,0 +1,48 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/gcs/store_client/in_memory_store_client.h"
+#include "ray/gcs/store_client/test/store_client_test_base.h"
+
+namespace ray {
+
+namespace gcs {
+
+class InMemoryStoreClientTest : public StoreClientTestBase {
+ public:
+  void InitStoreClient() override {
+    store_client_ = std::make_shared<InMemoryStoreClient>(*(io_service_pool_->Get()));
+  }
+
+  void DisconnectStoreClient() override {}
+};
+
+TEST_F(InMemoryStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); }
+
+TEST_F(InMemoryStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); }
+
+TEST_F(InMemoryStoreClientTest, AsyncGetAllTest) { TestAsyncGetAll(); }
+
+TEST_F(InMemoryStoreClientTest, AsyncPutAndDeleteWithIndexTest) {
+  TestAsyncPutAndDeleteWithIndex();
+}
+
+}  // namespace gcs
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

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

@@ -26,6 +26,10 @@ class RedisStoreClientTest : public StoreClientTestBase {
 
   virtual ~RedisStoreClientTest() {}
 
+  static void SetUpTestCase() { RedisServiceManagerForTest::SetUpTestCase(); }
+
+  static void TearDownTestCase() { RedisServiceManagerForTest::TearDownTestCase(); }
+
   void InitStoreClient() override {
     RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true);
     redis_client_ = std::make_shared<RedisClient>(options);
@@ -40,113 +44,11 @@ class RedisStoreClientTest : public StoreClientTestBase {
   std::shared_ptr<RedisClient> redis_client_;
 };
 
-TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) {
-  // AsyncPut without index.
-  auto put_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_->AsyncPut(table_name_, elem.first.Binary(),
-                                         elem.second.SerializeAsString(), put_calllback));
-  }
-  WaitPendingDone();
+TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); }
 
-  // AsyncGet
-  auto get_callback = [this](const Status &status,
-                             const boost::optional<std::string> &result) {
-    RAY_CHECK_OK(status);
-    RAY_CHECK(result);
-    rpc::ActorTableData data;
-    RAY_CHECK(data.ParseFromString(*result));
-    ActorID actor_id = ActorID::FromBinary(data.actor_id());
-    auto it = key_to_value_.find(actor_id);
-    RAY_CHECK(it != key_to_value_.end());
-    --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();
-}
+TEST_F(RedisStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); }
 
-TEST_F(RedisStoreClientTest, AsyncDeleteTest) {
-  // AsyncPut
-  auto put_calllback = [this](const Status &status) { --pending_count_; };
-  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();
-
-  // 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));
-  }
-  WaitPendingDone();
-}
-
-TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) {
-  // 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
-  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);
-        static std::unordered_set<ActorID> received_keys;
-        for (const auto &item : result) {
-          const ActorID &actor_id = ActorID::FromBinary(item.first);
-          auto it = received_keys.find(actor_id);
-          RAY_CHECK(it == received_keys.end());
-          received_keys.emplace(actor_id);
-
-          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());
-        }
-        pending_count_ -= result.size();
-      };
-
-  pending_count_ += key_to_value_.size();
-  RAY_CHECK_OK(store_client_->AsyncGetAll(table_name_, get_all_callback));
-  WaitPendingDone();
-}
+TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { TestAsyncGetAll(); }
 
 }  // namespace gcs
 

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

@@ -29,7 +29,7 @@ namespace ray {
 
 namespace gcs {
 
-class StoreClientTestBase : public RedisServiceManagerForTest {
+class StoreClientTestBase : public ::testing::Test {
  public:
   StoreClientTestBase() {}
 
@@ -57,6 +57,156 @@ class StoreClientTestBase : public RedisServiceManagerForTest {
   virtual void DisconnectStoreClient() = 0;
 
  protected:
+  void TestAsyncPutAndAsyncGet() {
+    // AsyncPut without index.
+    auto put_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_->AsyncPut(table_name_, elem.first.Binary(),
+                                           elem.second.SerializeAsString(),
+                                           put_calllback));
+    }
+    WaitPendingDone();
+
+    // AsyncGet
+    auto get_callback = [this](const Status &status,
+                               const boost::optional<std::string> &result) {
+      RAY_CHECK_OK(status);
+      RAY_CHECK(result);
+      rpc::ActorTableData data;
+      RAY_CHECK(data.ParseFromString(*result));
+      ActorID actor_id = ActorID::FromBinary(data.actor_id());
+      auto it = key_to_value_.find(actor_id);
+      RAY_CHECK(it != key_to_value_.end());
+      --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 TestAsyncDelete() {
+    // AsyncPut
+    auto put_calllback = [this](const Status &status) { --pending_count_; };
+    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();
+
+    // 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));
+    }
+    WaitPendingDone();
+  }
+
+  void TestAsyncPutAndDeleteWithIndex() {
+    // AsyncPut with index
+    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(),
+          elem.second.SerializeAsString(), put_calllback));
+    }
+    WaitPendingDone();
+
+    // AsyncDelete by index
+    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(),
+                                                     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
+    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);
+          static std::unordered_set<ActorID> received_keys;
+          for (const auto &item : result) {
+            const ActorID &actor_id = ActorID::FromBinary(item.first);
+            auto it = received_keys.find(actor_id);
+            RAY_CHECK(it == received_keys.end());
+            received_keys.emplace(actor_id);
+
+            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());
+          }
+          pending_count_ -= result.size();
+        };
+
+    pending_count_ += key_to_value_.size();
+    RAY_CHECK_OK(store_client_->AsyncGetAll(table_name_, get_all_callback));
+    WaitPendingDone();
+  }
+
   void GenTestData() {
     for (size_t i = 0; i < key_count_; i++) {
       rpc::ActorTableData actor;