Fix actor table Delete bug (#9499)

src/ray/gcs/gcs_server/gcs_table_storage.cc

@@ -106,6 +106,26 @@ Status GcsTableWithJobId<Key, Data>::DeleteByJobId(const JobID &job_id,
                                                  callback);
 }
 
+template <typename Key, typename Data>
+Status GcsTableWithJobId<Key, Data>::Delete(const Key &key,
+                                            const StatusCallback &callback) {
+  return this->store_client_->AsyncDeleteWithIndex(
+      this->table_name_, key.Binary(), GetJobIdFromKey(key).Binary(), callback);
+}
+
+template <typename Key, typename Data>
+Status GcsTableWithJobId<Key, Data>::BatchDelete(const std::vector<Key> &keys,
+                                                 const StatusCallback &callback) {
+  std::vector<std::string> keys_to_delete;
+  std::vector<std::string> indexs_to_delete;
+  for (auto key : keys) {
+    keys_to_delete.push_back(key.Binary());
+    indexs_to_delete.push_back(GetJobIdFromKey(key).Binary());
+  }
+  return this->store_client_->AsyncBatchDeleteWithIndex(this->table_name_, keys_to_delete,
+                                                        indexs_to_delete, callback);
+}
+
 template class GcsTable<JobID, JobTableData>;
 template class GcsTable<ClientID, GcsNodeInfo>;
 template class GcsTable<ClientID, ResourceMap>;

src/ray/gcs/gcs_server/gcs_table_storage.h

