From 3bc8df676ee2d8444898a0a8fd6323b46eb25791 Mon Sep 17 00:00:00 2001 From: eric_song Date: Fri, 10 Jul 2026 16:26:08 +0800 Subject: [PATCH] [Store] Prevent drain success with local disk replicas --- mooncake-store/include/master_service.h | 8 ++ mooncake-store/src/master_service.cpp | 81 +++++++++++++++++++- mooncake-store/tests/master_service_test.cpp | 51 ++++++++++++ 3 files changed, 136 insertions(+), 4 deletions(-) diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 42fd27efa3..cbb615e8f4 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -2052,6 +2052,9 @@ class MasterService { // Task manager ClientTaskManager task_manager_; + using DrainSourceSegmentsByClient = + std::unordered_map, boost::hash>; + struct ActiveDrainTask { UUID task_id; std::string tenant_id; @@ -2071,10 +2074,12 @@ class MasterService { std::chrono::system_clock::time_point created_at; std::chrono::system_clock::time_point last_updated_at; std::string message; + DrainSourceSegmentsByClient source_segments_by_client; uint64_t succeeded_units{0}; uint64_t failed_units{0}; uint64_t blocked_units{0}; uint64_t migrated_bytes{0}; + bool has_unsupported_local_disk_replicas{false}; std::unordered_map> active_tasks; std::unordered_set completed_unit_keys; @@ -2099,6 +2104,9 @@ class MasterService { std::string MakeDrainUnitKey(const std::string& tenant_id, const std::string& key, const std::string& source_segment) const; + std::string MakeDrainLocalDiskUnitKey(const std::string& tenant_id, + const std::string& key, + const UUID& client_id) const; std::thread job_dispatch_thread_; std::atomic job_dispatch_running_{false}; diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index bf55656980..b3cb925b37 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -8323,6 +8323,7 @@ tl::expected MasterService::ValidateDrainRequestLocked( tl::expected MasterService::CreateDrainJob( const CreateDrainJobRequest& request) { std::vector draining_segments; + DrainSourceSegmentsByClient source_segments_by_client; { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); @@ -8333,7 +8334,18 @@ tl::expected MasterService::CreateDrainJob( draining_segments.reserve(request.segments.size()); for (const auto& segment_name : request.segments) { - auto err = segment_access.SetSegmentStatusByName( + UUID client_id; + auto err = segment_access.GetClientIdBySegmentName(segment_name, + client_id); + if (err != ErrorCode::OK) { + for (const auto& updated_segment : draining_segments) { + (void)segment_access.SetSegmentStatusByName( + updated_segment, SegmentStatus::OK); + } + return tl::make_unexpected(err); + } + + err = segment_access.SetSegmentStatusByName( segment_name, SegmentStatus::DRAINING); if (err != ErrorCode::OK) { for (const auto& updated_segment : draining_segments) { @@ -8343,12 +8355,14 @@ tl::expected MasterService::CreateDrainJob( return tl::make_unexpected(err); } draining_segments.push_back(segment_name); + source_segments_by_client[client_id].push_back(segment_name); } } auto job = std::make_shared(); job->id = generate_uuid(); job->request = request; + job->source_segments_by_client = std::move(source_segments_by_client); job->created_at = std::chrono::system_clock::now(); job->last_updated_at = job->created_at; job->status = JobStatus::CREATED; @@ -8446,6 +8460,13 @@ std::string MasterService::MakeDrainUnitKey( ":" + std::to_string(key.size()) + ":" + key + ":" + source_segment; } +std::string MasterService::MakeDrainLocalDiskUnitKey( + const std::string& tenant_id, const std::string& key, + const UUID& client_id) const { + return "local-disk:" + UuidToString(client_id) + ":" + + MakeDrainUnitKey(tenant_id, key, ""); +} + std::optional MasterService::SelectDrainTargetForKey( const ObjectMetadata& metadata, const std::string& source_segment, const std::vector& requested_targets) { @@ -8553,7 +8574,6 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { for (const auto& [_, task] : job.active_tasks) { active_unit_keys.insert(task.unit_key); } - std::unordered_set blocked_unit_keys; { std::shared_lock shared_lock(snapshot_mutex_); @@ -8561,6 +8581,30 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { MetadataShardAccessorRO shard(this, i); for (const auto& [tenant_id, tenant_state] : shard->tenants) { for (const auto& [key, metadata] : tenant_state.metadata) { + // LOCAL_DISK replicas are client-scoped and have no + // segment name. Associate them with every requested + // source segment owned by that client and fail closed: + // REPLICA_MOVE cannot migrate them yet. + metadata.VisitReplicas( + &Replica::fn_is_local_disk_replica, + [&](const Replica& replica) { + const auto holder_id = + replica.get_local_disk_client_id(); + if (!holder_id.has_value() || + !job.source_segments_by_client.contains( + holder_id.value())) { + return; + } + + const auto unit_key = MakeDrainLocalDiskUnitKey( + tenant_id, key, holder_id.value()); + if (job.terminal_failed_unit_keys.insert(unit_key) + .second) { + job.failed_units++; + job.has_unsupported_local_disk_replicas = true; + } + }); + for (const auto& source_segment : job.request.segments) { const auto unit_key = MakeDrainUnitKey(tenant_id, key, source_segment); @@ -8636,7 +8680,9 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { job.status = JobStatus::RUNNING; job.last_updated_at = std::chrono::system_clock::now(); - job.message = "Drain job running"; + job.message = job.has_unsupported_local_disk_replicas + ? "Drain job cannot migrate LOCAL_DISK replicas" + : "Drain job running"; } bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { @@ -8646,12 +8692,36 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { std::unordered_set remaining_segments; std::unordered_set remaining_unit_keys; + bool has_remaining_local_disk_replicas = false; { std::shared_lock shared_lock(snapshot_mutex_); for (size_t i = 0; i < kNumShards; ++i) { MetadataShardAccessorRO shard(this, i); for (const auto& [tenant_id, tenant_state] : shard->tenants) { for (const auto& [key, metadata] : tenant_state.metadata) { + metadata.VisitReplicas( + &Replica::fn_is_local_disk_replica, + [&](const Replica& replica) { + const auto holder_id = + replica.get_local_disk_client_id(); + if (!holder_id.has_value()) { + return; + } + auto source_it = job.source_segments_by_client.find( + holder_id.value()); + if (source_it == + job.source_segments_by_client.end()) { + return; + } + + has_remaining_local_disk_replicas = true; + remaining_segments.insert(source_it->second.begin(), + source_it->second.end()); + remaining_unit_keys.insert( + MakeDrainLocalDiskUnitKey(tenant_id, key, + holder_id.value())); + }); + const auto replica_segments = metadata.GetReplicaSegmentNames(); for (const auto& source_segment : job.request.segments) { @@ -8713,7 +8783,10 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { job.status = JobStatus::FAILED; job.last_updated_at = std::chrono::system_clock::now(); - job.message = "Drain job failed: unrecoverable units remain"; + job.message = has_remaining_local_disk_replicas + ? "Drain job failed: LOCAL_DISK replicas remain on " + "source clients" + : "Drain job failed: unrecoverable units remain"; return true; } diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 46512a5749..ade78df067 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -6289,6 +6289,57 @@ TEST_F(MasterServiceTest, DrainJobSchedulesMoveTaskAndConvergesToDrained) { EXPECT_FALSE(segment_names.contains("segment_0")); } +TEST_F(MasterServiceTest, + DrainJobFailsWhenLocalDiskReplicaRemainsOnSourceClient) { + auto service_config = MasterServiceConfig::builder() + .set_enable_offload(true) + .set_default_kv_lease_ttl(0) + .build(); + auto service_ = std::make_unique(service_config); + + const auto ctx0 = PrepareSimpleSegment(*service_, "segment_0", 0x300000000, + kDefaultSegmentSize); + [[maybe_unused]] const auto ctx1 = PrepareSimpleSegment( + *service_, "segment_1", 0x400000000, kDefaultSegmentSize); + ASSERT_TRUE( + service_->MountLocalDiskSegment(ctx0.client_id, true).has_value()); + + const std::string key = "drain_local_disk_only_key"; + constexpr int64_t object_size = 1024; + OffloadTaskItem task{ + .tenant_id = "default", .key = key, .size = object_size}; + StorageObjectMetadata metadata; + metadata.data_size = object_size; + metadata.transport_endpoint = "local-disk-endpoint"; + ASSERT_TRUE( + service_->NotifyOffloadSuccess(ctx0.client_id, {task}, {metadata}) + .has_value()); + + CreateDrainJobRequest request; + request.segments = {"segment_0"}; + request.target_segments = {"segment_1"}; + request.max_concurrency = 1; + + auto job_id = service_->CreateDrainJob(request); + ASSERT_TRUE(job_id.has_value()); + + WaitUntil([&] { + auto query = service_->QueryDrainJob(job_id.value()); + return query.has_value() && query->status == JobStatus::FAILED; + }); + + auto query = service_->QueryDrainJob(job_id.value()); + ASSERT_TRUE(query.has_value()); + EXPECT_EQ(query->status, JobStatus::FAILED); + EXPECT_EQ(query->active_units, 0u); + EXPECT_EQ(query->failed_units, 1u); + EXPECT_NE(query->message.find("LOCAL_DISK"), std::string::npos); + + auto segment_status = service_->QuerySegmentStatus("segment_0"); + ASSERT_TRUE(segment_status.has_value()); + EXPECT_EQ(segment_status.value(), SegmentStatus::OK); +} + TEST_F(MasterServiceTest, CancelDrainJobRestoresSegmentStatus) { auto service_ = std::make_unique();