Skip to content

[Store] Prevent drain success with local disk replicas #2834

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
feichai0017 wants to merge 1 commit into
base: main
Choose a base branch
from
+136 −4
Open

[Store] Prevent drain success with local disk replicas #2834

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions mooncake-store/include/master_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2052,6 +2052,9 @@ class MasterService {
// Task manager
ClientTaskManager task_manager_;

using DrainSourceSegmentsByClient =
std::unordered_map<UUID, std::vector<std::string>, boost::hash<UUID>>;

struct ActiveDrainTask {
UUID task_id;
std::string tenant_id;
Expand All @@ -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<UUID, ActiveDrainTask, boost::hash<UUID>>
active_tasks;
std::unordered_set<std::string> completed_unit_keys;
Expand All @@ -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<bool> job_dispatch_running_{false};
Expand Down
81 changes: 77 additions & 4 deletions mooncake-store/src/master_service.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8323,6 +8323,7 @@ tl::expected<void, ErrorCode> MasterService::ValidateDrainRequestLocked(
tl::expected<UUID, ErrorCode> MasterService::CreateDrainJob(
const CreateDrainJobRequest& request) {
std::vector<std::string> draining_segments;
DrainSourceSegmentsByClient source_segments_by_client;
{
ScopedSegmentAccess segment_access =
segment_manager_.getSegmentAccess();
Expand All @@ -8333,7 +8334,18 @@ tl::expected<UUID, ErrorCode> 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) {
Expand All @@ -8343,12 +8355,14 @@ tl::expected<UUID, ErrorCode> 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<DrainJob>();
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;
Expand Down Expand Up @@ -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<std::string> MasterService::SelectDrainTargetForKey(
const ObjectMetadata& metadata, const std::string& source_segment,
const std::vector<std::string>& requested_targets) {
Expand Down Expand Up @@ -8553,14 +8574,37 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) {
for (const auto& [_, task] : job.active_tasks) {
active_unit_keys.insert(task.unit_key);
}

std::unordered_set<std::string> blocked_unit_keys;
{
std::shared_lock<std::shared_mutex> 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) {
// 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);
Expand Down Expand Up @@ -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) {
Expand All @@ -8646,12 +8692,36 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) {

std::unordered_set<std::string> remaining_segments;
std::unordered_set<std::string> remaining_unit_keys;
bool has_remaining_local_disk_replicas = false;
{
std::shared_lock<std::shared_mutex> 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) {
Expand Down Expand Up @@ -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;
}

Expand Down
51 changes: 51 additions & 0 deletions mooncake-store/tests/master_service_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<MasterService>(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<MasterService>();

Expand Down
Loading