From c59aa2a54cf55c9f6134fc4a47e7aaa0a6fe6955 Mon Sep 17 00:00:00 2001 From: lzw <2994587066@qq.com> Date: Fri, 10 Jul 2026 14:51:59 +0800 Subject: [PATCH] [Store] feat: add workload-aware KV cache hints --- mooncake-integration/store/store_py.cpp | 7 + .../store/store_py_internal.h | 5 +- mooncake-store/include/master_service.h | 6 +- mooncake-store/include/replica.h | 23 +++ mooncake-store/include/rpc_service.h | 6 + .../catalog_backed_snapshot_provider.cpp | 6 +- mooncake-store/src/master_service.cpp | 55 ++++- .../catalog_backed_snapshot_provider_test.cpp | 12 +- .../snapshot/snapshot_child_process_test.cpp | 93 +++++++++ .../tests/ha/snapshot/snapshot_test_utils.h | 30 ++- mooncake-store/tests/master_service_test.cpp | 195 ++++++++++++++++++ mooncake-wheel/tests/test_import_structure.py | 21 +- 12 files changed, 432 insertions(+), 27 deletions(-) diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 15ca9bac56..f562d91474 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -1826,6 +1826,12 @@ PYBIND11_MODULE(store, m) { .value("GENERAL", ObjectDataType::GENERAL) .export_values(); + py::class_(m, "WorkloadHints") + .def(py::init<>()) + .def_readwrite("session_id", &WorkloadHints::session_id) + .def_readwrite("retention_priority", + &WorkloadHints::retention_priority); + // Define the ReplicateConfig class py::class_(m, "ReplicateConfig") .def(py::init<>()) @@ -1842,6 +1848,7 @@ PYBIND11_MODULE(store, m) { &ReplicateConfig::prefer_alloc_in_same_node) .def_readwrite("data_type", &ReplicateConfig::data_type) .def_readwrite("group_ids", &ReplicateConfig::group_ids) + .def_readwrite("workload_hints", &ReplicateConfig::workload_hints) .def("__str__", [](const ReplicateConfig &config) { std::ostringstream oss; oss << config; diff --git a/mooncake-integration/store/store_py_internal.h b/mooncake-integration/store/store_py_internal.h index a7d030a4d7..30644fb8ff 100644 --- a/mooncake-integration/store/store_py_internal.h +++ b/mooncake-integration/store/store_py_internal.h @@ -861,10 +861,7 @@ bool parallelism_specs_equal_by_kind(const TensorParallelismSpec &lhs, } bool is_default_replicate_config(const ReplicateConfig &config) { - return config.replica_num == 1 && !config.with_soft_pin && - !config.with_hard_pin && config.preferred_segments.empty() && - config.preferred_segment.empty() && - !config.prefer_alloc_in_same_node && !config.group_ids.has_value(); + return config.IsDefault(); } std::optional parse_parallel_axis_spec( diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 661109fd51..a01f6b5456 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -55,6 +55,7 @@ struct MetadataStoragePlugin; namespace test { class MasterServiceSnapshotTestBase; class SnapshotChildProcessTest; +class MasterServiceTest; // Friended so the promotion-on-hit tests can drive a serialize/reset/ // deserialize cycle directly via the otherwise-private // MetadataSerializer, and inspect private clamp fields. This avoids @@ -85,6 +86,7 @@ class MasterService { // Test friend class for snapshot/restore testing friend class test::MasterServiceSnapshotTestBase; friend class test::SnapshotChildProcessTest; + friend class test::MasterServiceTest; friend class test::PromotionOnHitTest; friend class benchmarks::BatchEvictBench; friend class test::MasterServiceTenantQuotaTest; @@ -894,7 +896,7 @@ class MasterService { bool enable_soft_pin, bool enable_hard_pin = false, ObjectDataType data_type_ = ObjectDataType::UNKNOWN, std::string group_id_ = "", std::string tenant_id_ = "default", - std::string user_key_ = {}) + std::string user_key_ = {}, WorkloadHints workload_hints_ = {}) : client_id(client_id_), put_start_time(put_start_time_), size(value_length), @@ -902,6 +904,7 @@ class MasterService { group_id(std::move(group_id_)), tenant_id(std::move(tenant_id_)), user_key(std::move(user_key_)), + workload_hints(std::move(workload_hints_)), lease_timeout(), soft_pin_timeout(std::nullopt), hard_pinned(enable_hard_pin), @@ -928,6 +931,7 @@ class MasterService { const std::string group_id; const std::string tenant_id; const std::string user_key; + const WorkloadHints workload_hints; mutable SpinLock lock; // Default constructor, creates a time_point representing diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index c3acf46dcf..71b0d4c652 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -75,6 +75,23 @@ inline std::ostream& operator<<(std::ostream& os, return os; } +/** + * @brief Workload-level hints associated with an object. + */ +struct WorkloadHints { + std::string session_id{}; // Empty means unset. + int32_t retention_priority{0}; // Zero means unset. + + bool operator==(const WorkloadHints&) const = default; + + friend std::ostream& operator<<(std::ostream& os, + const WorkloadHints& hints) noexcept { + return os << "{ session_id: " << hints.session_id + << ", retention_priority: " << hints.retention_priority + << " }"; + } +}; + /** * @brief Configuration for replica management */ @@ -96,6 +113,11 @@ struct ReplicateConfig { // ungrouped. Grouped keys share metadata routing, coalesced lease refresh, // and memory eviction behavior. std::optional> group_ids{}; + WorkloadHints workload_hints{}; + + bool operator==(const ReplicateConfig&) const = default; + + bool IsDefault() const { return *this == ReplicateConfig{}; } ReplicateConfig ForSingleKey(size_t key_index) const { ReplicateConfig key_config = *this; @@ -142,6 +164,7 @@ struct ReplicateConfig { } os << "]"; } + os << ", workload_hints: " << config.workload_hints; os << " }"; return os; } diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 8e2cfa4895..6a67479196 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -17,9 +17,15 @@ namespace mooncake { +namespace test { +class MasterServiceTest; +} + // Forward declaration class HttpMetadataServer; class WrappedMasterService { + friend class test::MasterServiceTest; + public: // Constructor with optional metadata-cleanup-on-timeout configuration. // - http_metadata_server: in-process pointer used when the HTTP metadata diff --git a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp index dc5c8f6fc1..4b21451046 100644 --- a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp +++ b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp @@ -135,17 +135,19 @@ DeserializeStandbyObjectMetadata( // compatibility with MasterService::MetadataSerializer, which appends // them over time: // data_type (positive int) appears before the replicas; - // hard_pinned (bool) and group_id (str) trail them. + // hard_pinned (bool), group_id (str), and workload_hints (array) + // trail them. // v1: 7 + replica_count, no optional fields // v2: 8 + replica_count, either data_type or hard_pinned // v3: 9 + replica_count, data_type + hard_pinned or // hard_pinned + group_id // v4: 10 + replica_count, data_type + hard_pinned + group_id + // v5: 11 + replica_count, v4 + workload_hints // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX // replica_count from wrapping the bounds and slipping an out-of-bounds // index through. constexpr uint64_t kBaseFieldCount = 7; - constexpr uint64_t kMaxOptionalFieldCount = 3; + constexpr uint64_t kMaxOptionalFieldCount = 4; const uint64_t total_elements = object.via.array.size; const uint64_t min_elements = kBaseFieldCount + replica_count; if (total_elements < min_elements || diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index edeeefe184..fb07924964 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -3012,7 +3012,8 @@ auto MasterService::AllocateAndInsertMetadata( std::piecewise_construct, std::forward_as_tuple(key), std::forward_as_tuple(client_id, now, value_length, std::move(replicas), config.with_soft_pin, config.with_hard_pin, - config.data_type, group_id, tenant_id, key)); + config.data_type, group_id, tenant_id, key, + config.workload_hints)); if (!inserted) { LOG(INFO) << "key=" << key << ", info=object_already_exists"; abort_reserved_quota(); @@ -3713,6 +3714,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, merged_config.with_hard_pin || metadata.IsHardPinned(); merged_config.with_soft_pin = merged_config.with_soft_pin || metadata.IsSoftPinned(); + merged_config.workload_hints = metadata.workload_hints; const std::string existing_group_id = metadata.group_id; const uint64_t old_quota_charge = @@ -8557,7 +8559,8 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj, metadata_ptr->size, metadata_ptr->PopReplicas(), metadata_ptr->soft_pin_timeout.has_value(), metadata_ptr->IsHardPinned(), metadata_ptr->data_type, - metadata_ptr->group_id, tenant_id, user_key)); + metadata_ptr->group_id, tenant_id, user_key, + metadata_ptr->workload_hints)); it->second.lease_timeout = metadata_ptr->lease_timeout; it->second.soft_pin_timeout = metadata_ptr->soft_pin_timeout; @@ -8580,11 +8583,13 @@ MasterService::MetadataSerializer::SerializeMetadata( // Pack ObjectMetadata using array structure for efficiency // Format: [client_id, put_start_time, size, lease_timeout, // has_soft_pin_timeout, soft_pin_timeout, replicas_count, data_type, - // replicas..., hard_pinned, group_id] + // replicas..., hard_pinned, group_id, + // [workload_session_id, workload_retention_priority]] - size_t array_size = 10; // client_id, put_start_time, size, lease_timeout, + size_t array_size = 11; // client_id, put_start_time, size, lease_timeout, // has_soft_pin_timeout, soft_pin_timeout, - // replicas_count, data_type, hard_pinned, group_id + // replicas_count, data_type, hard_pinned, + // group_id, workload_hints array_size += metadata.CountReplicas(); // One element per replica packer.pack_array(array_size); @@ -8638,6 +8643,9 @@ MasterService::MetadataSerializer::SerializeMetadata( packer.pack(metadata.IsHardPinned()); packer.pack(metadata.group_id); + packer.pack_array(2); + packer.pack(metadata.workload_hints.session_id); + packer.pack(metadata.workload_hints.retention_priority); return {}; } @@ -8690,12 +8698,14 @@ MasterService::MetadataSerializer::DeserializeMetadata( // v1: 7 + replicas_count, no optional fields // v2: 8 + replicas_count, either data_type or hard_pinned // v3: 9 + replicas_count, data_type + hard_pinned or hard_pinned + - // group_id v4: 10 + replicas_count, data_type + hard_pinned + group_id + // group_id + // v4: 10 + replicas_count, data_type + hard_pinned + group_id + // v5: 11 + replicas_count, v4 + nested workload_hints // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX // replicas_count from wrapping the bounds and slipping an out-of-bounds // index past the size check. constexpr uint64_t kBaseFieldCount = 7; - constexpr uint64_t kMaxOptionalFieldCount = 3; + constexpr uint64_t kMaxOptionalFieldCount = 4; const uint64_t total_elements = obj.via.array.size; const uint64_t min_elements = kBaseFieldCount + replicas_count; if (total_elements < min_elements || @@ -8745,6 +8755,35 @@ MasterService::MetadataSerializer::DeserializeMetadata( group_id = array[index++].as(); } + WorkloadHints workload_hints; + if (index < total_elements) { + const auto& hints_object = array[index++]; + if (hints_object.type != msgpack::type::ARRAY || + hints_object.via.array.size != 2) { + return tl::unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + "deserialize workload_hints must be a two-element array")); + } + + const auto* hints_array = hints_object.via.array.ptr; + if (hints_array[0].type != msgpack::type::STR || + (hints_array[1].type != msgpack::type::POSITIVE_INTEGER && + hints_array[1].type != msgpack::type::NEGATIVE_INTEGER)) { + return tl::unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + "deserialize workload_hints has invalid field types")); + } + + try { + workload_hints.session_id = hints_array[0].as(); + workload_hints.retention_priority = hints_array[1].as(); + } catch (const std::exception& e) { + return tl::unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + "deserialize workload_hints failed: " + std::string(e.what()))); + } + } + // Create ObjectMetadata instance bool enable_soft_pin = has_soft_pin_timeout; auto metadata = std::make_unique( @@ -8752,7 +8791,7 @@ MasterService::MetadataSerializer::DeserializeMetadata( std::chrono::system_clock::time_point( std::chrono::milliseconds(put_start_time_timestamp)), size, std::move(replicas), enable_soft_pin, is_hard_pinned, data_type, - group_id); + group_id, "default", std::string{}, std::move(workload_hints)); metadata->lease_timeout = std::chrono::system_clock::time_point( std::chrono::milliseconds(lease_timestamp)); diff --git a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp index 071940c826..240791c95b 100644 --- a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp +++ b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp @@ -177,13 +177,19 @@ TEST_P(CatalogBackedSnapshotProviderTest, } TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithGroupId) { - // 10 + replica_count: current writer format (data_type + hard_pinned + - // trailing group_id). Regression test for the live snapshot restore - // failure against the latest metadata layout. + // 10 + replica_count: data_type + hard_pinned + trailing group_id. PublishSnapshotPayload(SnapshotMetadataFormat::kWithGroupId); ExpectLoadsDefaultObject(); } +TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithWorkloadHints) { + // 11 + replica_count: current writer format, with workload hints encoded + // as one trailing [session_id, retention_priority] array. The standby does + // not use the hints, but it must accept snapshots that contain them. + PublishSnapshotPayload(SnapshotMetadataFormat::kWithWorkloadHints); + ExpectLoadsDefaultObject(); +} + TEST_P(CatalogBackedSnapshotProviderTest, RejectsOverflowingReplicaCount) { // A near-UINT32_MAX replica_count must not wrap the format-detection // arithmetic into a valid-looking total and slip an out-of-bounds index diff --git a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp index 8902329a39..ca9d9aed75 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp +++ b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -194,6 +195,22 @@ class SnapshotChildProcessTest : public ::testing::Test { return false; } + std::optional GetObjectWorkloadHintsFromMetadata( + const std::string& key) { + for (auto& shard : service_->metadata_shards_) { + SharedMutexLocker lock(&shard.mutex, shared_lock_t{}); + auto tenant_it = shard.tenants.find("default"); + if (tenant_it == shard.tenants.end()) { + continue; + } + auto it = tenant_it->second.metadata.find(key); + if (it != tenant_it->second.metadata.end()) { + return it->second.workload_hints; + } + } + return std::nullopt; + } + std::string FindGroupIdOnDifferentShard(MasterService* svc, const std::string& key) { const size_t key_shard = svc->getShardIndex(key); @@ -508,6 +525,9 @@ TEST_F(SnapshotChildProcessTest, RestoreRebuildsGroupedObjectRouting) { replicate_config.replica_num = 1; replicate_config.group_ids = std::vector{ FindGroupIdOnDifferentShard(service_.get(), key)}; + replicate_config.workload_hints.session_id = kDefaultTestWorkloadSessionId; + replicate_config.workload_hints.retention_priority = + kDefaultTestRetentionPriority; auto put_start = service_->PutStart(client_id, key, "default", 1024, replicate_config); @@ -526,6 +546,11 @@ TEST_F(SnapshotChildProcessTest, RestoreRebuildsGroupedObjectRouting) { auto restored_replicas = service_->GetReplicaList(key, "default"); ASSERT_TRUE(restored_replicas.has_value()) << "Grouped key should remain reachable by key after restore"; + auto restored_hints = GetObjectWorkloadHintsFromMetadata(key); + ASSERT_TRUE(restored_hints.has_value()); + EXPECT_EQ(restored_hints->session_id, kDefaultTestWorkloadSessionId); + EXPECT_EQ(restored_hints->retention_priority, + kDefaultTestRetentionPriority); ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value()); EXPECT_FALSE(service_->ExistKey(key, "default").value_or(true)); } @@ -581,6 +606,74 @@ TEST_F(SnapshotChildProcessTest, << deserialize_result.error().message; EXPECT_FALSE(ObjectIsGroupedInMetadata(key, shard_idx)); + // This v1 payload predates workload hints; the new reader must supply + // their API defaults rather than rejecting the snapshot. + auto restored_hints = GetObjectWorkloadHintsFromMetadata(key); + ASSERT_TRUE(restored_hints.has_value()); + EXPECT_TRUE(restored_hints->session_id.empty()); + EXPECT_EQ(restored_hints->retention_priority, 0); +} + +TEST_F(SnapshotChildProcessTest, + DeserializePreviousMetadataWithoutWorkloadHintsUsesDefaults) { + CreateDefaultService(); + const std::string key = "snapshot_v4_without_workload_hints"; + const uint32_t shard_idx = GetShardIndexForTest(key); + const UUID client_id = generate_uuid(); + + msgpack::sbuffer shard_buffer; + MsgpackPacker shard_packer(&shard_buffer); + shard_packer.pack_map(1); + shard_packer.pack(std::string("metadata")); + shard_packer.pack_array(1); + shard_packer.pack_array(3); + shard_packer.pack(std::string("default")); + shard_packer.pack(key); + + // v4: 7 base fields + data_type + one replica + hard_pinned + group_id. + shard_packer.pack_array(11); + shard_packer.pack(UuidToString(client_id)); + shard_packer.pack(kDefaultTestPutStartTimeMs); + shard_packer.pack(kDefaultTestObjectSize); + shard_packer.pack(kDefaultTestLeaseTimeoutMs); + shard_packer.pack(false); + shard_packer.pack(uint64_t{0}); + shard_packer.pack(uint32_t{1}); + shard_packer.pack(static_cast(ObjectDataType::TENSOR)); + PackDiskReplica(shard_packer, kDefaultTestDiskFilePath, + kDefaultTestObjectSize); + shard_packer.pack(true); + shard_packer.pack(std::string("snapshot-v4-group")); + + auto compressed_shard = + zstd_compress(reinterpret_cast(shard_buffer.data()), + shard_buffer.size(), 3); + + msgpack::sbuffer root_buffer; + MsgpackPacker root_packer(&root_buffer); + root_packer.pack_map(3); + root_packer.pack(std::string("shards")); + root_packer.pack_map(1); + root_packer.pack(shard_idx); + root_packer.pack_bin(compressed_shard.size()); + root_packer.pack_bin_body( + reinterpret_cast(compressed_shard.data()), + compressed_shard.size()); + root_packer.pack(std::string("discarded_replicas")); + root_packer.pack_array(0); + root_packer.pack(std::string("replica_next_id")); + root_packer.pack(uint64_t{11}); + + auto deserialize_result = + DeserializeMetadataForTest(ToByteVector(root_buffer)); + ASSERT_TRUE(deserialize_result.has_value()) + << deserialize_result.error().message; + + EXPECT_TRUE(ObjectIsGroupedInMetadata(key, shard_idx)); + auto restored_hints = GetObjectWorkloadHintsFromMetadata(key); + ASSERT_TRUE(restored_hints.has_value()); + EXPECT_TRUE(restored_hints->session_id.empty()); + EXPECT_EQ(restored_hints->retention_priority, 0); } TEST_F(SnapshotChildProcessTest, LegacyEtcdConnstringFallbackIsPreserved) { diff --git a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h index 01262f671c..b0fff355d6 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h +++ b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -37,6 +38,8 @@ constexpr uint64_t kDefaultTestLeaseTimeoutMs = 4102444800000ULL; constexpr uint64_t kDefaultTestSnapshotSeq = 42; constexpr uint64_t kDefaultTestProducerViewVersion = 7; constexpr int64_t kDefaultTestCreatedAtMs = 1700000000000LL; +constexpr char kDefaultTestWorkloadSessionId[] = "test-session"; +constexpr int32_t kDefaultTestRetentionPriority = -17; struct CatalogBackendParam { std::string name; @@ -53,13 +56,16 @@ struct CatalogBackendParam { // kDataTypeAndHardPinned: // 9 + replica_count, data_type plus trailing hard_pinned // kWithGroupId: 10 + replica_count, data_type + hard_pinned + group_id -// (the current writer format) +// kWithWorkloadHints: +// 11 + replica_count, kWithGroupId plus trailing +// [session_id, retention_priority] array (current writer) enum class SnapshotMetadataFormat { kLegacy, kDataTypeOnly, kHardPinnedOnly, kDataTypeAndHardPinned, kWithGroupId, + kWithWorkloadHints, }; class ScopedEnvVar { @@ -177,18 +183,25 @@ inline std::vector BuildMetadataPayload( const bool include_data_type = format == SnapshotMetadataFormat::kDataTypeOnly || format == SnapshotMetadataFormat::kDataTypeAndHardPinned || - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithWorkloadHints; const bool include_hard_pinned = format == SnapshotMetadataFormat::kHardPinnedOnly || format == SnapshotMetadataFormat::kDataTypeAndHardPinned || - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithWorkloadHints; const bool include_group_id = - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithWorkloadHints; + const bool include_workload_hints = + format == SnapshotMetadataFormat::kWithWorkloadHints; constexpr uint32_t kReplicaCount = 1; - // 7 leading fields + replicas + optional data_type/hard_pinned/group_id. + // 7 leading fields + replicas + optional data_type/hard_pinned/group_id/ + // workload_hints. const size_t array_size = 7 + kReplicaCount + (include_data_type ? 1 : 0) + (include_hard_pinned ? 1 : 0) + - (include_group_id ? 1 : 0); + (include_group_id ? 1 : 0) + + (include_workload_hints ? 1 : 0); msgpack::sbuffer shard_buffer; MsgpackPacker shard_packer(&shard_buffer); @@ -216,6 +229,11 @@ inline std::vector BuildMetadataPayload( if (include_group_id) { shard_packer.pack(std::string("test-group")); } + if (include_workload_hints) { + shard_packer.pack_array(2); + shard_packer.pack(std::string(kDefaultTestWorkloadSessionId)); + shard_packer.pack(kDefaultTestRetentionPriority); + } return WrapShardIntoMetadataRoot(shard_buffer); } diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 46512a5749..ac99fea771 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -1,4 +1,5 @@ #include "master_service.h" +#include "master_client.h" #include "rpc_service.h" #include @@ -64,6 +65,11 @@ class ScopedEnvVar { class MasterServiceTest : public ::testing::Test { protected: + struct ObjectHintState { + WorkloadHints workload_hints; + std::string group_id; + }; + void SetUp() override { google::InitGoogleLogging("MasterServiceTest"); FLAGS_logtostderr = true; @@ -307,6 +313,32 @@ class MasterServiceTest : public ::testing::Test { EXPECT_TRUE(predicate()); } + std::optional GetObjectHintState( + MasterService& service, const std::string& key, + const std::string& tenant_id = "default") const { + const size_t shard_idx = service.getMetadataShardIndex(tenant_id, key); + auto& shard = service.metadata_shards_[shard_idx]; + SharedMutexLocker lock(&shard.mutex, shared_lock_t{}); + auto tenant_it = shard.tenants.find(tenant_id); + if (tenant_it == shard.tenants.end()) { + return std::nullopt; + } + auto metadata_it = tenant_it->second.metadata.find(key); + if (metadata_it == tenant_it->second.metadata.end()) { + return std::nullopt; + } + return ObjectHintState{ + .workload_hints = metadata_it->second.workload_hints, + .group_id = metadata_it->second.group_id, + }; + } + + std::optional GetObjectHintState( + WrappedMasterService& service, const std::string& key, + const std::string& tenant_id = "default") const { + return GetObjectHintState(service.master_service_, key, tenant_id); + } + std::vector replica_list; std::vector policy_files_; size_t next_policy_file_ = 0; @@ -341,6 +373,169 @@ TEST(TenantScopedStorageKeyTest, RoundTripsAndParsesLegacyKeys) { EXPECT_EQ(legacy_key, "legacy_key"); } +TEST(WorkloadHintsTest, DefaultsAndForSingleKeyPreserveSharedHints) { + ReplicateConfig config; + EXPECT_TRUE(config.IsDefault()); + EXPECT_TRUE(config.workload_hints.session_id.empty()); + EXPECT_EQ(config.workload_hints.retention_priority, 0); + + config.workload_hints.session_id = "session-shared-by-batch"; + config.workload_hints.retention_priority = -17; + config.group_ids = std::vector{"group-a", "group-b"}; + EXPECT_FALSE(config.IsDefault()); + + ReplicateConfig key_config = config.ForSingleKey(1); + EXPECT_EQ(key_config.workload_hints.session_id, "session-shared-by-batch"); + EXPECT_EQ(key_config.workload_hints.retention_priority, -17); + ASSERT_TRUE(key_config.group_ids.has_value()); + EXPECT_EQ(*key_config.group_ids, (std::vector{"group-b"})); +} + +TEST_F(MasterServiceTest, PutStoresHintsAndDuplicatePutKeepsOriginalHints) { + MasterService service; + const auto mounted = PrepareSimpleSegment(service); + + ReplicateConfig config; + config.workload_hints.session_id = "put-session"; + config.workload_hints.retention_priority = -9; + PutCompletedObject(service, mounted.client_id, "hinted-put", config); + + auto stored = GetObjectHintState(service, "hinted-put"); + ASSERT_TRUE(stored.has_value()); + EXPECT_EQ(stored->workload_hints.session_id, "put-session"); + EXPECT_EQ(stored->workload_hints.retention_priority, -9); + + ReplicateConfig duplicate_config = config; + duplicate_config.workload_hints.session_id = "replacement-session"; + duplicate_config.workload_hints.retention_priority = 99; + auto duplicate = service.PutStart(mounted.client_id, "hinted-put", + "default", 1024, duplicate_config); + ASSERT_FALSE(duplicate.has_value()); + EXPECT_EQ(duplicate.error(), ErrorCode::OBJECT_ALREADY_EXISTS); + + stored = GetObjectHintState(service, "hinted-put"); + ASSERT_TRUE(stored.has_value()); + EXPECT_EQ(stored->workload_hints.session_id, "put-session"); + EXPECT_EQ(stored->workload_hints.retention_priority, -9); +} + +TEST_F(MasterServiceTest, UpsertPreservesExistingHintsAndSetsHintsOnInsert) { + MasterService service; + const auto mounted = PrepareSimpleSegment(service); + + ReplicateConfig original_config; + original_config.workload_hints.session_id = "original-session"; + original_config.workload_hints.retention_priority = 7; + PutCompletedObject(service, mounted.client_id, "hinted-upsert", + original_config); + + ReplicateConfig incoming_config; + incoming_config.workload_hints.session_id = "incoming-session"; + incoming_config.workload_hints.retention_priority = 100; + + auto same_size = service.UpsertStart(mounted.client_id, "hinted-upsert", + "default", 1024, incoming_config); + ASSERT_TRUE(same_size.has_value()) << toString(same_size.error()); + auto stored = GetObjectHintState(service, "hinted-upsert"); + ASSERT_TRUE(stored.has_value()); + EXPECT_EQ(stored->workload_hints.session_id, "original-session"); + EXPECT_EQ(stored->workload_hints.retention_priority, 7); + ASSERT_TRUE(service + .UpsertEnd(mounted.client_id, "hinted-upsert", "default", + ReplicaType::MEMORY) + .has_value()); + + incoming_config.workload_hints.session_id = "new-size-session"; + incoming_config.workload_hints.retention_priority = -100; + auto different_size = service.UpsertStart( + mounted.client_id, "hinted-upsert", "default", 2048, incoming_config); + ASSERT_TRUE(different_size.has_value()) << toString(different_size.error()); + stored = GetObjectHintState(service, "hinted-upsert"); + ASSERT_TRUE(stored.has_value()); + EXPECT_EQ(stored->workload_hints.session_id, "original-session"); + EXPECT_EQ(stored->workload_hints.retention_priority, 7); + ASSERT_TRUE(service + .UpsertEnd(mounted.client_id, "hinted-upsert", "default", + ReplicaType::MEMORY) + .has_value()); + + auto insert = service.UpsertStart(mounted.client_id, "hinted-upsert-new", + "default", 1024, incoming_config); + ASSERT_TRUE(insert.has_value()) << toString(insert.error()); + stored = GetObjectHintState(service, "hinted-upsert-new"); + ASSERT_TRUE(stored.has_value()); + EXPECT_EQ(stored->workload_hints.session_id, "new-size-session"); + EXPECT_EQ(stored->workload_hints.retention_priority, -100); + ASSERT_TRUE(service + .UpsertEnd(mounted.client_id, "hinted-upsert-new", + "default", ReplicaType::MEMORY) + .has_value()); +} + +TEST_F(MasterServiceTest, BatchPutRpcPreservesHintsAlongsideGroupIds) { + const int rpc_port = getFreeTcpPort(); + ASSERT_GT(rpc_port, 0); + + WrappedMasterServiceConfig service_config; + service_config.default_kv_lease_ttl = 100; + service_config.enable_metric_reporting = false; + WrappedMasterService service(service_config); + coro_rpc::coro_rpc_server server(/*thread_num=*/2, rpc_port, + /*address=*/"127.0.0.1"); + RegisterRpcService(server, service); + auto start_error = server.async_start(); + ASSERT_FALSE(start_error.hasResult()); + + const UUID client_id = generate_uuid(); + const Segment segment = MakeSegment("workload-hints-rpc-segment"); + ASSERT_TRUE(service.MountSegment(segment, client_id).has_value()); + + { + MasterClient client(client_id); + ErrorCode connect_result = ErrorCode::INTERNAL_ERROR; + WaitUntil( + [&] { + connect_result = + client.Connect("127.0.0.1:" + std::to_string(rpc_port)); + return connect_result == ErrorCode::OK; + }, + std::chrono::seconds(3)); + ASSERT_EQ(connect_result, ErrorCode::OK); + + const std::vector keys = {"rpc-hint-grouped", + "rpc-hint-ungrouped"}; + ReplicateConfig config; + config.workload_hints.session_id = "rpc-session"; + config.workload_hints.retention_priority = -31; + config.group_ids = std::vector{"rpc-workload-group", ""}; + + auto results = client.BatchPutStart( + keys, std::vector>{{1024}, {2048}}, config); + ASSERT_EQ(results.size(), keys.size()); + ASSERT_TRUE(results[0].has_value()) << toString(results[0].error()); + ASSERT_TRUE(results[1].has_value()) << toString(results[1].error()); + + auto grouped = GetObjectHintState(service, keys[0]); + ASSERT_TRUE(grouped.has_value()); + EXPECT_EQ(grouped->workload_hints.session_id, "rpc-session"); + EXPECT_EQ(grouped->workload_hints.retention_priority, -31); + EXPECT_EQ(grouped->group_id, "rpc-workload-group"); + + auto ungrouped = GetObjectHintState(service, keys[1]); + ASSERT_TRUE(ungrouped.has_value()); + EXPECT_EQ(ungrouped->workload_hints.session_id, "rpc-session"); + EXPECT_EQ(ungrouped->workload_hints.retention_priority, -31); + EXPECT_TRUE(ungrouped->group_id.empty()); + + auto end_results = client.BatchPutEnd(keys, ReplicaType::MEMORY); + ASSERT_EQ(end_results.size(), keys.size()); + EXPECT_TRUE(end_results[0].has_value()); + EXPECT_TRUE(end_results[1].has_value()); + } + + server.stop(); +} + std::string GenerateKeyForSegment(const UUID& client_id, const std::unique_ptr& service, const std::string& segment_name) { diff --git a/mooncake-wheel/tests/test_import_structure.py b/mooncake-wheel/tests/test_import_structure.py index e05c980f30..06510b8e6f 100644 --- a/mooncake-wheel/tests/test_import_structure.py +++ b/mooncake-wheel/tests/test_import_structure.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 import unittest -class TestImportStructure(unittest.TestCase): +class TestImportStructure(unittest.TestCase): def test_new_import_structure(self): """Test that the new import structure works.""" import mooncake.engine @@ -16,15 +16,29 @@ def test_new_import_structure(self): # Verify direct access to TransferOpcode self.assertIsNotNone(mooncake.engine.TransferOpcode) - from mooncake.store import MooncakeDistributedStore, ReplicateConfig + from mooncake.store import ( + MooncakeDistributedStore, + ReplicateConfig, + WorkloadHints, + ) # Just verify we can create instances store = MooncakeDistributedStore() config = ReplicateConfig() config.group_ids = ["group-a", ""] + hints = WorkloadHints() + self.assertEqual(hints.session_id, "") + self.assertEqual(hints.retention_priority, 0) + + hints.session_id = "session-a" + hints.retention_priority = 7 + config.workload_hints = hints + self.assertIsNotNone(store) self.assertEqual(config.group_ids, ["group-a", ""]) + self.assertEqual(config.workload_hints.session_id, "session-a") + self.assertEqual(config.workload_hints.retention_priority, 7) config.group_ids = None self.assertIsNone(config.group_ids) @@ -43,5 +57,6 @@ def test_top_level_buffer_pool_alias(self): self.assertIs(BufferPool, RegisteredBufferPool) -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main()