@@ -84,14 +84,15 @@ class GcsTable {
   /// \param key The key that will be deleted from the table.
   /// \param callback Callback that will be called after delete finishes.
   /// \return Status
-  Status Delete(const Key &key, const StatusCallback &callback);
+  virtual Status Delete(const Key &key, const StatusCallback &callback);
 
   /// Delete a batch of data from the table asynchronously.
   ///
   /// \param keys The batch key that will be deleted from the table.
   /// \param callback Callback that will be called after delete finishes.
   /// \return Status
-  Status BatchDelete(const std::vector<Key> &keys, const StatusCallback &callback);
+  virtual Status BatchDelete(const std::vector<Key> &keys,
+                             const StatusCallback &callback);
 
  protected:
   std::string table_name_;
@@ -132,6 +133,21 @@ class GcsTableWithJobId : public GcsTable<Key, Data> {
   /// \return Status
   Status DeleteByJobId(const JobID &job_id, const StatusCallback &callback);
 
+  /// Delete data and index from the table asynchronously.
+  ///
+  /// \param key The key that will be deleted from the table.
+  /// \param callback Callback that will be called after delete finishes.
+  /// \return Status
+  Status Delete(const Key &key, const StatusCallback &callback) override;
+
+  /// Delete a batch of data and index from the table asynchronously.
+  ///
+  /// \param keys The batch key that will be deleted from the table.
+  /// \param callback Callback that will be called after delete finishes.
+  /// \return Status
+  Status BatchDelete(const std::vector<Key> &keys,
+                     const StatusCallback &callback) override;
+
  protected:
   virtual JobID GetJobIdFromKey(const Key &key) = 0;
 };

src/ray/gcs/gcs_server/test/gcs_table_storage_test_base.h

@@ -52,29 +52,50 @@ class GcsTableStorageTestBase : public ::testing::Test {
 
     // Delete.
     Delete(table, job1_id);
+    Delete(table, job2_id);
     ASSERT_EQ(Get(table, job1_id, values), 0);
-    ASSERT_EQ(Get(table, job2_id, values), 1);
+    ASSERT_EQ(Get(table, job2_id, values), 0);
   }
 
   void TestGcsTableWithJobIdApi() {
     auto table = gcs_table_storage_->ActorTable();
-    JobID job_id = JobID::FromInt(3);
-    auto actor_table_data = Mocker::GenActorTableData(job_id);
-    ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
+    JobID job_id1 = JobID::FromInt(1);
+    JobID job_id2 = JobID::FromInt(2);
+    JobID job_id3 = JobID::FromInt(3);
+    auto actor_table_data1 = Mocker::GenActorTableData(job_id1);
+    auto actor_table_data2 = Mocker::GenActorTableData(job_id2);
+    auto actor_table_data3 = Mocker::GenActorTableData(job_id3);
+    ActorID actor_id1 = ActorID::FromBinary(actor_table_data1->actor_id());
+    ActorID actor_id2 = ActorID::FromBinary(actor_table_data2->actor_id());
+    ActorID actor_id3 = ActorID::FromBinary(actor_table_data3->actor_id());
 
     // Put.
-    Put(table, actor_id, *actor_table_data);
+    Put(table, actor_id1, *actor_table_data1);
+    Put(table, actor_id2, *actor_table_data2);
+    Put(table, actor_id3, *actor_table_data3);
 
     // Get.
     std::vector<rpc::ActorTableData> values;
-    ASSERT_EQ(Get(table, actor_id, values), 1);
+    ASSERT_EQ(Get(table, actor_id1, values), 1);
 
     // Get by job id.
-    ASSERT_EQ(GetByJobId(table, job_id, actor_id, values), 1);
+    ASSERT_EQ(GetByJobId(table, job_id1, actor_id1, values), 1);
 
     // Delete.
-    Delete(table, actor_id);
-    ASSERT_EQ(Get(table, actor_id, values), 0);
+    Delete(table, actor_id1);
+    ASSERT_EQ(Get(table, actor_id1, values), 0);
+    ASSERT_EQ(GetByJobId(table, job_id1, actor_id1, values), 0);
+
+    std::vector<ActorID> keys;
+    keys.push_back(actor_id2);
+    keys.push_back(actor_id3);
+    BatchDelete(table, keys);
+
+    ASSERT_EQ(Get(table, actor_id2, values), 0);
+    ASSERT_EQ(GetByJobId(table, job_id2, actor_id2, values), 0);
+
+    ASSERT_EQ(Get(table, actor_id3, values), 0);
+    ASSERT_EQ(GetByJobId(table, job_id3, actor_id3, values), 0);
   }
 
   template <typename TABLE, typename KEY, typename VALUE>
@@ -137,6 +158,17 @@ class GcsTableStorageTestBase : public ::testing::Test {
     WaitPendingDone();
   }
 
+  template <typename TABLE, typename KEY>
+  void BatchDelete(TABLE &table, const std::vector<KEY> &keys) {
+    auto on_done = [this](const Status &status) {
+      RAY_CHECK_OK(status);
+      --pending_count_;
+    };
+    ++pending_count_;
+    RAY_CHECK_OK(table.BatchDelete(keys, on_done));
+    WaitPendingDone();
+  }
+
   void WaitPendingDone() { WaitPendingDone(pending_count_); }
 
   void WaitPendingDone(std::atomic<int> &pending_count) {

src/ray/gcs/store_client/in_memory_store_client.cc

@@ -77,6 +77,36 @@ Status InMemoryStoreClient::AsyncDelete(const std::string &table_name,
   return Status::OK();
 }
 
+Status InMemoryStoreClient::AsyncDeleteWithIndex(const std::string &table_name,
+                                                 const std::string &key,
+                                                 const std::string &index_key,
+                                                 const StatusCallback &callback) {
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+  // Remove key-value data.
+  table->records_.erase(key);
+
+  // Remove index-key data.
+  auto iter = table->index_keys_.find(index_key);
+  if (iter != table->index_keys_.end()) {
+    auto it = std::find(iter->second.begin(), iter->second.end(), key);
+    if (it != iter->second.end()) {
+      iter->second.erase(it);
+      if (iter->second.size() == 0) {
+        table->index_keys_.erase(iter);
+      }
+    }
+  }
+
+  main_io_service_.post([callback]() {
+    if (callback) {
+      callback(Status::OK());
+    }
+  });
+
+  return Status::OK();
+}
+
 Status InMemoryStoreClient::AsyncBatchDelete(const std::string &table_name,
                                              const std::vector<std::string> &keys,
                                              const StatusCallback &callback) {
@@ -89,6 +119,40 @@ Status InMemoryStoreClient::AsyncBatchDelete(const std::string &table_name,
   return Status::OK();
 }
 
+Status InMemoryStoreClient::AsyncBatchDeleteWithIndex(
+    const std::string &table_name, const std::vector<std::string> &keys,
+    const std::vector<std::string> &index_keys, const StatusCallback &callback) {
+  RAY_CHECK(keys.size() == index_keys.size());
+
+  auto table = GetOrCreateTable(table_name);
+  absl::MutexLock lock(&(table->mutex_));
+
+  for (size_t i = 0; i < keys.size(); ++i) {
+    const std::string &key = keys[i];
+    const std::string &index_key = index_keys[i];
+    table->records_.erase(key);
+
+    auto iter = table->index_keys_.find(index_key);
+    if (iter != table->index_keys_.end()) {
+      auto it = std::find(iter->second.begin(), iter->second.end(), key);
+      if (it != iter->second.end()) {
+        iter->second.erase(it);
+        if (iter->second.size() == 0) {
+          table->index_keys_.erase(iter);
+        }
+      }
+    }
+  }
+
+  main_io_service_.post([callback]() {
+    if (callback) {
+      callback(Status::OK());
+    }
+  });
+
+  return Status::OK();
+}
+
 Status InMemoryStoreClient::AsyncGetByIndex(
     const std::string &table_name, const std::string &index_key,
     const MapCallback<std::string, std::string> &callback) {

src/ray/gcs/store_client/in_memory_store_client.h

@@ -50,10 +50,19 @@ class InMemoryStoreClient : public StoreClient {
   Status AsyncDelete(const std::string &table_name, const std::string &key,
                      const StatusCallback &callback) override;
 
+  Status AsyncDeleteWithIndex(const std::string &table_name, const std::string &key,
+                              const std::string &index_key,
+                              const StatusCallback &callback) override;
+
   Status AsyncBatchDelete(const std::string &table_name,
                           const std::vector<std::string> &keys,
                           const StatusCallback &callback) override;
 
+  Status AsyncBatchDeleteWithIndex(const std::string &table_name,
+                                   const std::vector<std::string> &keys,
+                                   const std::vector<std::string> &index_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

@@ -114,6 +114,18 @@ Status RedisStoreClient::AsyncDelete(const std::string &table_name,
   return shard_context->RunArgvAsync(args, delete_callback);
 }
 
+Status RedisStoreClient::AsyncDeleteWithIndex(const std::string &table_name,
+                                              const std::string &key,
+                                              const std::string &index_key,
+                                              const StatusCallback &callback) {
+  std::vector<std::string> redis_keys;
+  redis_keys.reserve(20);
+  redis_keys.push_back(GenRedisKey(table_name, key));
+  redis_keys.push_back(GenRedisKey(table_name, key, index_key));
+
+  return DeleteByKeys(redis_keys, callback);
+}
+
 Status RedisStoreClient::AsyncBatchDelete(const std::string &table_name,
                                           const std::vector<std::string> &keys,
                                           const StatusCallback &callback) {
@@ -125,6 +137,21 @@ Status RedisStoreClient::AsyncBatchDelete(const std::string &table_name,
   return DeleteByKeys(redis_keys, callback);
 }
 
+Status RedisStoreClient::AsyncBatchDeleteWithIndex(
+    const std::string &table_name, const std::vector<std::string> &keys,
+    const std::vector<std::string> &index_keys, const StatusCallback &callback) {
+  RAY_CHECK(keys.size() == index_keys.size());
+
+  std::vector<std::string> redis_keys;
+  redis_keys.reserve(2 * keys.size());
+  for (size_t i = 0; i < keys.size(); ++i) {
+    redis_keys.push_back(GenRedisKey(table_name, keys[i]));
+    redis_keys.push_back(GenRedisKey(table_name, keys[i], index_keys[i]));
+  }
+
+  return DeleteByKeys(redis_keys, callback);
+}
+
 Status RedisStoreClient::AsyncGetByIndex(
     const std::string &table_name, const std::string &index_key,
     const MapCallback<std::string, std::string> &callback) {

src/ray/gcs/store_client/redis_store_client.h

@@ -48,10 +48,19 @@ class RedisStoreClient : public StoreClient {
   Status AsyncDelete(const std::string &table_name, const std::string &key,
                      const StatusCallback &callback) override;
 
+  Status AsyncDeleteWithIndex(const std::string &table_name, const std::string &key,
+                              const std::string &index_key,
+                              const StatusCallback &callback) override;
+
   Status AsyncBatchDelete(const std::string &table_name,
                           const std::vector<std::string> &keys,
                           const StatusCallback &callback) override;
 
+  Status AsyncBatchDeleteWithIndex(const std::string &table_name,
+                                   const std::vector<std::string> &keys,
+                                   const std::vector<std::string> &index_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/store_client.h

@@ -92,6 +92,19 @@ class StoreClient {
   virtual Status AsyncDelete(const std::string &table_name, const std::string &key,
                              const StatusCallback &callback) = 0;
 
+  /// Delete data from the given table asynchronously, this can delete
+  /// key--value and index--key.
+  ///
+  /// \param table_name The name of the table from which data is to be deleted.
+  /// \param key The key that will be deleted from the table.
+  /// \param index_key The index key of the given key.
+  /// \param callback Callback that will be called after delete finishes.
+  /// \return Status
+  virtual Status AsyncDeleteWithIndex(const std::string &table_name,
+                                      const std::string &key,
+                                      const std::string &index_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.
@@ -102,6 +115,20 @@ class StoreClient {
                                   const std::vector<std::string> &keys,
                                   const StatusCallback &callback) = 0;
 
+  /// Batch delete data from the given table asynchronously, this can delete all
+  /// key--value data and index--key data.
+  ///
+  /// \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 index_keys The index keys of the given keys, they are in one-to-one
+  ///                   correspondence
+  /// \param callback Callback that will be called after delete finishes.
+  /// \return Status
+  virtual Status AsyncBatchDeleteWithIndex(const std::string &table_name,
+                                           const std::vector<std::string> &keys,
+                                           const std::vector<std::string> &index_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

@@ -38,6 +38,12 @@ TEST_F(InMemoryStoreClientTest, AsyncGetAllAndBatchDeleteTest) {
   TestAsyncGetAllAndBatchDelete();
 }
 
+TEST_F(InMemoryStoreClientTest, TestAsyncDeleteWithIndex) { TestAsyncDeleteWithIndex(); }
+
+TEST_F(InMemoryStoreClientTest, TestAsyncBatchDeleteWithIndex) {
+  TestAsyncBatchDeleteWithIndex();
+}
+
 }  // namespace gcs
 
 }  // namespace ray

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

@@ -55,6 +55,12 @@ TEST_F(RedisStoreClientTest, AsyncGetAllAndBatchDeleteTest) {
   TestAsyncGetAllAndBatchDelete();
 }
 
+TEST_F(RedisStoreClientTest, TestAsyncDeleteWithIndex) { TestAsyncDeleteWithIndex(); }
+
+TEST_F(RedisStoreClientTest, TestAsyncBatchDeleteWithIndex) {
+  TestAsyncBatchDeleteWithIndex();
+}
+
 }  // namespace gcs
 
 }  // namespace ray

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

@@ -201,6 +201,37 @@ class StoreClientTestBase : public ::testing::Test {
     WaitPendingDone();
   }
 
+  void DeleteWithIndex() {
+    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_->AsyncDeleteWithIndex(table_name_, elem.first.Binary(),
+                                                       key_to_index_[elem.first].Hex(),
+                                                       delete_calllback));
+    }
+    WaitPendingDone();
+  }
+
+  void BatchDeleteWithIndex() {
+    auto delete_calllback = [this](const Status &status) {
+      RAY_CHECK_OK(status);
+      --pending_count_;
+    };
+    ++pending_count_;
+    std::vector<std::string> keys;
+    std::vector<std::string> index_keys;
+    for (auto &elem : key_to_value_) {
+      keys.push_back(elem.first.Binary());
+      index_keys.push_back(key_to_index_[elem.first].Hex());
+    }
+    RAY_CHECK_OK(store_client_->AsyncBatchDeleteWithIndex(table_name_, keys, index_keys,
+                                                          delete_calllback));
+    WaitPendingDone();
+  }
+
   void TestAsyncPutAndAsyncGet() {
     // AsyncPut without index.
     Put();
@@ -242,6 +273,34 @@ class StoreClientTestBase : public ::testing::Test {
     GetEmpty();
   }
 
+  void TestAsyncDeleteWithIndex() {
+    // AsyncPut with index
+    PutWithIndex();
+
+    // AsyncGet with index
+    GetByIndex();
+
+    // AsyncDelete key-value and index-key
+    DeleteWithIndex();
+
+    // AsyncGet
+    GetEmpty();
+  }
+
+  void TestAsyncBatchDeleteWithIndex() {
+    // AsyncPut with index
+    PutWithIndex();
+
+    // AsyncGetAll
+    GetByIndex();
+
+    // AsyncBatchDeleteWithIndex
+    BatchDeleteWithIndex();
+
+    // AsyncGet
+    GetEmpty();
+  }
+
   void GenTestData() {
     for (size_t i = 0; i < key_count_; i++) {
       rpc::ActorTableData actor;