From ddc1e483fbd3b38f83f79a2a959b9cb5c5f79b3c Mon Sep 17 00:00:00 2001 From: kisuke95 <2522134184@qq.com> Date: Wed, 5 Aug 2020 18:05:51 +0800 Subject: [PATCH] Fix actor table Delete bug (#9499) --- src/ray/gcs/gcs_server/gcs_table_storage.cc | 20 ++++++ src/ray/gcs/gcs_server/gcs_table_storage.h | 20 +++++- .../test/gcs_table_storage_test_base.h | 50 ++++++++++++--- .../store_client/in_memory_store_client.cc | 64 +++++++++++++++++++ .../gcs/store_client/in_memory_store_client.h | 9 +++ .../gcs/store_client/redis_store_client.cc | 27 ++++++++ src/ray/gcs/store_client/redis_store_client.h | 9 +++ src/ray/gcs/store_client/store_client.h | 27 ++++++++ .../test/in_memory_store_client_test.cc | 6 ++ .../test/redis_store_client_test.cc | 6 ++ .../test/store_client_test_base.h | 59 +++++++++++++++++ 11 files changed, 286 insertions(+), 11 deletions(-) diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc index 541f5a0fa..0f4dfbb81 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.cc +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -106,6 +106,26 @@ Status GcsTableWithJobId::DeleteByJobId(const JobID &job_id, callback); } +template +Status GcsTableWithJobId::Delete(const Key &key, + const StatusCallback &callback) { + return this->store_client_->AsyncDeleteWithIndex( + this->table_name_, key.Binary(), GetJobIdFromKey(key).Binary(), callback); +} + +template +Status GcsTableWithJobId::BatchDelete(const std::vector &keys, + const StatusCallback &callback) { + std::vector keys_to_delete; + std::vector 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; template class GcsTable; template class GcsTable; diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.h b/src/ray/gcs/gcs_server/gcs_table_storage.h index 726144bc6..b596e9a9d 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.h +++ b/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 &keys, const StatusCallback &callback); + virtual Status BatchDelete(const std::vector &keys, + const StatusCallback &callback); protected: std::string table_name_; @@ -132,6 +133,21 @@ class GcsTableWithJobId : public GcsTable { /// \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 &keys, + const StatusCallback &callback) override; + protected: virtual JobID GetJobIdFromKey(const Key &key) = 0; }; diff --git a/src/ray/gcs/gcs_server/test/gcs_table_storage_test_base.h b/src/ray/gcs/gcs_server/test/gcs_table_storage_test_base.h index 9c471a578..49168b0f3 100644 --- a/src/ray/gcs/gcs_server/test/gcs_table_storage_test_base.h +++ b/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 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 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 @@ -137,6 +158,17 @@ class GcsTableStorageTestBase : public ::testing::Test { WaitPendingDone(); } + template + void BatchDelete(TABLE &table, const std::vector &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 &pending_count) { diff --git a/src/ray/gcs/store_client/in_memory_store_client.cc b/src/ray/gcs/store_client/in_memory_store_client.cc index 494b472a6..789b130a8 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.cc +++ b/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 &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 &keys, + const std::vector &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 &callback) { diff --git a/src/ray/gcs/store_client/in_memory_store_client.h b/src/ray/gcs/store_client/in_memory_store_client.h index a75ca0723..a26ab5e69 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.h +++ b/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 &keys, const StatusCallback &callback) override; + Status AsyncBatchDeleteWithIndex(const std::string &table_name, + const std::vector &keys, + const std::vector &index_keys, + const StatusCallback &callback) override; + Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) override; diff --git a/src/ray/gcs/store_client/redis_store_client.cc b/src/ray/gcs/store_client/redis_store_client.cc index ce3c18c62..f20270fd3 100644 --- a/src/ray/gcs/store_client/redis_store_client.cc +++ b/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 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 &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 &keys, + const std::vector &index_keys, const StatusCallback &callback) { + RAY_CHECK(keys.size() == index_keys.size()); + + std::vector 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 &callback) { diff --git a/src/ray/gcs/store_client/redis_store_client.h b/src/ray/gcs/store_client/redis_store_client.h index ccf606cba..14d77b1b6 100644 --- a/src/ray/gcs/store_client/redis_store_client.h +++ b/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 &keys, const StatusCallback &callback) override; + Status AsyncBatchDeleteWithIndex(const std::string &table_name, + const std::vector &keys, + const std::vector &index_keys, + const StatusCallback &callback) override; + Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) override; diff --git a/src/ray/gcs/store_client/store_client.h b/src/ray/gcs/store_client/store_client.h index 38e6e8ebd..c993f1dca 100644 --- a/src/ray/gcs/store_client/store_client.h +++ b/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 &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 &keys, + const std::vector &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. diff --git a/src/ray/gcs/store_client/test/in_memory_store_client_test.cc b/src/ray/gcs/store_client/test/in_memory_store_client_test.cc index b259b02af..3f84b8e40 100644 --- a/src/ray/gcs/store_client/test/in_memory_store_client_test.cc +++ b/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 diff --git a/src/ray/gcs/store_client/test/redis_store_client_test.cc b/src/ray/gcs/store_client/test/redis_store_client_test.cc index ba7790fba..afb570a57 100644 --- a/src/ray/gcs/store_client/test/redis_store_client_test.cc +++ b/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 diff --git a/src/ray/gcs/store_client/test/store_client_test_base.h b/src/ray/gcs/store_client/test/store_client_test_base.h index 0b062d976..a2a0a0fb6 100644 --- a/src/ray/gcs/store_client/test/store_client_test_base.h +++ b/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 keys; + std::vector 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;