diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h index 323b2148c..d9f7c84a0 100644 --- a/src/ray/common/ray_config_def.h +++ b/src/ray/common/ray_config_def.h @@ -248,6 +248,9 @@ RAY_CONFIG(int32_t, num_actor_checkpoints_to_keep, 20) /// Maximum number of ids in one batch to send to GCS to delete keys. RAY_CONFIG(uint32_t, maximum_gcs_deletion_batch_size, 1000) +/// Maximum number of items in one batch to scan from GCS storage. +RAY_CONFIG(uint32_t, maximum_gcs_scan_batch_size, 1000) + /// When getting objects from object store, print a warning every this number of attempts. RAY_CONFIG(uint32_t, object_store_get_warn_per_num_attempts, 50) diff --git a/src/ray/gcs/callback.h b/src/ray/gcs/callback.h index af805f239..3d1ecf15a 100644 --- a/src/ray/gcs/callback.h +++ b/src/ray/gcs/callback.h @@ -16,6 +16,7 @@ #define RAY_GCS_CALLBACK_H #include +#include #include #include "ray/common/status.h" @@ -53,15 +54,10 @@ using SubscribeCallback = std::function; template using ItemCallback = std::function; -/// This callback is used to receive a large amount of results. -/// \param status Status indicates whether the scan was successful. -/// \param has_more Whether more data will be called back. -/// If `has_more == true`, there are more data to be received. This callback will -/// be called again. -/// \param result The items returned by storage. -template -using SegmentedCallback = - std::function &result)>; +/// This callback is used to receive multiple key-value items from GCS. +/// \param result The key-value items returned by GCS. +template +using MapCallback = std::function &result)>; } // namespace gcs diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc index ed0931306..50061ce3e 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.cc +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -44,18 +44,15 @@ Status GcsTable::Get(const Key &key, } template -Status GcsTable::GetAll( - const SegmentedCallback> &callback) { - auto on_done = [callback]( - const Status &status, bool has_more, - const std::vector> &result) { - std::vector> values; +Status GcsTable::GetAll(const MapCallback &callback) { + auto on_done = [callback](const std::unordered_map &result) { + std::unordered_map values; for (auto &item : result) { Data data; data.ParseFromString(item.second); - values.emplace_back(std::move(std::make_pair(Key::FromBinary(item.first), data))); + values[Key::FromBinary(item.first)] = data; } - callback(status, has_more, values); + callback(values); }; return store_client_->AsyncGetAll(table_name_, on_done); } @@ -68,19 +65,13 @@ Status GcsTable::Delete(const Key &key, const StatusCallback &callbac template Status GcsTable::BatchDelete(const std::vector &keys, const StatusCallback &callback) { - // TODO(ffbin): We will use redis store client batch delete interface directly later. - auto finished_count = std::make_shared(0); - int size = keys.size(); - for (Key key : keys) { - auto done = [finished_count, size, callback](const Status &status) { - ++(*finished_count); - if (*finished_count == size) { - callback(Status::OK()); - } - }; - RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key.Binary(), done)); + std::vector keys_to_delete; + keys_to_delete.reserve(keys.size()); + for (auto &key : keys) { + keys_to_delete.emplace_back(std::move(key.Binary())); } - return Status::OK(); + return this->store_client_->AsyncBatchDelete(this->table_name_, keys_to_delete, + callback); } template @@ -92,8 +83,8 @@ Status GcsTableWithJobId::Put(const Key &key, const Data &value, } template -Status GcsTableWithJobId::GetByJobId( - const JobID &job_id, const SegmentedCallback> &callback) { +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"); diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.h b/src/ray/gcs/gcs_server/gcs_table_storage.h index b94014abf..7f1c1e622 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.h +++ b/src/ray/gcs/gcs_server/gcs_table_storage.h @@ -74,9 +74,8 @@ class GcsTable { /// Get all data from the table asynchronously. /// /// \param callback Callback that will be called after data has been received. - /// If the callback return `has_more == true` mean there's more data to be received. /// \return Status - Status GetAll(const SegmentedCallback> &callback); + Status GetAll(const MapCallback &callback); /// Delete data from the table asynchronously. /// @@ -122,8 +121,7 @@ class GcsTableWithJobId : public GcsTable { /// \param job_id The key to lookup from the table. /// \param callback Callback that will be called after read finishes. /// \return Status - Status GetByJobId(const JobID &job_id, - const SegmentedCallback> &callback); + Status GetByJobId(const JobID &job_id, const MapCallback &callback); /// Delete all the data of the specified job id from the table asynchronously. /// 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 2a001141d..b805276f9 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.cc +++ b/src/ray/gcs/store_client/in_memory_store_client.cc @@ -58,14 +58,12 @@ Status InMemoryStoreClient::AsyncGet(const std::string &table_name, Status InMemoryStoreClient::AsyncGetAll( const std::string &table_name, - const SegmentedCallback> &callback) { + const MapCallback &callback) { auto table = GetOrCreateTable(table_name); absl::MutexLock lock(&(table->mutex_)); - std::vector> result; - for (auto &record : table->records_) { - result.emplace_back(std::make_pair(record.first, record.second)); - } - main_io_service_.post([result, callback]() { callback(Status::OK(), false, result); }); + std::unordered_map result; + result.insert(table->records_.begin(), table->records_.end()); + main_io_service_.post([result, callback]() { callback(result); }); return Status::OK(); } @@ -79,6 +77,18 @@ Status InMemoryStoreClient::AsyncDelete(const std::string &table_name, return Status::OK(); } +Status InMemoryStoreClient::AsyncBatchDelete(const std::string &table_name, + const std::vector &keys, + const StatusCallback &callback) { + auto table = GetOrCreateTable(table_name); + absl::MutexLock lock(&(table->mutex_)); + for (auto &key : keys) { + table->records_.erase(key); + } + main_io_service_.post([callback]() { callback(Status::OK()); }); + 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 b07676af0..f9f893cc1 100644 --- a/src/ray/gcs/store_client/in_memory_store_client.h +++ b/src/ray/gcs/store_client/in_memory_store_client.h @@ -42,13 +42,16 @@ class InMemoryStoreClient : public StoreClient { Status AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) override; - Status AsyncGetAll( - const std::string &table_name, - const SegmentedCallback> &callback) override; + Status AsyncGetAll(const std::string &table_name, + const MapCallback &callback) override; Status AsyncDelete(const std::string &table_name, const std::string &key, const StatusCallback &callback) override; + Status AsyncBatchDelete(const std::string &table_name, + const std::vector &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 cb20dc59f..99e390ffb 100644 --- a/src/ray/gcs/store_client/redis_store_client.cc +++ b/src/ray/gcs/store_client/redis_store_client.cc @@ -23,11 +23,12 @@ namespace ray { namespace gcs { +std::string RedisStoreClient::separator_ = ":"; + Status RedisStoreClient::AsyncPut(const std::string &table_name, const std::string &key, const std::string &data, const StatusCallback &callback) { - std::string full_key = table_name + key; - return DoPut(full_key, data, callback); + return DoPut(GenRedisKey(table_name, key), data, callback); } Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name, @@ -45,8 +46,7 @@ Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name, } // Write data to Redis. - std::string full_key = table_name + key; - status = DoPut(full_key, data, callback); + status = DoPut(GenRedisKey(table_name, key), data, callback); if (!status.ok()) { // Run callback if failed. @@ -57,25 +57,10 @@ Status RedisStoreClient::AsyncPutWithIndex(const std::string &table_name, }; // Write index to Redis. - std::string index_table_key = index_key + table_name + key; + std::string index_table_key = GenRedisKey(table_name, key, index_key); return DoPut(index_table_key, key, write_callback); } -Status RedisStoreClient::DoPut(const std::string &key, const std::string &data, - const StatusCallback &callback) { - std::vector args = {"SET", key, data}; - RedisCallback write_callback = nullptr; - if (callback) { - write_callback = [callback](const std::shared_ptr &reply) { - auto status = reply->ReadAsStatus(); - callback(status); - }; - } - - auto shard_context = redis_client_->GetShardContext(key); - return shard_context->RunArgvAsync(args, write_callback); -} - Status RedisStoreClient::AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) { RAY_CHECK(callback != nullptr); @@ -91,18 +76,24 @@ Status RedisStoreClient::AsyncGet(const std::string &table_name, const std::stri callback(Status::OK(), result); }; - std::string full_key = table_name + key; - std::vector args = {"GET", full_key}; + std::string redis_key = GenRedisKey(table_name, key); + std::vector args = {"GET", redis_key}; - auto shard_context = redis_client_->GetShardContext(full_key); + auto shard_context = redis_client_->GetShardContext(redis_key); return shard_context->RunArgvAsync(args, redis_callback); } Status RedisStoreClient::AsyncGetAll( const std::string &table_name, - const SegmentedCallback> &callback) { - RAY_CHECK(0) << "Not implemented! Will implement this function in next PR."; - return Status::OK(); + 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 on_done = [callback, + scanner](const std::unordered_map &result) { + callback(result); + }; + return scanner->ScanKeysAndValues(on_done); } Status RedisStoreClient::AsyncDelete(const std::string &table_name, @@ -110,27 +101,279 @@ Status RedisStoreClient::AsyncDelete(const std::string &table_name, const StatusCallback &callback) { RedisCallback delete_callback = nullptr; if (callback) { - delete_callback = [callback](std::shared_ptr reply) { - int64_t deleted_count = reply->ReadAsInteger(); - RAY_LOG(DEBUG) << "Delete done, total delete count " << deleted_count; + delete_callback = [callback](const std::shared_ptr &reply) { callback(Status::OK()); }; } - std::string full_key = table_name + key; - std::vector args = {"DEL", full_key}; + std::string redis_key = GenRedisKey(table_name, key); + std::vector args = {"DEL", redis_key}; - auto shard_context = redis_client_->GetShardContext(full_key); + auto shard_context = redis_client_->GetShardContext(redis_key); return shard_context->RunArgvAsync(args, delete_callback); } +Status RedisStoreClient::AsyncBatchDelete(const std::string &table_name, + const std::vector &keys, + const StatusCallback &callback) { + std::vector redis_keys; + redis_keys.reserve(keys.size()); + for (auto &key : keys) { + redis_keys.push_back(GenRedisKey(table_name, key)); + } + return DeleteByKeys(redis_keys, callback); +} + Status RedisStoreClient::AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) { - RAY_CHECK(0) << "Not implemented! Will implement this function in next PR."; + std::string match_pattern = GenRedisMatchPattern(table_name, index_key); + auto scanner = std::make_shared(redis_client_, table_name, match_pattern); + auto on_done = [this, table_name, index_key, callback, scanner]( + const Status &status, const std::vector &result) { + if (!result.empty()) { + std::vector keys; + keys.reserve(result.size()); + for (auto &item : result) { + keys.push_back(GetKeyFromRedisKey(item, table_name, index_key)); + } + auto batch_delete_callback = [this, result, callback](const Status &status) { + RAY_CHECK_OK(status); + // Delete index keys. + RAY_CHECK_OK(DeleteByKeys(result, callback)); + }; + RAY_CHECK_OK(AsyncBatchDelete(table_name, keys, batch_delete_callback)); + } else { + callback(status); + } + }; + + return scanner->ScanKeys(on_done); +} + +Status RedisStoreClient::DoPut(const std::string &key, const std::string &data, + const StatusCallback &callback) { + std::vector args = {"SET", key, data}; + RedisCallback write_callback = nullptr; + if (callback) { + write_callback = [callback](const std::shared_ptr &reply) { + auto status = reply->ReadAsStatus(); + callback(status); + }; + } + + auto shard_context = redis_client_->GetShardContext(key); + return shard_context->RunArgvAsync(args, write_callback); +} + +Status RedisStoreClient::DeleteByKeys(const std::vector &keys, + const StatusCallback &callback) { + // The `DEL` command for each shard. + auto del_commands_by_shards = GenCommandsByShards(redis_client_, "DEL", keys); + + auto finished_count = std::make_shared(0); + int size = del_commands_by_shards.size(); + for (auto &item : del_commands_by_shards) { + auto delete_callback = [finished_count, size, + callback](const std::shared_ptr &reply) { + ++(*finished_count); + if (*finished_count == size) { + callback(Status::OK()); + } + }; + RAY_CHECK_OK(item.first->RunArgvAsync(item.second, delete_callback)); + } return Status::OK(); } +std::unordered_map> +RedisStoreClient::GenCommandsByShards(const std::shared_ptr &redis_client, + const std::string &command, + const std::vector &keys) { + std::unordered_map> commands_by_shards; + for (auto &key : keys) { + auto shard_context = redis_client->GetShardContext(key).get(); + auto it = commands_by_shards.find(shard_context); + if (it == commands_by_shards.end()) { + commands_by_shards[shard_context].push_back(command); + commands_by_shards[shard_context].push_back(key); + } else { + it->second.push_back(key); + } + } + return commands_by_shards; +} + +std::string RedisStoreClient::GenRedisKey(const std::string &table_name, + const std::string &key) { + std::stringstream ss; + ss << table_name << separator_ << key; + return ss.str(); +} + +std::string RedisStoreClient::GenRedisKey(const std::string &table_name, + const std::string &key, + const std::string &index_key) { + std::stringstream ss; + ss << table_name << separator_ << index_key << separator_ << key; + return ss.str(); +} + +std::string RedisStoreClient::GenRedisMatchPattern(const std::string &table_name) { + std::stringstream ss; + ss << table_name << separator_ << "*"; + return ss.str(); +} + +std::string RedisStoreClient::GenRedisMatchPattern(const std::string &table_name, + const std::string &index_key) { + std::stringstream ss; + ss << table_name << separator_ << index_key << separator_ << "*"; + return ss.str(); +} + +std::string RedisStoreClient::GetKeyFromRedisKey(const std::string &redis_key, + const std::string &table_name) { + auto pos = table_name.size() + separator_.size(); + return redis_key.substr(pos, redis_key.size() - pos); +} + +std::string RedisStoreClient::GetKeyFromRedisKey(const std::string &redis_key, + const std::string &table_name, + const std::string &index_key) { + auto pos = table_name.size() + separator_.size() * 2 + index_key.size(); + return redis_key.substr(pos, redis_key.size() - pos); +} + +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)) { + for (size_t index = 0; index < redis_client_->GetShardContexts().size(); ++index) { + shard_to_cursor_[index] = 0; + } +} + +Status RedisStoreClient::RedisScanner::ScanKeysAndValues( + 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); + } + }; + return ScanKeys(on_done); +} + +Status RedisStoreClient::RedisScanner::ScanKeys( + 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); + return Status::OK(); +} + +void RedisStoreClient::RedisScanner::Scan(const StatusCallback &callback) { + if (shard_to_cursor_.empty()) { + callback(Status::OK()); + return; + } + + size_t batch_count = RayConfig::instance().maximum_gcs_scan_batch_size(); + for (const auto &item : shard_to_cursor_) { + ++pending_request_count_; + + size_t shard_index = item.first; + size_t cursor = item.second; + + auto scan_callback = [this, shard_index, + callback](const std::shared_ptr &reply) { + OnScanCallback(shard_index, reply, callback); + }; + + // Scan by prefix from Redis. + std::vector args = {"SCAN", std::to_string(cursor), + "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(); + } + } +} + +void RedisStoreClient::RedisScanner::OnScanCallback( + 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_); + auto shard_it = shard_to_cursor_.find(shard_index); + RAY_CHECK(shard_it != shard_to_cursor_.end()); + // If cursor is equal to 0, it means that the scan of this shard is finished, so we + // erase it from shard_to_cursor_. + if (cursor == 0) { + shard_to_cursor_.erase(shard_it); + } else { + shard_it->second = cursor; + } + + keys_.insert(scan_result.begin(), scan_result.end()); + } + + // 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 = 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(item.second, mget_callback)); + } +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/store_client/redis_store_client.h b/src/ray/gcs/store_client/redis_store_client.h index 200288995..67f1c5f80 100644 --- a/src/ray/gcs/store_client/redis_store_client.h +++ b/src/ray/gcs/store_client/redis_store_client.h @@ -15,8 +15,7 @@ #ifndef RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H #define RAY_GCS_STORE_CLIENT_REDIS_STORE_CLIENT_H -#include -#include +#include "absl/container/flat_hash_set.h" #include "ray/gcs/redis_client.h" #include "ray/gcs/redis_context.h" #include "ray/gcs/store_client/store_client.h" @@ -28,7 +27,7 @@ namespace gcs { class RedisStoreClient : public StoreClient { public: - RedisStoreClient(std::shared_ptr redis_client) + explicit RedisStoreClient(std::shared_ptr redis_client) : redis_client_(std::move(redis_client)) {} Status AsyncPut(const std::string &table_name, const std::string &key, @@ -41,20 +40,97 @@ class RedisStoreClient : public StoreClient { Status AsyncGet(const std::string &table_name, const std::string &key, const OptionalItemCallback &callback) override; - Status AsyncGetAll( - const std::string &table_name, - const SegmentedCallback> &callback) override; + Status AsyncGetAll(const std::string &table_name, + const MapCallback &callback) override; Status AsyncDelete(const std::string &table_name, const std::string &key, const StatusCallback &callback) override; + Status AsyncBatchDelete(const std::string &table_name, + const std::vector &keys, + const StatusCallback &callback) override; + Status AsyncDeleteByIndex(const std::string &table_name, const std::string &index_key, const StatusCallback &callback) override; private: + /// \class RedisScanner + /// This class is used to scan data from Redis. + /// + /// If you called one method, should never call the other methods. + /// Otherwise it will disturb the status of the RedisScanner. + class RedisScanner { + public: + explicit RedisScanner(std::shared_ptr redis_client, + std::string table_name, std::string match_pattern); + + Status ScanKeysAndValues(const MapCallback &callback); + + Status ScanKeys(const MultiItemCallback &callback); + + private: + void Scan(const StatusCallback &callback); + + void OnScanCallback(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_; + + /// 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_; + + /// The pending shard scan count. + std::atomic pending_request_count_{0}; + + std::shared_ptr redis_client_; + }; + Status DoPut(const std::string &key, const std::string &data, const StatusCallback &callback); + Status DeleteByKeys(const std::vector &keys, + const StatusCallback &callback); + + static std::unordered_map> GenCommandsByShards( + const std::shared_ptr &redis_client, const std::string &command, + const std::vector &keys); + + /// The separator is used when building redis key. + static std::string separator_; + + static std::string GenRedisKey(const std::string &table_name, const std::string &key); + + static std::string GenRedisKey(const std::string &table_name, const std::string &key, + const std::string &index_key); + + static std::string GenRedisMatchPattern(const std::string &table_name); + + static std::string GenRedisMatchPattern(const std::string &table_name, + const std::string &index_key); + + static std::string GetKeyFromRedisKey(const std::string &redis_key, + const std::string &table_name); + + static std::string GetKeyFromRedisKey(const std::string &redis_key, + const std::string &table_name, + const std::string &index_key); + 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 d7e22ddf0..1c90418a9 100644 --- a/src/ray/gcs/store_client/store_client.h +++ b/src/ray/gcs/store_client/store_client.h @@ -69,11 +69,9 @@ class StoreClient { /// /// \param table_name The name of the table to be read. /// \param callback Callback that will be called after data has been received. - /// If the callback return `has_more == true` mean there's more data to be received. /// \return Status - virtual Status AsyncGetAll( - const std::string &table_name, - const SegmentedCallback> &callback) = 0; + virtual Status AsyncGetAll(const std::string &table_name, + const MapCallback &callback) = 0; /// Delete data from the given table asynchronously. /// @@ -84,6 +82,16 @@ class StoreClient { virtual Status AsyncDelete(const std::string &table_name, const std::string &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. + /// \param keys The keys that will be deleted from the table. + /// \param callback Callback that will be called after delete finishes. + /// \return Status + virtual Status AsyncBatchDelete(const std::string &table_name, + const std::vector &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 f302d3d80..b259b02af 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 @@ -30,14 +30,14 @@ class InMemoryStoreClientTest : public StoreClientTestBase { TEST_F(InMemoryStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); } -TEST_F(InMemoryStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); } - -TEST_F(InMemoryStoreClientTest, AsyncGetAllTest) { TestAsyncGetAll(); } - TEST_F(InMemoryStoreClientTest, AsyncPutAndDeleteWithIndexTest) { TestAsyncPutAndDeleteWithIndex(); } +TEST_F(InMemoryStoreClientTest, AsyncGetAllAndBatchDeleteTest) { + TestAsyncGetAllAndBatchDelete(); +} + } // 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 a4f7408de..ed8b347ce 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 @@ -46,9 +46,13 @@ class RedisStoreClientTest : public StoreClientTestBase { TEST_F(RedisStoreClientTest, AsyncPutAndAsyncGetTest) { TestAsyncPutAndAsyncGet(); } -TEST_F(RedisStoreClientTest, AsyncDeleteTest) { TestAsyncDelete(); } +TEST_F(RedisStoreClientTest, AsyncPutAndDeleteWithIndexTest) { + TestAsyncPutAndDeleteWithIndex(); +} -TEST_F(RedisStoreClientTest, DISABLED_AsyncGetAllTest) { TestAsyncGetAll(); } +TEST_F(RedisStoreClientTest, AsyncGetAllAndBatchDeleteTest) { + TestAsyncGetAllAndBatchDelete(); +} } // namespace gcs 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 0a6739ce7..265a28e86 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 @@ -31,7 +31,7 @@ namespace gcs { class StoreClientTestBase : public ::testing::Test { public: - StoreClientTestBase() {} + StoreClientTestBase() = default; virtual ~StoreClientTestBase() {} @@ -57,8 +57,7 @@ class StoreClientTestBase : public ::testing::Test { virtual void DisconnectStoreClient() = 0; protected: - void TestAsyncPutAndAsyncGet() { - // AsyncPut without index. + void Put() { auto put_calllback = [this](const Status &status) { RAY_CHECK_OK(status); --pending_count_; @@ -70,8 +69,22 @@ class StoreClientTestBase : public ::testing::Test { put_calllback)); } WaitPendingDone(); + } - // AsyncGet + void Delete() { + 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_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback)); + } + WaitPendingDone(); + } + + void Get() { auto get_callback = [this](const Status &status, const boost::optional &result) { RAY_CHECK_OK(status); @@ -91,101 +104,49 @@ class StoreClientTestBase : public ::testing::Test { WaitPendingDone(); } - void TestAsyncDelete() { - // AsyncPut - auto put_calllback = [this](const Status &status) { --pending_count_; }; + void GetEmpty() { for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK(store_client_->AsyncPut(table_name_, elem.first.Binary(), - elem.second.SerializeAsString(), - put_calllback)); - } - WaitPendingDone(); + auto key = elem.first.Binary(); + auto get_callback = [this, key](const Status &status, + const boost::optional &result) { + RAY_CHECK_OK(status); + RAY_CHECK(!result); + --pending_count_; + }; - // AsyncDelete - 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_->AsyncDelete(table_name_, elem.first.Binary(), delete_calllback)); - } - WaitPendingDone(); - - // AsyncGet - auto get_callback = [this](const Status &status, - const boost::optional &result) { - RAY_CHECK_OK(status); - RAY_CHECK(!result); - --pending_count_; - }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK( - store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); + RAY_CHECK_OK(store_client_->AsyncGet(table_name_, key, get_callback)); } WaitPendingDone(); } - void TestAsyncPutAndDeleteWithIndex() { - // AsyncPut with index + void PutWithIndex() { auto put_calllback = [this](const Status &status) { --pending_count_; }; for (const auto &elem : key_to_value_) { ++pending_count_; RAY_CHECK_OK(store_client_->AsyncPutWithIndex( - table_name_, elem.first.Binary(), key_to_index_[elem.first].Binary(), + table_name_, elem.first.Binary(), key_to_index_[elem.first].Hex(), elem.second.SerializeAsString(), put_calllback)); } WaitPendingDone(); + } - // AsyncDelete by index + void DeleteByIndex() { auto delete_calllback = [this](const Status &status) { RAY_CHECK_OK(status); --pending_count_; }; for (const auto &elem : index_to_keys_) { ++pending_count_; - RAY_CHECK_OK(store_client_->AsyncDeleteByIndex(table_name_, elem.first.Binary(), + RAY_CHECK_OK(store_client_->AsyncDeleteByIndex(table_name_, elem.first.Hex(), delete_calllback)); } WaitPendingDone(); - - // AsyncGet - auto get_callback = [this](const Status &status, - const boost::optional &result) { - RAY_CHECK_OK(status); - RAY_CHECK(!result); - --pending_count_; - }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - RAY_CHECK_OK( - store_client_->AsyncGet(table_name_, elem.first.Binary(), get_callback)); - } - WaitPendingDone(); } - void TestAsyncGetAll() { - // AsyncPut - auto put_calllback = [this](const Status &status) { --pending_count_; }; - for (const auto &elem : key_to_value_) { - ++pending_count_; - // Get index - auto it = key_to_index_.find(elem.first); - const JobID &index = it->second; - RAY_CHECK_OK(store_client_->AsyncPutWithIndex( - table_name_, elem.first.Binary(), index.Binary(), - elem.second.SerializeAsString(), put_calllback)); - } - WaitPendingDone(); - - // AsyncGetAll + void GetAll() { auto get_all_callback = - [this](const Status &status, bool has_more, - const std::vector> &result) { - RAY_CHECK_OK(status); + [this](const std::unordered_map &result) { static std::unordered_set received_keys; for (const auto &item : result) { const ActorID &actor_id = ActorID::FromBinary(item.first); @@ -196,9 +157,7 @@ class StoreClientTestBase : public ::testing::Test { auto map_it = key_to_value_.find(actor_id); RAY_CHECK(map_it != key_to_value_.end()); } - if (!has_more) { - RAY_CHECK(received_keys.size() == key_to_value_.size()); - } + RAY_CHECK(received_keys.size() == key_to_value_.size()); pending_count_ -= result.size(); }; @@ -207,6 +166,58 @@ class StoreClientTestBase : public ::testing::Test { WaitPendingDone(); } + void BatchDelete() { + auto delete_calllback = [this](const Status &status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + ++pending_count_; + std::vector keys; + for (auto &elem : key_to_value_) { + keys.push_back(elem.first.Binary()); + } + RAY_CHECK_OK(store_client_->AsyncBatchDelete(table_name_, keys, delete_calllback)); + WaitPendingDone(); + } + + void TestAsyncPutAndAsyncGet() { + // AsyncPut without index. + Put(); + + // AsyncGet + Get(); + + // AsyncDelete + Delete(); + + GetEmpty(); + } + + void TestAsyncPutAndDeleteWithIndex() { + // AsyncPut with index + PutWithIndex(); + + // AsyncDelete by index + DeleteByIndex(); + + // AsyncGet + GetEmpty(); + } + + void TestAsyncGetAllAndBatchDelete() { + // AsyncPut + Put(); + + // AsyncGetAll + GetAll(); + + // AsyncBatchDelete + BatchDelete(); + + // AsyncGet + GetEmpty(); + } + void GenTestData() { for (size_t i = 0; i < key_count_; i++) { rpc::ActorTableData actor;