[core] Introduces Postable for InternalKVInterface.

python/ray/_raylet.pyx

@@ -5249,7 +5249,7 @@ cdef void async_callback(shared_ptr[CRayObject] obj,
         user_callback = <object>user_callback_ptr
         user_callback(result)
     except Exception:
-        # Only log the error here because this calllback is called from Cpp
+        # Only log the error here because this callback is called from Cpp
         # and Cython will ignore the exception anyway
         logger.exception(f"failed to run async callback (user func)")
     finally:

python/ray/includes/global_state_accessor.pxd

@@ -102,7 +102,7 @@ cdef extern from * namespace "ray::gcs" nogil:
         std::make_unique<RedisStoreClient>(std::move(redis_client)));
 
       bool ret_val = false;
-      cli->Get("session", key, [&](std::optional<std::string> result) {
+      cli->Get("session", key, {[&](std::optional<std::string> result) {
         if (result.has_value()) {
           *data = result.value();
           ret_val = true;
@@ -111,7 +111,7 @@ cdef extern from * namespace "ray::gcs" nogil:
                         << " from persistent storage.";
           ret_val = false;
         }
-      });
+      }, io_service});
       io_service.run_for(std::chrono::milliseconds(1000));
 
       return ret_val;

python/ray/serve/tests/test_callback.py

@@ -216,7 +216,7 @@ def test_http_proxy_return_aribitary_objects(ray_instance):
     ],
     indirect=True,
 )
-def test_http_proxy_calllback_failures(ray_instance, capsys):
+def test_http_proxy_callback_failures(ray_instance, capsys):
     """Test http proxy keeps restarting when callback function fails"""
 
     try:

src/mock/ray/gcs/gcs_server/gcs_actor_manager.h

