From 41072fbcc8b3aea9fe726e960ca695f6365e9e40 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Thu, 4 Jun 2020 20:51:43 +0800 Subject: [PATCH] Implement GetByJobId in gcs table storage (#8727) --- src/ray/gcs/gcs_server/gcs_table_storage.cc | 14 +- .../test/gcs_table_storage_test_base.h | 21 +++ .../store_client/in_memory_store_client.cc | 20 +++ .../gcs/store_client/in_memory_store_client.h | 3 + .../gcs/store_client/redis_store_client.cc | 132 ++++++++++-------- src/ray/gcs/store_client/redis_store_client.h | 29 ++-- src/ray/gcs/store_client/store_client.h | 10 ++ .../test/store_client_test_base.h | 21 +++ 8 files changed, 177 insertions(+), 73 deletions(-) diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc index 50061ce3e..627605a6c 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.cc +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -85,9 +85,17 @@ Status GcsTableWithJobId::Put(const Key &key, const Data &value, template Status GcsTableWithJobId::GetByJobId(const JobID &job_id, const MapCallback &callback) { - // TODO(ffbin): We will add this function after redis store client support - // AsyncGetByIndex interface. - return Status::NotImplemented("GetByJobId not implemented"); + auto on_done = [callback](const std::unordered_map &result) { + std::unordered_map values; + for (auto &item : result) { + Data data; + data.ParseFromString(item.second); + values[Key::FromBinary(item.first)] = std::move(data); + } + callback(values); + }; + return this->store_client_->AsyncGetByIndex(this->table_name_, job_id.Binary(), + on_done); } template 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 ddd6c59c7..2457373f1 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 @@ -67,6 +67,9 @@ class GcsTableStorageTestBase : public ::testing::Test { std::vector values; ASSERT_EQ(Get(table, actor_id, values), 1); + // Get by job id. + ASSERT_EQ(GetByJobId(table, job_id, actor_id, values), 1); + // Delete. Delete(table, actor_id); ASSERT_EQ(Get(table, actor_id, values), 0); @@ -96,6 +99,24 @@ class GcsTableStorageTestBase : public ::testing::Test { return values.size(); } + template + int GetByJobId(TABLE &table, const JobID &job_id, const KEY &key, + std::vector &values) { + auto on_done = [this, &values](const std::unordered_map &result) { + --pending_count_; + values.clear(); + if (!result.empty()) { + for (auto &item : result) { + values.push_back(item.second); + } + } + }; + ++pending_count_; + RAY_CHECK_OK(table.GetByJobId(job_id, on_done)); + WaitPendingDone(); + return values.size(); + } + template void Delete(TABLE &table, const KEY &key) { auto on_done = [this](Status status) { 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 b805276f9..86bb619bb 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.cc +++ b/src/ray/gcs/store_client/in_memory_store_client.cc @@ -89,6 +89,26 @@ Status InMemoryStoreClient::AsyncBatchDelete(const std::string &table_name, return Status::OK(); } +Status InMemoryStoreClient::AsyncGetByIndex( + const std::string &table_name, const std::string &index_key, + const MapCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + auto iter = table->index_keys_.find(index_key); + std::unordered_map result; + if (iter != table->index_keys_.end()) { + for (auto &key : iter->second) { + auto kv_iter = table->records_.find(key); + if (kv_iter != table->records_.end()) { + result[kv_iter->first] = kv_iter->second; + } + } + } + main_io_service_.post([result, callback]() { callback(result); }); + + return Status::OK(); +} + Status InMemoryStoreClient::AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &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 f9f893cc1..f51480c15 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.h +++ b/src/ray/gcs/store_client/in_memory_store_client.h @@ -42,6 +42,9 @@ class InMemoryStoreClient : public StoreClient { Status AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) override; + Status AsyncGetByIndex(const std::string &table_name, const std::string &index_key, + const MapCallback &callback) override; + Status AsyncGetAll(const std::string &table_name, const MapCallback &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 cd77251bc..2836e7b7b 100644 --- a/src/ray/gcs/store_client/redis_store_client.cc +++ b/src/ray/gcs/store_client/redis_store_client.cc @@ -89,12 +89,12 @@ Status RedisStoreClient::AsyncGetAll( const MapCallback &callback) { RAY_CHECK(callback); std::string match_pattern = GenRedisMatchPattern(table_name); - auto scanner = std::make_shared(redis_client_, table_name, match_pattern); + auto scanner = std::make_shared(redis_client_, table_name); auto on_done = [callback, scanner](const std::unordered_map &result) { callback(result); }; - return scanner->ScanKeysAndValues(on_done); + return scanner->ScanKeysAndValues(match_pattern, on_done); } Status RedisStoreClient::AsyncDelete(const std::string &table_name, @@ -125,11 +125,35 @@ Status RedisStoreClient::AsyncBatchDelete(const std::string &table_name, return DeleteByKeys(redis_keys, callback); } +Status RedisStoreClient::AsyncGetByIndex( + const std::string &table_name, const std::string &index_key, + const MapCallback &callback) { + RAY_CHECK(callback); + std::string match_pattern = GenRedisMatchPattern(table_name, index_key); + auto scanner = std::make_shared(redis_client_, table_name); + auto on_done = [this, callback, scanner, table_name, index_key]( + const Status &status, const std::vector &result) { + if (!result.empty()) { + std::vector keys; + keys.reserve(result.size()); + for (auto &item : result) { + keys.push_back( + GenRedisKey(table_name, GetKeyFromRedisKey(item, table_name, index_key))); + } + + RAY_CHECK_OK(MGetValues(redis_client_, table_name, keys, callback)); + } else { + callback(std::unordered_map()); + } + }; + return scanner->ScanKeys(match_pattern, on_done); +} + Status RedisStoreClient::AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) { std::string match_pattern = GenRedisMatchPattern(table_name, index_key); - auto scanner = std::make_shared(redis_client_, table_name, match_pattern); + auto scanner = std::make_shared(redis_client_, table_name); auto on_done = [this, table_name, index_key, callback, scanner]( const Status &status, const std::vector &result) { if (!result.empty()) { @@ -149,7 +173,7 @@ Status RedisStoreClient::AsyncDeleteByIndex(const std::string &table_name, } }; - return scanner->ScanKeys(on_done); + return scanner->ScanKeys(match_pattern, on_done); } Status RedisStoreClient::DoPut(const std::string &key, const std::string &data, @@ -248,42 +272,74 @@ std::string RedisStoreClient::GetKeyFromRedisKey(const std::string &redis_key, return redis_key.substr(pos, redis_key.size() - pos); } +Status RedisStoreClient::MGetValues( + std::shared_ptr redis_client, std::string table_name, + const std::vector &keys, + const ItemCallback> &callback) { + // The `MGET` command for each shard. + auto mget_commands_by_shards = GenCommandsByShards(redis_client, "MGET", keys); + + auto finished_count = std::make_shared(0); + int size = mget_commands_by_shards.size(); + for (auto &item : mget_commands_by_shards) { + auto mget_keys = std::move(item.second); + auto mget_callback = [table_name, finished_count, size, mget_keys, + callback](const std::shared_ptr &reply) { + std::unordered_map key_value_map; + if (!reply->IsNil()) { + auto value = reply->ReadAsStringArray(); + // The 0 th element of mget_keys is "MGET", so we start from the 1 th element. + for (int index = 0; index < (int)value.size(); ++index) { + key_value_map[GetKeyFromRedisKey(mget_keys[index + 1], table_name)] = + value[index]; + } + } + + ++(*finished_count); + if (*finished_count == size) { + callback(key_value_map); + } + }; + RAY_CHECK_OK(item.first->RunArgvAsync(mget_keys, mget_callback)); + } + return Status::OK(); +} + RedisStoreClient::RedisScanner::RedisScanner(std::shared_ptr redis_client, - std::string table_name, - std::string match_pattern) - : table_name_(std::move(table_name)), - match_pattern_(std::move(match_pattern)), - redis_client_(std::move(redis_client)) { + std::string table_name) + : table_name_(std::move(table_name)), redis_client_(std::move(redis_client)) { for (size_t index = 0; index < redis_client_->GetShardContexts().size(); ++index) { shard_to_cursor_[index] = 0; } } Status RedisStoreClient::RedisScanner::ScanKeysAndValues( + std::string match_pattern, const ItemCallback> &callback) { auto on_done = [this, callback](const Status &status, const std::vector &result) { if (result.empty()) { callback(std::unordered_map()); } else { - MGetValues(result, callback); + RAY_CHECK_OK(MGetValues(redis_client_, table_name_, result, callback)); } }; - return ScanKeys(on_done); + return ScanKeys(match_pattern, on_done); } Status RedisStoreClient::RedisScanner::ScanKeys( - const MultiItemCallback &callback) { + std::string match_pattern, const MultiItemCallback &callback) { auto on_done = [this, callback](const Status &status) { std::vector result; result.insert(result.begin(), keys_.begin(), keys_.end()); callback(status, result); }; - Scan(on_done); + Scan(match_pattern, on_done); return Status::OK(); } -void RedisStoreClient::RedisScanner::Scan(const StatusCallback &callback) { +void RedisStoreClient::RedisScanner::Scan(std::string match_pattern, + const StatusCallback &callback) { if (shard_to_cursor_.empty()) { callback(Status::OK()); return; @@ -296,18 +352,16 @@ void RedisStoreClient::RedisScanner::Scan(const StatusCallback &callback) { size_t shard_index = item.first; size_t cursor = item.second; - auto scan_callback = [this, shard_index, + auto scan_callback = [this, match_pattern, shard_index, callback](const std::shared_ptr &reply) { - OnScanCallback(shard_index, reply, callback); + OnScanCallback(match_pattern, shard_index, reply, callback); }; - // Scan by prefix from Redis. std::vector args = {"SCAN", std::to_string(cursor), - "MATCH", match_pattern_, + "MATCH", match_pattern, "COUNT", std::to_string(batch_count)}; auto shard_context = redis_client_->GetShardContexts()[shard_index]; Status status = shard_context->RunArgvAsync(args, scan_callback); - if (!status.ok()) { RAY_LOG(FATAL) << "Scan failed, status " << status.ToString(); } @@ -315,12 +369,11 @@ void RedisStoreClient::RedisScanner::Scan(const StatusCallback &callback) { } void RedisStoreClient::RedisScanner::OnScanCallback( - size_t shard_index, const std::shared_ptr &reply, - const StatusCallback &callback) { + std::string match_pattern, size_t shard_index, + const std::shared_ptr &reply, const StatusCallback &callback) { RAY_CHECK(reply); std::vector scan_result; size_t cursor = reply->ReadAsScanArray(&scan_result); - // Update shard cursors and keys_. { absl::MutexLock lock(&mutex_); @@ -340,40 +393,7 @@ void RedisStoreClient::RedisScanner::OnScanCallback( // If pending_request_count_ is equal to 0, it means that the scan of this batch is // completed and the next batch is started if any. if (--pending_request_count_ == 0) { - Scan(callback); - } -} - -void RedisStoreClient::RedisScanner::MGetValues( - const std::vector &keys, - const ItemCallback> &callback) { - // The `MGET` command for each shard. - auto mget_commands_by_shards = GenCommandsByShards(redis_client_, "MGET", keys); - - auto finished_count = std::make_shared(0); - int size = mget_commands_by_shards.size(); - for (auto &item : mget_commands_by_shards) { - auto mget_keys = std::move(item.second); - auto mget_callback = [this, finished_count, size, mget_keys, - callback](const std::shared_ptr &reply) { - if (!reply->IsNil()) { - auto value = reply->ReadAsStringArray(); - { - absl::MutexLock lock(&mutex_); - // The 0 th element of mget_keys is "MGET", so we start from the 1 th element. - for (int index = 0; index < (int)value.size(); ++index) { - key_value_map_[GetKeyFromRedisKey(mget_keys[index + 1], table_name_)] = - value[index]; - } - } - } - - ++(*finished_count); - if (*finished_count == size) { - callback(key_value_map_); - } - }; - RAY_CHECK_OK(item.first->RunArgvAsync(mget_keys, mget_callback)); + Scan(match_pattern, callback); } } diff --git a/src/ray/gcs/store_client/redis_store_client.h b/src/ray/gcs/store_client/redis_store_client.h index b2cada940..8f8eb748c 100644 --- a/src/ray/gcs/store_client/redis_store_client.h +++ b/src/ray/gcs/store_client/redis_store_client.h @@ -40,6 +40,9 @@ class RedisStoreClient : public StoreClient { Status AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) override; + Status AsyncGetByIndex(const std::string &table_name, const std::string &index_key, + const MapCallback &callback) override; + Status AsyncGetAll(const std::string &table_name, const MapCallback &callback) override; @@ -62,26 +65,23 @@ class RedisStoreClient : public StoreClient { class RedisScanner { public: explicit RedisScanner(std::shared_ptr redis_client, - std::string table_name, std::string match_pattern); + std::string table_name); - Status ScanKeysAndValues(const MapCallback &callback); + Status ScanKeysAndValues(std::string match_pattern, + const MapCallback &callback); - Status ScanKeys(const MultiItemCallback &callback); + Status ScanKeys(std::string match_pattern, + const MultiItemCallback &callback); private: - void Scan(const StatusCallback &callback); + void Scan(std::string match_pattern, const StatusCallback &callback); - void OnScanCallback(size_t shard_index, const std::shared_ptr &reply, + void OnScanCallback(std::string match_pattern, size_t shard_index, + const std::shared_ptr &reply, const StatusCallback &callback); - void MGetValues(const std::vector &keys, - const MapCallback &callback); - std::string table_name_; - /// The scan match pattern. - std::string match_pattern_; - /// Mutex to protect the shard_to_cursor_ field and the keys_ field and the /// key_value_map_ field. absl::Mutex mutex_; @@ -89,9 +89,6 @@ class RedisStoreClient : public StoreClient { /// All keys that scanned from redis. absl::flat_hash_set keys_; - /// Key-Value pairs that scanned from redis. - std::unordered_map key_value_map_; - /// The scan cursor for each shard. std::unordered_map shard_to_cursor_; @@ -132,6 +129,10 @@ class RedisStoreClient : public StoreClient { const std::string &table_name, const std::string &index_key); + static Status MGetValues(std::shared_ptr redis_client, + std::string table_name, const std::vector &keys, + const MapCallback &callback); + std::shared_ptr redis_client_; }; diff --git a/src/ray/gcs/store_client/store_client.h b/src/ray/gcs/store_client/store_client.h index 1c90418a9..a59fe87ec 100644 --- a/src/ray/gcs/store_client/store_client.h +++ b/src/ray/gcs/store_client/store_client.h @@ -65,6 +65,16 @@ class StoreClient { virtual Status AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) = 0; + /// Get data by index from the given table asynchronously. + /// + /// \param table_name The name of the table to be read. + /// \param index_key The secondary key that will be used to get the indexed data. + /// \param callback Callback that will be called after read finishes. + /// \return Status + virtual Status AsyncGetByIndex( + const std::string &table_name, const std::string &index_key, + const MapCallback &callback) = 0; + /// Get all data from the given table asynchronously. /// /// \param table_name The name of the table to be read. 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 fff318c54..1705e49dc 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 @@ -132,6 +132,24 @@ class StoreClientTestBase : public ::testing::Test { WaitPendingDone(); } + void GetByIndex() { + auto get_calllback = + [this](const std::unordered_map &result) { + if (!result.empty()) { + auto key = ActorID::FromBinary(result.begin()->first); + auto it = key_to_index_.find(key); + RAY_CHECK(it != key_to_index_.end()); + RAY_CHECK(index_to_keys_[it->second].size() == result.size()); + } + pending_count_ -= result.size(); + }; + auto iter = index_to_keys_.begin(); + pending_count_ += iter->second.size(); + RAY_CHECK_OK( + store_client_->AsyncGetByIndex(table_name_, iter->first.Hex(), get_calllback)); + WaitPendingDone(); + } + void DeleteByIndex() { auto delete_calllback = [this](const Status &status) { RAY_CHECK_OK(status); @@ -198,6 +216,9 @@ class StoreClientTestBase : public ::testing::Test { // AsyncPut with index PutWithIndex(); + // AsyncGet with index + GetByIndex(); + // AsyncDelete by index DeleteByIndex();