[core][bugfix] Fix the C++ GcsClient Del not respecting del_by_prefix (…

…ray-project#45604) Signed-off-by: Ruiyang Wang <[email protected]>
nemo9cby · Jun 2, 2024 · b4d8083 · b4d8083
1 parent a6a9170
commit b4d8083
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/src/ray/gcs/gcs_client/accessor.cc b/src/ray/gcs/gcs_client/accessor.cc
@@ -1151,6 +1151,7 @@ Status InternalKVAccessor::AsyncInternalKVDel(const std::string &ns,
   rpc::InternalKVDelRequest req;
   req.set_namespace_(ns);
   req.set_key(key);
+  req.set_del_by_prefix(del_by_prefix);
   client_impl_->GetGcsRpcClient().InternalKVDel(
       req,
       [callback](const Status &status, const rpc::InternalKVDelReply &reply) {

diff --git a/src/ray/gcs/gcs_client/test/gcs_client_test.cc b/src/ray/gcs/gcs_client/test/gcs_client_test.cc
@@ -1034,6 +1034,43 @@ TEST_P(GcsClientTest, TestRegisterHeadNode) {
   }
 }
 
+TEST_P(GcsClientTest, TestInternalKVDelByPrefix) {
+  // Test Del can del by prefix
+  bool added;
+  RAY_CHECK_OK(gcs_client_->InternalKV().Put("test_ns",
+                                             "test_key1",
+                                             "test_value1",
+                                             /*overwrite=*/false,
+                                             /*timeout_ms=*/-1,
+                                             added));
+  ASSERT_TRUE(added);
+  RAY_CHECK_OK(gcs_client_->InternalKV().Put("test_ns",
+                                             "test_key2",
+                                             "test_value2",
+                                             /*overwrite=*/false,
+                                             /*timeout_ms=*/-1,
+                                             added));
+  ASSERT_TRUE(added);
+  RAY_CHECK_OK(gcs_client_->InternalKV().Put("test_ns",
+                                             "other_key",
+                                             "test_value3",
+                                             /*overwrite=*/false,
+                                             /*timeout_ms=*/-1,
+                                             added));
+  ASSERT_TRUE(added);
+
+  int num_deleted;
+  RAY_CHECK_OK(gcs_client_->InternalKV().Del(
+      "test_ns", "test_key", /*del_by_prefix=*/true, /*timeout_ms=*/-1, num_deleted));
+  ASSERT_EQ(num_deleted, 2);
+
+  // ... and the other key should still be there
+  std::string value;
+  RAY_CHECK_OK(
+      gcs_client_->InternalKV().Get("test_ns", "other_key", /*timeout_ms=*/-1, value));
+  ASSERT_EQ(value, "test_value3");
+}
+
 // TODO(sang): Add tests after adding asyncAdd
 
 }  // namespace ray