@@ -28,14 +28,16 @@ namespace gcs {
 class MockGcsActorManager : public GcsActorManager {
  public:
   MockGcsActorManager(RuntimeEnvManager &runtime_env_manager,
-                      GcsFunctionManager &function_manager)
+                      GcsFunctionManager &function_manager,
+                      instrumented_io_context &io_service)
       : GcsActorManager(
             nullptr,
             nullptr,
             nullptr,
             runtime_env_manager,
             function_manager,
             [](const ActorID &) {},
+            io_service,
             [](const rpc::Address &) { return nullptr; }) {}
 
   MOCK_METHOD(void,

src/mock/ray/gcs/gcs_server/gcs_kv_manager.h

@@ -20,78 +20,78 @@ namespace gcs {
 
 class MockInternalKVInterface : public ray::gcs::InternalKVInterface {
  public:
-  MockInternalKVInterface() {}
+  explicit MockInternalKVInterface() = default;
 
   MOCK_METHOD(void,
               Get,
               (const std::string &ns,
                const std::string &key,
-               std::function<void(std::optional<std::string>)> callback),
+               Postable<void(std::optional<std::string>)> callback),
+              (override));
+  MOCK_METHOD(void,
+              MultiGet,
+              (const std::string &ns,
+               const std::vector<std::string> &keys,
+               Postable<void(std::unordered_map<std::string, std::string>)> callback),
               (override));
-  MOCK_METHOD(
-      void,
-      MultiGet,
-      (const std::string &ns,
-       const std::vector<std::string> &keys,
-       std::function<void(std::unordered_map<std::string, std::string>)> callback),
-      (override));
   MOCK_METHOD(void,
               Put,
               (const std::string &ns,
                const std::string &key,
                const std::string &value,
                bool overwrite,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
   MOCK_METHOD(void,
               Del,
               (const std::string &ns,
                const std::string &key,
                bool del_by_prefix,
-               std::function<void(int64_t)> callback),
+               Postable<void(int64_t)> callback),
               (override));
   MOCK_METHOD(void,
               Exists,
               (const std::string &ns,
                const std::string &key,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
   MOCK_METHOD(void,
               Keys,
               (const std::string &ns,
                const std::string &prefix,
-               std::function<void(std::vector<std::string>)> callback),
+               Postable<void(std::vector<std::string>)> callback),
               (override));
 };
 
 // Fake internal KV interface that simply stores keys and values in a C++ map.
 // Only supports Put and Get.
 // Warning: Naively prepends the namespace to the key, so e.g.
 // the (namespace, key) pairs ("a", "bc") and ("ab", "c") will collide which is a bug.
-
+// TODO(ryw): DO NOT SUBMIT. Get and Put used to be sync, now it's async. We need to use
+// promise wait.
 class FakeInternalKVInterface : public ray::gcs::InternalKVInterface {
  public:
-  FakeInternalKVInterface() {}
+  explicit FakeInternalKVInterface() = default;
 
   // The C++ map.
   std::unordered_map<std::string, std::string> kv_store_ = {};
 
   void Get(const std::string &ns,
            const std::string &key,
-           std::function<void(std::optional<std::string>)> callback) override {
+           Postable<void(std::optional<std::string>)> callback) override {
     std::string full_key = ns + key;
     auto it = kv_store_.find(full_key);
     if (it == kv_store_.end()) {
-      callback(std::nullopt);
+      std::move(callback).Post("FakeInternalKVInterface::Get", std::nullopt);
     } else {
-      callback(it->second);
+      std::move(callback).Post("FakeInternalKVInterface::Get", it->second);
     }
   }
 
-  void MultiGet(const std::string &ns,
-                const std::vector<std::string> &keys,
-                std::function<void(std::unordered_map<std::string, std::string>)>
-                    callback) override {
+  void MultiGet(
+      const std::string &ns,
+      const std::vector<std::string> &keys,
+      Postable<void(std::unordered_map<std::string, std::string>)> callback) override {
     std::unordered_map<std::string, std::string> result;
     for (const auto &key : keys) {
       std::string full_key = ns + key;
@@ -100,20 +100,20 @@ class FakeInternalKVInterface : public ray::gcs::InternalKVInterface {
         result[key] = it->second;
       }
     }
-    callback(result);
+    std::move(callback).Post("FakeInternalKVInterface::MultiGet", result);
   }
 
   void Put(const std::string &ns,
            const std::string &key,
            const std::string &value,
            bool overwrite,
-           std::function<void(bool)> callback) override {
+           Postable<void(bool)> callback) override {
     std::string full_key = ns + key;
     if (kv_store_.find(full_key) != kv_store_.end() && !overwrite) {
-      callback(false);
+      std::move(callback).Post("FakeInternalKVInterface::Put", false);
     } else {
       kv_store_[full_key] = value;
-      callback(true);
+      std::move(callback).Post("FakeInternalKVInterface::Put", true);
     }
   }
 
@@ -122,19 +122,19 @@ class FakeInternalKVInterface : public ray::gcs::InternalKVInterface {
               (const std::string &ns,
                const std::string &key,
                bool del_by_prefix,
-               std::function<void(int64_t)> callback),
+               Postable<void(int64_t)> callback),
               (override));
   MOCK_METHOD(void,
               Exists,
               (const std::string &ns,
                const std::string &key,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
   MOCK_METHOD(void,
               Keys,
               (const std::string &ns,
                const std::string &prefix,
-               std::function<void(std::vector<std::string>)> callback),
+               Postable<void(std::vector<std::string>)> callback),
               (override));
 };
 

src/mock/ray/gcs/gcs_server/gcs_node_manager.h

@@ -18,10 +18,11 @@ namespace gcs {
 
 class MockGcsNodeManager : public GcsNodeManager {
  public:
-  MockGcsNodeManager()
+  MockGcsNodeManager(instrumented_io_context &io_context)
       : GcsNodeManager(/*gcs_publisher=*/nullptr,
                        /*gcs_table_storage=*/nullptr,
                        /*raylet_client_pool=*/nullptr,
+                       io_context,
                        /*cluster_id=*/ClusterID::Nil()) {}
   MOCK_METHOD(void,
               HandleRegisterNode,

src/mock/ray/gcs/gcs_server/gcs_resource_manager.h

@@ -18,10 +18,8 @@ namespace ray {
 namespace gcs {
 static instrumented_io_context __mock_io_context_;
 static ClusterResourceManager __mock_cluster_resource_manager_(__mock_io_context_);
-static GcsNodeManager __mock_gcs_node_manager_(nullptr,
-                                               nullptr,
-                                               nullptr,
-                                               ClusterID::Nil());
+static GcsNodeManager __mock_gcs_node_manager_(
+    nullptr, nullptr, nullptr, __mock_io_context_, ClusterID::Nil());
 
 class MockGcsResourceManager : public GcsResourceManager {
  public:

src/mock/ray/gcs/store_client/in_memory_store_client.h

@@ -22,65 +22,58 @@ class MockInMemoryStoreClient : public InMemoryStoreClient {
               (const std::string &table_name,
                const std::string &key,
                const std::string &data,
-               const StatusCallback &callback),
-              (override));
-  MOCK_METHOD(Status,
-              AsyncPutWithIndex,
-              (const std::string &table_name,
-               const std::string &key,
-               const std::string &index_key,
-               const std::string &data,
-               const StatusCallback &callback),
+               bool overwrite,
+               Postable<void(bool)> callback),
               (override));
+
   MOCK_METHOD(Status,
               AsyncGet,
               (const std::string &table_name,
                const std::string &key,
-               const OptionalItemCallback<std::string> &callback),
-              (override));
-  MOCK_METHOD(Status,
-              AsyncGetByIndex,
-              (const std::string &table_name,
-               const std::string &index_key,
-               (const MapCallback<std::string, std::string> &callback)),
+               ToPostable<OptionalItemCallback<std::string>> callback),
               (override));
+
   MOCK_METHOD(Status,
               AsyncGetAll,
               (const std::string &table_name,
-               (const MapCallback<std::string, std::string> &callback)),
+               Postable<void(absl::flat_hash_map<std::string, std::string>)> callback),
               (override));
+
   MOCK_METHOD(Status,
-              AsyncDelete,
+              AsyncMultiGet,
               (const std::string &table_name,
-               const std::string &key,
-               const StatusCallback &callback),
+               const std::vector<std::string> &keys,
+               Postable<void(absl::flat_hash_map<std::string, std::string>)> callback),
               (override));
+
   MOCK_METHOD(Status,
-              AsyncDeleteWithIndex,
+              AsyncDelete,
               (const std::string &table_name,
                const std::string &key,
-               const std::string &index_key,
-               const StatusCallback &callback),
+               Postable<void(bool)> callback),
               (override));
+
   MOCK_METHOD(Status,
               AsyncBatchDelete,
               (const std::string &table_name,
                const std::vector<std::string> &keys,
-               const StatusCallback &callback),
+               Postable<void(int64_t)> callback),
               (override));
+
   MOCK_METHOD(Status,
-              AsyncBatchDeleteWithIndex,
+              AsyncGetKeys,
               (const std::string &table_name,
-               const std::vector<std::string> &keys,
-               const std::vector<std::string> &index_keys,
-               const StatusCallback &callback),
+               const std::string &prefix,
+               Postable<void(std::vector<std::string>)> callback),
               (override));
+
   MOCK_METHOD(Status,
-              AsyncDeleteByIndex,
+              AsyncExists,
               (const std::string &table_name,
-               const std::string &index_key,
-               const StatusCallback &callback),
+               const std::string &key,
+               Postable<void(bool)> callback),
               (override));
+
   MOCK_METHOD(int, GetNextJobID, (), (override));
 };
 

src/mock/ray/gcs/store_client/redis_store_client.h

@@ -25,14 +25,7 @@ class MockRedisStoreClient : public RedisStoreClient {
                const std::string &data,
                const StatusCallback &callback),
               (override));
-  MOCK_METHOD(Status,
-              AsyncPutWithIndex,
-              (const std::string &table_name,
-               const std::string &key,
-               const std::string &index_key,
-               const std::string &data,
-               const StatusCallback &callback),
-              (override));
+
   MOCK_METHOD(Status,
               AsyncGet,
               (const std::string &table_name,

src/mock/ray/gcs/store_client/store_client.h

@@ -23,50 +23,50 @@ class MockStoreClient : public StoreClient {
                const std::string &key,
                const std::string &data,
                bool overwrite,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
   MOCK_METHOD(Status,
               AsyncGet,
               (const std::string &table_name,
                const std::string &key,
-               const OptionalItemCallback<std::string> &callback),
+               ToPostable<OptionalItemCallback<std::string>> callback),
               (override));
   MOCK_METHOD(Status,
               AsyncGetAll,
               (const std::string &table_name,
-               (const MapCallback<std::string, std::string> &callback)),
+               Postable<void(absl::flat_hash_map<std::string, std::string>)> callback),
               (override));
   MOCK_METHOD(Status,
               AsyncMultiGet,
               (const std::string &table_name,
                const std::vector<std::string> &key,
-               (const MapCallback<std::string, std::string> &callback)),
+               Postable<void(absl::flat_hash_map<std::string, std::string>)> callback),
               (override));
   MOCK_METHOD(Status,
               AsyncDelete,
               (const std::string &table_name,
                const std::string &key,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
   MOCK_METHOD(Status,
               AsyncBatchDelete,
               (const std::string &table_name,
                const std::vector<std::string> &keys,
-               std::function<void(int64_t)> callback),
+               Postable<void(int64_t)> callback),
               (override));
   MOCK_METHOD(int, GetNextJobID, (), (override));
   MOCK_METHOD(Status,
               AsyncGetKeys,
               (const std::string &table_name,
                const std::string &prefix,
-               std::function<void(std::vector<std::string>)> callback),
+               Postable<void(std::vector<std::string>)> callback),
               (override));
 
   MOCK_METHOD(Status,
               AsyncExists,
               (const std::string &table_name,
                const std::string &key,
-               std::function<void(bool)> callback),
+               Postable<void(bool)> callback),
               (override));
 };