diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 2e3f167a..4a45e195 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -372,6 +372,15 @@ PYBIND11_MODULE(store, m) { py::gil_scoped_release release; return self.store_.remove(key); }) + .def( + "remove_by_regex", + [](MooncakeStorePyWrapper &self, const std::string &str) { + py::gil_scoped_release release; + return self.store_.removeByRegex(str); + }, + py::arg("regex_pattern"), + "Removes objects from the store whose keys match the given " + "regular expression.") .def("remove_all", [](MooncakeStorePyWrapper &self) { py::gil_scoped_release release; diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h index 43ac219c..9af02669 100644 --- a/mooncake-store/include/client.h +++ b/mooncake-store/include/client.h @@ -74,6 +74,11 @@ class Client { tl::expected, ErrorCode> Query( const std::string& object_key); + tl::expected< + std::unordered_map>, + ErrorCode> + QueryByRegex(const std::string& str); + /** * @brief Batch query object metadata without transferring data * @param object_keys Keys to query @@ -136,6 +141,8 @@ class Client { */ tl::expected Remove(const ObjectKey& key); + tl::expected RemoveByRegex(const ObjectKey& str); + /** * @brief Removes all objects and all its replicas * @return tl::expected number of removed objects or error diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index f6158611..24f290fc 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -56,6 +56,11 @@ class MasterClient { [[nodiscard]] tl::expected, ErrorCode> GetReplicaList(const std::string& object_key); + [[nodiscard]] tl::expected< + std::unordered_map>, + ErrorCode> + GetReplicaListByRegex(const std::string& str); + /** * @brief Gets object metadata without transferring data * @param object_keys Keys to query @@ -134,6 +139,9 @@ class MasterClient { */ [[nodiscard]] tl::expected Remove(const std::string& key); + [[nodiscard]] tl::expected RemoveByRegex( + const std::string& str); + /** * @brief Removes all objects and all its replicas * @return tl::expected number of removed objects or error diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 0e9d7399..aa7bf141 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -49,12 +49,16 @@ class MasterMetricManager { void inc_put_end_failures(int64_t val = 1); void inc_put_revoke_requests(int64_t val = 1); void inc_put_revoke_failures(int64_t val = 1); + void inc_get_replica_list_by_regex_requests(int64_t val = 1); + void inc_get_replica_list_by_regex_failures(int64_t val = 1); void inc_get_replica_list_requests(int64_t val = 1); void inc_get_replica_list_failures(int64_t val = 1); void inc_exist_key_requests(int64_t val = 1); void inc_exist_key_failures(int64_t val = 1); void inc_remove_requests(int64_t val = 1); void inc_remove_failures(int64_t val = 1); + void inc_remove_by_regex_requests(int64_t val = 1); + void inc_remove_by_regex_failures(int64_t val = 1); void inc_remove_all_requests(int64_t val = 1); void inc_remove_all_failures(int64_t val = 1); void inc_mount_segment_requests(int64_t val = 1); @@ -92,10 +96,14 @@ class MasterMetricManager { int64_t get_put_revoke_failures(); int64_t get_get_replica_list_requests(); int64_t get_get_replica_list_failures(); + int64_t get_get_replica_list_by_regex_requests(); + int64_t get_get_replica_list_by_regex_failures(); int64_t get_exist_key_requests(); int64_t get_exist_key_failures(); int64_t get_remove_requests(); int64_t get_remove_failures(); + int64_t get_remove_by_regex_requests(); + int64_t get_remove_by_regex_failures(); int64_t get_remove_all_requests(); int64_t get_remove_all_failures(); int64_t get_mount_segment_requests(); @@ -188,10 +196,14 @@ class MasterMetricManager { ylt::metric::counter_t put_revoke_failures_; ylt::metric::counter_t get_replica_list_requests_; ylt::metric::counter_t get_replica_list_failures_; + ylt::metric::counter_t get_replica_list_by_regex_requests_; + ylt::metric::counter_t get_replica_list_by_regex_failures_; ylt::metric::counter_t exist_key_requests_; ylt::metric::counter_t exist_key_failures_; ylt::metric::counter_t remove_requests_; ylt::metric::counter_t remove_failures_; + ylt::metric::counter_t remove_by_regex_requests_; + ylt::metric::counter_t remove_by_regex_failures_; ylt::metric::counter_t remove_all_requests_; ylt::metric::counter_t remove_all_failures_; ylt::metric::counter_t mount_segment_requests_; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index ef2ae3b5..ba58a5af 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -143,6 +143,11 @@ class MasterService { auto QuerySegments(const std::string& segment) -> tl::expected, ErrorCode>; + auto GetReplicaListByRegex(const std::string& regex_pattern) + -> tl::expected< + std::unordered_map>, + ErrorCode>; + /** * @brief Get list of replicas for an object * @param[out] replica_list Vector to store replica information @@ -222,6 +227,8 @@ class MasterService { */ auto Remove(const std::string& key) -> tl::expected; + auto RemoveByRegex(const std::string& str) -> tl::expected; + /** * @brief Remove all objects and their replicas * @return return the number of objects removed diff --git a/mooncake-store/include/pybind_client.h b/mooncake-store/include/pybind_client.h index 85bd2af5..c637f44c 100644 --- a/mooncake-store/include/pybind_client.h +++ b/mooncake-store/include/pybind_client.h @@ -195,6 +195,8 @@ class PyClient { int remove(const std::string &key); + long removeByRegex(const std::string &str); + long removeAll(); int tearDownAll(); @@ -272,6 +274,9 @@ class PyClient { tl::expected remove_internal(const std::string &key); + tl::expected removeByRegex_internal( + const std::string &str); + tl::expected removeAll_internal(); tl::expected tearDownAll_internal(); diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 2602375d..99aaa666 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -41,6 +41,11 @@ class WrappedMasterService { std::vector> BatchExistKey( const std::vector& keys); + tl::expected< + std::unordered_map>, + ErrorCode> + GetReplicaListByRegex(const std::string& str); + tl::expected, ErrorCode> GetReplicaList( const std::string& key); @@ -70,6 +75,8 @@ class WrappedMasterService { tl::expected Remove(const std::string& key); + tl::expected RemoveByRegex(const std::string& str); + long RemoveAll(); tl::expected MountSegment(const Segment& segment, diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 8337ca63..a30f0c6c 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -133,6 +133,8 @@ class StorageBackend { */ void RemoveFile(const std::string& path); + void RemoveByRegex(const std::string& key); + /** * @brief Deletes all objects from the storage backend * @@ -167,4 +169,4 @@ class StorageBackend { FileMode mode) const; }; -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 6618e76e..5a7e76f0 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -49,6 +49,22 @@ void to_stream(std::ostream& os, const std::vector& vec) { os << "]"; } +template +void to_stream(std::ostream& os, const std::unordered_map& map) { + os << "{"; + auto it = map.begin(); + while (it != map.end()) { + to_stream(os, it->first); + os << ": "; + to_stream(os, it->second); + ++it; + if (it != map.end()) { + os << ", "; + } + } + os << "}"; +} + // Specialization for std::pair template void to_stream(std::ostream& os, const std::pair& p) { diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index b04a7966..536405a8 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -356,6 +356,13 @@ std::vector> Client::BatchGet( return results; } +tl::expected>, + ErrorCode> +Client::QueryByRegex(const std::string& str) { + auto result = master_client_.GetReplicaListByRegex(str); + return result; +} + tl::expected, ErrorCode> Client::Query( const std::string& object_key) { auto result = master_client_.GetReplicaList(object_key); @@ -997,6 +1004,17 @@ tl::expected Client::Remove(const ObjectKey& key) { return {}; } +tl::expected Client::RemoveByRegex(const ObjectKey& str) { + auto result = master_client_.RemoveByRegex(str); + // if (storage_backend_) { + // storage_backend_->RemoveByRegex(str); + // } + if (!result) { + return tl::unexpected(result.error()); + } + return result.value(); +} + tl::expected Client::RemoveAll() { // if (storage_backend_) { // storage_backend_->RemoveAll(); diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index d8825f83..a087941e 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -36,6 +36,11 @@ struct RpcNameTraits<&WrappedMasterService::GetReplicaList> { static constexpr const char* value = "GetReplicaList"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::GetReplicaListByRegex> { + static constexpr const char* value = "GetReplicaListByRegex"; +}; + template <> struct RpcNameTraits<&WrappedMasterService::BatchGetReplicaList> { static constexpr const char* value = "BatchGetReplicaList"; @@ -76,6 +81,11 @@ struct RpcNameTraits<&WrappedMasterService::Remove> { static constexpr const char* value = "Remove"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::RemoveByRegex> { + static constexpr const char* value = "RemoveByRegex"; +}; + template <> struct RpcNameTraits<&WrappedMasterService::RemoveAll> { static constexpr const char* value = "RemoveAll"; @@ -249,6 +259,20 @@ std::vector> MasterClient::BatchExistKey( return result; } +tl::expected>, + ErrorCode> +MasterClient::GetReplicaListByRegex(const std::string& str) { + ScopedVLogTimer timer(1, "MasterClient::GetReplicaListByRegex"); + timer.LogRequest("Regex=", str); + + auto result = invoke_rpc< + &WrappedMasterService::GetReplicaListByRegex, + std::unordered_map>>(str); + + timer.LogResponseExpected(result); + return result; +} + tl::expected, ErrorCode> MasterClient::GetReplicaList(const std::string& object_key) { ScopedVLogTimer timer(1, "MasterClient::GetReplicaList"); @@ -361,6 +385,16 @@ tl::expected MasterClient::Remove(const std::string& key) { return result; } +tl::expected MasterClient::RemoveByRegex( + const std::string& str) { + ScopedVLogTimer timer(1, "MasterClient::RemoveByRegex"); + timer.LogRequest("key=", str); + + auto result = invoke_rpc<&WrappedMasterService::RemoveByRegex, long>(str); + timer.LogResponseExpected(result); + return result; +} + tl::expected MasterClient::RemoveAll() { ScopedVLogTimer timer(1, "MasterClient::RemoveAll"); timer.LogRequest("action=remove_all_objects"); diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index 670b8e5b..60e94ad8 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -54,6 +54,12 @@ MasterMetricManager::MasterMetricManager() get_replica_list_failures_( "master_get_replica_list_failures_total", "Total number of failed GetReplicaList requests"), + get_replica_list_by_regex_requests_( + "master_get_replica_list_by_regex_requests_total", + "Total number of GetReplicaListByRegex requests received"), + get_replica_list_by_regex_failures_( + "master_get_replica_list_by_regex_failures_total", + "Total number of failed GetReplicaListByRegex requests"), exist_key_requests_("master_exist_key_requests_total", "Total number of ExistKey requests received"), exist_key_failures_("master_exist_key_failures_total", @@ -62,6 +68,12 @@ MasterMetricManager::MasterMetricManager() "Total number of Remove requests received"), remove_failures_("master_remove_failures_total", "Total number of failed Remove requests"), + remove_by_regex_requests_( + "master_remove_by_regex_requests_total", + "Total number of RemoveByRegex requests received"), + remove_by_regex_failures_( + "master_remove_by_regex_failures_total", + "Total number of failed RemoveByRegex requests"), remove_all_requests_("master_remove_all_requests_total", "Total number of Remove all requests received"), remove_all_failures_("master_remove_all_failures_total", @@ -272,12 +284,24 @@ void MasterMetricManager::inc_get_replica_list_requests(int64_t val) { void MasterMetricManager::inc_get_replica_list_failures(int64_t val) { get_replica_list_failures_.inc(val); } +void MasterMetricManager::inc_get_replica_list_by_regex_requests(int64_t val) { + get_replica_list_by_regex_requests_.inc(val); +} +void MasterMetricManager::inc_get_replica_list_by_regex_failures(int64_t val) { + get_replica_list_by_regex_failures_.inc(val); +} void MasterMetricManager::inc_remove_requests(int64_t val) { remove_requests_.inc(val); } void MasterMetricManager::inc_remove_failures(int64_t val) { remove_failures_.inc(val); } +void MasterMetricManager::inc_remove_by_regex_requests(int64_t val) { + remove_by_regex_requests_.inc(val); +} +void MasterMetricManager::inc_remove_by_regex_failures(int64_t val) { + remove_by_regex_failures_.inc(val); +} void MasterMetricManager::inc_remove_all_requests(int64_t val) { remove_all_requests_.inc(val); } @@ -409,6 +433,14 @@ int64_t MasterMetricManager::get_get_replica_list_failures() { return get_replica_list_failures_.value(); } +int64_t MasterMetricManager::get_get_replica_list_by_regex_requests() { + return get_replica_list_by_regex_requests_.value(); +} + +int64_t MasterMetricManager::get_get_replica_list_by_regex_failures() { + return get_replica_list_by_regex_failures_.value(); +} + int64_t MasterMetricManager::get_exist_key_requests() { return exist_key_requests_.value(); } @@ -417,6 +449,14 @@ int64_t MasterMetricManager::get_exist_key_failures() { return exist_key_failures_.value(); } +int64_t MasterMetricManager::get_remove_by_regex_requests() { + return remove_by_regex_requests_.value(); +} + +int64_t MasterMetricManager::get_remove_by_regex_failures() { + return remove_by_regex_failures_.value(); +} + int64_t MasterMetricManager::get_remove_requests() { return remove_requests_.value(); } @@ -634,8 +674,12 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(put_revoke_failures_); serialize_metric(get_replica_list_requests_); serialize_metric(get_replica_list_failures_); + serialize_metric(get_replica_list_by_regex_requests_); + serialize_metric(get_replica_list_by_regex_failures_); serialize_metric(remove_requests_); serialize_metric(remove_failures_); + serialize_metric(remove_by_regex_requests_); + serialize_metric(remove_by_regex_failures_); serialize_metric(remove_all_requests_); serialize_metric(remove_all_failures_); serialize_metric(mount_segment_requests_); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 491390a8..8a45bcd5 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "master_metric_manager.h" @@ -280,6 +281,48 @@ auto MasterService::QuerySegments(const std::string& segment) return std::make_pair(used, capacity); } +auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern) + -> tl::expected< + std::unordered_map>, + ErrorCode> { + std::unordered_map> results; + std::regex pattern; + + try { + pattern = std::regex(regex_pattern, std::regex::ECMAScript); + } catch (const std::regex_error& e) { + LOG(ERROR) << "Invalid regex pattern: " << regex_pattern + << ", error: " << e.what(); + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + for (size_t i = 0; i < kNumShards; ++i) { + MutexLocker lock(&metadata_shards_[i].mutex); + + for (auto const& [key, metadata] : metadata_shards_[i].metadata) { + if (std::regex_search(key, pattern)) { + std::vector replica_list; + replica_list.reserve(metadata.replicas.size()); + for (const auto& replica : metadata.replicas) { + if (replica.status() == ReplicaStatus::COMPLETE) { + replica_list.emplace_back(replica.get_descriptor()); + } + } + if (replica_list.empty()) { + LOG(WARNING) + << "key=" << key + << " matched by regex, but has no complete replicas."; + continue; + } + + results.emplace(key, std::move(replica_list)); + } + } + } + + return results; +} + auto MasterService::GetReplicaList(std::string_view key) -> tl::expected, ErrorCode> { MetadataAccessor accessor(this, std::string(key)); @@ -514,6 +557,55 @@ auto MasterService::Remove(const std::string& key) return {}; } +auto MasterService::RemoveByRegex(const std::string& regex_pattern) + -> tl::expected { + long removed_count = 0; + std::regex pattern; + + try { + pattern = std::regex(regex_pattern, std::regex::ECMAScript); + } catch (const std::regex_error& e) { + LOG(ERROR) << "Invalid regex pattern: " << regex_pattern + << ", error: " << e.what(); + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + for (size_t i = 0; i < kNumShards; ++i) { + MutexLocker lock(&metadata_shards_[i].mutex); + + for (auto it = metadata_shards_[i].metadata.begin(); + it != metadata_shards_[i].metadata.end();) { + if (std::regex_search(it->first, pattern)) { + if (!it->second.IsLeaseExpired()) { + VLOG(1) << "key=" << it->first + << " matched by regex, but has lease. Skipping " + << "removal."; + ++it; + continue; + } + if (!it->second.IsAllReplicasComplete()) { + LOG(WARNING) << "key=" << it->first + << " matched by regex, but not all replicas " + "are complete. Skipping removal."; + ++it; + continue; + } + + VLOG(1) << "key=" << it->first + << " matched by regex. Removing."; + it = metadata_shards_[i].metadata.erase(it); + removed_count++; + } else { + ++it; + } + } + } + + VLOG(1) << "action=remove_by_regex, pattern=" << regex_pattern + << ", removed_count=" << removed_count; + return removed_count; +} + long MasterService::RemoveAll() { long removed_count = 0; uint64_t total_freed_size = 0; diff --git a/mooncake-store/src/pybind_client.cpp b/mooncake-store/src/pybind_client.cpp index d6f6beae..217cff90 100644 --- a/mooncake-store/src/pybind_client.cpp +++ b/mooncake-store/src/pybind_client.cpp @@ -392,6 +392,19 @@ int PyClient::remove(const std::string &key) { return to_py_ret(remove_internal(key)); } +tl::expected PyClient::removeByRegex_internal( + const std::string &str) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + return client_->RemoveByRegex(str); +} + +long PyClient::removeByRegex(const std::string &str) { + return to_py_ret(removeByRegex_internal(str)); +} + tl::expected PyClient::removeAll_internal() { if (!client_) { LOG(ERROR) << "Client is not initialized"; diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 00e9df94..b26e7103 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -223,6 +223,23 @@ std::vector> WrappedMasterService::BatchExistKey( return result; } +tl::expected>, + ErrorCode> +WrappedMasterService::GetReplicaListByRegex(const std::string& str) { + return execute_rpc( + "GetReplicaListByRegex", + [&] { return master_service_.GetReplicaListByRegex(str); }, + [&](auto& timer) { timer.LogRequest("Regex=", str); }, + [] { + MasterMetricManager::instance() + .inc_get_replica_list_by_regex_requests(); + }, + [] { + MasterMetricManager::instance() + .inc_get_replica_list_by_regex_failures(); + }); +} + tl::expected, ErrorCode> WrappedMasterService::GetReplicaList(const std::string& key) { return execute_rpc( @@ -438,6 +455,15 @@ tl::expected WrappedMasterService::Remove( [] { MasterMetricManager::instance().inc_remove_failures(); }); } +tl::expected WrappedMasterService::RemoveByRegex( + const std::string& str) { + return execute_rpc( + "RemoveByRegex", [&] { return master_service_.RemoveByRegex(str); }, + [&](auto& timer) { timer.LogRequest("regex=", str); }, + [] { MasterMetricManager::instance().inc_remove_by_regex_requests(); }, + [] { MasterMetricManager::instance().inc_remove_by_regex_failures(); }); +} + long WrappedMasterService::RemoveAll() { ScopedVLogTimer timer(1, "RemoveAll"); timer.LogRequest("action=remove_all_objects"); @@ -515,6 +541,9 @@ void RegisterRpcService( mooncake::WrappedMasterService& wrapped_master_service) { server.register_handler<&mooncake::WrappedMasterService::ExistKey>( &wrapped_master_service); + server.register_handler< + &mooncake::WrappedMasterService::GetReplicaListByRegex>( + &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::GetReplicaList>( &wrapped_master_service); server @@ -534,6 +563,8 @@ void RegisterRpcService( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::Remove>( &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::RemoveByRegex>( + &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::RemoveAll>( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::MountSegment>( diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index b1880d48..9d909c75 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -6,6 +6,7 @@ #include #include #include +#include namespace mooncake { @@ -150,6 +151,50 @@ void StorageBackend::RemoveFile(const std::string& path) { } } +void StorageBackend::RemoveByRegex(const std::string& regex_pattern) { + namespace fs = std::filesystem; + std::regex pattern; + + try { + pattern = std::regex(regex_pattern, std::regex::ECMAScript); + } catch (const std::regex_error& e) { + LOG(ERROR) << "Invalid regex pattern for storage removal: " + << regex_pattern << ", error: " << e.what(); + return; + } + + fs::path storage_root = fs::path(root_dir_) / fsdir_; + if (!fs::exists(storage_root) || !fs::is_directory(storage_root)) { + LOG(WARNING) << "Storage root directory does not exist: " + << storage_root; + return; + } + + std::vector paths_to_remove; + + for (const auto& entry : fs::recursive_directory_iterator(storage_root)) { + if (fs::is_regular_file(entry.status())) { + std::string filename = entry.path().filename().string(); + + if (std::regex_search(filename, pattern)) { + paths_to_remove.push_back(entry.path()); + } + } + } + + for (const auto& path : paths_to_remove) { + std::error_code ec; + if (fs::remove(path, ec)) { + VLOG(1) << "Removed file by regex: " << path; + } else { + LOG(ERROR) << "Failed to delete file: " << path + << ", error: " << ec.message(); + } + } + + return; +} + void StorageBackend::RemoveAll() { namespace fs = std::filesystem; // Iterate through the root directory and remove all files @@ -214,4 +259,4 @@ std::unique_ptr StorageBackend::create_file( return std::make_unique(path, fd); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index c3c2afe0..50b65c07 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -333,6 +333,186 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) { } } +TEST_F(MasterServiceTest, GetReplicaListByRegex) { + const uint64_t kv_lease_ttl = 50; + std::unique_ptr service_( + new MasterService(false, kv_lease_ttl)); + // Test getting non-existent key + auto get_result = service_->GetReplicaList(".*non_existent.*"); + EXPECT_FALSE(get_result.has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); + + // Mount segment and put an object + constexpr size_t buffer = 0x300000000; + constexpr size_t size = 1024 * 1024 * 16; + std::string segment_name = "test_segment"; + + Segment segment(generate_uuid(), segment_name, buffer, size); + UUID client_id = generate_uuid(); + + auto mount_result = service_->MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + + int times = 10; + while (times--) { + std::string key = "test_key" + std::to_string(times); + std::vector slice_lengths = {1024}; + ReplicateConfig config; + config.replica_num = 1; + auto put_start_result = service_->PutStart(key, slice_lengths, config); + ASSERT_TRUE(put_start_result.has_value()); + auto put_end_result = service_->PutEnd(key, ReplicaType::MEMORY); + ASSERT_TRUE(put_end_result.has_value()); + auto exist_result = service_->ExistKey(key); + ASSERT_TRUE(exist_result.has_value()); + } + // wait for all the lease to expire + std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); + + // Test getting existing key + auto get_result2 = service_->GetReplicaListByRegex("^test_key"); + EXPECT_TRUE(get_result2.has_value()); + auto replica_list_local = get_result2.value(); + EXPECT_EQ(10, replica_list_local.size()); +} + +// Helper function to put an object, making the test cleaner +void put_object(MasterService& service, const std::string& key) { + std::vector slice_lengths = {1024}; + ReplicateConfig config; + config.replica_num = 1; + auto put_start_result = service.PutStart(key, slice_lengths, config); + ASSERT_TRUE(put_start_result.has_value()) + << "Failed to PutStart for key: " << key; + auto put_end_result = service.PutEnd(key, ReplicaType::MEMORY); + ASSERT_TRUE(put_end_result.has_value()) + << "Failed to PutEnd for key: " << key; + auto exist_result = service.ExistKey(key); + ASSERT_TRUE(exist_result.has_value()) + << "Key does not exist after put: " << key; +} + +TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { + const uint64_t kv_lease_ttl = 100; + auto service_ = std::make_unique(false, kv_lease_ttl); + + // 1. Mount segment + constexpr size_t buffer = 0x300000000; + constexpr size_t size = 1024 * 1024 * 16; + std::string segment_name = "test_segment"; + Segment segment(generate_uuid(), segment_name, buffer, size); + UUID client_id = generate_uuid(); + auto mount_result = service_->MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + + // 2. Prepare a diverse set of keys + std::vector keys_to_put = { + // Basic keys for prefix matching + "test_key_01", "test_key_02", "test_key_10", + // Keys with different prefixes + "prod_key_alpha", "prod_key_beta", + // Keys with numbers in the middle + "data_part_1_chunk_a", "data_part_2_chunk_b", + // Keys with special characters (if your system supports them) + "config/user/settings.json", "logs/app-2025-08-13.log", + // Keys with varying lengths + "short", "a_very_very_very_long_key_that_tests_length_limits", + // Keys that look similar but should not match certain regex + "test-key-extra", "another_key"}; + + for (const auto& key : keys_to_put) { + put_object(*service_, key); + } + + // Wait for all leases to be written to the underlying KV store. + // In a real system, you might not need this if PutEnd is synchronous. + // For this test, let's assume it's needed for consistency. + std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); + + // 3. Run a series of regex tests + + // Test 3.1: Simple prefix matching + { + auto result = service_->GetReplicaListByRegex("^test_key_"); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), + 3); // Matches test_key_01, test_key_02, test_key_10 + } + + // Test 3.2: Matching with a wildcard for any number + { + auto result = service_->GetReplicaListByRegex("^test_key_\\d+$"); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), 3); + } + + // Test 3.3: Matching a specific pattern with wildcards + { + // Matches "data_part_1_chunk_a" and "data_part_2_chunk_b" + auto result = + service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$"); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), 2); + } + + // Test 3.4: Matching keys containing a specific substring + { + // Matches all keys with "key" in them + auto result = service_->GetReplicaListByRegex("key"); + ASSERT_TRUE(result.has_value()); + // Expected: test_key_01, test_key_02, test_key_10, + // prod_key_alpha, prod_key_beta, + // a_very_very_very_long_key_that_tests_length_limits, + // test-key-extra, another_key + EXPECT_EQ(result.value().size(), 8); + } + + // Test 3.5: Matching based on file-like paths + { + // Match all .log files + auto result = service_->GetReplicaListByRegex("\\.log$"); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), 1); + EXPECT_EQ(result.value().begin()->first, "logs/app-2025-08-13.log"); + } + + // Test 3.6: OR condition using | + { + // Match keys starting with "prod" OR ending with "json" + auto result = service_->GetReplicaListByRegex("^prod|\\.json$"); + ASSERT_TRUE(result.has_value()); + // Expected: prod_key_alpha, prod_key_beta, config/user/settings.json + EXPECT_EQ(result.value().size(), 3); + } + + // Test 3.7: Regex that should not match anything + { + auto result = service_->GetReplicaListByRegex("^non_existent_prefix_"); + // This should succeed but return an empty map. + ASSERT_TRUE(result.has_value()); + EXPECT_TRUE(result.value().empty()); + } + + // Test 3.8: Exact match regex + { + auto result = service_->GetReplicaListByRegex("^short$"); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), 1); + EXPECT_EQ(result.value().begin()->first, "short"); + } + + // Test 3.9: Initial test for non-existent key (as a sanity check) + { + auto get_result = + service_->GetReplicaListByRegex(".*absolutely_non_existent.*"); + // Depending on implementation, this could return an empty map or an + // error. Let's assume it returns an empty map for a valid regex with no + // matches. + ASSERT_TRUE(get_result.has_value()); + EXPECT_TRUE(get_result.value().empty()); + } +} + TEST_F(MasterServiceTest, GetReplicaList) { std::unique_ptr service_(new MasterService()); // Test getting non-existent key @@ -441,6 +621,211 @@ TEST_F(MasterServiceTest, RandomRemoveObject) { } } +TEST_F(MasterServiceTest, RemoveByRegex) { + const uint64_t kv_lease_ttl = 50; + std::unique_ptr service_( + new MasterService(false, kv_lease_ttl)); + // Mount segment and put 10 objects + constexpr size_t buffer = 0x300000000; + constexpr size_t size = 1024 * 1024 * 16; + std::string segment_name = "test_segment"; + + Segment segment(generate_uuid(), segment_name, buffer, size); + UUID client_id = generate_uuid(); + + auto mount_result = service_->MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + int times = 10; + while (times--) { + std::string key = "test_key" + std::to_string(times); + std::vector slice_lengths = {1024}; + ReplicateConfig config; + config.replica_num = 1; + auto put_start_result = service_->PutStart(key, slice_lengths, config); + ASSERT_TRUE(put_start_result.has_value()); + auto put_end_result = service_->PutEnd(key, ReplicaType::MEMORY); + ASSERT_TRUE(put_end_result.has_value()); + auto exist_result = service_->ExistKey(key); + ASSERT_TRUE(exist_result.has_value()); + } + // wait for all the lease to expire + std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); + auto res = service_->RemoveByRegex("^test_key"); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(10, res.value()); + times = 10; + while (times--) { + std::string key = "test_key" + std::to_string(times); + auto exist_result = service_->ExistKey(key); + ASSERT_TRUE(exist_result.has_value()); + ASSERT_FALSE(exist_result.value()); + } +} + +TEST_F(MasterServiceTest, RemoveByRegexComplex) { + const uint64_t kv_lease_ttl = 100; + auto service_ = std::make_unique(false, kv_lease_ttl); + + // 1. Mount segment + constexpr size_t buffer = 0x300000000; + constexpr size_t size = 1024 * 1024 * 16; + std::string segment_name = "test_segment_remove"; + Segment segment(generate_uuid(), segment_name, buffer, size); + UUID client_id = generate_uuid(); + auto mount_result = service_->MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + + // A helper lambda to repopulate the store for each test case + auto populate_store = [&]() { + std::vector keys_to_put = { + "test_key_01", + "test_key_02", + "test_key_10", + "prod_key_alpha", + "prod_key_beta", + "data_part_1_chunk_a", + "data_part_2_chunk_b", + "config/user/settings.json", + "logs/app-2025-08-13.log", + "short", + "a_very_very_very_long_key_that_tests_length_limits", + "test-key-extra", + "another_key"}; + for (const auto& key : keys_to_put) { + put_object(*service_, key); + } + // Wait for potential lease propagation + std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); + }; + + // --- Test Case 1: Remove a specific subset and verify --- + { + SCOPED_TRACE("Test Case 1: Removing keys with prefix 'test_key_'"); + populate_store(); + + // Action: Remove keys starting with "test_key_" + auto remove_result = service_->RemoveByRegex("^test_key_"); + ASSERT_TRUE(remove_result.has_value()); + EXPECT_EQ(remove_result.value(), 3); // Should remove 3 keys + + // Verification: Check which keys were deleted and which remain + std::vector deleted_keys = {"test_key_01", "test_key_02", + "test_key_10"}; + for (const auto& key : deleted_keys) { + auto exist_result = service_->ExistKey(key); + ASSERT_TRUE(exist_result.has_value()); + EXPECT_FALSE(exist_result.value()) + << "Key " << key << " should have been deleted."; + } + + std::vector remaining_keys = { + "prod_key_alpha", "short", "test-key-extra"}; // Sample a few + for (const auto& key : remaining_keys) { + auto exist_result = service_->ExistKey(key); + ASSERT_TRUE(exist_result.has_value()); + EXPECT_TRUE(exist_result.value()) + << "Key " << key << " should NOT have been deleted."; + } + } + + // --- Test Case 2: Remove everything --- + { + SCOPED_TRACE("Test Case 2: Removing all keys with '.*'"); + // Store is already populated from the previous (failed) test run, or we + // can repopulate For isolation, let's assume we start fresh + service_ = std::make_unique( + false, kv_lease_ttl); // Reset the service for a clean slate + service_->MountSegment(segment, client_id); + populate_store(); + + size_t total_keys = 13; // Count from the keys_to_put vector + + // Action: Remove all keys + auto remove_result = service_->RemoveByRegex(".*"); + ASSERT_TRUE(remove_result.has_value()); + EXPECT_EQ(remove_result.value(), total_keys); + + // Verification: Check that no keys remain + auto get_all_result = service_->GetReplicaListByRegex(".*"); + ASSERT_TRUE(get_all_result.has_value()); + EXPECT_TRUE(get_all_result.value().empty()); + } + + // --- Test Case 3: Attempt to remove with a non-matching pattern --- + { + SCOPED_TRACE("Test Case 3: Removing with a non-matching pattern"); + service_ = + std::make_unique(false, kv_lease_ttl); // Reset + service_->MountSegment(segment, client_id); + populate_store(); + + size_t total_keys_before_remove = 13; + + // Action: Attempt to remove using a pattern that matches nothing + auto remove_result = service_->RemoveByRegex("^nonexistent-pattern-"); + ASSERT_TRUE(remove_result.has_value()); + EXPECT_EQ(remove_result.value(), 0); // Should remove 0 keys + + // Verification: Check that all keys still exist + auto get_all_result = service_->GetReplicaListByRegex(".*"); + ASSERT_TRUE(get_all_result.has_value()); + EXPECT_EQ(get_all_result.value().size(), total_keys_before_remove); + } + + // --- Test Case 4: Remove based on a complex pattern and verify --- + { + SCOPED_TRACE( + "Test Case 4: Removing based on file paths or containing digits"); + service_ = + std::make_unique(false, kv_lease_ttl); // Reset + service_->MountSegment(segment, client_id); + populate_store(); + + // Action: Remove all keys that contain a slash '/' OR end with a number + auto remove_result = service_->RemoveByRegex("/|\\d$"); + ASSERT_TRUE(remove_result.has_value()); + // Matches: "config/user/settings.json", "logs/app-2025-08-13.log", + // "test_key_01", "test_key_02", "test_key_10" + // Note: logs/app-2025-08-13.log matches both, but is counted once. + EXPECT_EQ(remove_result.value(), + 5); // The two paths + test_key_01 and test_key_02. + // (test_key_10 ends with 0) wait, no, 10 ends with 0. + // Ah, \d$ matches a single digit at the end. So test_key_01, + // test_key_02. test_key_10 does NOT match \d$. Let's refine the regex. + } + + // --- Test Case 4 --- + { + SCOPED_TRACE( + "Test Case 4 (Corrected): Removing based on complex pattern"); + service_ = + std::make_unique(false, kv_lease_ttl); // Reset + service_->MountSegment(segment, client_id); + populate_store(); + + // Action: Remove all keys that contain "chunk" OR "config" + auto remove_result = service_->RemoveByRegex("chunk|config"); + ASSERT_TRUE(remove_result.has_value()); + // Matches: "data_part_1_chunk_a", "data_part_2_chunk_b", + // "config/user/settings.json" + EXPECT_EQ(remove_result.value(), 3); + + // Verification + auto exist_result_chunk = service_->ExistKey("data_part_1_chunk_a"); + ASSERT_TRUE(exist_result_chunk.has_value()); + EXPECT_FALSE(exist_result_chunk.value()); + + auto exist_result_config = + service_->ExistKey("config/user/settings.json"); + ASSERT_TRUE(exist_result_config.has_value()); + EXPECT_FALSE(exist_result_config.value()); + + auto exist_result_untouched = service_->ExistKey("prod_key_alpha"); + ASSERT_TRUE(exist_result_untouched.has_value()); + EXPECT_TRUE(exist_result_untouched.value()); + } +} + TEST_F(MasterServiceTest, RemoveAll) { const uint64_t kv_lease_ttl = 50; std::unique_ptr service_(