zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add interface for fuzz match (#734) 

---------

Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>
Co-authored-by: Teng Ma <805522925@qq.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Xuchun Shang 2025-08-14 17:09:06 +08:00 committed by GitHub
parent 81f492c033
commit 70ac0b611f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 737 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mooncake-integration/store/store_py.cpp
View File

@ -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;

7
mooncake-store/include/client.h
View File

@ -74,6 +74,11 @@ class Client {
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> Query(
const std::string& object_key);
tl::expected<
std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
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<void, ErrorCode> Remove(const ObjectKey& key);
tl::expected<long, ErrorCode> RemoveByRegex(const ObjectKey& str);
/**
* @brief Removes all objects and all its replicas
* @return tl::expected<long, ErrorCode> number of removed objects or error

8
mooncake-store/include/master_client.h
View File

@ -56,6 +56,11 @@ class MasterClient {
[[nodiscard]] tl::expected<std::vector<Replica::Descriptor>, ErrorCode>
GetReplicaList(const std::string& object_key);
[[nodiscard]] tl::expected<
std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
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<void, ErrorCode> Remove(const std::string& key);
[[nodiscard]] tl::expected<long, ErrorCode> RemoveByRegex(
const std::string& str);
/**
* @brief Removes all objects and all its replicas
* @return tl::expected<long, ErrorCode> number of removed objects or error

12
mooncake-store/include/master_metric_manager.h
View File

@ -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_;

7
mooncake-store/include/master_service.h
View File

@ -143,6 +143,11 @@ class MasterService {
auto QuerySegments(const std::string& segment)
-> tl::expected<std::pair<size_t, size_t>, ErrorCode>;
auto GetReplicaListByRegex(const std::string& regex_pattern)
-> tl::expected<
std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
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<void, ErrorCode>;
auto RemoveByRegex(const std::string& str) -> tl::expected<long, ErrorCode>;
/**
* @brief Remove all objects and their replicas
* @return return the number of objects removed

5
mooncake-store/include/pybind_client.h
View File

@ -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<void, ErrorCode> remove_internal(const std::string &key);
tl::expected<long, ErrorCode> removeByRegex_internal(
const std::string &str);
tl::expected<int64_t, ErrorCode> removeAll_internal();
tl::expected<void, ErrorCode> tearDownAll_internal();

7
mooncake-store/include/rpc_service.h
View File

@ -41,6 +41,11 @@ class WrappedMasterService {
std::vector<tl::expected<bool, ErrorCode>> BatchExistKey(
const std::vector<std::string>& keys);
tl::expected<
std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
ErrorCode>
GetReplicaListByRegex(const std::string& str);
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> GetReplicaList(
const std::string& key);
@ -70,6 +75,8 @@ class WrappedMasterService {
tl::expected<void, ErrorCode> Remove(const std::string& key);
tl::expected<long, ErrorCode> RemoveByRegex(const std::string& str);
long RemoveAll();
tl::expected<void, ErrorCode> MountSegment(const Segment& segment,

4
mooncake-store/include/storage_backend.h
View File

@ -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
} // namespace mooncake

16
mooncake-store/include/utils.h
View File

@ -49,6 +49,22 @@ void to_stream(std::ostream& os, const std::vector<T>& vec) {
os << "]";
}
template <typename K, typename V>
void to_stream(std::ostream& os, const std::unordered_map<K, V>& 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 <typename T1, typename T2>
void to_stream(std::ostream& os, const std::pair<T1, T2>& p) {

18
mooncake-store/src/client.cpp
View File

@ -356,6 +356,13 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchGet(
return results;
}
tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
ErrorCode>
Client::QueryByRegex(const std::string& str) {
auto result = master_client_.GetReplicaListByRegex(str);
return result;
}
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> Client::Query(
const std::string& object_key) {
auto result = master_client_.GetReplicaList(object_key);
@ -997,6 +1004,17 @@ tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key) {
return {};
}
tl::expected<long, ErrorCode> 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<long, ErrorCode> Client::RemoveAll() {
// if (storage_backend_) {
// storage_backend_->RemoveAll();

34
mooncake-store/src/master_client.cpp
View File

@ -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<tl::expected<bool, ErrorCode>> MasterClient::BatchExistKey(
return result;
}
tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
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<std::string, std::vector<Replica::Descriptor>>>(str);
timer.LogResponseExpected(result);
return result;
}
tl::expected<std::vector<Replica::Descriptor>, ErrorCode>
MasterClient::GetReplicaList(const std::string& object_key) {
ScopedVLogTimer timer(1, "MasterClient::GetReplicaList");
@ -361,6 +385,16 @@ tl::expected<void, ErrorCode> MasterClient::Remove(const std::string& key) {
return result;
}
tl::expected<long, ErrorCode> 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<long, ErrorCode> MasterClient::RemoveAll() {
ScopedVLogTimer timer(1, "MasterClient::RemoveAll");
timer.LogRequest("action=remove_all_objects");

44
mooncake-store/src/master_metric_manager.cpp
View File

@ -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_);

92
mooncake-store/src/master_service.cpp
View File

@ -4,6 +4,7 @@
#include <cstdint>
#include <queue>
#include <shared_mutex>
#include <regex>
#include <ylt/util/tl/expected.hpp>
#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<std::string, std::vector<Replica::Descriptor>>,
ErrorCode> {
std::unordered_map<std::string, std::vector<Replica::Descriptor>> 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::Descriptor> 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<std::vector<Replica::Descriptor>, 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, ErrorCode> {
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;

13
mooncake-store/src/pybind_client.cpp
View File

@ -392,6 +392,19 @@ int PyClient::remove(const std::string &key) {
return to_py_ret(remove_internal(key));
}
tl::expected<long, ErrorCode> 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<int64_t, ErrorCode> PyClient::removeAll_internal() {
if (!client_) {
LOG(ERROR) << "Client is not initialized";

31
mooncake-store/src/rpc_service.cpp
View File

@ -223,6 +223,23 @@ std::vector<tl::expected<bool, ErrorCode>> WrappedMasterService::BatchExistKey(
return result;
}
tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
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<std::vector<Replica::Descriptor>, ErrorCode>
WrappedMasterService::GetReplicaList(const std::string& key) {
return execute_rpc(
@ -438,6 +455,15 @@ tl::expected<void, ErrorCode> WrappedMasterService::Remove(
[] { MasterMetricManager::instance().inc_remove_failures(); });
}
tl::expected<long, ErrorCode> 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>(

47
mooncake-store/src/storage_backend.cpp
View File

@ -6,6 +6,7 @@
#include <sys/uio.h>
#include <string>
#include <vector>
#include <regex>
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<fs::path> 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<StorageFile> StorageBackend::create_file(
return std::make_unique<PosixFile>(path, fd);
}
} // namespace mooncake
} // namespace mooncake

385
mooncake-store/tests/master_service_test.cpp
View File

@ -333,6 +333,186 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) {
}
}
TEST_F(MasterServiceTest, GetReplicaListByRegex) {
const uint64_t kv_lease_ttl = 50;
std::unique_ptr<MasterService> 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<uint64_t> 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<uint64_t> 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<MasterService>(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<std::string> 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<MasterService> 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<MasterService> 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<uint64_t> 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<MasterService>(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<std::string> 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<std::string> 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<std::string> 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<MasterService>(
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<MasterService>(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<MasterService>(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<MasterService>(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<MasterService> service_(