From d13ba4a3f514633d4e62b420b6dbb57f44696b40 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Tue, 20 Oct 2026 13:07:24 +0800 Subject: [PATCH] [Store] Recover drain jobs through batch OpLog --- mooncake-store/include/client_offboarding.h | 1 + mooncake-store/include/drain_recovery.h | 162 +++ .../include/ha/oplog/oplog_applier.h | 2 +- mooncake-store/include/ha/oplog/oplog_types.h | 3 + .../include/ha/snapshot/batch_oplog/capture.h | 9 + .../include/ha/snapshot/batch_oplog/codec.h | 12 + .../ha/snapshot/batch_oplog/metadata.h | 7 + mooncake-store/include/hot_standby_service.h | 3 + .../include/master_metric_manager.h | 2 + mooncake-store/include/master_service.h | 60 +- mooncake-store/include/metadata_store.h | 21 +- mooncake-store/include/task_manager.h | 35 +- mooncake-store/src/CMakeLists.txt | 2 + mooncake-store/src/drain_recovery.cpp | 318 ++++++ mooncake-store/src/ha/oplog/oplog_applier.cpp | 96 +- .../batch_oplog_pruning_coordinator.cpp | 9 + .../batch_oplog/batch_oplog_snapshot_gc.cpp | 10 + .../batch_oplog_snapshot_provider.cpp | 27 +- .../src/ha/snapshot/batch_oplog/codec.cpp | 18 + .../src/ha/snapshot/batch_oplog/metadata.cpp | 43 +- .../src/ha/snapshot/batch_oplog/writer.cpp | 37 +- mooncake-store/src/ha/standby_controller.cpp | 2 +- .../src/ha/standby_metadata_store.cpp | 10 + mooncake-store/src/hot_standby_service.cpp | 73 +- mooncake-store/src/master_metric_manager.cpp | 9 + mooncake-store/src/master_service.cpp | 1004 ++++++++++++----- .../src/master_service_drain_recovery.cpp | 479 ++++++++ mooncake-store/src/segment.cpp | 8 + mooncake-store/src/task_manager.cpp | 156 ++- .../tests/ha/master_service_ha_test.cpp | 755 ++++++++++++- .../tests/ha/oplog/oplog_applier_test.cpp | 20 + .../ha/snapshot/batch_oplog/codec_test.cpp | 16 + .../tests/ha/snapshot/batch_oplog/gc_test.cpp | 24 +- .../ha/snapshot/batch_oplog/metadata_test.cpp | 47 + .../snapshot/batch_oplog/promotion_test.cpp | 50 + .../batch_oplog/pruning_coordinator_test.cpp | 35 +- .../ha/snapshot/batch_oplog/writer_test.cpp | 215 +++- .../master_service/master_service_test_peer.h | 9 + mooncake-store/tests/task_manager_test.cpp | 90 ++ 39 files changed, 3461 insertions(+), 418 deletions(-) create mode 100644 mooncake-store/include/drain_recovery.h create mode 100644 mooncake-store/src/drain_recovery.cpp create mode 100644 mooncake-store/src/master_service_drain_recovery.cpp diff --git a/mooncake-store/include/client_offboarding.h b/mooncake-store/include/client_offboarding.h index 6aad5347ab..b988aa7580 100644 --- a/mooncake-store/include/client_offboarding.h +++ b/mooncake-store/include/client_offboarding.h @@ -32,6 +32,7 @@ struct PreparedSegmentOffboarding { std::string segment_name; std::string transport_endpoint; size_t metrics_dec_capacity{0}; + bool recovered{false}; }; // Process-local residual work for one terminal Client incarnation. The job is diff --git a/mooncake-store/include/drain_recovery.h b/mooncake-store/include/drain_recovery.h new file mode 100644 index 0000000000..a0b037db8f --- /dev/null +++ b/mooncake-store/include/drain_recovery.h @@ -0,0 +1,162 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "replica.h" +#include "types.h" + +namespace mooncake { + +// Persist values only. Runtime locks, callbacks and allocator pointers stay +// local. +struct DrainJobRecord { + UUID id{}; + int32_t status{0}; + std::vector sources, targets; + uint32_t concurrency{0}; + int64_t created_ms{0}, updated_ms{0}; + std::string message; + uint64_t succeeded{0}, failed{0}, blocked{0}, migrated_bytes{0}; + YLT_REFL(DrainJobRecord, id, status, sources, targets, concurrency, + created_ms, updated_ms, message, succeeded, failed, blocked, + migrated_bytes); +}; + +struct DrainUnitRecord { + UUID job_id{}, task_id{}; + std::string unit_key, tenant, key, source, target; + uint64_t bytes{0}; + uint32_t retries{0}; + bool completed{false}, terminal_failed{false}; + YLT_REFL(DrainUnitRecord, job_id, task_id, unit_key, tenant, key, source, + target, bytes, retries, completed, terminal_failed); +}; + +struct DrainSegmentRecord { + UUID id{}, client_id{}; + std::string name, endpoint; + uint64_t base{0}, size{0}; + // UNDEFINED updates registration identity without changing durable status. + int32_t status{0}; + YLT_REFL(DrainSegmentRecord, id, client_id, name, endpoint, base, size, + status); +}; + +struct DrainJobUpdate { + std::optional job; + std::optional unit; + std::vector segments; + YLT_REFL(DrainJobUpdate, job, unit, segments); +}; + +struct DurableTaskRecord { + UUID id{}, client_id{}; + int32_t type{0}, status{0}; + std::string payload, message; + int64_t created_ms{0}, updated_ms{0}; + uint32_t max_retries{0}; + std::optional drain; + bool removed{false}; + YLT_REFL(DurableTaskRecord, id, client_id, type, status, payload, message, + created_ms, updated_ms, max_retries, drain, removed); +}; + +struct ReplicationRecoveryRecord { + std::string tenant, key; + UUID client_id{}; + int32_t type{0}; + int64_t started_ms{0}; + ReplicaID source_id{0}; + std::vector targets; + uint64_t pending_bytes{0}; + UUID dynamic_lease{}; + uint64_t dynamic_version{0}; + YLT_REFL(ReplicationRecoveryRecord, tenant, key, client_id, type, + started_ms, source_id, targets, pending_bytes, dynamic_lease, + dynamic_version); +}; + +struct ReplicationStartRecord { + ReplicationRecoveryRecord runtime; + std::string metadata_payload; + std::vector segments; + YLT_REFL(ReplicationStartRecord, runtime, metadata_payload, segments); +}; + +enum class RecoveryRecordKind : uint8_t { + Job, + Unit, + Task, + Segment, + Replication, + RetiredReplica, +}; + +struct RetiredReplicaRecord { + Replica::Descriptor replica; + int64_t release_ms{0}; + YLT_REFL(RetiredReplicaRecord, replica, release_ms); +}; + +struct DrainRecoveryRecord { + RecoveryRecordKind kind; + std::string key, payload; + YLT_REFL(DrainRecoveryRecord, kind, key, payload); +}; + +using DrainRecoveryRecords = std::map; + +inline constexpr uint32_t kMaxDrainUnitRetries = 3; + +std::string DrainRecoveryKey(RecoveryRecordKind kind, const std::string& key); +std::string DrainUnitKey(const std::string& tenant, const std::string& key, + const std::string& source); +bool ValidateDrainJob(const DrainJobRecord& job); +bool ValidateDrainUnit(const DrainUnitRecord& unit); +bool ValidateDurableTask(const DurableTaskRecord& task); +bool ValidateReplicationRecovery( + const ReplicationRecoveryRecord& runtime, + const std::vector& replicas, uint64_t object_size, + bool allow_missing_replicas = false); + +// Updated by standby replay; snapshot capture and promotion freeze replay. +// Ordered keys permit bounded traversal without a second full state copy. +class DrainRecoveryState { + public: + bool Apply(const DrainJobUpdate& update); + bool Apply(const DurableTaskRecord& task); + bool Apply(const ReplicationRecoveryRecord& replication); + bool Restore(const DrainRecoveryRecord& record); + void EraseReplication(const std::string& tenant, const std::string& key); + void EraseSegment(const std::string& name, const std::string& endpoint); + bool Retire(const RetiredReplicaRecord& replica); + void PruneRetired(int64_t now_ms); + const DrainRecoveryRecords& records() const { return records_; } + void Clear() { records_.clear(); } + + private: + DrainRecoveryRecords records_; +}; + +template +std::string EncodeRecoveryValue(const T& value) { + const auto encoded = struct_pack::serialize(value); + return {encoded.begin(), encoded.end()}; +} + +template +bool DecodeRecoveryValue(const std::string& encoded, T& value) { + return struct_pack::deserialize_to(value, encoded) == struct_pack::errc::ok; +} + +template +DrainRecoveryRecord MakeRecoveryRecord(RecoveryRecordKind kind, + const std::string& key, const T& value) { + return {kind, DrainRecoveryKey(kind, key), EncodeRecoveryValue(value)}; +} + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_applier.h b/mooncake-store/include/ha/oplog/oplog_applier.h index 4bb49185a9..01b818f6a3 100644 --- a/mooncake-store/include/ha/oplog/oplog_applier.h +++ b/mooncake-store/include/ha/oplog/oplog_applier.h @@ -70,7 +70,7 @@ class OpLogApplier { * @brief Apply PUT_END operation * @param entry OpLog entry */ - void ApplyPutEnd(const OpLogEntry& entry); + bool ApplyPutEnd(const OpLogEntry& entry); /** * @brief Apply PUT_REVOKE operation diff --git a/mooncake-store/include/ha/oplog/oplog_types.h b/mooncake-store/include/ha/oplog/oplog_types.h index f766895e89..5bf15fc92a 100644 --- a/mooncake-store/include/ha/oplog/oplog_types.h +++ b/mooncake-store/include/ha/oplog/oplog_types.h @@ -22,6 +22,9 @@ enum class OpType : uint8_t { WEIGHT_METADATA_DELETE = 9, WEIGHT_LEASE_UPSERT = 10, WEIGHT_LEASE_DELETE = 11, + DRAIN_JOB_UPDATE = 12, + TASK_UPDATE = 13, + REPLICATION_START = 14, OP_TYPE_MAX, }; diff --git a/mooncake-store/include/ha/snapshot/batch_oplog/capture.h b/mooncake-store/include/ha/snapshot/batch_oplog/capture.h index ecefeb3530..43ce4ad796 100644 --- a/mooncake-store/include/ha/snapshot/batch_oplog/capture.h +++ b/mooncake-store/include/ha/snapshot/batch_oplog/capture.h @@ -22,6 +22,9 @@ class BatchOpLogSnapshotCapture { segments(std::move(other.segments)), weight_metadata(std::move(other.weight_metadata)), cursor_(std::move(other.cursor_)), + recovery_cursor_(other.recovery_cursor_), + recovery_end_(other.recovery_end_), + recovery_done_(other.recovery_done_), generation_(std::exchange(other.generation_, 0)), lease_state_(std::move(other.lease_state_)) {} BatchOpLogSnapshotCapture& operator=( @@ -34,6 +37,9 @@ class BatchOpLogSnapshotCapture { segments = std::move(other.segments); weight_metadata = std::move(other.weight_metadata); cursor_ = std::move(other.cursor_); + recovery_cursor_ = other.recovery_cursor_; + recovery_end_ = other.recovery_end_; + recovery_done_ = other.recovery_done_; generation_ = std::exchange(other.generation_, 0); lease_state_ = std::move(other.lease_state_); } @@ -45,6 +51,7 @@ class BatchOpLogSnapshotCapture { ~BatchOpLogSnapshotCapture() { Release(); } bool done() const { return !lease_state_ || cursor_.done(); } + bool recovery_done() const { return !lease_state_ || recovery_done_; } uint64_t last_included_seq; uint64_t last_included_batch_id; @@ -91,6 +98,8 @@ class BatchOpLogSnapshotCapture { } StandbyMetadataStore::SnapshotCursor cursor_; + DrainRecoveryRecords::const_iterator recovery_cursor_, recovery_end_; + bool recovery_done_{true}; uint64_t generation_; std::shared_ptr lease_state_; }; diff --git a/mooncake-store/include/ha/snapshot/batch_oplog/codec.h b/mooncake-store/include/ha/snapshot/batch_oplog/codec.h index ef8b8cf96a..fca3ebacc4 100644 --- a/mooncake-store/include/ha/snapshot/batch_oplog/codec.h +++ b/mooncake-store/include/ha/snapshot/batch_oplog/codec.h @@ -18,6 +18,18 @@ struct BatchOpLogSnapshotObjectChunk { YLT_REFL(BatchOpLogSnapshotObjectChunk, chunk_index, objects); }; +struct BatchOpLogRecoveryChunk { + uint64_t chunk_index{0}; + std::vector records; + YLT_REFL(BatchOpLogRecoveryChunk, chunk_index, records); +}; + +std::vector EncodeBatchOpLogRecoveryChunk( + uint64_t index, std::vector records); +tl::expected +DecodeBatchOpLogRecoveryChunk(const std::vector& encoded, + uint64_t index, uint64_t count); + std::vector EncodeBatchOpLogSnapshotSegments( const std::vector& segments); tl::expected, std::string> diff --git a/mooncake-store/include/ha/snapshot/batch_oplog/metadata.h b/mooncake-store/include/ha/snapshot/batch_oplog/metadata.h index a3e0453f31..fd3c1e6f0c 100644 --- a/mooncake-store/include/ha/snapshot/batch_oplog/metadata.h +++ b/mooncake-store/include/ha/snapshot/batch_oplog/metadata.h @@ -18,6 +18,9 @@ inline constexpr char kBatchOpLogSnapshotFormat[] = inline constexpr uint32_t kBatchOpLogWeightSnapshotSchemaVersion = 2; inline constexpr char kBatchOpLogWeightSnapshotFormat[] = "standby-oplog-materialized/v2"; +inline constexpr uint32_t kBatchOpLogRecoverySnapshotSchemaVersion = 3; +inline constexpr char kBatchOpLogRecoverySnapshotFormat[] = + "standby-oplog-materialized/v3"; struct BatchOpLogSnapshotDescriptor { uint32_t schema_version{kBatchOpLogSnapshotSchemaVersion}; @@ -56,6 +59,7 @@ struct BatchOpLogSnapshotManifest { BatchOpLogSnapshotObjectDescriptor segments; std::vector object_chunks; std::optional weight_metadata; + std::vector recovery_chunks; }; std::string EncodeBatchOpLogSnapshotDescriptor( @@ -71,6 +75,9 @@ DecodeBatchOpLogSnapshotManifest(std::string_view value); std::string BuildBatchOpLogSnapshotId(uint64_t last_included_batch_id, int64_t maintenance_lease_id); std::string BuildBatchOpLogSnapshotRoot(const std::string& cluster_id); +std::string BuildBatchOpLogSnapshotRecoveryChunkKey( + const std::string& snapshot_root, std::string_view snapshot_id, + uint64_t chunk_index); std::string BuildBatchOpLogSnapshotMaintenanceKey( const std::string& cluster_id); diff --git a/mooncake-store/include/hot_standby_service.h b/mooncake-store/include/hot_standby_service.h index eddfbd70f4..945c560489 100644 --- a/mooncake-store/include/hot_standby_service.h +++ b/mooncake-store/include/hot_standby_service.h @@ -153,6 +153,9 @@ class HotStandbyService { ErrorCode PromoteAndExportSnapshot(StandbySnapshot& out); bool IsBatchOpLogSnapshotMode() const; + bool CopyNextBatchOpLogRecoveryChunk(size_t count, + BatchOpLogSnapshotCapture& capture, + std::vector& out); tl::expected PromoteAndDetachBatchOpLogStore(); diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 4c00ab1be8..3448ed8bee 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -390,6 +390,7 @@ class MasterMetricManager { // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics void inc_copy_start_requests(int64_t val = 1); + void inc_orphaned_draining_restore(); void inc_copy_start_failures(int64_t val = 1); void inc_copy_end_requests(int64_t val = 1); void inc_copy_end_failures(int64_t val = 1); @@ -785,6 +786,7 @@ class MasterMetricManager { // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics ylt::metric::counter_t copy_start_requests_; + ylt::metric::counter_t orphaned_draining_restore_; ylt::metric::counter_t copy_start_failures_; ylt::metric::counter_t copy_end_requests_; ylt::metric::counter_t copy_end_failures_; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 60222c45cd..c7fb21a7d2 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -1435,13 +1435,16 @@ class MasterService { void FinalizeRemovedReplicasAfterDurable( const OpLogEntry& durable_entry, const std::vector& replica_ids, QuotaEraseMode quota_mode, - const std::vector& previous_media_hint = {}); + const std::vector& previous_media_hint = {}, + std::optional retire_at = + std::nullopt); // The durable finalizers of a decision taken under a key's lock that was // let go of before the OpLog entry became durable. Those that name the // object by `ObjectRef` act only on that same publication, so a same-key // recreate that landed meanwhile is left alone. void FinalizeMetadataEraseAfterDurable(route::ObjectRef object, const TenantId& tenant_id, + QuotaEraseMode quota_mode); void FinalizeExpiredProcessingReplicasAfterDurable( route::ObjectRef object, const OpLogEntry& durable_entry, @@ -2175,6 +2178,8 @@ class MasterService { const std::chrono::seconds put_start_discard_timeout_sec_; const std::chrono::seconds put_start_release_timeout_sec_; class DiscardedReplicas { + friend class MasterService; + public: DiscardedReplicas() = delete; @@ -2245,7 +2250,47 @@ class MasterService { std::unordered_set terminal_failed_unit_keys; }; - static constexpr uint32_t kMaxDrainUnitRetries = 3; + static DrainJobRecord DescribeDrainJob(const DrainJob& job); + static void PublishDrainJob(DrainJob& job, const DrainJobRecord& record); + std::vector DescribeDrainSegments( + ScopedSegmentAccess& access, const std::vector& names, + std::optional status); + ErrorCode SetDrainSegmentStatus(ScopedSegmentAccess& access, + const std::string& name, + SegmentStatus status); + ErrorCode GetDrainSegmentStatus(ScopedSegmentAccess& access, + const std::string& name, + SegmentStatus& status) const; + ErrorCode PersistRecoveryEntry(OpType type, const std::string& payload, + const std::string& key, + const std::string& tenant = "default"); + ErrorCode AwaitRecoveryDurability(uint64_t sequence); + ErrorCode PersistRecoveryTasks(std::span tasks); + ErrorCode PersistDrainJobUpdate(const DrainJobUpdate& update); + ErrorCode PersistReplicationStart( + MetadataAccessorRW& accessor, + std::optional& reservation); + ErrorCode PersistReplicationRevoke(MetadataAccessorRW& accessor); + tl::expected CreateMoveTaskImpl( + const std::string& key, const TenantId& tenant_id, + const std::string& source, const std::string& target, + std::optional drain); + tl::expected RestoreDrainRecovery( + const DrainRecoveryState& state, int64_t restore_epoch_ms); + tl::expected RestoreReplicationRecovery( + const StandbyMetadataStore& metadata_store); + struct RestoredDrainSegment : DrainSegmentRecord { + explicit RestoredDrainSegment(DrainSegmentRecord record) + : DrainSegmentRecord(std::move(record)) {} + SegmentLifetime lifetime; + std::shared_ptr unmount_mutex{ + std::make_shared()}; + }; + // Accessed under the segment-manager access lock, like mounted segments. + std::unordered_map + restored_drain_segments_; + void ForgetRestoredSegments( + const std::vector& segments); tl::expected ValidateDrainRequest( const CreateDrainJobRequest& request); @@ -2253,15 +2298,13 @@ class MasterService { ScopedSegmentAccess& segment_access, const CreateDrainJobRequest& request); void ProcessDrainJobs(); + void WarnOrphanedDrainingSegments(); void RefreshDrainJobTasks(DrainJob& job); void ScheduleDrainJobTasks(DrainJob& job); bool MaybeCompleteDrainJob(DrainJob& job); std::optional SelectDrainTargetForKey( const ObjectMetadata& metadata, const std::string& source_segment, const std::vector& requested_targets); - std::string MakeDrainUnitKey(const TenantId& tenant_id, - const std::string& key, - const std::string& source_segment) const; std::thread job_dispatch_thread_; std::atomic job_dispatch_running_{false}; @@ -2388,7 +2431,9 @@ class MasterService { const ObjectMetadata& metadata) const; std::string SerializeMetadataForOpLogFromReplicaDescriptors( const ObjectMetadata& metadata, - const std::vector& replicas) const; + const std::vector& replicas, + const std::vector& retired = {}, + bool replication_finished = false) const; ErrorCode InitializeBatchOpLogWriter(std::shared_ptr backend, bool require_fenced_writer); tl::expected AppendOpLogVisibleBeforeDurable( @@ -2420,7 +2465,8 @@ class MasterService { std::vector standby_memory_segments_; std::unordered_map standby_accounted_memory_bytes_; - ErrorCode ValidateStandbyRemountSegment(const Segment& segment) const; + ErrorCode ValidateStandbyRemountSegment(const Segment& segment, + const UUID& client_id) const; // True while the replica sits on an endpoint in // invalid_replica_endpoints_. The read path asks only this, never the set. diff --git a/mooncake-store/include/metadata_store.h b/mooncake-store/include/metadata_store.h index 08ff3c5598..f6720ae1cc 100644 --- a/mooncake-store/include/metadata_store.h +++ b/mooncake-store/include/metadata_store.h @@ -1,5 +1,7 @@ #pragma once +#include "drain_recovery.h" + #include #include #include @@ -110,8 +112,18 @@ struct MetadataPayload { struct_pack::compatible session_id; + struct_pack::compatible, 3> + retired_replicas; + + // Ordinary metadata updates can project only COMPLETE replicas. Preserve + // their live COPY/MOVE targets separately until an explicit settlement. + struct_pack::compatible, 4> + processing_replicas; + struct_pack::compatible replication_finished; + YLT_REFL(MetadataPayload, client_id, size, replicas, group_id, data_type, - hard_pinned, session_id); + hard_pinned, session_id, retired_replicas, processing_replicas, + replication_finished); // Convert to StandbyObjectMetadata StandbyObjectMetadata ToStandbyMetadata() const { @@ -164,6 +176,10 @@ class MetadataStore { public: virtual ~MetadataStore() = default; + // Written by the ordered applier; snapshot capture pauses that writer. + DrainRecoveryState& DrainRecovery() { return drain_recovery_; } + const DrainRecoveryState& DrainRecovery() const { return drain_recovery_; } + // NEW: tenant-aware methods (primary API) virtual bool PutMetadata(const std::string& tenant_id, const std::string& key, @@ -216,6 +232,9 @@ class MetadataStore { virtual bool RemoveWeightLease(uint64_t lease_id, const WeightRevisionIdentity& identity, uint64_t fenced_metadata_generation) = 0; + + private: + DrainRecoveryState drain_recovery_; }; } // namespace mooncake diff --git a/mooncake-store/include/task_manager.h b/mooncake-store/include/task_manager.h index b5cf8c69ba..dfb853a161 100644 --- a/mooncake-store/include/task_manager.h +++ b/mooncake-store/include/task_manager.h @@ -8,6 +8,9 @@ #include #include #include +#include +#include +#include "drain_recovery.h" #include "types.h" #include "mutex.h" #include "master_config.h" @@ -75,7 +78,8 @@ struct Task { std::string message; UUID assigned_client; - uint32_t max_retry_attempts; + uint32_t max_retry_attempts{0}; + std::optional drain; bool is_finished() const { return is_finished_status(status); } @@ -91,6 +95,9 @@ struct Task { } }; +DurableTaskRecord ToDurableTask(const Task& task, bool removed = false); +Task FromDurableTask(const DurableTaskRecord& record); + struct ReplicaCopyPayload { std::string tenant_id = "default"; std::string key; @@ -171,9 +178,10 @@ class ScopedTaskWriteAccess { template tl::expected submit_task_typed( const UUID& client_id, - const typename TaskPayloadTraits::type& payload) { + const typename TaskPayloadTraits::type& payload, + std::optional drain = std::nullopt) { std::string json = serialize_payload(payload); - return submit_task(client_id, Type, json); + return submit_task(client_id, Type, json, std::move(drain)); } std::vector pop_tasks(const UUID& client_id, size_t batch_size); @@ -193,9 +201,9 @@ class ScopedTaskWriteAccess { void prune_expired_tasks(); private: - tl::expected submit_task(const UUID& client_id, - TaskType type, - const std::string& payload); + tl::expected submit_task( + const UUID& client_id, TaskType type, const std::string& payload, + std::optional drain); ClientTaskManager* manager_; SharedMutexLocker lock_; @@ -206,13 +214,18 @@ class ClientTaskManager { friend class ScopedTaskReadAccess; friend class ScopedTaskWriteAccess; - explicit ClientTaskManager(const TaskManagerConfig& config) + using Persistence = + std::function)>; + + explicit ClientTaskManager(const TaskManagerConfig& config, + Persistence persistence = {}) : max_total_finished_tasks_(config.max_total_finished_tasks), max_total_pending_tasks_(config.max_total_pending_tasks), max_total_processing_tasks_(config.max_total_processing_tasks), pending_task_timeout_sec_(config.pending_task_timeout_sec), processing_task_timeout_sec_(config.processing_task_timeout_sec), - max_retry_attempts_(config.max_retry_attempts) {} + max_retry_attempts_(config.max_retry_attempts), + persistence_(std::move(persistence)) {} ~ClientTaskManager() = default; @@ -235,6 +248,12 @@ class ClientTaskManager { uint64_t pending_task_timeout_sec_; uint64_t processing_task_timeout_sec_; uint32_t max_retry_attempts_; + Persistence persistence_; + ErrorCode Persist(const Task& task, bool removed = false) const { + if (!persistence_) return ErrorCode::OK; + const auto record = ToDurableTask(task, removed); + return persistence_({&record, 1}); + } size_t total_pending_tasks_ GUARDED_BY(mutex_) = 0; size_t total_processing_tasks_ GUARDED_BY(mutex_) = 0; diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 91fa180b66..6239ebaab0 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -31,6 +31,8 @@ set(MOONCAKE_STORE_MASTER_SOURCES client_offboarding.cpp dynamic_replication_controller.cpp master_service.cpp + master_service_drain_recovery.cpp + drain_recovery.cpp promotion_candidate_tracker.cpp master_store_backend.cpp weight_store_manager.cpp diff --git a/mooncake-store/src/drain_recovery.cpp b/mooncake-store/src/drain_recovery.cpp new file mode 100644 index 0000000000..c2484d60cb --- /dev/null +++ b/mooncake-store/src/drain_recovery.cpp @@ -0,0 +1,318 @@ +#include "drain_recovery.h" + +#include +#include +#include + +#include "rpc_types.h" +#include "segment/status.h" +#include "task_manager.h" + +namespace mooncake { +namespace { +bool ValidTime(int64_t value) { + return value >= 0 && + value <= std::chrono::duration_cast( + std::chrono::system_clock::time_point::max() + .time_since_epoch()) + .count(); +} +} // namespace + +std::string DrainRecoveryKey(RecoveryRecordKind kind, const std::string& key) { + return std::to_string(static_cast(kind)) + ":" + key; +} + +std::string DrainUnitKey(const std::string& tenant, const std::string& key, + const std::string& source) { + return std::to_string(tenant.size()) + ":" + tenant + ":" + + std::to_string(key.size()) + ":" + key + ":" + source; +} + +bool ValidateDrainJob(const DrainJobRecord& job) { + if (job.id == UUID{} || job.sources.empty() || job.concurrency == 0 || + job.status < static_cast(JobStatus::CREATED) || + job.status > static_cast(JobStatus::CANCELED) || + !ValidTime(job.created_ms) || !ValidTime(job.updated_ms)) { + return false; + } + const std::unordered_set sources(job.sources.begin(), + job.sources.end()); + if (sources.size() != job.sources.size()) return false; + return std::none_of( + job.targets.begin(), job.targets.end(), + [&](const auto& target) { return sources.contains(target); }); +} + +bool ValidateDrainUnit(const DrainUnitRecord& unit) { + return unit.job_id != UUID{} && TenantId(unit.tenant).IsValid() && + unit.source != unit.target && + unit.unit_key == DrainUnitKey(unit.tenant, unit.key, unit.source) && + unit.retries <= kMaxDrainUnitRetries && + !(unit.completed && unit.terminal_failed) && + (!(unit.completed || unit.terminal_failed) || + unit.task_id == UUID{}); +} + +bool ValidateDurableTask(const DurableTaskRecord& task) { + if (task.id == UUID{}) return false; + if (task.removed) return true; + if (!ValidTime(task.created_ms) || !ValidTime(task.updated_ms) || + task.type < static_cast(TaskType::REPLICA_COPY) || + task.type > static_cast(TaskType::REPLICA_MOVE) || + task.status < static_cast(TaskStatus::PENDING) || + task.status > static_cast(TaskStatus::SUCCESS)) + return false; + if (!task.drain) return true; + const auto& unit = *task.drain; + if (!ValidateDrainUnit(unit) || unit.task_id != task.id || + task.type != static_cast(TaskType::REPLICA_MOVE)) + return false; + ReplicaMovePayload payload; + try { + struct_json::from_json(payload, task.payload); + } catch (const std::runtime_error&) { + return false; + } + return payload.tenant_id == unit.tenant && payload.key == unit.key && + payload.source == unit.source && payload.target == unit.target; +} + +bool ValidateReplicationRecovery( + const ReplicationRecoveryRecord& runtime, + const std::vector& replicas, uint64_t object_size, + bool allow_missing_replicas) { + if (!TenantId(runtime.tenant).IsValid() || object_size == 0 || + !ValidTime(runtime.started_ms) || runtime.source_id == 0 || + runtime.type < 0 || runtime.type > 1 || + (runtime.type == 1 && runtime.targets.size() > 1)) + return false; + std::unordered_set targets(runtime.targets.begin(), + runtime.targets.end()); + if (targets.size() != runtime.targets.size() || targets.contains(0) || + targets.contains(runtime.source_id)) + return false; + if (!targets.empty() && object_size > UINT64_MAX / targets.size()) + return false; + // Full target allocation is charged until COPY/MOVE settles. Changing + // this billing model requires translating these versioned records. + if (runtime.pending_bytes != object_size * targets.size()) return false; + bool source_found = false; + size_t targets_found = 0; + for (const auto& replica : replicas) { + if (replica.id == runtime.source_id) { + if (replica.status != ReplicaStatus::COMPLETE) return false; + source_found = true; + } + if (targets.contains(replica.id)) { + if (!replica.is_memory_replica() || + replica.status != ReplicaStatus::PROCESSING || + replica.get_memory_descriptor().buffer_descriptor.size_ != + object_size) + return false; + ++targets_found; + } + } + return allow_missing_replicas || + (source_found && targets_found == targets.size()); +} + +bool DrainRecoveryState::Restore(const DrainRecoveryRecord& record) { + bool valid = false; + std::string identity; + switch (record.kind) { + case RecoveryRecordKind::Job: { + DrainJobRecord job; + valid = DecodeRecoveryValue(record.payload, job) && + ValidateDrainJob(job); + identity = UuidToString(job.id); + break; + } + case RecoveryRecordKind::Unit: { + DrainUnitRecord unit; + valid = DecodeRecoveryValue(record.payload, unit) && + ValidateDrainUnit(unit); + identity = UuidToString(unit.job_id) + ":" + unit.unit_key; + break; + } + case RecoveryRecordKind::Task: { + DurableTaskRecord task; + valid = DecodeRecoveryValue(record.payload, task) && + !task.removed && ValidateDurableTask(task); + identity = UuidToString(task.id); + break; + } + case RecoveryRecordKind::Segment: { + DrainSegmentRecord segment; + valid = + DecodeRecoveryValue(record.payload, segment) && + !segment.endpoint.empty() && segment.size > 0 && + segment.status >= static_cast(SegmentStatus::OK) && + segment.status <= static_cast(SegmentStatus::DRAINED); + identity = segment.name; + break; + } + case RecoveryRecordKind::Replication: { + ReplicationRecoveryRecord replication; + valid = DecodeRecoveryValue(record.payload, replication) && + TenantId(replication.tenant).IsValid() && + replication.source_id != 0 && replication.type >= 0 && + replication.type <= 1 && ValidTime(replication.started_ms); + identity = + TenantId(replication.tenant).MakeScopedKey(replication.key); + break; + } + case RecoveryRecordKind::RetiredReplica: { + RetiredReplicaRecord retired; + valid = DecodeRecoveryValue(record.payload, retired) && + retired.replica.id != 0 && ValidTime(retired.release_ms); + identity = std::to_string(retired.replica.id); + break; + } + } + return valid && record.key == DrainRecoveryKey(record.kind, identity) && + records_.emplace(record.key, record).second; +} + +bool DrainRecoveryState::Apply(const DrainJobUpdate& update) { + // Validate every value before publishing any part of this event. + DrainRecoveryState prepared; + if (update.job && !prepared.Restore(MakeRecoveryRecord( + RecoveryRecordKind::Job, UuidToString(update.job->id), + *update.job))) + return false; + if (update.unit && + (!prepared.Restore(MakeRecoveryRecord( + RecoveryRecordKind::Unit, + UuidToString(update.unit->job_id) + ":" + update.unit->unit_key, + *update.unit)) || + (update.job && update.job->id != update.unit->job_id))) + return false; + for (auto segment : update.segments) { + if (segment.status == static_cast(SegmentStatus::UNDEFINED)) { + segment.status = static_cast(SegmentStatus::OK); + const auto key = + DrainRecoveryKey(RecoveryRecordKind::Segment, segment.name); + if (auto previous = records_.find(key); + previous != records_.end()) { + DrainSegmentRecord old; + if (!DecodeRecoveryValue(previous->second.payload, old)) + return false; + segment.status = old.status; + } + } + if (!prepared.Restore(MakeRecoveryRecord(RecoveryRecordKind::Segment, + segment.name, segment))) { + return false; + } + } + for (auto& [key, record] : prepared.records_) { + records_.insert_or_assign(key, std::move(record)); + } + return true; +} + +bool DrainRecoveryState::Apply(const DurableTaskRecord& task) { + if (!ValidateDurableTask(task)) return false; + const auto key = + DrainRecoveryKey(RecoveryRecordKind::Task, UuidToString(task.id)); + if (task.removed) { + records_.erase(key); + return true; + } + const bool created = !records_.contains(key); + if (created && task.drain) { + const auto job_key = DrainRecoveryKey(RecoveryRecordKind::Job, + UuidToString(task.drain->job_id)); + if (!records_.contains(job_key)) return false; + DrainJobRecord job; + if (!DecodeRecoveryValue(records_.at(job_key).payload, job) || + job.status >= static_cast(JobStatus::SUCCEEDED)) + return false; + auto unit = MakeRecoveryRecord( + RecoveryRecordKind::Unit, + UuidToString(task.drain->job_id) + ":" + task.drain->unit_key, + *task.drain); + if (auto previous = records_.find(unit.key); + previous != records_.end()) { + DrainUnitRecord old; + if (!DecodeRecoveryValue(previous->second.payload, old) || + old.task_id != UUID{} || old.completed || old.terminal_failed || + old.retries != task.drain->retries) + return false; + } + if (job.status < static_cast(JobStatus::RUNNING)) { + job.status = static_cast(JobStatus::RUNNING); + job.updated_ms = task.created_ms; + job.message = "Drain job running"; + records_.at(job_key) = MakeRecoveryRecord( + RecoveryRecordKind::Job, UuidToString(job.id), job); + } + records_.insert_or_assign(unit.key, std::move(unit)); + } + records_.insert_or_assign( + key, MakeRecoveryRecord(RecoveryRecordKind::Task, UuidToString(task.id), + task)); + return true; +} + +bool DrainRecoveryState::Apply(const ReplicationRecoveryRecord& replication) { + auto record = MakeRecoveryRecord( + RecoveryRecordKind::Replication, + TenantId(replication.tenant).MakeScopedKey(replication.key), + replication); + DrainRecoveryState prepared; + if (!prepared.Restore(record)) return false; + records_.insert_or_assign(record.key, std::move(record)); + return true; +} + +void DrainRecoveryState::EraseReplication(const std::string& tenant, + const std::string& key) { + records_.erase(DrainRecoveryKey(RecoveryRecordKind::Replication, + TenantId(tenant).MakeScopedKey(key))); +} + +void DrainRecoveryState::EraseSegment(const std::string& name, + const std::string& endpoint) { + records_.erase(DrainRecoveryKey(RecoveryRecordKind::Segment, name)); + const auto prefix = + DrainRecoveryKey(RecoveryRecordKind::RetiredReplica, ""); + for (auto it = records_.lower_bound(prefix); + it != records_.end() && it->first.starts_with(prefix);) { + RetiredReplicaRecord retired; + if (DecodeRecoveryValue(it->second.payload, retired) && + retired.replica.is_memory_replica() && + retired.replica.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_ == endpoint) { + it = records_.erase(it); + } else + ++it; + } +} + +bool DrainRecoveryState::Retire(const RetiredReplicaRecord& replica) { + auto record = + MakeRecoveryRecord(RecoveryRecordKind::RetiredReplica, + std::to_string(replica.replica.id), replica); + DrainRecoveryState prepared; + if (!prepared.Restore(record)) return false; + records_.insert_or_assign(record.key, std::move(record)); + return true; +} + +void DrainRecoveryState::PruneRetired(int64_t now_ms) { + const auto prefix = + DrainRecoveryKey(RecoveryRecordKind::RetiredReplica, ""); + for (auto it = records_.lower_bound(prefix); + it != records_.end() && it->first.starts_with(prefix);) { + RetiredReplicaRecord retired; + if (DecodeRecoveryValue(it->second.payload, retired) && + retired.release_ms <= now_ms) + it = records_.erase(it); + else + ++it; + } +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_applier.cpp b/mooncake-store/src/ha/oplog/oplog_applier.cpp index 1dd5aff657..b2a8edf526 100644 --- a/mooncake-store/src/ha/oplog/oplog_applier.cpp +++ b/mooncake-store/src/ha/oplog/oplog_applier.cpp @@ -1,5 +1,7 @@ #include "ha/oplog/oplog_applier.h" +#include +#include #include #include "ha_metric_manager.h" @@ -85,7 +87,7 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { bool applied = true; switch (entry.op_type) { case OpType::PUT_END: - ApplyPutEnd(entry); + applied = ApplyPutEnd(entry); break; case OpType::PUT_REVOKE: ApplyPutRevoke(entry); @@ -114,6 +116,42 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { case OpType::WEIGHT_LEASE_DELETE: applied = ApplyWeightLeaseDelete(entry); break; + case OpType::DRAIN_JOB_UPDATE: { + DrainJobUpdate update; + applied = DecodeRecoveryValue(entry.payload, update) && + metadata_store_->DrainRecovery().Apply(update); + break; + } + case OpType::TASK_UPDATE: { + DurableTaskRecord task; + applied = DecodeRecoveryValue(entry.payload, task) && + metadata_store_->DrainRecovery().Apply(task); + break; + } + case OpType::REPLICATION_START: { + ReplicationStartRecord start; + MetadataPayload payload; + applied = DecodeRecoveryValue(entry.payload, start) && + start.runtime.tenant == entry.tenant_id && + start.runtime.key == entry.object_key && + DecodeRecoveryValue(start.metadata_payload, payload); + if (!applied) break; + applied = ValidateReplicationRecovery( + start.runtime, payload.replicas, payload.size); + if (!applied) break; + DrainRecoveryState prepared; + applied = prepared.Apply(start.runtime) && + prepared.Apply(DrainJobUpdate{std::nullopt, std::nullopt, + start.segments}); + if (!applied) break; + auto metadata = payload.ToStandbyMetadata(); + applied = metadata_store_->PutMetadata( + entry.tenant_id, entry.object_key, metadata) && + metadata_store_->DrainRecovery().Apply(start.runtime) && + metadata_store_->DrainRecovery().Apply(DrainJobUpdate{ + std::nullopt, std::nullopt, start.segments}); + break; + } default: LOG(ERROR) << "OpLogApplier: unsupported op_type=" << static_cast(entry.op_type) @@ -125,6 +163,11 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { return false; } + metadata_store_->DrainRecovery().PruneRetired( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + // Update expected sequence ID expected_sequence_id_.store(entry.sequence_id + 1); @@ -157,7 +200,7 @@ void OpLogApplier::Recover(uint64_t last_applied_sequence_id) { << expected_sequence_id_.load(); } -void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { +bool OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { // Payload contains serialized metadata (replicas, size, etc.) in JSON // format. Deserialize the payload immediately and store structured // metadata. This allows Standby to serve requests immediately after @@ -175,7 +218,7 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { << entry.object_key << ", sequence_id=" << entry.sequence_id; } - return; + return true; } // Deserialize payload using struct_pack (msgpack binary format) @@ -190,7 +233,43 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { StandbyObjectMetadata empty_metadata; metadata_store_->PutMetadata(entry.tenant_id, entry.object_key, empty_metadata); - return; + return true; + } + + if (payload.replication_finished.value_or(false)) { + metadata_store_->DrainRecovery().EraseReplication(entry.tenant_id, + entry.object_key); + } else { + const auto& records = metadata_store_->DrainRecovery().records(); + const auto runtime_key = DrainRecoveryKey( + RecoveryRecordKind::Replication, + TenantId(entry.tenant_id).MakeScopedKey(entry.object_key)); + if (auto it = records.find(runtime_key); it != records.end()) { + ReplicationRecoveryRecord runtime; + if (!DecodeRecoveryValue(it->second.payload, runtime)) return false; + for (const auto& target : payload.processing_replicas.value_or( + std::vector{})) { + if (target.status != ReplicaStatus::PROCESSING) return false; + if (std::find(runtime.targets.begin(), runtime.targets.end(), + target.id) != runtime.targets.end() && + std::none_of(payload.replicas.begin(), + payload.replicas.end(), + [&](const auto& existing) { + return existing.id == target.id; + })) { + payload.replicas.push_back(target); + } + } + } + } + + for (const auto& retired : payload.retired_replicas.value_or( + std::vector{})) { + if (!metadata_store_->DrainRecovery().Retire(retired)) { + LOG(ERROR) << "Invalid retired MOVE source, key=" + << entry.object_key; + return false; + } } // Convert to StandbyObjectMetadata and store @@ -206,9 +285,12 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { << ", replicas=" << metadata.replicas.size() << ", size=" << metadata.size; } + return true; } void OpLogApplier::ApplyPutRevoke(const OpLogEntry& entry) { + metadata_store_->DrainRecovery().EraseReplication(entry.tenant_id, + entry.object_key); // PUT_REVOKE means the object should be removed from metadata store // (but the key itself may still exist if there are other replicas). // Current implementation removes the entire key; if we later support @@ -225,6 +307,8 @@ void OpLogApplier::ApplyPutRevoke(const OpLogEntry& entry) { } void OpLogApplier::ApplyRemove(const OpLogEntry& entry) { + metadata_store_->DrainRecovery().EraseReplication(entry.tenant_id, + entry.object_key); if (!metadata_store_->Remove(entry.tenant_id, entry.object_key)) { LOG(WARNING) << "OpLogApplier: failed to Remove key=" << entry.object_key @@ -443,6 +527,10 @@ void OpLogApplier::ApplySegmentUnmount(const OpLogEntry& entry) { << entry.object_key; return; } + if (auto segment = segment_registry_.GetSegment(op.transport_endpoint)) { + metadata_store_->DrainRecovery().EraseSegment( + segment->segment_name, segment->transport_endpoint); + } segment_registry_.OnSegmentUnmount(op.transport_endpoint); HAMetricManager::instance().inc_oplog_applied_entries(); } diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_pruning_coordinator.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_pruning_coordinator.cpp index 3a60f1fae1..820cd7c275 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_pruning_coordinator.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_pruning_coordinator.cpp @@ -74,6 +74,15 @@ bool ValidateSnapshot(SnapshotObjectStore& store, const std::string& root, weights.crc32c)) return false; } + for (size_t i = 0; i < manifest->recovery_chunks.size(); ++i) { + if (cancelled()) return false; + const auto& chunk = manifest->recovery_chunks[i]; + if (chunk.chunk_index != i || + chunk.key != ha::BuildBatchOpLogSnapshotRecoveryChunkKey( + root, descriptor.snapshot_id, i) || + !InspectArtifact(store, chunk.key, chunk.stored_size, chunk.crc32c)) + return false; + } return true; } diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_gc.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_gc.cpp index 546e7fd88e..a9f4dd6132 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_gc.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_gc.cpp @@ -109,6 +109,16 @@ ErrorCode BatchOpLogSnapshotGc::Run( chunk.crc32c)) return ErrorCode::INTERNAL_ERROR; } + for (size_t i = 0; i < manifest->recovery_chunks.size(); ++i) { + const auto& chunk = manifest->recovery_chunks[i]; + if (chunk.chunk_index != i || + chunk.key != ha::BuildBatchOpLogSnapshotRecoveryChunkKey( + root_, descriptor->snapshot_id, i) || + !VerifyObject(store_, chunk.key, chunk.stored_size, + chunk.crc32c)) { + return ErrorCode::INTERNAL_ERROR; + } + } if (manifest->weight_metadata) { const auto& weights = *manifest->weight_metadata; if (weights.key != ha::BuildBatchOpLogSnapshotWeightMetadataKey( diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_provider.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_provider.cpp index 2533fdef3b..ef06fb43f7 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_provider.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_provider.cpp @@ -323,6 +323,28 @@ AttemptResult RestorePointer( } } + for (size_t i = 0; i < manifest.recovery_chunks.size(); ++i) { + if (cancelled && cancelled()) + return Invalid(ErrorCode::ETCD_CTX_CANCELLED); + const auto& chunk = manifest.recovery_chunks[i]; + if (chunk.chunk_index != i || + chunk.key != ha::BuildBatchOpLogSnapshotRecoveryChunkKey( + snapshot_root, descriptor.snapshot_id, i)) + return Invalid(); + std::vector bytes; + auto read = ReadVerifiedObject(object_store, chunk.key, + chunk.stored_size, chunk.crc32c, bytes); + if (read.disposition != AttemptDisposition::kSuccess) return read; + auto decoded = + DecodeBatchOpLogRecoveryChunk(bytes, i, chunk.object_count); + if (!decoded) return Invalid(ErrorCode::DESERIALIZE_FAIL); + for (const auto& record : decoded->records) { + if (!metadata.DrainRecovery().Restore(record)) { + return Invalid(ErrorCode::DESERIALIZE_FAIL); + } + } + } + auto suffix = ReplaySuffix(cluster_id, backend, *applier, &baseline, cancelled); if (suffix.disposition != AttemptDisposition::kSuccess) { @@ -336,8 +358,11 @@ AttemptResult RestorePointer( uint64_t chunk_bytes = 0; for (const auto& chunk : manifest.object_chunks) chunk_bytes += chunk.stored_size; + for (const auto& chunk : manifest.recovery_chunks) + chunk_bytes += chunk.stored_size; HAMetricManager::instance().update_snapshot_runtime([&](auto& metrics) { - metrics.chunk_count = manifest.object_chunks.size(); + metrics.chunk_count = + manifest.object_chunks.size() + manifest.recovery_chunks.size(); metrics.chunk_bytes = chunk_bytes; metrics.snapshot_bytes = metrics.chunk_bytes + manifest.segments.stored_size + diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/codec.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/codec.cpp index 766dd3369b..850b5a552c 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/codec.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/codec.cpp @@ -73,4 +73,22 @@ DecodeBatchOpLogSnapshotObjectChunk(const std::vector& encoded, return decoded; } +std::vector EncodeBatchOpLogRecoveryChunk( + uint64_t index, std::vector records) { + return Encode(BatchOpLogRecoveryChunk{index, std::move(records)}); +} + +tl::expected +DecodeBatchOpLogRecoveryChunk(const std::vector& encoded, + uint64_t index, uint64_t count) { + auto decoded = + Decode(encoded, "snapshot recovery chunk"); + if (!decoded) return decoded; + if (decoded->chunk_index != index || decoded->records.size() != count) { + return tl::make_unexpected( + std::string("Snapshot recovery chunk index/count mismatch")); + } + return decoded; +} + } // namespace mooncake diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/metadata.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/metadata.cpp index 605f824c33..26fcf9bb73 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/metadata.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/metadata.cpp @@ -128,7 +128,10 @@ bool ValidateIdentity(const SnapshotMetadata& metadata, std::string* reason) { const bool with_weights = metadata.schema_version == kBatchOpLogWeightSnapshotSchemaVersion && metadata.snapshot_format == kBatchOpLogWeightSnapshotFormat; - if (!legacy && !with_weights) { + const bool with_recovery = + metadata.schema_version == kBatchOpLogRecoverySnapshotSchemaVersion && + metadata.snapshot_format == kBatchOpLogRecoverySnapshotFormat; + if (!legacy && !with_weights && !with_recovery) { SetReason(reason, "unsupported snapshot version and format"); return false; } @@ -323,6 +326,13 @@ std::string EncodeBatchOpLogSnapshotManifest( chunks.append(EncodeChunkDescriptor(descriptor)); } root["object_chunks"] = std::move(chunks); + if (manifest.schema_version == kBatchOpLogRecoverySnapshotSchemaVersion) { + Json::Value recovery(Json::arrayValue); + for (const auto& chunk : manifest.recovery_chunks) { + recovery.append(EncodeChunkDescriptor(chunk)); + } + root["recovery_chunks"] = std::move(recovery); + } return WriteJson(root); } @@ -342,7 +352,9 @@ DecodeBatchOpLogSnapshotManifest(std::string_view value) { if (!DecodeObjectDescriptor(root["segments"], decoded.segments, &reason)) { return tl::make_unexpected(std::move(reason)); } - if (decoded.schema_version == kBatchOpLogWeightSnapshotSchemaVersion) { + if (decoded.schema_version == kBatchOpLogWeightSnapshotSchemaVersion || + (decoded.schema_version == kBatchOpLogRecoverySnapshotSchemaVersion && + root.isMember("weight_metadata"))) { BatchOpLogSnapshotObjectDescriptor weights; if (!root.isMember("weight_metadata") || !DecodeObjectDescriptor(root["weight_metadata"], weights, @@ -368,9 +380,36 @@ DecodeBatchOpLogSnapshotManifest(std::string_view value) { } decoded.object_chunks.push_back(std::move(descriptor)); } + if (decoded.schema_version == kBatchOpLogRecoverySnapshotSchemaVersion) { + if (!root.isMember("recovery_chunks") || + !root["recovery_chunks"].isArray()) { + return tl::make_unexpected( + std::string("recovery_chunks must be an array")); + } + const auto& recovery = root["recovery_chunks"]; + for (Json::ArrayIndex i = 0; i < recovery.size(); ++i) { + BatchOpLogSnapshotChunkDescriptor chunk; + if (!DecodeChunkDescriptor(recovery[i], i, chunk, &reason)) { + return tl::make_unexpected( + std::string("Invalid recovery chunk: ") + reason); + } + decoded.recovery_chunks.push_back(std::move(chunk)); + } + } else if (root.isMember("recovery_chunks")) { + return tl::make_unexpected( + std::string("Recovery chunks require snapshot v3")); + } return decoded; } +std::string BuildBatchOpLogSnapshotRecoveryChunkKey( + const std::string& root, std::string_view snapshot_id, uint64_t index) { + const auto prefix = BuildArtifactPrefix(root, snapshot_id); + return prefix.empty() + ? std::string() + : prefix + "recovery/" + std::to_string(index) + ".bin"; +} + std::string BuildBatchOpLogSnapshotId(uint64_t last_included_batch_id, int64_t maintenance_lease_id) { if (maintenance_lease_id <= 0) { diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/writer.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/writer.cpp index 7f1d9e7ce1..7ac5bbaae5 100644 --- a/mooncake-store/src/ha/snapshot/batch_oplog/writer.cpp +++ b/mooncake-store/src/ha/snapshot/batch_oplog/writer.cpp @@ -109,6 +109,10 @@ tl::expected BatchOpLogSnapshotWriter::Write( manifest.schema_version = ha::kBatchOpLogWeightSnapshotSchemaVersion; manifest.snapshot_format = ha::kBatchOpLogWeightSnapshotFormat; } + if (!capture.recovery_done()) { + manifest.schema_version = ha::kBatchOpLogRecoverySnapshotSchemaVersion; + manifest.snapshot_format = ha::kBatchOpLogRecoverySnapshotFormat; + } manifest.snapshot_id = snapshot_id; manifest.last_included_seq = capture.last_included_seq; manifest.last_included_batch_id = capture.last_included_batch_id; @@ -179,6 +183,28 @@ tl::expected BatchOpLogSnapshotWriter::Write( manifest.object_chunks.push_back(std::move(descriptor)); ++chunk_index; } + chunk_index = 0; + while (!capture.recovery_done()) { + std::vector records; + if (!standby.CopyNextBatchOpLogRecoveryChunk(chunk_object_count, + capture, records) || + records.empty()) { + return fail("Recovery snapshot capture was cancelled or empty"); + } + const auto count = records.size(); + auto encoded = + EncodeBatchOpLogRecoveryChunk(chunk_index, std::move(records)); + const auto key = ha::BuildBatchOpLogSnapshotRecoveryChunkKey( + snapshot_root, snapshot_id, chunk_index); + ha::BatchOpLogSnapshotChunkDescriptor chunk{ + chunk_index, key, count, encoded.size(), + Crc32cValue(encoded.data(), encoded.size())}; + upload = object_store_.UploadBuffer(key, encoded); + if (!upload) + return fail("Failed to upload recovery chunk: " + upload.error()); + manifest.recovery_chunks.push_back(std::move(chunk)); + ++chunk_index; + } end_capture(); auto verify = @@ -203,6 +229,12 @@ tl::expected BatchOpLogSnapshotWriter::Write( } } + for (const auto& chunk : manifest.recovery_chunks) { + verify = VerifyObject(object_store_, chunk.key, chunk.stored_size, + chunk.crc32c); + if (!verify) return fail(std::move(verify.error())); + } + const std::string manifest_key = ha::BuildBatchOpLogSnapshotManifestKey(snapshot_root, snapshot_id); const std::string manifest_json = @@ -248,8 +280,11 @@ tl::expected BatchOpLogSnapshotWriter::Write( uint64_t chunk_bytes = 0; for (const auto& chunk : manifest.object_chunks) chunk_bytes += chunk.stored_size; + for (const auto& chunk : manifest.recovery_chunks) + chunk_bytes += chunk.stored_size; HAMetricManager::instance().update_snapshot_runtime([&](auto& metrics) { - metrics.chunk_count = manifest.object_chunks.size(); + metrics.chunk_count = + manifest.object_chunks.size() + manifest.recovery_chunks.size(); metrics.chunk_bytes = chunk_bytes; metrics.snapshot_bytes = metrics.chunk_bytes + manifest.segments.stored_size + diff --git a/mooncake-store/src/ha/standby_controller.cpp b/mooncake-store/src/ha/standby_controller.cpp index 55153f2a25..641854ad8d 100644 --- a/mooncake-store/src/ha/standby_controller.cpp +++ b/mooncake-store/src/ha/standby_controller.cpp @@ -310,7 +310,7 @@ class CapabilityDrivenStandbyController final : public StandbyController { return tl::unexpected(promote_error); } - if (standby_service_->IsBatchOpLogSnapshotMode()) { + if (capabilities_.has_oplog_following) { auto handoff = standby_service_->PromoteAndDetachBatchOpLogStore(); if (!handoff) { standby_service_->MarkAppliedCursorForRestart(); diff --git a/mooncake-store/src/ha/standby_metadata_store.cpp b/mooncake-store/src/ha/standby_metadata_store.cpp index 91af56c690..d6e23d7178 100644 --- a/mooncake-store/src/ha/standby_metadata_store.cpp +++ b/mooncake-store/src/ha/standby_metadata_store.cpp @@ -80,6 +80,7 @@ size_t StandbyMetadataStore::GetKeyCount() const { void StandbyMetadataStore::Clear() { std::lock_guard lock(mutex_); store_.clear(); + DrainRecovery().Clear(); weight_metadata_.clear(); weight_metadata_tombstones_.clear(); weight_leases_.clear(); @@ -266,6 +267,15 @@ bool StandbyMetadataStore::ValidateReplicaIds(ReplicaID& max_replica_id) const { } } } + for (const auto& [key, value] : DrainRecovery().records()) { + if (value.kind != RecoveryRecordKind::RetiredReplica) continue; + RetiredReplicaRecord retired; + if (!DecodeRecoveryValue(value.payload, retired) || + retired.replica.id == 0 || + retired.replica.id == std::numeric_limits::max()) + return false; + max_replica_id = std::max(max_replica_id, retired.replica.id); + } return true; } diff --git a/mooncake-store/src/hot_standby_service.cpp b/mooncake-store/src/hot_standby_service.cpp index 2f028fe5f3..cc0976a287 100644 --- a/mooncake-store/src/hot_standby_service.cpp +++ b/mooncake-store/src/hot_standby_service.cpp @@ -172,7 +172,9 @@ ErrorCode HotStandbyService::PrepareBootstrapBaselineLocked( uint64_t& baseline_seq_id) { const uint64_t local_last_seq_id = GetLocalLastAppliedSequenceIdLocked(); const bool has_local_metadata = - metadata_store_ && metadata_store_->GetKeyCount() > 0; + metadata_store_ && + (metadata_store_->GetKeyCount() > 0 || + !metadata_store_->DrainRecovery().records().empty()); const bool has_recoverable_local_state = has_local_metadata && local_last_seq_id > 0; @@ -185,29 +187,29 @@ ErrorCode HotStandbyService::PrepareBootstrapBaselineLocked( oplog_applier_ = std::make_unique(metadata_store_.get(), cluster_id_); - if (batch_oplog_snapshot_provider_ && config_.enable_snapshot_bootstrap) { - if (preserve_promotion_baseline_ && metadata_store_ && - batch_snapshot_baseline_) { - const uint64_t preserved_seq = batch_snapshot_baseline_->last_seq; - oplog_applier_->Recover(preserved_seq); - if (!preserved_segments_.empty()) { - oplog_applier_->LoadSegmentRegistry(preserved_segments_); - } - preserved_segments_.clear(); - preserve_promotion_baseline_ = false; - applied_seq_id_.store(preserved_seq, std::memory_order_release); - primary_seq_id_.store(preserved_seq, std::memory_order_release); - baseline_seq_id = preserved_seq; - { - std::lock_guard cursor_lock( - batch_snapshot_cursor_mutex_); - last_applied_batch_snapshot_prefix_ = *batch_snapshot_baseline_; - } - LOG(INFO) << "Standby restart resumes preserved promotion cursor, " - << "applied_seq=" << preserved_seq - << ", batch_id=" << batch_snapshot_baseline_->batch_id; - return ErrorCode::OK; + if (preserve_promotion_baseline_ && metadata_store_ && + batch_snapshot_baseline_) { + const uint64_t preserved_seq = batch_snapshot_baseline_->last_seq; + oplog_applier_->Recover(preserved_seq); + if (!preserved_segments_.empty()) { + oplog_applier_->LoadSegmentRegistry(preserved_segments_); } + preserved_segments_.clear(); + preserve_promotion_baseline_ = false; + applied_seq_id_.store(preserved_seq, std::memory_order_release); + primary_seq_id_.store(preserved_seq, std::memory_order_release); + baseline_seq_id = preserved_seq; + { + std::lock_guard cursor_lock( + batch_snapshot_cursor_mutex_); + last_applied_batch_snapshot_prefix_ = *batch_snapshot_baseline_; + } + LOG(INFO) << "Standby restart resumes preserved promotion cursor, " + << "applied_seq=" << preserved_seq + << ", batch_id=" << batch_snapshot_baseline_->batch_id; + return ErrorCode::OK; + } + if (batch_oplog_snapshot_provider_ && config_.enable_snapshot_bootstrap) { preserve_promotion_baseline_ = false; return LoadBatchOpLogSnapshotBaselineLocked(baseline_seq_id); } @@ -805,8 +807,7 @@ bool HotStandbyService::IsBatchOpLogSnapshotMode() const { tl::expected HotStandbyService::PromoteAndDetachBatchOpLogStore() { std::unique_lock lock(mutex_); - if (!batch_oplog_snapshot_provider_ || !config_.enable_oplog_following || - !batch_standby_reader_) { + if (!config_.enable_oplog_following || !batch_standby_reader_) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } if (!IsReadyForPromotion()) { @@ -972,6 +973,24 @@ bool HotStandbyService::CopyNextBatchOpLogSnapshotChunk( return metadata_store_->CopyNextSnapshotChunk(count, capture.cursor_, out); } +bool HotStandbyService::CopyNextBatchOpLogRecoveryChunk( + size_t count, BatchOpLogSnapshotCapture& capture, + std::vector& out) { + auto state = snapshot_capture_state_; + std::lock_guard lock(state->mutex); + out.clear(); + if (count == 0 || capture.lease_state_ != state || !state->active || + capture.generation_ != state->generation || !metadata_store_) + return false; + while (out.size() < count && !capture.recovery_done()) { + out.push_back(capture.recovery_cursor_->second); + ++capture.recovery_cursor_; + capture.recovery_done_ = + capture.recovery_cursor_ == capture.recovery_end_; + } + return true; +} + void HotStandbyService::EndBatchOpLogSnapshotCapture( BatchOpLogSnapshotCapture& capture) { if (capture.lease_state_ == snapshot_capture_state_) { @@ -1017,6 +1036,10 @@ void HotStandbyService::HandleSnapshotCaptureRequest( metadata_store_->BeginSnapshotTraversal(), state->generation, state); capture.weight_metadata = metadata_store_->SnapshotWeightMetadata(); + const auto& recovery = metadata_store_->DrainRecovery().records(); + capture.recovery_cursor_ = recovery.begin(); + capture.recovery_end_ = recovery.end(); + capture.recovery_done_ = recovery.empty(); ready_snapshot_capture_ = std::move(capture); state->active = true; } diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index db710d76e7..58a279b6a6 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -485,6 +485,9 @@ MasterMetricManager::MasterMetricManager() // MoveRevoke Counters copy_start_requests_("master_copy_start_requests_total", "Total number of CopyStart requests received"), + orphaned_draining_restore_( + "master_orphaned_draining_restore_total", + "Orphaned draining segments found during restore"), copy_start_failures_("master_copy_start_failures_total", "Total number of failed CopyStart requests"), copy_end_requests_("master_copy_end_requests_total", @@ -640,6 +643,7 @@ void MasterMetricManager::update_metrics_for_zero_output() { // Update CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke // counters copy_start_requests_.inc(0); + orphaned_draining_restore_.inc(0); copy_start_failures_.inc(0); copy_end_requests_.inc(0); copy_end_failures_.inc(0); @@ -1861,6 +1865,10 @@ int64_t MasterMetricManager::get_promotion_candidate_dropped_limit() { void MasterMetricManager::inc_copy_start_requests(int64_t val) { copy_start_requests_.inc(val); } + +void MasterMetricManager::inc_orphaned_draining_restore() { + orphaned_draining_restore_.inc(); +} void MasterMetricManager::inc_copy_start_failures(int64_t val) { copy_start_failures_.inc(val); } @@ -2106,6 +2114,7 @@ std::string MasterMetricManager::serialize_metrics() { // Serialize CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke // Counters serialize_metric(copy_start_requests_); + serialize_metric(orphaned_draining_restore_); serialize_metric(copy_start_failures_); serialize_metric(copy_end_requests_); serialize_metric(copy_end_failures_); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 488c018e3f..3337eeb2f9 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -216,7 +216,13 @@ MasterService::MasterService(const MasterServiceConfig& config) put_start_release_timeout_sec_(config.put_start_release_timeout_sec), offloading_queue_limit_(config.offloading_queue_limit), offload_cap_ratio_(config.offload_cap_ratio), - task_manager_(config.task_manager_config), + task_manager_(config.task_manager_config, + config.enable_ha && config.enable_oplog + ? ClientTaskManager::Persistence( + [this](std::span tasks) { + return PersistRecoveryTasks(tasks); + }) + : ClientTaskManager::Persistence{}), batch_oplog_writer_factory_( [](OrderedOpLogWriterConfig writer_config, OrderedOpLogWriter::WriteBatchFn write_batch) { @@ -675,6 +681,14 @@ MasterService::DeleteTenantQuotaPolicy(const TenantId& tenant_id) { auto MasterService::MountSegment(const Segment& segment, const UUID& client_id) -> tl::expected { + { + std::shared_lock lock(master_state_mutex_); + if (invalid_replica_endpoints_.contains(segment.name) || + invalid_replica_endpoints_.contains(segment.te_endpoint)) { + lock.unlock(); + return ReMountSegment({segment}, client_id); + } + } ErrorCode mount_result = ErrorCode::INTERNAL_ERROR; { std::unique_lock client_lock(client_mutex_); @@ -768,7 +782,17 @@ auto MasterService::MountNoFSegment(const NoFSegment& segment, } ErrorCode MasterService::ValidateStandbyRemountSegment( - const Segment& segment) const { + const Segment& segment, const UUID& client_id) const { + if (auto recovered = restored_drain_segments_.find(segment.name); + recovered != restored_drain_segments_.end()) { + const auto& expected = recovered->second; + if (!expected.lifetime.isAvailable() || expected.id != segment.id || + expected.client_id != client_id || expected.base != segment.base || + expected.size != segment.size || + expected.endpoint != segment.te_endpoint) { + return ErrorCode::INVALID_PARAMS; + } + } const StandbySegmentInfo* match = nullptr; for (const auto& standby : standby_memory_segments_) { if (standby.transport_endpoint == segment.te_endpoint || @@ -827,7 +851,7 @@ auto MasterService::ReMountSegment(const std::vector& segments, auto segment_access = segment_manager_.getSegmentAccess(); for (const auto& segment : segments) { auto standby_validation = - ValidateStandbyRemountSegment(segment); + ValidateStandbyRemountSegment(segment, client_id); if (standby_validation != ErrorCode::OK) { discard_provisional_record(); return tl::make_unexpected(standby_validation); @@ -840,7 +864,12 @@ auto MasterService::ReMountSegment(const std::vector& segments, } } } - if (ok_client_.contains(client_id)) { + const bool pending_remount = std::any_of( + segments.begin(), segments.end(), [&](const Segment& segment) { + return invalid_replica_endpoints_.contains(segment.name) || + invalid_replica_endpoints_.contains(segment.te_endpoint); + }); + if (ok_client_.contains(client_id) && !pending_remount) { LOG(WARNING) << "client_id=" << client_id << ", warn=client_already_remounted"; // Return OK because this is an idempotent operation @@ -855,6 +884,7 @@ auto MasterService::ReMountSegment(const std::vector& segments, std::vector descriptors; std::vector> buffers; uint64_t imported_size{0}; + uint64_t retired_size{0}; }; std::vector restores; restores.reserve(segments.size()); @@ -943,7 +973,8 @@ auto MasterService::ReMountSegment(const std::vector& segments, [](const Replica& replica) { return replica.is_memory_replica() && replica.status() != ReplicaStatus::REMOVED && - replica.status() != ReplicaStatus::FAILED; + replica.status() != ReplicaStatus::FAILED && + !replica.has_invalid_mem_handle(); }, [&](Replica& replica) { auto descriptor = replica.get_descriptor() @@ -951,6 +982,11 @@ auto MasterService::ReMountSegment(const std::vector& segments, .buffer_descriptor; SegmentRestore* match = nullptr; for (auto& restore : restores) { + if (!invalid_replica_endpoints_.contains( + restore.segment.name) && + !invalid_replica_endpoints_.contains( + restore.segment.te_endpoint)) + continue; if (descriptor.transport_endpoint_ == restore.segment.te_endpoint || descriptor.transport_endpoint_ == @@ -977,6 +1013,33 @@ auto MasterService::ReMountSegment(const std::vector& segments, } } } + { + std::lock_guard lock(discarded_replicas_mutex_); + for (auto& item : discarded_replicas_) { + for (auto& replica : item.replicas_) { + if (!replica.is_memory_replica() || + replica.has_invalid_mem_handle() || + !IsAwaitingRemount(replica)) + continue; + auto descriptor = replica.get_descriptor() + .get_memory_descriptor() + .buffer_descriptor; + for (auto& restore : restores) { + if (descriptor.transport_endpoint_ == + restore.segment.te_endpoint || + descriptor.transport_endpoint_ == + restore.segment.name) { + descriptor.transport_endpoint_ = + restore.segment.te_endpoint; + restore.replicas.push_back(&replica); + restore.descriptors.push_back(descriptor); + restore.retired_size += descriptor.size_; + break; + } + } + } + } + } if (ambiguous_endpoint) { return fail_remount(ErrorCode::INVALID_PARAMS); } @@ -1037,10 +1100,12 @@ auto MasterService::ReMountSegment(const std::vector& segments, uint64_t{0}, [](uint64_t sum, const auto& descriptor) { return sum + descriptor.size_; }); + restore.imported_size -= restore.retired_size; auto accounted = standby_accounted_memory_bytes_.find(restore.segment.name); - if (accounted == standby_accounted_memory_bytes_.end() || - accounted->second < restore.imported_size) { + if (restore.imported_size != 0 && + (accounted == standby_accounted_memory_bytes_.end() || + accounted->second < restore.imported_size)) { return fail_remount(ErrorCode::INTERNAL_ERROR); } allocator_replacements.push_back({restore.segment.id, @@ -1066,19 +1131,24 @@ auto MasterService::ReMountSegment(const std::vector& segments, segment_access.BindBufferToSegment(restore.segment.id, *buffer); } + if (auto recovered = + restored_drain_segments_.find(restore.segment.name); + recovered != restored_drain_segments_.end()) { + (void)segment_access.SetSegmentStatusByName( + restore.segment.name, + static_cast(recovered->second.status)); + restored_drain_segments_.erase(recovered); + } } } for (auto& restore : restores) { - if (restore.imported_size != 0) { + if (auto accounted = + standby_accounted_memory_bytes_.find(restore.segment.name); + accounted != standby_accounted_memory_bytes_.end()) { MasterMetricManager::instance().dec_allocated_mem_size( restore.segment.name, - static_cast(restore.imported_size)); - auto accounted = - standby_accounted_memory_bytes_.find(restore.segment.name); - accounted->second -= restore.imported_size; - if (accounted->second == 0) { - standby_accounted_memory_bytes_.erase(accounted); - } + static_cast(accounted->second)); + standby_accounted_memory_bytes_.erase(accounted); } for (size_t i = 0; i < restore.replicas.size(); ++i) { (void)restore.replicas[i]->replace_memory_buffer( @@ -1095,8 +1165,9 @@ auto MasterService::ReMountSegment(const std::vector& segments, } // Change the client status to OK - ok_client_.insert(client_id); - MasterMetricManager::instance().inc_active_clients(); + if (ok_client_.insert(client_id).second) { + MasterMetricManager::instance().inc_active_clients(); + } } if (enable_oplog_ && ordered_oplog_writer_) { @@ -1352,6 +1423,17 @@ void MasterService::RebuildTenantQuotaUsageFromMetadata() { for (const auto& [_, record] : tenant_state.metadata) { const auto& metadata = record.metadata; auto& charged_bytes = usage[tenant_id]; + const uint64_t pending = record.state.replication_task + ? record.state.replication_task + ->pending_quota_charge_bytes + : 0; + if (pending > TenantQuotaAccount::kMaxChargedBytes || + charged_bytes > + TenantQuotaAccount::kMaxChargedBytes - pending) { + throw std::overflow_error( + "rebuilt pending quota exceeds 2^63 - 1 bytes"); + } + charged_bytes += pending; const uint64_t charge = CompletedMemoryQuotaCharge(metadata); if (charge > TenantQuotaAccount::kMaxChargedBytes || charged_bytes > @@ -1660,7 +1742,8 @@ tl::expected MasterService::SettlePrimaryWriteQuotaIfReady( void MasterService::FinalizeRemovedReplicasAfterDurable( const OpLogEntry& durable_entry, const std::vector& replica_ids, QuotaEraseMode quota_mode, - const std::vector& previous_media_hint) { + const std::vector& previous_media_hint, + std::optional retire_at) { if (replica_ids.empty()) { return; } @@ -1748,6 +1831,11 @@ void MasterService::FinalizeRemovedReplicasAfterDurable( accessor.EraseFromProcessing(); } } + if (retire_at && *retire_at > std::chrono::system_clock::now()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back(std::move(erased_replicas), + *retire_at); + } } void MasterService::FinalizeMetadataEraseAfterDurable( @@ -1922,11 +2010,6 @@ tl::expected MasterService::PersistStaleHandleCleanupForHA( const auto op_type = plan.would_invalidate ? OpType::REMOVE : OpType::PUT_END; - const std::string payload = - plan.would_invalidate ? std::string{} - : SerializeMetadataForOpLogFromReplicaDescriptors( - metadata, plan.remaining); - auto reservation = ReserveBatchOpLogSlot(); if (!reservation) { return tl::make_unexpected(reservation.error()); @@ -1936,6 +2019,11 @@ tl::expected MasterService::PersistStaleHandleCleanupForHA( metadata.VisitReplicas( [&ids](const Replica& replica) { return ids.contains(replica.id()); }, [](Replica& replica) { replica.mark_removed(); }); + const std::string payload = + plan.would_invalidate ? std::string{} + : SerializeMetadataForOpLogFromReplicaDescriptors( + metadata, plan.remaining); + auto result = AppendReservedOpLogWithDurableFinalize( std::move(reservation.value()), op_type, tenant_id.value(), key, payload, @@ -2632,6 +2720,7 @@ tl::expected MasterService::ClearStaleHandles( bool MasterService::ProcessClientOffboardingJob(ClientOffboardingJob& job) { bool quota_recompute_needed = false; + std::vector forgotten; { std::shared_lock snapshot_lock(master_state_mutex_); @@ -2680,6 +2769,25 @@ bool MasterService::ProcessClientOffboardingJob(ClientOffboardingJob& job) { } } + { + auto access = segment_manager_.getSegmentAccess(); + for (auto& [name, recovered] : restored_drain_segments_) { + if (recovered.client_id != job.client_id) continue; + const auto queued = std::any_of( + job.prepared_segments.begin(), job.prepared_segments.end(), + [&](const auto& segment) { + return segment.segment_id == recovered.id; + }); + if (!queued) { + recovered.lifetime.setAvailable(false); + recovered.status = + static_cast(SegmentStatus::UNMOUNTING); + job.prepared_segments.push_back( + {recovered.id, name, recovered.endpoint, 0, true}); + } + } + } + if (!job.metadata_cleanup_accepted) { const auto cleanup_result = ClearStaleHandles([&job](const Replica& replica) { @@ -2745,6 +2853,23 @@ bool MasterService::ProcessClientOffboardingJob(ClientOffboardingJob& job) { it != job.prepared_segments.end();) { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); + auto recovered = restored_drain_segments_.find(it->segment_name); + if (it->recovered && (recovered == restored_drain_segments_.end() || + recovered->second.id != it->segment_id)) { + it = job.prepared_segments.erase(it); + continue; + } + std::shared_ptr recovery_mutex; + std::unique_lock recovery_claim; + if (recovered != restored_drain_segments_.end()) { + recovery_mutex = recovered->second.unmount_mutex; + recovery_claim = + std::unique_lock(*recovery_mutex, std::try_to_lock); + if (!recovery_claim.owns_lock()) { + ++it; + continue; + } + } const auto commit_result = segment_access.CommitUnmountSegment( it->segment_id, job.client_id, it->metrics_dec_capacity, /*retain_name_registration=*/true); @@ -2777,6 +2902,13 @@ bool MasterService::ProcessClientOffboardingJob(ClientOffboardingJob& job) { } } + if (auto recovered = + restored_drain_segments_.find(it->segment_name); + recovered != restored_drain_segments_.end()) { + recovered->second.lifetime.setAvailable(false); + forgotten.push_back(recovered->second); + restored_drain_segments_.erase(recovered); + } cleanupHttpMetadata(it->segment_name); segment_access.ReleaseUnmountedSegmentName(it->segment_id, it->segment_name); @@ -2788,6 +2920,7 @@ bool MasterService::ProcessClientOffboardingJob(ClientOffboardingJob& job) { } } + ForgetRestoredSegments(forgotten); if (quota_recompute_needed) { RecomputeTenantEffectiveQuotas(); } @@ -2837,62 +2970,102 @@ void MasterService::TaskCleanupThreadFunc() { auto MasterService::UnmountSegment(const UUID& segment_id, const UUID& client_id) -> tl::expected { - size_t metrics_dec_capacity = 0; // to update the metrics - + size_t metrics_dec_capacity = 0; std::shared_lock client_lock(client_mutex_); std::shared_lock shared_lock(master_state_mutex_); auto retaining_clients = GetRetainingClientIdsLocked(); client_lock.unlock(); - // 1. Prepare to unmount the segment by deleting its allocator + std::string segment_name, te_endpoint; + std::optional recovered; + std::shared_ptr recovery_mutex; + std::unique_lock recovery_claim; + bool mounted = true; { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - ErrorCode err = segment_access.PrepareUnmountSegment( - segment_id, metrics_dec_capacity); - if (err == ErrorCode::SEGMENT_NOT_FOUND) { - // Return OK because this is an idempotent operation - return {}; - } - if (err != ErrorCode::OK) { - return tl::make_unexpected(err); - } + auto access = segment_manager_.getSegmentAccess(); + auto previous = std::find_if( + restored_drain_segments_.begin(), restored_drain_segments_.end(), + [&](const auto& entry) { return entry.second.id == segment_id; }); + if (previous != restored_drain_segments_.end()) { + if (previous->second.client_id != client_id) { + return tl::make_unexpected(ErrorCode::ILLEGAL_CLIENT); + } + recovery_mutex = previous->second.unmount_mutex; + recovery_claim = + std::unique_lock(*recovery_mutex, std::try_to_lock); + if (!recovery_claim.owns_lock()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + recovered = previous->second; + previous->second.lifetime.setAvailable(false); + previous->second.status = + static_cast(SegmentStatus::UNMOUNTING); + segment_name = recovered->name; + te_endpoint = recovered->endpoint; + } + Segment registered; + if (access.GetSegment(segment_id, registered)) { + segment_name = registered.name; + te_endpoint = registered.te_endpoint; + } + auto error = + access.PrepareUnmountSegment(segment_id, metrics_dec_capacity); + if (error == ErrorCode::SEGMENT_NOT_FOUND) { + if (!recovered) return {}; + mounted = false; + } else if (error != ErrorCode::OK) + return tl::make_unexpected(error); } - - // Keep HA, snapshot, and CXL behavior unchanged. Regular memory segments - // become unreadable as soon as PrepareUnmountSegment releases their - // allocator; only the physical metadata sweep is deferred. - if (enable_async_segment_cleanup_) { + if (enable_async_segment_cleanup_) replica_cleanup_worker_.Schedule(); - } else { + else ClearInvalidHandles(retaining_clients); - } - - // Cache endpoint before commit removes segment from registry. - std::string segment_name; - std::string te_endpoint; - if (!segment_manager_.GetSegmentBasicInfo(segment_id, segment_name, - te_endpoint)) { - return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); - } - - // 3. Commit the unmount operation + if (mounted) { + auto access = segment_manager_.getSegmentAccess(); + auto error = access.CommitUnmountSegment(segment_id, client_id, + metrics_dec_capacity); + if (error != ErrorCode::OK) return tl::make_unexpected(error); + } + if (enable_oplog_ && !te_endpoint.empty()) { + const auto payload = EncodeRecoveryValue(SegmentUnmountOp{te_endpoint}); + if (recovered) { + auto error = PersistRecoveryEntry(OpType::SEGMENT_UNMOUNT, payload, + te_endpoint); + if (error != ErrorCode::OK) { + auto access = segment_manager_.getSegmentAccess(); + if (auto previous = restored_drain_segments_.find(segment_name); + previous != restored_drain_segments_.end()) { + previous->second.status = recovered->status; + previous->second.lifetime.setAvailable(true); + } + return tl::make_unexpected(error); + } + } else + PersistSegmentOpForHAOrEnqueue("UnmountSegment", + OpType::SEGMENT_UNMOUNT, te_endpoint, + payload); + } + std::vector forgotten; { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - auto err = segment_access.CommitUnmountSegment(segment_id, client_id, - metrics_dec_capacity); - if (err != ErrorCode::OK) { - return tl::make_unexpected(err); - } - } - - if (enable_oplog_ && ordered_oplog_writer_ && !te_endpoint.empty()) { - SegmentUnmountOp op{te_endpoint}; - auto bytes = struct_pack::serialize(op); - PersistSegmentOpForHAOrEnqueue("UnmountSegment", - OpType::SEGMENT_UNMOUNT, te_endpoint, - std::string(bytes.begin(), bytes.end())); - } + auto access = segment_manager_.getSegmentAccess(); + if (auto previous = restored_drain_segments_.find(segment_name); + previous != restored_drain_segments_.end()) { + forgotten.push_back(previous->second); + restored_drain_segments_.erase(previous); + } else if (std::any_of(standby_memory_segments_.begin(), + standby_memory_segments_.end(), + [&](const auto& segment) { + return segment.segment_name == segment_name; + })) { + DrainSegmentRecord previous; + previous.name = segment_name; + previous.endpoint = te_endpoint; + forgotten.push_back(std::move(previous)); + } + } + if (recovery_claim.owns_lock()) recovery_claim.unlock(); + shared_lock.unlock(); + ForgetRestoredSegments(forgotten); RecomputeTenantEffectiveQuotas(); return {}; } @@ -3248,7 +3421,7 @@ auto MasterService::QuerySegmentStatus(const std::string& segment_name) -> tl::expected { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); SegmentStatus status = SegmentStatus::UNDEFINED; - auto err = segment_access.GetSegmentStatusByName(segment_name, status); + auto err = GetDrainSegmentStatus(segment_access, segment_name, status); if (err != ErrorCode::OK) { return tl::make_unexpected(err); } @@ -3261,6 +3434,10 @@ auto MasterService::QuerySegmentStatusById(const UUID& segment_id) SegmentStatus status = SegmentStatus::UNDEFINED; auto err = segment_access.GetSegmentStatusById(segment_id, status); if (err != ErrorCode::OK) { + for (const auto& [name, recovered] : restored_drain_segments_) { + if (recovered.id == segment_id) + return static_cast(recovered.status); + } return tl::make_unexpected(err); } return status; @@ -3706,14 +3883,36 @@ tl::expected MasterService::SetSegmentStatus( ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); SegmentStatus current = SegmentStatus::UNDEFINED; - auto err = segment_access.GetSegmentStatusByName(segment_name, current); + auto err = GetDrainSegmentStatus(segment_access, segment_name, current); if (err != ErrorCode::OK) { return tl::make_unexpected(err); } if (current != SegmentStatus::OK && current != SegmentStatus::DRAINING) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } - err = segment_access.SetSegmentStatusByName(segment_name, status); + if (current == SegmentStatus::DRAINING && status == SegmentStatus::OK) { + auto tasks = task_manager_.get_read_access(); + for (const auto& [id, task] : tasks) { + if (task.type != TaskType::REPLICA_MOVE || task.is_finished()) + continue; + ReplicaMovePayload payload; + try { + struct_json::from_json(payload, task.payload); + } catch (const std::exception&) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + if (payload.source == segment_name) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + } + } + auto persisted = PersistDrainJobUpdate( + {std::nullopt, std::nullopt, + DescribeDrainSegments(segment_access, {segment_name}, status)}); + if (persisted != ErrorCode::OK) return tl::make_unexpected(persisted); + err = SetDrainSegmentStatus(segment_access, segment_name, status); if (err != ErrorCode::OK) { return tl::make_unexpected(err); } @@ -3786,6 +3985,48 @@ tl::expected MasterService::RestoreFromStandbyState( return RecordRestoreKnownOwner(new_known_owner_records, owner); }; + decltype(restored_drain_segments_) recovery_segments; + std::unordered_set> recovery_segment_ids; + if (metadata_store) { + for (const auto& [key, value] : + metadata_store->DrainRecovery().records()) { + if (value.kind != RecoveryRecordKind::Segment) continue; + DrainSegmentRecord recovered; + if (!DecodeRecoveryValue(value.payload, recovered)) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto segment = std::find_if( + segments.begin(), segments.end(), [&](const auto& candidate) { + return candidate.segment_name == recovered.name && + candidate.transport_endpoint == recovered.endpoint && + candidate.capacity == recovered.size && + candidate.is_memory_segment; + }); + if (segment == segments.end()) + continue; // An authoritative unmount wins. + if (!recovery_segment_ids.insert(recovered.id).second) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + if (recovered.base > + std::numeric_limits::max() - recovered.size) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + if (!record_for_known_owner(recovered.client_id)) { + new_known_owner_records.emplace( + recovered.client_id, + std::make_shared( + ClientLivenessRecord::Clock::now())); + } + const auto name = recovered.name; + recovery_segments.emplace( + name, RestoredDrainSegment(std::move(recovered))); + } + } + + const int64_t restore_epoch_ms = + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); const auto& resolve_standby_object = ResolveStandbyObject; ReplicaID max_live_id = 0; @@ -3825,8 +4066,7 @@ tl::expected MasterService::RestoreFromStandbyState( if (!seg.is_memory_segment) { continue; } - if (seg.segment_name.empty() || seg.transport_endpoint.empty() || - seg.capacity == 0) { + if (seg.transport_endpoint.empty() || seg.capacity == 0) { LOG(ERROR) << "RestoreFromStandbySnapshot: invalid memory segment " << "name=" << seg.segment_name << ", endpoint=" << seg.transport_endpoint @@ -3903,6 +4143,41 @@ tl::expected MasterService::RestoreFromStandbyState( // failed restore can roll its own state back before reporting the error. std::vector> installed_keys; + if (metadata_store) { + for (const auto& [key, value] : + metadata_store->DrainRecovery().records()) { + if (value.kind != RecoveryRecordKind::RetiredReplica) continue; + RetiredReplicaRecord retired; + if (!DecodeRecoveryValue(value.payload, retired) || + !retired.replica.is_memory_replica()) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + if (retired.release_ms <= restore_epoch_ms) continue; + const auto& buffer = + retired.replica.get_memory_descriptor().buffer_descriptor; + auto segment = + memory_segments_by_alias.find(buffer.transport_endpoint_); + if (segment == memory_segments_by_alias.end()) continue; + auto owner = recovery_segments.find(segment->second->segment_name); + if (owner == recovery_segments.end() || buffer.size_ == 0 || + buffer.buffer_address_ < owner->second.base || + buffer.size_ > owner->second.size || + buffer.buffer_address_ - owner->second.base > + owner->second.size - buffer.size_) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto& ranges = bounded_memory_ranges[segment->second]; + auto next = ranges.lower_bound(buffer.buffer_address_); + if ((next != ranges.end() && + buffer.buffer_address_ + buffer.size_ > next->first) || + (next != ranges.begin() && + std::prev(next)->first + std::prev(next)->second > + buffer.buffer_address_)) + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + ranges.emplace(buffer.buffer_address_, buffer.size_); + } + } + const auto restore_chunk = [&](const std::vector& objects) -> tl::expected { @@ -3977,6 +4252,17 @@ tl::expected MasterService::RestoreFromStandbyState( } const auto* segment = segment_it->second; + if (auto owner = + recovery_segments.find(segment->segment_name); + owner != recovery_segments.end()) { + if (buffer.buffer_address_ < owner->second.base || + buffer.size_ > owner->second.size || + buffer.buffer_address_ - owner->second.base > + owner->second.size - buffer.size_) { + return tl::make_unexpected( + ErrorCode::DESERIALIZE_FAIL); + } + } if (desc.status != ReplicaStatus::REMOVED && desc.status != ReplicaStatus::FAILED) { auto& bytes = @@ -4022,6 +4308,15 @@ tl::expected MasterService::RestoreFromStandbyState( restored_allocators.at(buffer.transport_endpoint_); auto restored_buffer = std::make_unique(alloc, buffer); + if (auto owner = + recovery_segments.find(segment->segment_name); + owner != recovery_segments.end()) { + restored_buffer->bindClientLiveness( + new_known_owner_records.at( + owner->second.client_id)); + restored_buffer->bindSegmentLifetime( + owner->second.lifetime); + } replicas.push_back(Replica( desc.id, std::move(restored_buffer), desc.status)); } else if (desc.is_nof_replica()) { @@ -4198,6 +4493,11 @@ tl::expected MasterService::RestoreFromStandbyState( const auto prev_standby_accounted_memory_bytes = standby_accounted_memory_bytes_; const auto prev_standby_memory_segments = standby_memory_segments_; + decltype(restored_drain_segments_) prev_recovered_segments; + { + auto access = segment_manager_.getSegmentAccess(); + prev_recovered_segments = std::move(restored_drain_segments_); + } const auto prev_standby_allocator_keepalive = standby_allocator_keepalive_; const auto prev_invalid_replica_endpoints = invalid_replica_endpoints_; @@ -4217,10 +4517,22 @@ tl::expected MasterService::RestoreFromStandbyState( if (meta_it == tenant_it->second.metadata.end()) { continue; } + if (auto& runtime = meta_it->second.state.replication_task; + runtime) { + if (auto* source = meta_it->second.metadata.GetReplicaByID( + runtime->source_id)) { + source->dec_refcnt(); + } + runtime.reset(); + } EraseMetadata(tenant_it->second, meta_it, tenant); } standby_accounted_memory_bytes_ = prev_standby_accounted_memory_bytes; standby_memory_segments_ = prev_standby_memory_segments; + { + auto segment_access = segment_manager_.getSegmentAccess(); + restored_drain_segments_ = std::move(prev_recovered_segments); + } standby_allocator_keepalive_ = prev_standby_allocator_keepalive; invalid_replica_endpoints_ = prev_invalid_replica_endpoints; for (const auto& [client_id, record] : new_known_owner_records) { @@ -4236,6 +4548,10 @@ tl::expected MasterService::RestoreFromStandbyState( standby_accounted_memory_bytes_ = std::move(restored_accounted_memory_bytes); standby_memory_segments_ = std::move(restored_memory_segments); + { + auto segment_access = segment_manager_.getSegmentAccess(); + restored_drain_segments_ = std::move(recovery_segments); + } standby_allocator_keepalive_ = std::move(restored_allocators); invalid_replica_endpoints_ = std::move(restored_invalid_endpoints); for (auto& [client_id, record] : new_known_owner_records) { @@ -4255,8 +4571,19 @@ tl::expected MasterService::RestoreFromStandbyState( } } - if (enable_multi_tenants_) { - RebuildTenantQuotaUsageFromMetadata(); + if (metadata_store) { + auto recovery = RestoreReplicationRecovery(*metadata_store); + if (!recovery) { + rollback_restored_state(); + return recovery; + } + } + try { + if (enable_multi_tenants_) RebuildTenantQuotaUsageFromMetadata(); + } catch (const std::exception& error) { + LOG(ERROR) << "Recovered quota is invalid: " << error.what(); + rollback_restored_state(); + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); } auto restored_weight_metadata = @@ -4371,6 +4698,15 @@ tl::expected MasterService::RestoreFromStandbyState( << total; } + if (metadata_store) { + auto recovery = RestoreDrainRecovery(metadata_store->DrainRecovery(), + restore_epoch_ms); + if (!recovery) { + rollback_restored_state(); + return recovery; + } + } + // Publish the accounting swap only after nothing can fail anymore: the // fail-closed paths above must stay a no-op for the exported gauges, and // the client/snapshot mutexes held here mean no concurrent reader can see @@ -6846,6 +7182,7 @@ tl::expected MasterService::CopyStart( return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } std::shared_lock shared_lock(master_state_mutex_); + std::optional start_reservation; if (!dynamic_copy) { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); @@ -6931,6 +7268,12 @@ tl::expected MasterService::CopyStart( return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } + if (enable_oplog_) { + auto slot = ReserveBatchOpLogSlot(); + if (!slot) return tl::make_unexpected(slot.error()); + start_reservation.emplace(std::move(*slot)); + } + const uint64_t pending_quota_charge = SaturatingMultiply(static_cast(metadata.TotalLength()), static_cast(new_replica_count)); @@ -7024,6 +7367,9 @@ tl::expected MasterService::CopyStart( // Add replicas to the object. // DO NOT ACCESS source AFTER THIS !!! metadata.AddReplicas(std::move(replicas)); + if (auto error = PersistReplicationStart(accessor, start_reservation); + error != ErrorCode::OK) + return tl::make_unexpected(error); return response; } @@ -7083,6 +7429,10 @@ tl::expected MasterService::CopyEnd( LOG(ERROR) << "key=" << key << ", source_id=" << source_id << ", status=" << (source == nullptr ? "nullptr" : "invalid") << ", copy source becomes invalid during data transfer"; + if (auto error = PersistReplicationRevoke(accessor); + error != ErrorCode::OK) { + return tl::make_unexpected(error); + } // Release the refcnt taken in CopyStart. The success path below does // this once the copy completes; this error path must do it too, or the // source replica stays pinned and can never be evicted. @@ -7129,13 +7479,29 @@ tl::expected MasterService::CopyEnd( } } - std::optional batch_reservation; - if (enable_ha_ && enable_oplog_) { + if (enable_oplog_) { auto reservation = ReserveBatchOpLogSlot(); - if (!reservation) { - return tl::make_unexpected(reservation.error()); + if (!reservation) return tl::make_unexpected(reservation.error()); + std::vector post; + for (const auto& replica : metadata.GetAllReplicas()) { + if (std::find(failed_target_ids.begin(), failed_target_ids.end(), + replica.id()) != failed_target_ids.end()) + continue; + auto descriptor = replica.get_descriptor(); + if (std::find(commit_target_ids.begin(), commit_target_ids.end(), + replica.id()) != commit_target_ids.end()) { + descriptor.status = ReplicaStatus::COMPLETE; + } else if (!replica.is_completed()) + continue; + post.push_back(std::move(descriptor)); } - batch_reservation = std::move(reservation.value()); + auto persisted = AppendReservedOpLogWithDurableFinalize( + std::move(*reservation), OpType::PUT_END, + metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors(metadata, post, {}, + true), + nullptr); + if (!persisted) return tl::make_unexpected(persisted.error()); } source->dec_refcnt(); @@ -7163,34 +7529,6 @@ tl::expected MasterService::CopyEnd( dynamic_replication_.ClearPendingLocked(accessor.GetTenantId(), key, accessor.GetState()); - if (enable_oplog_ && ordered_oplog_writer_) { - std::vector post; - metadata.VisitReplicas(&Replica::fn_is_completed, - [&post](const Replica& replica) { - post.push_back(replica.get_descriptor()); - }); - auto payload = - SerializeMetadataForOpLogFromReplicaDescriptors(metadata, post); - if (batch_reservation) { - auto persist_result = AppendReservedOpLogWithDurableFinalize( - std::move(*batch_reservation), OpType::PUT_END, - tenant_id.value(), key, payload, nullptr); - if (!persist_result) { - LOG(WARNING) - << "CopyEnd: PUT_END persist failed for key=" << key - << ", err=" << static_cast(persist_result.error()); - } - } else { - auto persist_result = AppendOpLogVisibleBeforeDurable( - OpType::PUT_END, tenant_id.value(), key, payload); - if (!persist_result) { - LOG(WARNING) - << "CopyEnd: PUT_END persist failed for key=" << key - << ", err=" << static_cast(persist_result.error()); - } - } - } - SyncCacheTotalAccounting(metadata); if (enable_multi_tenants_) { @@ -7261,6 +7599,10 @@ tl::expected MasterService::CopyRevoke( auto& metadata = accessor.Get(); const auto source_id = task.source_id; const auto replica_ids = task.replica_ids; + if (auto error = PersistReplicationRevoke(accessor); + error != ErrorCode::OK) { + return tl::make_unexpected(error); + } auto source = metadata.GetReplicaByID(source_id); if (source == nullptr) { LOG(WARNING) << "key=" << key << ", source_id=" << source_id @@ -7303,6 +7645,7 @@ tl::expected MasterService::MoveStart( return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } std::shared_lock shared_lock(master_state_mutex_); + std::optional start_reservation; if (src_segment == tgt_segment) { LOG(ERROR) << "key=" << key << ", move_tgt=" << tgt_segment << " cannot be the same as move_src=" << src_segment; @@ -7350,6 +7693,17 @@ tl::expected MasterService::MoveStart( return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } + if (enable_oplog_) { + auto slot = ReserveBatchOpLogSlot(); + if (!slot) return tl::make_unexpected(slot.error()); + start_reservation.emplace(std::move(*slot)); + } + + if (const auto* target = metadata.GetReplicaBySegmentName(tgt_segment); + target && !target->is_completed()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + std::vector replicas; // See CopyStart: target allocation owns its own Client guard and must not // be nested under the source executor's guard. The metadata shard keeps @@ -7433,6 +7787,9 @@ tl::expected MasterService::MoveStart( // Add replicas to the object. // DO NOT ACCESS source AFTER THIS !!! metadata.AddReplicas(std::move(replicas)); + if (auto error = PersistReplicationStart(accessor, start_reservation); + error != ErrorCode::OK) + return tl::make_unexpected(error); return response; } @@ -7480,6 +7837,10 @@ tl::expected MasterService::MoveEnd( LOG(ERROR) << "key=" << key << ", source_id=" << source_id << ", status=" << (source == nullptr ? "nullptr" : "invalid") << ", move source becomes invalid during data transfer"; + if (auto error = PersistReplicationRevoke(accessor); + error != ErrorCode::OK) { + return tl::make_unexpected(error); + } // Release the refcnt taken in MoveStart. The success path below does // this once the move completes; this error path must do it too, or the // source replica stays pinned and can never be evicted. @@ -7513,6 +7874,10 @@ tl::expected MasterService::MoveEnd( LOG(WARNING) << "key=" << key << ", replica_id=" << target_id << ", move target becomes invalid during data transfer"; + if (auto error = PersistReplicationRevoke(accessor); + error != ErrorCode::OK) { + return tl::make_unexpected(error); + } // Release the refcnt taken in MoveStart. The success path below // does this once the move completes; this error path must do it // too, or the source replica stays pinned. @@ -7531,49 +7896,44 @@ tl::expected MasterService::MoveEnd( } } - if (enable_oplog_ && ordered_oplog_writer_) { - // Build post-mutation descriptors: - // - existing COMPLETE replicas, except the source (about to be - // popped) - // - target (if any) flipped to COMPLETE + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) return tl::make_unexpected(reservation.error()); std::vector post; - for (const auto& rep : metadata.GetAllReplicas()) { - if (rep.id() == source_id) continue; - if (rep.status() == ReplicaStatus::COMPLETE) { - post.push_back(rep.get_descriptor()); + for (const auto& replica : metadata.GetAllReplicas()) { + if (replica.id() == source_id) continue; + auto descriptor = replica.get_descriptor(); + if (has_target && replica.id() == target_id) { + descriptor.status = ReplicaStatus::COMPLETE; + } else if (!replica.is_completed()) continue; - } - if (has_target && rep.id() == target_id) { - Replica::Descriptor desc = rep.get_descriptor(); - desc.status = ReplicaStatus::COMPLETE; - post.push_back(std::move(desc)); - } - } - - tl::expected persist_result; - if (enable_oplog_) { - auto reservation = ReserveBatchOpLogSlot(); - if (!reservation) { - return tl::make_unexpected(reservation.error()); - } - source->mark_removed(); - persist_result = AppendReservedOpLogWithDurableFinalize( - std::move(reservation.value()), OpType::PUT_END, - metadata.tenant_id.value(), key, - SerializeMetadataForOpLogFromReplicaDescriptors(metadata, post), - [this, removed_ids = std::vector{source_id}]( - const OpLogEntry& durable_entry) { - FinalizeRemovedReplicasAfterDurable( - durable_entry, removed_ids, QuotaEraseMode::kFull); - }); - } else { - persist_result = AppendOpLogWithDurableFinalize( - OpType::PUT_END, metadata.tenant_id.value(), key, - SerializeMetadataForOpLogFromReplicaDescriptors(metadata, post), - nullptr); - } - if (!persist_result) { - return tl::make_unexpected(persist_result.error()); + post.push_back(std::move(descriptor)); + } + const auto release_at = metadata.EvictionDeadline(); + std::vector retired; + std::optional retire_at; + if (source->is_memory_replica() && + release_at > std::chrono::system_clock::now()) { + retired.push_back( + {source->get_descriptor(), + std::chrono::duration_cast( + release_at.time_since_epoch()) + .count()}); + retire_at = release_at; + } + source->mark_removed(); + auto persisted = AppendReservedOpLogWithDurableFinalize( + std::move(*reservation), OpType::PUT_END, + metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors(metadata, post, + retired, true), + [this, source_id, retire_at](const OpLogEntry& entry) { + FinalizeRemovedReplicasAfterDurable( + entry, {source_id}, QuotaEraseMode::kFull, {}, retire_at); + }); + if (!persisted) { + source->mark_complete(); + return tl::make_unexpected(persisted.error()); } } @@ -7665,6 +8025,10 @@ tl::expected MasterService::MoveRevoke( auto& metadata = accessor.Get(); auto source_id = task.source_id; + if (auto error = PersistReplicationRevoke(accessor); + error != ErrorCode::OK) { + return tl::make_unexpected(error); + } auto source = metadata.GetReplicaByID(source_id); if (source == nullptr) { LOG(WARNING) << "key=" << key << ", source_id=" << source_id @@ -10484,7 +10848,7 @@ void MasterService::DiscardExpiredProcessingReplicas( OpType::PUT_END, tenant_it->first.value(), metadata_it->first, SerializeMetadataForOpLogFromReplicaDescriptors( - metadata, post_descriptors), + metadata, post_descriptors, {}, true), enable_oplog_ ? [this, source_id, target_ids = std::move(target_ids), @@ -12047,7 +12411,10 @@ void MasterService::ClientMonitorFunc() { .transport_endpoint = segment.te_endpoint, .metrics_dec_capacity = - metrics_dec_capacity}); + metrics_dec_capacity, + .recovered = + restored_drain_segments_.contains( + segment.name)}); } else if (prepare_result != ErrorCode::SEGMENT_NOT_FOUND) { job.pending_prepare_segments.push_back( @@ -12474,6 +12841,13 @@ tl::expected MasterService::CreateCopyTask( tl::expected MasterService::CreateMoveTask( const std::string& key, const TenantId& tenant_id, const std::string& source, const std::string& target) { + return CreateMoveTaskImpl(key, tenant_id, source, target, std::nullopt); +} + +tl::expected MasterService::CreateMoveTaskImpl( + const std::string& key, const TenantId& tenant_id, + const std::string& source, const std::string& target, + std::optional drain) { auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); if (!normalized_tenant_result) { return tl::make_unexpected(normalized_tenant_result.error()); @@ -12588,10 +12962,12 @@ tl::expected MasterService::CreateMoveTask( return task_manager_.get_write_access() .submit_task_typed( - select_client, {.tenant_id = object_id.tenant_id.value(), - .key = object_id.user_key, - .source = source, - .target = target}); + select_client, + {.tenant_id = object_id.tenant_id.value(), + .key = object_id.user_key, + .source = source, + .target = target}, + std::move(drain)); } tl::expected MasterService::QueryTask( @@ -12616,6 +12992,9 @@ tl::expected, ErrorCode> MasterService::FetchTasks( return std::vector{}; } std::shared_lock shared_lock(master_state_mutex_); + if (enable_oplog_) + batch_size = + std::min(batch_size, static_cast(oplog_batch_max_entries_)); const auto& tasks = task_manager_.get_write_access().pop_tasks(client_id, batch_size); std::vector assignments; @@ -12697,40 +13076,26 @@ tl::expected MasterService::ValidateDrainRequestLocked( tl::expected MasterService::CreateDrainJob( const CreateDrainJobRequest& request) { - // Held until the job is registered so SetSegmentStatus never sees a - // segment this job has marked DRAINING without also seeing the job. std::lock_guard lock(job_mutex_); - std::vector draining_segments; - { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - auto valid = ValidateDrainRequestLocked(segment_access, request); - if (!valid.has_value()) { - return tl::make_unexpected(valid.error()); - } - - draining_segments.reserve(request.segments.size()); - for (const auto& segment_name : request.segments) { - auto err = segment_access.SetSegmentStatusByName( - segment_name, SegmentStatus::DRAINING); - if (err != ErrorCode::OK) { - for (const auto& updated_segment : draining_segments) { - (void)segment_access.SetSegmentStatusByName( - updated_segment, SegmentStatus::OK); - } - return tl::make_unexpected(err); - } - draining_segments.push_back(segment_name); - } - } - auto job = std::make_shared(); job->id = generate_uuid(); job->request = request; job->created_at = std::chrono::system_clock::now(); job->last_updated_at = job->created_at; - job->status = JobStatus::CREATED; job->message = "Drain job created"; + auto access = segment_manager_.getSegmentAccess(); + auto valid = ValidateDrainRequestLocked(access, request); + if (!valid) return tl::make_unexpected(valid.error()); + auto persisted = + PersistDrainJobUpdate({DescribeDrainJob(*job), std::nullopt, + DescribeDrainSegments(access, request.segments, + SegmentStatus::DRAINING)}); + if (persisted != ErrorCode::OK) return tl::make_unexpected(persisted); + for (const auto& source : request.segments) { + auto error = + SetDrainSegmentStatus(access, source, SegmentStatus::DRAINING); + if (error != ErrorCode::OK) return tl::make_unexpected(error); + } drain_jobs_.emplace(job->id, job); return job->id; } @@ -12792,32 +13157,36 @@ tl::expected MasterService::CancelDrainJob( ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } - job->status = JobStatus::CANCELED; - job->last_updated_at = std::chrono::system_clock::now(); - job->message = "Drain job canceled"; + auto next = DescribeDrainJob(*job); + next.status = static_cast(JobStatus::CANCELED); + next.updated_ms = + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); + next.message = "Drain job canceled"; + auto access = segment_manager_.getSegmentAccess(); + auto persisted = PersistDrainJobUpdate( + {next, std::nullopt, + DescribeDrainSegments(access, job->request.segments, + SegmentStatus::OK)}); + if (persisted != ErrorCode::OK) return tl::make_unexpected(persisted); + PublishDrainJob(*job, next); segments_to_restore = job->request.segments; } ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); for (const auto& segment_name : segments_to_restore) { SegmentStatus status = SegmentStatus::UNDEFINED; - if (segment_access.GetSegmentStatusByName(segment_name, status) == + if (GetDrainSegmentStatus(segment_access, segment_name, status) == ErrorCode::OK && status != SegmentStatus::UNMOUNTING) { - (void)segment_access.SetSegmentStatusByName(segment_name, - SegmentStatus::OK); + (void)SetDrainSegmentStatus(segment_access, segment_name, + SegmentStatus::OK); } } return {}; } -std::string MasterService::MakeDrainUnitKey( - const TenantId& tenant_id, const std::string& key, - const std::string& source_segment) const { - return std::to_string(tenant_id.value().size()) + ":" + tenant_id.value() + - ":" + std::to_string(key.size()) + ":" + key + ":" + source_segment; -} - std::optional MasterService::SelectDrainTargetForKey( const ObjectMetadata& metadata, const std::string& source_segment, const std::vector& requested_targets) { @@ -12861,45 +13230,54 @@ std::optional MasterService::SelectDrainTargetForKey( } void MasterService::RefreshDrainJobTasks(DrainJob& job) { - auto read_access = task_manager_.get_read_access(); - std::vector finished_task_ids; - finished_task_ids.reserve(job.active_tasks.size()); - - for (const auto& [task_id, active_task] : job.active_tasks) { - auto task_opt = read_access.find_task_by_id(task_id); - if (!task_opt.has_value()) { - finished_task_ids.push_back(task_id); - job.failed_units++; - job.terminal_failed_unit_keys.insert(active_task.unit_key); - continue; - } - if (!task_opt->is_finished()) { + auto access = task_manager_.get_read_access(); + for (auto it = job.active_tasks.begin(); it != job.active_tasks.end();) { + const auto& active = it->second; + auto task = access.find_task_by_id(it->first); + if (task && !task->is_finished()) { + ++it; continue; } - - finished_task_ids.push_back(task_id); - if (task_opt->status == TaskStatus::SUCCESS) { - job.succeeded_units++; - job.migrated_bytes += active_task.bytes; - job.completed_unit_keys.insert(active_task.unit_key); + auto next = DescribeDrainJob(job); + DrainUnitRecord unit{job.id, + {}, + active.unit_key, + active.tenant_id.value(), + active.key, + active.source_segment, + active.target_segment, + active.bytes}; + if (auto retry = job.retry_counts.find(unit.unit_key); + retry != job.retry_counts.end()) + unit.retries = retry->second; + if (!task) { + ++next.failed; + unit.terminal_failed = true; + } else if (task->status == TaskStatus::SUCCESS) { + ++next.succeeded; + next.migrated_bytes += active.bytes; + unit.completed = true; } else { - job.failed_units++; - auto& retry_count = job.retry_counts[active_task.unit_key]; - retry_count++; - if (retry_count >= kMaxDrainUnitRetries) { - job.terminal_failed_unit_keys.insert(active_task.unit_key); - } + ++next.failed; + ++unit.retries; + unit.terminal_failed = unit.retries >= kMaxDrainUnitRetries; } - } - - for (const auto& task_id : finished_task_ids) { - job.active_tasks.erase(task_id); + if (PersistDrainJobUpdate({next, unit, {}}) != ErrorCode::OK) { + ++it; + continue; + } + PublishDrainJob(job, next); + if (unit.completed) job.completed_unit_keys.insert(unit.unit_key); + if (unit.retries) job.retry_counts[unit.unit_key] = unit.retries; + if (unit.terminal_failed) + job.terminal_failed_unit_keys.insert(unit.unit_key); + it = job.active_tasks.erase(it); } } void MasterService::ScheduleDrainJobTasks(DrainJob& job) { - if (job.status == JobStatus::CREATED) { - job.status = JobStatus::PLANNING; + if (job.status < JobStatus::RUNNING) { + job.status = JobStatus::RUNNING; } const uint32_t max_concurrency = @@ -12935,8 +13313,8 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { for (const auto& [key, record] : tenant_state.metadata) { const auto& metadata = record.metadata; for (const auto& source_segment : job.request.segments) { - const auto unit_key = - MakeDrainUnitKey(tenant_id, key, source_segment); + const auto unit_key = DrainUnitKey(tenant_id.value(), + key, source_segment); if (job.completed_unit_keys.contains(unit_key) || active_unit_keys.contains(unit_key) || job.terminal_failed_unit_keys.contains(unit_key)) { @@ -13002,36 +13380,54 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { job.blocked_units = blocked_unit_keys.size(); for (const auto& plan : plans) { - auto task_id = CreateMoveTask(plan.key, plan.tenant_id, - plan.source_segment, plan.target_segment); - if (task_id.has_value()) { - ActiveDrainTask active_task; - active_task.task_id = task_id.value(); - active_task.tenant_id = plan.tenant_id; - active_task.key = plan.key; - active_task.source_segment = plan.source_segment; - active_task.target_segment = plan.target_segment; - active_task.bytes = plan.bytes; - active_task.unit_key = plan.unit_key; - job.active_tasks.emplace(task_id.value(), std::move(active_task)); + DrainUnitRecord unit{job.id, + {}, + plan.unit_key, + plan.tenant_id.value(), + plan.key, + plan.source_segment, + plan.target_segment, + plan.bytes}; + if (auto retry = job.retry_counts.find(plan.unit_key); + retry != job.retry_counts.end()) + unit.retries = retry->second; + auto task_id = + CreateMoveTaskImpl(plan.key, plan.tenant_id, plan.source_segment, + plan.target_segment, unit); + if (task_id) { + job.active_tasks.emplace( + *task_id, + ActiveDrainTask{*task_id, plan.tenant_id, plan.key, + plan.source_segment, plan.target_segment, + plan.bytes, plan.unit_key}); } else if (task_id.error() == ErrorCode::NO_AVAILABLE_HANDLE || task_id.error() == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS || task_id.error() == ErrorCode::OBJECT_HAS_REPLICATION_TASK) { - job.blocked_units++; + ++job.blocked_units; } else { - job.failed_units++; - auto& retry_count = job.retry_counts[plan.unit_key]; - retry_count++; - if (retry_count >= kMaxDrainUnitRetries) { + auto next = DescribeDrainJob(job); + ++next.failed; + ++unit.retries; + unit.terminal_failed = unit.retries >= kMaxDrainUnitRetries; + if (PersistDrainJobUpdate({next, unit, {}}) != ErrorCode::OK) + continue; + PublishDrainJob(job, next); + job.retry_counts[plan.unit_key] = unit.retries; + if (unit.terminal_failed) job.terminal_failed_unit_keys.insert(plan.unit_key); - } } } - job.status = JobStatus::RUNNING; - job.last_updated_at = std::chrono::system_clock::now(); - job.message = "Drain job running"; + auto next = DescribeDrainJob(job); + next.status = static_cast(JobStatus::RUNNING); + next.updated_ms = std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); + next.message = "Drain job running"; + if (PersistDrainJobUpdate({next, std::nullopt, {}}) == ErrorCode::OK) { + PublishDrainJob(job, next); + } } bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { @@ -13054,8 +13450,8 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { replica_segments.end(), source_segment) != replica_segments.end()) { remaining_segments.insert(source_segment); - remaining_unit_keys.insert(MakeDrainUnitKey( - tenant_id, key, source_segment)); + remaining_unit_keys.insert(DrainUnitKey( + tenant_id.value(), key, source_segment)); } } } @@ -13063,53 +13459,45 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { } } - { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - for (const auto& segment_name : job.request.segments) { - if (!remaining_segments.contains(segment_name)) { - (void)segment_access.SetSegmentStatusByName( - segment_name, SegmentStatus::DRAINED); - } - } + bool all_failed = !remaining_unit_keys.empty(); + for (const auto& key : remaining_unit_keys) { + if (!job.terminal_failed_unit_keys.contains(key)) all_failed = false; } - - if (remaining_segments.empty()) { - job.status = JobStatus::SUCCEEDED; - job.last_updated_at = std::chrono::system_clock::now(); - job.message = "Drain job finished successfully"; - return true; + auto next = DescribeDrainJob(job); + if (remaining_segments.empty() || all_failed) { + next.status = static_cast(remaining_segments.empty() + ? JobStatus::SUCCEEDED + : JobStatus::FAILED); + next.updated_ms = + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); + next.message = remaining_segments.empty() + ? "Drain job finished successfully" + : "Drain job failed: unrecoverable units remain"; } - - bool all_remaining_terminal_failed = !remaining_unit_keys.empty(); - for (const auto& unit_key : remaining_unit_keys) { - if (!job.terminal_failed_unit_keys.contains(unit_key)) { - all_remaining_terminal_failed = false; - break; - } + auto access = segment_manager_.getSegmentAccess(); + std::vector changes; + for (const auto& source : job.request.segments) { + auto desired = all_failed ? SegmentStatus::OK : SegmentStatus::DRAINED; + if (!all_failed && remaining_segments.contains(source)) continue; + SegmentStatus current; + if (GetDrainSegmentStatus(access, source, current) != ErrorCode::OK || + current == desired || current == SegmentStatus::UNMOUNTING) + continue; + auto records = DescribeDrainSegments(access, {source}, desired); + changes.insert(changes.end(), records.begin(), records.end()); } - if (!all_remaining_terminal_failed) { + if ((next.status != static_cast(job.status) || !changes.empty()) && + PersistDrainJobUpdate({next, std::nullopt, changes}) != ErrorCode::OK) return false; + for (const auto& change : changes) { + (void)SetDrainSegmentStatus(access, change.name, + static_cast(change.status)); } - - { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - for (const auto& segment_name : job.request.segments) { - SegmentStatus status = SegmentStatus::UNDEFINED; - if (segment_access.GetSegmentStatusByName(segment_name, status) == - ErrorCode::OK && - status != SegmentStatus::UNMOUNTING) { - (void)segment_access.SetSegmentStatusByName(segment_name, - SegmentStatus::OK); - } - } - } - - job.status = JobStatus::FAILED; - job.last_updated_at = std::chrono::system_clock::now(); - job.message = "Drain job failed: unrecoverable units remain"; - return true; + PublishDrainJob(job, next); + return job.status == JobStatus::SUCCEEDED || + job.status == JobStatus::FAILED; } void MasterService::ProcessDrainJobs() { @@ -13141,8 +13529,15 @@ void MasterService::ProcessDrainJobs() { } void MasterService::JobDispatchThreadFunc() { + auto next_warning = + std::chrono::steady_clock::now() + std::chrono::minutes(1); while (job_dispatch_running_) { ProcessDrainJobs(); + if (std::chrono::steady_clock::now() >= next_warning) { + WarnOrphanedDrainingSegments(); + next_warning = + std::chrono::steady_clock::now() + std::chrono::minutes(1); + } std::this_thread::sleep_for( std::chrono::milliseconds(kJobDispatchThreadSleepMs)); } @@ -13516,11 +13911,28 @@ std::string MasterService::SerializeMetadataForOpLogWithoutMemReplicas( std::string MasterService::SerializeMetadataForOpLogFromReplicaDescriptors( const ObjectMetadata& metadata, - const std::vector& replicas) const { + const std::vector& replicas, + const std::vector& retired, + bool replication_finished) const { MetadataPayload payload; payload.client_id = metadata.client_id; payload.size = metadata.TotalLength(); payload.replicas = replicas; + if (!retired.empty()) payload.retired_replicas = retired; + payload.replication_finished = replication_finished; + if (!replication_finished) { + std::vector processing; + metadata.VisitReplicas( + [](const Replica& replica) { + return replica.is_processing() && + !replica.has_invalid_mem_handle() && + !replica.has_invalid_nof_handle(); + }, + [&](const Replica& replica) { + processing.push_back(replica.get_descriptor()); + }); + payload.processing_replicas = std::move(processing); + } payload.group_id = metadata.group_id; payload.session_id = metadata.session_id; payload.data_type = metadata.data_type; diff --git a/mooncake-store/src/master_service_drain_recovery.cpp b/mooncake-store/src/master_service_drain_recovery.cpp new file mode 100644 index 0000000000..289bbe9abd --- /dev/null +++ b/mooncake-store/src/master_service_drain_recovery.cpp @@ -0,0 +1,479 @@ +#include "master_service.h" + +#include +#include +#include +#include + +#include "ha/standby_metadata_store.h" +#include "master_metric_manager.h" + +namespace mooncake { +ErrorCode MasterService::AwaitRecoveryDurability(uint64_t sequence) { + return async_simple::coro::syncAwait( + [this, sequence]() -> async_simple::coro::Lazy { + co_return co_await ordered_oplog_writer_->AwaitDurable(sequence); + }()); +} + +ErrorCode MasterService::PersistRecoveryEntry(OpType type, + const std::string& payload, + const std::string& key, + const std::string& tenant) { + if (!enable_oplog_) return ErrorCode::OK; + auto sequence = AppendOpLogVisibleBeforeDurable(type, tenant, key, payload); + if (!sequence) return sequence.error(); + // AwaitDurable completes independently of finalizer callbacks. Callers may + // hold a task/job/shard lock that an earlier finalizer still needs. + return AwaitRecoveryDurability(*sequence); +} + +ErrorCode MasterService::PersistDrainJobUpdate(const DrainJobUpdate& update) { + const auto key = update.job ? UuidToString(update.job->id) + : update.segments.front().name; + return PersistRecoveryEntry(OpType::DRAIN_JOB_UPDATE, + EncodeRecoveryValue(update), key); +} + +ErrorCode MasterService::PersistRecoveryTasks( + std::span tasks) { + if (!enable_oplog_ || tasks.empty()) return ErrorCode::OK; + std::vector entries; + entries.reserve(tasks.size()); + for (const auto& task : tasks) { + OpLogEntry entry; + entry.op_type = OpType::TASK_UPDATE; + entry.object_key = UuidToString(task.id); + entry.payload = EncodeRecoveryValue(task); + entries.push_back(std::move(entry)); + } + auto sequence = + AppendOpLogBatchWithDurableFinalize(std::move(entries), nullptr); + return sequence ? AwaitRecoveryDurability(*sequence) : sequence.error(); +} + +ErrorCode MasterService::PersistReplicationStart( + MetadataAccessorRW& accessor, + std::optional& reservation) { + if (!enable_oplog_) return ErrorCode::OK; + auto& metadata = accessor.Get(); + const auto& task = accessor.GetReplicationTask(); + ReplicationRecoveryRecord record{ + metadata.tenant_id.value(), + metadata.user_key, + task.client_id, + static_cast(task.type), + std::chrono::duration_cast( + task.start_time.time_since_epoch()) + .count(), + task.source_id, + task.replica_ids, + task.pending_quota_charge_bytes, + task.dynamic_replication_lease_id, + task.dynamic_replication_version_epoch}; + std::vector segments; + { + auto access = segment_manager_.getSegmentAccess(); + segments = DescribeDrainSegments( + access, metadata.GetReplicaSegmentNames(), std::nullopt); + } + auto persisted = AppendReservedOpLogWithDurableFinalize( + std::move(*reservation), OpType::REPLICATION_START, + metadata.tenant_id.value(), metadata.user_key, + EncodeRecoveryValue(ReplicationStartRecord{ + record, SerializeMetadataForOpLog(metadata), std::move(segments)}), + nullptr); + if (!persisted) { + // Submission was rejected: no client has received the new targets. + const auto targets = task.replica_ids; + const auto pending = task.pending_quota_charge_bytes; + if (auto* source = metadata.GetReplicaByID(task.source_id)) + source->dec_refcnt(); + metadata.ForgetDynamicReplicas(targets); + EraseReplicasWithCacheTotalAccounting( + metadata, [&](const Replica& replica) { + return std::find(targets.begin(), targets.end(), + replica.id()) != targets.end(); + }); + ReleaseTenantQuota(accessor.QuotaAccount(), pending); + accessor.EraseReplicationTask(); + dynamic_replication_.ClearPendingLocked( + accessor.GetTenantId(), metadata.user_key, accessor.GetState()); + return persisted.error(); + } + // A failed admitted write can have an uncertain outcome. Retain its + // buffers on the failed leader rather than allowing them to be reused. + return AwaitRecoveryDurability(persisted->sequence_id); +} + +ErrorCode MasterService::PersistReplicationRevoke( + MetadataAccessorRW& accessor) { + if (!enable_oplog_) return ErrorCode::OK; + auto& metadata = accessor.Get(); + const auto& task = accessor.GetReplicationTask(); + std::vector remaining; + metadata.VisitReplicas( + [&](const Replica& replica) { + return replica.is_completed() && + !replica.has_invalid_mem_handle() && + std::find(task.replica_ids.begin(), task.replica_ids.end(), + replica.id()) == task.replica_ids.end(); + }, + [&](const Replica& replica) { + remaining.push_back(replica.get_descriptor()); + }); + return PersistRecoveryEntry( + remaining.empty() ? OpType::REMOVE : OpType::PUT_END, + remaining.empty() ? std::string() + : SerializeMetadataForOpLogFromReplicaDescriptors( + metadata, remaining, {}, true), + metadata.user_key, metadata.tenant_id.value()); +} + +DrainJobRecord MasterService::DescribeDrainJob(const DrainJob& job) { + const auto milliseconds = [](auto time) { + return std::chrono::duration_cast( + time.time_since_epoch()) + .count(); + }; + return {job.id, + static_cast(job.status), + job.request.segments, + job.request.target_segments, + job.request.max_concurrency, + milliseconds(job.created_at), + milliseconds(job.last_updated_at), + job.message, + job.succeeded_units, + job.failed_units, + job.blocked_units, + job.migrated_bytes}; +} + +void MasterService::PublishDrainJob(DrainJob& job, + const DrainJobRecord& record) { + job.id = record.id; + job.status = static_cast(record.status); + job.request = {record.sources, record.targets, record.concurrency}; + job.created_at = std::chrono::system_clock::time_point( + std::chrono::milliseconds(record.created_ms)); + job.last_updated_at = std::chrono::system_clock::time_point( + std::chrono::milliseconds(record.updated_ms)); + job.message = record.message; + job.succeeded_units = record.succeeded; + job.failed_units = record.failed; + job.blocked_units = record.blocked; + job.migrated_bytes = record.migrated_bytes; +} + +std::vector MasterService::DescribeDrainSegments( + ScopedSegmentAccess& access, const std::vector& names, + std::optional status) { + std::vector> segments; + access.GetAllSegments(segments); + std::vector records; + for (const auto& [segment, client] : segments) { + if (std::find(names.begin(), names.end(), segment.name) != + names.end()) { + SegmentStatus current; + access.GetSegmentStatusById(segment.id, current); + if (status == SegmentStatus::OK && + current >= SegmentStatus::GRACEFULLY_UNMOUNTING) + continue; + records.push_back({segment.id, client, segment.name, + segment.te_endpoint, segment.base, segment.size, + static_cast( + status.value_or(SegmentStatus::UNDEFINED))}); + } + } + for (const auto& name : names) { + if (std::any_of( + records.begin(), records.end(), + [&](const auto& record) { return record.name == name; })) + continue; + if (auto previous = restored_drain_segments_.find(name); + previous != restored_drain_segments_.end()) { + if (status == SegmentStatus::OK && + previous->second.status >= + static_cast(SegmentStatus::GRACEFULLY_UNMOUNTING)) + continue; + DrainSegmentRecord record = previous->second; + record.status = + static_cast(status.value_or(SegmentStatus::UNDEFINED)); + records.push_back(std::move(record)); + } + } + return records; +} + +ErrorCode MasterService::SetDrainSegmentStatus(ScopedSegmentAccess& access, + const std::string& name, + SegmentStatus status) { + const auto error = access.SetSegmentStatusByName(name, status); + auto restored = restored_drain_segments_.find(name); + if (restored != restored_drain_segments_.end() && + (error == ErrorCode::OK || error == ErrorCode::SEGMENT_NOT_FOUND)) { + restored->second.status = static_cast(status); + return ErrorCode::OK; + } + return error; +} + +ErrorCode MasterService::GetDrainSegmentStatus(ScopedSegmentAccess& access, + const std::string& name, + SegmentStatus& status) const { + const auto error = access.GetSegmentStatusByName(name, status); + if (error != ErrorCode::OK) { + if (auto restored = restored_drain_segments_.find(name); + restored != restored_drain_segments_.end()) { + status = static_cast(restored->second.status); + return ErrorCode::OK; + } + } + return error; +} + +tl::expected MasterService::RestoreReplicationRecovery( + const StandbyMetadataStore& store) { + for (const auto& [key, value] : store.DrainRecovery().records()) { + if (value.kind != RecoveryRecordKind::Replication) continue; + ReplicationRecoveryRecord runtime; + if (!DecodeRecoveryValue(value.payload, runtime)) { + LOG(ERROR) << "Invalid replication recovery record: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto standby_metadata = store.GetMetadata(runtime.tenant, runtime.key); + if (!standby_metadata || + !ValidateReplicationRecovery(runtime, standby_metadata->replicas, + standby_metadata->size, true)) { + LOG(ERROR) << "Invalid replication source/targets/quota: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + MetadataAccessorRW accessor( + this, MakeObjectIdentityForRequest(runtime.key, + TenantId(runtime.tenant))); + // An authoritative unmount can remove every replica of this object. + // Its existing task history still drives normal job failure/retry. + if (!accessor.Exists()) continue; + auto& metadata = accessor.Get(); + accessor.GetState().replication_task.emplace( + runtime.client_id, + std::chrono::system_clock::time_point( + std::chrono::milliseconds(runtime.started_ms)), + static_cast(runtime.type), runtime.source_id, + runtime.targets, runtime.pending_bytes, runtime.dynamic_lease, + runtime.dynamic_version); + if (auto* source = metadata.GetReplicaByID(runtime.source_id)) { + source->inc_refcnt(); + } + } + return {}; +} + +tl::expected MasterService::RestoreDrainRecovery( + const DrainRecoveryState& state, int64_t restore_epoch_ms) { + decltype(drain_jobs_) jobs; + std::vector tasks; + std::vector< + std::pair, std::chrono::system_clock::time_point>> + retired; + for (const auto& [key, value] : state.records()) { + if (value.kind == RecoveryRecordKind::Job) { + DrainJobRecord record; + if (!DecodeRecoveryValue(value.payload, record) || + !ValidateDrainJob(record)) { + LOG(ERROR) << "Invalid drain job header: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto job = std::make_shared(); + PublishDrainJob(*job, record); + jobs.emplace(record.id, std::move(job)); + } else if (value.kind == RecoveryRecordKind::Task) { + DurableTaskRecord record; + if (!DecodeRecoveryValue(value.payload, record) || + !ValidateDurableTask(record)) { + LOG(ERROR) << "Invalid task history or drain payload: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + tasks.push_back(FromDurableTask(record)); + } else if (value.kind == RecoveryRecordKind::RetiredReplica) { + RetiredReplicaRecord record; + if (!DecodeRecoveryValue(value.payload, record) || + !record.replica.is_memory_replica()) { + LOG(ERROR) << "Invalid retired source replica: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + if (record.release_ms <= restore_epoch_ms) continue; + const auto& descriptor = + record.replica.get_memory_descriptor().buffer_descriptor; + auto allocator = standby_allocator_keepalive_.find( + descriptor.transport_endpoint_); + if (allocator == standby_allocator_keepalive_.end()) continue; + auto owner = std::find_if( + restored_drain_segments_.begin(), + restored_drain_segments_.end(), [&](const auto& segment) { + return segment.second.endpoint == + descriptor.transport_endpoint_ || + segment.second.name == + descriptor.transport_endpoint_; + }); + if (owner == restored_drain_segments_.end()) { + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto buffer = std::make_unique(allocator->second, + descriptor); + buffer->bindSegmentLifetime(owner->second.lifetime); + buffer->bindClientLiveness( + client_liveness_records_.at(owner->second.client_id)); + std::vector replicas; + replicas.push_back(Replica(record.replica.id, std::move(buffer), + record.replica.status)); + retired.emplace_back( + std::move(replicas), + std::chrono::system_clock::time_point( + std::chrono::milliseconds(record.release_ms))); + } + } + for (const auto& [key, value] : state.records()) { + if (value.kind != RecoveryRecordKind::Unit) continue; + DrainUnitRecord unit; + if (!DecodeRecoveryValue(value.payload, unit) || + !ValidateDrainUnit(unit) || !jobs.contains(unit.job_id)) { + LOG(ERROR) << "Invalid drain unit or missing parent job: " << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + auto& job = *jobs.at(unit.job_id); + if (std::find(job.request.segments.begin(), job.request.segments.end(), + unit.source) == job.request.segments.end() || + (!job.request.target_segments.empty() && + std::find(job.request.target_segments.begin(), + job.request.target_segments.end(), + unit.target) == job.request.target_segments.end())) { + LOG(ERROR) << "Drain unit source/target outside job request: " + << key; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + if (unit.completed) job.completed_unit_keys.insert(unit.unit_key); + if (unit.terminal_failed) { + job.terminal_failed_unit_keys.insert(unit.unit_key); + } + if (unit.retries != 0) + job.retry_counts.emplace(unit.unit_key, unit.retries); + if (unit.task_id != UUID{}) { + job.active_tasks.emplace( + unit.task_id, + ActiveDrainTask{unit.task_id, TenantId(unit.tenant), unit.key, + unit.source, unit.target, unit.bytes, + unit.unit_key}); + } + } + for (const auto& [id, job] : jobs) { + uint64_t retries = 0; + for (const auto& [key, count] : job->retry_counts) retries += count; + if (job->succeeded_units != job->completed_unit_keys.size() || + job->failed_units < job->terminal_failed_unit_keys.size() || + retries > job->failed_units || + job->active_tasks.size() > job->request.max_concurrency || + (job->status >= JobStatus::SUCCEEDED && + !job->active_tasks.empty())) { + LOG(ERROR) << "Inconsistent drain job counters or active tasks: " + << UuidToString(id); + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + } + std::unordered_set owned_sources; + for (const auto& [id, job] : jobs) { + if (job->status < JobStatus::SUCCEEDED) { + owned_sources.insert(job->request.segments.begin(), + job->request.segments.end()); + } + } + for (const auto& [name, segment] : restored_drain_segments_) { + if (segment.status == static_cast(SegmentStatus::DRAINING) && + !owned_sources.contains(name)) { + LOG(WARNING) + << "Recovered orphaned DRAINING segment=" << name + << "; kept non-allocatable. Finish outstanding moves before " + "PUT /api/v1/segments/status."; + MasterMetricManager::instance().inc_orphaned_draining_restore(); + } + } + std::sort(tasks.begin(), tasks.end(), [](const Task& a, const Task& b) { + return std::tie(a.last_updated_at, a.id) < + std::tie(b.last_updated_at, b.id); + }); + { + std::lock_guard lock(discarded_replicas_mutex_); + for (auto& [replicas, release_at] : retired) { + discarded_replicas_.emplace_back(std::move(replicas), release_at); + } + } + std::lock_guard lock(job_mutex_); + auto access = task_manager_.get_write_access(); + access.clear_all(); + for (auto& task : tasks) access.restore_task(std::move(task)); + drain_jobs_ = std::move(jobs); + return {}; +} + +void MasterService::WarnOrphanedDrainingSegments() { + std::unordered_set owned; + { + std::lock_guard lock(job_mutex_); + for (const auto& [id, job] : drain_jobs_) { + std::lock_guard job_lock(job->mutex); + if (job->status < JobStatus::SUCCEEDED) { + owned.insert(job->request.segments.begin(), + job->request.segments.end()); + } + } + } + const auto warn = [&](const std::string& name, uint64_t size) { + if (!owned.contains(name)) { + LOG(WARNING) << "Orphaned DRAINING segment=" << name + << " stranded_capacity_bytes=" << size + << "; finish outstanding transfers and use PUT " + "/api/v1/segments/status to resume allocation."; + } + }; + auto access = segment_manager_.getSegmentAccess(); + std::vector> mounted; + access.GetAllSegments(mounted); + std::unordered_set names; + for (const auto& [segment, client] : mounted) { + names.insert(segment.name); + SegmentStatus status; + access.GetSegmentStatusById(segment.id, status); + if (status == SegmentStatus::DRAINING) warn(segment.name, segment.size); + } + for (const auto& [name, segment] : restored_drain_segments_) { + if (!names.contains(name) && + segment.status == static_cast(SegmentStatus::DRAINING)) { + warn(name, segment.size); + } + } +} + +void MasterService::ForgetRestoredSegments( + const std::vector& segments) { + if (segments.empty()) return; + std::unique_lock lock(master_state_mutex_); + for (const auto& segment : segments) { + if (auto accounted = standby_accounted_memory_bytes_.find(segment.name); + accounted != standby_accounted_memory_bytes_.end()) { + MasterMetricManager::instance().dec_allocated_mem_size( + segment.name, static_cast(accounted->second)); + standby_accounted_memory_bytes_.erase(accounted); + } + standby_allocator_keepalive_.erase(segment.name); + standby_allocator_keepalive_.erase(segment.endpoint); + invalid_replica_endpoints_.erase(segment.name); + invalid_replica_endpoints_.erase(segment.endpoint); + std::erase_if(standby_memory_segments_, [&](const auto& candidate) { + return candidate.segment_name == segment.name && + candidate.transport_endpoint == segment.endpoint; + }); + } +} + +} // namespace mooncake diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp index 04a85e5b9e..f1775f8b12 100644 --- a/mooncake-store/src/segment.cpp +++ b/mooncake-store/src/segment.cpp @@ -317,6 +317,14 @@ ErrorCode ScopedSegmentAccess::ReMountSegment( if (validation != ErrorCode::OK) { return validation; } + const auto existing = + segment_manager_->mounted_segments_.find(segment.id); + if (existing != segment_manager_->mounted_segments_.end() && + (existing->second.status == SegmentStatus::DRAINING || + existing->second.status == SegmentStatus::DRAINED)) { + // Ownership and registration geometry were validated above. + continue; + } ErrorCode err = MountSegment(segment, client_id, client_liveness); if (err == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS || err == ErrorCode::INTERNAL_ERROR) { diff --git a/mooncake-store/src/task_manager.cpp b/mooncake-store/src/task_manager.cpp index 9ae2657045..54125a48c8 100644 --- a/mooncake-store/src/task_manager.cpp +++ b/mooncake-store/src/task_manager.cpp @@ -5,6 +5,40 @@ namespace mooncake { +DurableTaskRecord ToDurableTask(const Task& task, bool removed) { + const auto milliseconds = [](auto time) { + return std::chrono::duration_cast( + time.time_since_epoch()) + .count(); + }; + return {task.id, + task.assigned_client, + static_cast(task.type), + static_cast(task.status), + task.payload, + task.message, + milliseconds(task.created_at), + milliseconds(task.last_updated_at), + task.max_retry_attempts, + task.drain, + removed}; +} + +Task FromDurableTask(const DurableTaskRecord& record) { + return {.id = record.id, + .type = static_cast(record.type), + .status = static_cast(record.status), + .payload = record.payload, + .created_at = std::chrono::system_clock::time_point( + std::chrono::milliseconds(record.created_ms)), + .last_updated_at = std::chrono::system_clock::time_point( + std::chrono::milliseconds(record.updated_ms)), + .message = record.message, + .assigned_client = record.client_id, + .max_retry_attempts = record.max_retries, + .drain = record.drain}; +} + std::optional ScopedTaskReadAccess::find_task_by_id( const UUID& task_id) const { auto it = manager_->all_tasks_.find(task_id); @@ -27,7 +61,8 @@ size_t ScopedTaskReadAccess::size() const { } tl::expected ScopedTaskWriteAccess::submit_task( - const UUID& client_id, TaskType type, const std::string& payload) { + const UUID& client_id, TaskType type, const std::string& payload, + std::optional drain) { if (manager_->total_pending_tasks_ >= manager_->max_total_pending_tasks_) { LOG(ERROR) << "Cannot submit new task: pending task limit reached (" << manager_->total_pending_tasks_ << "/" @@ -47,7 +82,12 @@ tl::expected ScopedTaskWriteAccess::submit_task( .last_updated_at = now, .message = "", .assigned_client = client_id, - .max_retry_attempts = manager_->max_retry_attempts_}; + .max_retry_attempts = manager_->max_retry_attempts_, + .drain = std::move(drain)}; + if (task.drain) task.drain->task_id = task.id; + if (const auto error = manager_->Persist(task); error != ErrorCode::OK) { + return tl::make_unexpected(error); + } manager_->total_pending_tasks_++; manager_->all_tasks_[task.id] = task; manager_->pending_tasks_[client_id].push(task.id); @@ -58,54 +98,50 @@ tl::expected ScopedTaskWriteAccess::submit_task( std::vector ScopedTaskWriteAccess::pop_tasks(const UUID& client_id, size_t batch_size) { std::vector result; - - auto pit = manager_->pending_tasks_.find(client_id); - if (pit == manager_->pending_tasks_.end()) { - return result; - } - - auto& queue = pit->second; - auto& processing_set = manager_->processing_tasks_[client_id]; - + auto pending = manager_->pending_tasks_.find(client_id); + if (pending == manager_->pending_tasks_.end()) return result; + auto& queue = pending->second; + const size_t available = manager_->max_total_processing_tasks_ > + manager_->total_processing_tasks_ + ? manager_->max_total_processing_tasks_ - + manager_->total_processing_tasks_ + : 0; + batch_size = std::min(batch_size, available); while (!queue.empty() && result.size() < batch_size) { - if (manager_->total_processing_tasks_ >= - manager_->max_total_processing_tasks_) { - break; - } - - const UUID task_id = queue.front(); + const auto id = queue.front(); queue.pop(); - - auto it = manager_->all_tasks_.find(task_id); - if (it == manager_->all_tasks_.end()) { - // Terminal tasks can be pruned before their queued IDs are - // consumed. - VLOG(1) << "Skipping stale pending queue entry for task " - << task_id; + auto task = manager_->all_tasks_.find(id); + if (task == manager_->all_tasks_.end() || + task->second.status != TaskStatus::PENDING) continue; + auto next = task->second; + next.mark_processing(); + result.push_back(std::move(next)); + } + if (result.empty()) return result; + if (manager_->persistence_) { + std::vector records; + records.reserve(result.size()); + for (const auto& task : result) records.push_back(ToDurableTask(task)); + if (manager_->persistence_(records) != ErrorCode::OK) { + // Restore the selected prefix before the untouched queue tail. + std::queue retry; + for (const auto& task : result) retry.push(task.id); + while (!queue.empty()) { + retry.push(queue.front()); + queue.pop(); + } + queue.swap(retry); + return {}; } - - Task& task = it->second; - if (task.status != TaskStatus::PENDING) { - continue; - } - if (manager_->total_pending_tasks_ > 0) { - manager_->total_pending_tasks_--; - } - task.mark_processing(); - - const auto [_, inserted] = processing_set.insert(task_id); - if (!inserted) { - LOG(WARNING) << "Task " << task_id - << " is already in processing set for client " - << client_id; - } else { - manager_->total_processing_tasks_++; - } - - result.push_back(task); } - + auto& processing = manager_->processing_tasks_[client_id]; + for (const auto& task : result) { + manager_->all_tasks_.at(task.id) = task; + --manager_->total_pending_tasks_; + ++manager_->total_processing_tasks_; + processing.insert(task.id); + } return result; } @@ -139,11 +175,16 @@ ErrorCode ScopedTaskWriteAccess::complete_task(const UUID& client_id, return ErrorCode::OK; } + Task next = task; + next.mark_complete(status, message); + if (auto error = manager_->Persist(next); error != ErrorCode::OK) + return error; + if (task.status == TaskStatus::PENDING && manager_->total_pending_tasks_ > 0) { manager_->total_pending_tasks_--; } - task.mark_complete(status, message); + task = std::move(next); auto ps_it = manager_->processing_tasks_.find(client_id); if (ps_it != manager_->processing_tasks_.end()) { @@ -163,6 +204,10 @@ void ScopedTaskWriteAccess::prune_finished_tasks() { while (manager_->finished_task_history_.size() > manager_->max_total_finished_tasks_) { UUID oldest_task_id = manager_->finished_task_history_.front(); + auto task = manager_->all_tasks_.find(oldest_task_id); + if (task != manager_->all_tasks_.end() && + manager_->Persist(task->second, true) != ErrorCode::OK) + break; manager_->finished_task_history_.pop_front(); manager_->all_tasks_.erase(oldest_task_id); } @@ -197,7 +242,13 @@ void ScopedTaskWriteAccess::prune_expired_tasks() { } if (now - task.created_at > pending_timeout_duration) { - task.mark_complete(TaskStatus::FAILED, "pending timeout"); + Task next = task; + next.mark_complete(TaskStatus::FAILED, "pending timeout"); + if (manager_->Persist(next) != ErrorCode::OK) { + keep_queue.push(task_id); + continue; + } + task = std::move(next); if (manager_->total_pending_tasks_ > 0) { manager_->total_pending_tasks_--; } @@ -242,8 +293,11 @@ void ScopedTaskWriteAccess::prune_expired_tasks() { } if (now - task.last_updated_at > processing_timeout) { - task.mark_complete(TaskStatus::FAILED, + Task next = task; + next.mark_complete(TaskStatus::FAILED, "processing timeout"); + if (manager_->Persist(next) != ErrorCode::OK) continue; + task = std::move(next); to_remove.push_back(task_id); if (manager_->total_processing_tasks_ > 0) { manager_->total_processing_tasks_--; @@ -297,7 +351,11 @@ ErrorCode ScopedTaskWriteAccess::fail_task_if_pending( return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; } - task.mark_complete(TaskStatus::FAILED, message); + Task next = task; + next.mark_complete(TaskStatus::FAILED, message); + if (auto error = manager_->Persist(next); error != ErrorCode::OK) + return error; + task = std::move(next); if (manager_->total_pending_tasks_ > 0) { manager_->total_pending_tasks_--; } diff --git a/mooncake-store/tests/ha/master_service_ha_test.cpp b/mooncake-store/tests/ha/master_service_ha_test.cpp index 367c50d53c..4ebdd5c6c7 100644 --- a/mooncake-store/tests/ha/master_service_ha_test.cpp +++ b/mooncake-store/tests/ha/master_service_ha_test.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -15,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -32,6 +34,8 @@ #include "ha/oplog/oplog_batch_types.h" #include "ha/oplog/oplog_applier.h" #include "ha/oplog/ordered_oplog_writer.h" +#include "ha/standby_metadata_store.h" +#include "master_metric_manager.h" #include "types.h" namespace mooncake::test { @@ -43,6 +47,7 @@ class FakeBatchHaKvBackend : public HaKvBackend { } ErrorCode Get(std::string_view key, std::string& value) override { + std::lock_guard lock(kvs_mutex_); auto it = kvs_.find(std::string(key)); if (it == kvs_.end()) { return ErrorCode::ETCD_KEY_NOT_EXIST; @@ -52,12 +57,14 @@ class FakeBatchHaKvBackend : public HaKvBackend { } ErrorCode Put(std::string_view key, std::string_view value) override { + std::lock_guard lock(kvs_mutex_); kvs_[std::string(key)] = std::string(value); return ErrorCode::OK; } ErrorCode Range(std::string_view begin_key, std::string_view end_key, size_t limit, std::vector& kvs) override { + std::lock_guard lock(kvs_mutex_); kvs.clear(); for (auto it = kvs_.lower_bound(std::string(begin_key)); it != kvs_.end() && it->first < end_key; ++it) { @@ -72,6 +79,7 @@ class FakeBatchHaKvBackend : public HaKvBackend { bool SupportsTxn() const override { return true; } ErrorCode Txn(const KvTxn& txn) override { + std::lock_guard lock(kvs_mutex_); for (const auto& compare : txn.compares) { auto it = kvs_.find(compare.key); if (compare.kind == KvCompareKind::kKeyNotExists) { @@ -90,6 +98,7 @@ class FakeBatchHaKvBackend : public HaKvBackend { } private: + std::mutex kvs_mutex_; std::map kvs_; }; @@ -538,6 +547,69 @@ class MasterServiceHATest : public ::testing::Test { ASSERT_EQ(ErrorCode::OK, read_err); } + static MasterServiceConfig DrainRecoveryConfig(const std::string& cluster) { + return MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster) + .set_default_kv_lease_ttl(0) + .set_put_start_release_timeout_sec(3600) + .set_pending_task_timeout_sec(0) + .set_processing_task_timeout_sec(0) + .set_max_retry_attempts(5) + .build(); + } + + static void StopDrainWorkers(MasterService& service) { + MasterServiceTestPeer(service).StopJobDispatchForTesting(); + MasterServiceTestPeer::EvictionRunning(service) = false; + auto& eviction = MasterServiceTestPeer::EvictionThread(service); + if (eviction.joinable()) eviction.join(); + } + + static ::testing::AssertionResult ReplayAndPromoteDrain( + MasterService& primary, MasterService& promoted, + const std::shared_ptr& backend, + const std::string& cluster) { + primary.StopBatchOpLogWriter(); + StandbyMetadataStore metadata; + OpLogApplier applier(&metadata, cluster); + OpLogBatchStandbyReader reader(cluster, *backend, applier); + const auto replay = reader.PollOnce(); + if (replay.error != ErrorCode::OK || + !reader.GetLastAppliedDurablePrefix()) { + return ::testing::AssertionFailure() + << "batch replay: " << toString(replay.error); + } + ReplicaID max_replica_id = 0; + if (!metadata.ValidateReplicaIds(max_replica_id)) { + return ::testing::AssertionFailure() << "invalid replica IDs"; + } + const auto restored = promoted.RestoreFromBatchOpLogPromotion( + metadata, *reader.GetLastAppliedDurablePrefix(), + applier.GetSegmentRegistry().GetAllSegments(), max_replica_id, 1); + if (!restored) { + return ::testing::AssertionFailure() + << "promotion: " << toString(restored.error()); + } + const auto error = + MasterServiceTestPeer(promoted).SetBatchOpLogBackendForTesting( + backend); + return error == ErrorCode::OK + ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() << toString(error); + } + + static tl::expected CompleteDrainTask( + MasterService& service, const UUID& client, const UUID& task, + TaskStatus status) { + TaskCompleteRequest request; + request.id = task; + request.status = status; + request.message = status == TaskStatus::SUCCESS ? "copied" : "failed"; + return service.MarkTaskToComplete(client, request); + } + static GatedOrderedOpLogWriter* InstallGatedWriter( MasterService& service, std::shared_ptr backend) { MasterServiceTestPeer(service).SetBatchOpLogWriterFactoryForTesting( @@ -5416,6 +5488,9 @@ TEST_F(MasterServiceHATest, CopyEndVisibleBeforeBatchRecordDurable) { {"batch_copy_visible_dst"}) .has_value()); + ReadBatchEventually(storage, 4, batch); + ASSERT_EQ(OpType::REPLICATION_START, batch.entries.front().op_type); + backend->BlockTxn(); auto copy_future = std::async(std::launch::async, [&] { return service.CopyEnd(source.client_id, key, kDefaultTenant); @@ -5431,15 +5506,15 @@ TEST_F(MasterServiceHATest, CopyEndVisibleBeforeBatchRecordDurable) { DurablePrefix prefix; ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); - EXPECT_EQ(3u, prefix.batch_id); - EXPECT_EQ(3u, prefix.last_seq); + EXPECT_EQ(4u, prefix.batch_id); + EXPECT_EQ(4u, prefix.last_seq); auto replicas = service.GetReplicaList(key, kDefaultTenant); ASSERT_TRUE(replicas.has_value()); EXPECT_EQ(2u, replicas->replicas.size()); backend->AllowTxn(); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 5, batch); ASSERT_EQ(1u, batch.entries.size()); EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); EXPECT_EQ(key, batch.entries[0].object_key); @@ -5483,6 +5558,9 @@ TEST_F(MasterServiceHATest, "batch_move_visible_dst") .has_value()); + ReadBatchEventually(storage, 4, batch); + ASSERT_EQ(OpType::REPLICATION_START, batch.entries.front().op_type); + backend->BlockTxn(); auto move_future = std::async(std::launch::async, [&] { return service.MoveEnd(source.client_id, key, kDefaultTenant); @@ -5498,8 +5576,8 @@ TEST_F(MasterServiceHATest, DurablePrefix prefix; ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); - EXPECT_EQ(3u, prefix.batch_id); - EXPECT_EQ(3u, prefix.last_seq); + EXPECT_EQ(4u, prefix.batch_id); + EXPECT_EQ(4u, prefix.last_seq); auto replicas = service.GetReplicaList(key, kDefaultTenant); ASSERT_TRUE(replicas.has_value()); @@ -5514,7 +5592,7 @@ TEST_F(MasterServiceHATest, EXPECT_FALSE(before_finalize.has_value()); backend->AllowTxn(); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 5, batch); ASSERT_EQ(1u, batch.entries.size()); EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); EXPECT_EQ(key, batch.entries[0].object_key); @@ -5599,7 +5677,7 @@ TEST_F(MasterServiceHATest, ASSERT_TRUE(upsert.has_value()) << toString(upsert.error()); backend->AllowTxn(); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 5, batch); uint64_t charged_bytes = 2 * object_size; for (int i = 0; i < 50 && charged_bytes == 2 * object_size; ++i) { charged_bytes = @@ -6257,6 +6335,7 @@ TEST_F(MasterServiceHATest, BatchReplicaClearSegmentWritesBatchRecordOpLog) { ASSERT_TRUE(copy_start.has_value()); auto copy_end = service.CopyEnd(mounted.client_id, key, kDefaultTenant); ASSERT_TRUE(copy_end.has_value()); + ReadBatchEventually(storage, 5, batch); std::this_thread::sleep_for(std::chrono::milliseconds(60)); auto res = @@ -6264,14 +6343,20 @@ TEST_F(MasterServiceHATest, BatchReplicaClearSegmentWritesBatchRecordOpLog) { ASSERT_TRUE(res.has_value()); ASSERT_EQ(1u, res->size()); EXPECT_EQ(key, (*res)[0]); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 6, batch); ASSERT_EQ(1u, batch.entries.size()); EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); EXPECT_EQ(key, batch.entries[0].object_key); - EXPECT_EQ(4u, batch.entries[0].sequence_id); + EXPECT_EQ(6u, batch.entries[0].sequence_id); EXPECT_FALSE(batch.entries[0].payload.empty()); + MetadataPayload cleared; + ASSERT_TRUE(DecodeRecoveryValue(batch.entries[0].payload, cleared)); + ASSERT_EQ(1u, cleared.replicas.size()); + EXPECT_EQ("batch_clear_seg2", cleared.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_); } TEST_F(MasterServiceHATest, BatchReplicaClearAllReleasesAfterDurable) { @@ -6378,7 +6463,7 @@ TEST_F(MasterServiceHATest, BatchReplicaClearSegmentReleasesAfterDurable) { ASSERT_TRUE(copy_start.has_value()); ASSERT_TRUE( service.CopyEnd(mounted.client_id, key, kDefaultTenant).has_value()); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 5, batch); ASSERT_TRUE(writer->PauseCallbacksAfter(batch.last_seq)); std::this_thread::sleep_for(std::chrono::milliseconds(60)); @@ -6397,7 +6482,7 @@ TEST_F(MasterServiceHATest, BatchReplicaClearSegmentReleasesAfterDurable) { EXPECT_FALSE(before_finalize.has_value()); backend->AllowTxn(); - ReadBatchEventually(storage, 5, batch); + ReadBatchEventually(storage, 6, batch); EXPECT_EQ(2048, TenantUsedBytes(service)); ASSERT_TRUE(writer->RunCallbacksThrough(batch.last_seq)); EXPECT_EQ(1024, TenantUsedBytes(service)); @@ -6956,14 +7041,17 @@ TEST_F(MasterServiceHATest, MasterServiceTestPeer(service).RunBatchEvictForTesting( /*evict_ratio_target=*/1.0, /*evict_ratio_lowerbound=*/1.0); - ReadBatchEventually(storage, 4, batch); + ReadBatchEventually(storage, 5, batch); ASSERT_EQ(1u, batch.entries.size()); EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); EXPECT_EQ(key, batch.entries[0].object_key); - EXPECT_EQ(4u, batch.entries[0].sequence_id); + EXPECT_EQ(5u, batch.entries[0].sequence_id); EXPECT_FALSE(batch.entries[0].payload.empty()); + MetadataPayload settled; + ASSERT_TRUE(DecodeRecoveryValue(batch.entries[0].payload, settled)); + EXPECT_TRUE(settled.replication_finished.value_or(false)); } TEST_F(MasterServiceHATest, WeightLeaseShorterTtlRenewalReplaysThroughReader) { @@ -7155,6 +7243,647 @@ TEST_F(MasterServiceHATest, WeightLeaseRenewalRechecksExpiryAfterGroupWait) { EXPECT_EQ(WeightManagementError::LEASE_EXPIRED, renewed.error()); } +TEST_F(MasterServiceHATest, DrainJobSurvivesBatchReplayAndPromotion) { + const std::string cluster = "drain_recovery"; + auto backend = std::make_shared(); + MasterService primary(MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster) + .build()); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + PrepareSimpleSegment(primary, "drain-source"); + auto created = primary.CreateDrainJob({{"drain-source"}, {}, 1}); + ASSERT_TRUE(created.has_value()); + primary.StopBatchOpLogWriter(); + + StandbyMetadataStore metadata; + OpLogApplier applier(&metadata, cluster); + OpLogBatchStandbyReader reader(cluster, *backend, applier); + auto replay = reader.PollOnce(); + ASSERT_EQ(ErrorCode::OK, replay.error); + ASSERT_TRUE(reader.GetLastAppliedDurablePrefix()); + MasterService promoted; + ASSERT_TRUE(promoted.RestoreFromBatchOpLogPromotion( + metadata, *reader.GetLastAppliedDurablePrefix(), + applier.GetSegmentRegistry().GetAllSegments(), 0, 1)); + auto restored = promoted.QueryDrainJob(*created); + ASSERT_TRUE(restored.has_value()); + EXPECT_EQ(*created, restored->id); + EXPECT_EQ(std::vector{"drain-source"}, restored->segments); +} + +TEST_F(MasterServiceHATest, + DrainRecoveryResumesPendingAndProcessingTasksWithoutDoubleCounting) { + const std::string cluster = "drain_recovery_progress"; + const auto config = DrainRecoveryConfig(cluster); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + const auto source = MakeSegment("drain-source"); + const auto target = + MakeSegment("drain-target", kDefaultSegmentBase + kDefaultSegmentSize); + const auto source_client = generate_uuid(); + const auto target_client = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(source, source_client)); + ASSERT_TRUE(primary.MountSegment(target, target_client)); + PutObjectOnSegment(primary, source_client, "first", source.name, 1024); + PutObjectOnSegment(primary, source_client, "second", source.name, 2048); + const auto job = primary.CreateDrainJob({{source.name}, {target.name}, 2}); + ASSERT_TRUE(job); + MasterServiceTestPeer(primary).ProcessDrainJobsForTesting(); + std::vector original_tasks; + { + const auto tasks = + MasterServiceTestPeer::TaskManager(primary).get_read_access(); + for (const auto& [id, task] : tasks) original_tasks.push_back(task); + } + ASSERT_EQ(2u, original_tasks.size()); + const auto fetched = primary.FetchTasks(source_client, 1); + ASSERT_TRUE(fetched); + ASSERT_EQ(1u, fetched->size()); + const auto processing_id = fetched->front().id; + ReplicaMovePayload processing_payload; + struct_json::from_json(processing_payload, fetched->front().payload); + ASSERT_TRUE(primary.MoveStart(source_client, processing_payload.key, + kDefaultTenant, source.name, target.name)); + + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + auto progress = promoted.QueryDrainJob(*job); + ASSERT_TRUE(progress); + EXPECT_EQ(2u, progress->active_units); + EXPECT_EQ(0u, progress->succeeded_units); + EXPECT_EQ(0u, progress->migrated_bytes); + for (const auto& original : original_tasks) { + const auto task = promoted.QueryTask(original.id); + ASSERT_TRUE(task); + EXPECT_EQ(original.id, task->id); + EXPECT_EQ(source_client, task->assigned_client); + EXPECT_EQ(original.id == processing_id ? TaskStatus::PROCESSING + : TaskStatus::PENDING, + task->status); + const auto restored = MasterServiceTestPeer::TaskManager(promoted) + .get_read_access() + .find_task_by_id(original.id); + ASSERT_TRUE(restored); + EXPECT_EQ(original.payload, restored->payload); + EXPECT_EQ(original.max_retry_attempts, restored->max_retry_attempts); + } + ASSERT_TRUE(promoted.MountSegment(source, source_client)); + ASSERT_TRUE(promoted.MountSegment(target, target_client)); + ASSERT_EQ(SegmentStatus::DRAINING, + promoted.QuerySegmentStatus(source.name).value()); + const auto canceled = promoted.CancelDrainJob(*job); + ASSERT_FALSE(canceled); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, canceled.error()); + // The in-flight MOVE must finish with the same task and replica IDs. + ASSERT_TRUE(promoted.MoveEnd(source_client, processing_payload.key, + kDefaultTenant)); + ASSERT_TRUE(CompleteDrainTask(promoted, source_client, processing_id, + TaskStatus::SUCCESS)); + MasterServiceTestPeer(promoted).ProcessDrainJobsForTesting(); + const auto pending = promoted.FetchTasks(source_client, 8); + ASSERT_TRUE(pending); + ASSERT_EQ(1u, pending->size()); + EXPECT_NE(processing_id, pending->front().id); + ReplicaMovePayload pending_payload; + struct_json::from_json(pending_payload, pending->front().payload); + ASSERT_TRUE(promoted.MoveStart(source_client, pending_payload.key, + kDefaultTenant, source.name, target.name)); + ASSERT_TRUE( + promoted.MoveEnd(source_client, pending_payload.key, kDefaultTenant)); + ASSERT_TRUE(CompleteDrainTask(promoted, source_client, pending->front().id, + TaskStatus::SUCCESS)); + MasterServiceTestPeer(promoted).ProcessDrainJobsForTesting(); + progress = promoted.QueryDrainJob(*job); + ASSERT_TRUE(progress); + EXPECT_EQ(JobStatus::SUCCEEDED, progress->status); + EXPECT_EQ(2u, progress->succeeded_units); + EXPECT_EQ(3072u, progress->migrated_bytes); + EXPECT_EQ(0u, progress->active_units); + + MasterService promoted_again(config); + StopDrainWorkers(promoted_again); + ASSERT_TRUE( + ReplayAndPromoteDrain(promoted, promoted_again, backend, cluster)); + ASSERT_TRUE(promoted_again.MountSegment(source, source_client)); + ASSERT_TRUE(promoted_again.MountSegment(target, target_client)); + MasterServiceTestPeer(promoted_again).ProcessDrainJobsForTesting(); + const auto final = promoted_again.QueryDrainJob(*job); + ASSERT_TRUE(final); + EXPECT_EQ(JobStatus::SUCCEEDED, final->status); + EXPECT_EQ(2u, final->succeeded_units); + EXPECT_EQ(3072u, final->migrated_bytes); + EXPECT_EQ(0u, final->failed_units); + EXPECT_EQ(0u, final->active_units); + EXPECT_EQ(SegmentStatus::DRAINED, + promoted_again.QuerySegmentStatus(source.name).value()); + for (const auto& key : {"first", "second"}) { + const auto replicas = + promoted_again.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas); + ASSERT_EQ(1u, replicas->replicas.size()); + EXPECT_EQ(target.te_endpoint, + replicas->replicas.front() + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_); + } + for (const auto& original : original_tasks) { + const auto task = promoted_again.QueryTask(original.id); + ASSERT_TRUE(task); + EXPECT_EQ(TaskStatus::SUCCESS, task->status); + EXPECT_EQ("copied", task->message); + } + ASSERT_TRUE(promoted_again.FetchTasks(source_client, 8)->empty()); +} + +TEST_F(MasterServiceHATest, DrainRecoveryRetainsRetryBudgetAndCanceledJobs) { + const std::string cluster = "drain_recovery_retries"; + const auto config = DrainRecoveryConfig(cluster); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + const auto source = MakeSegment("retry-source"); + const auto target = + MakeSegment("retry-target", kDefaultSegmentBase + kDefaultSegmentSize); + const auto owner = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(source, owner)); + ASSERT_TRUE(primary.MountSegment(target, owner)); + PutObjectOnSegment(primary, owner, "retry-key", source.name); + const auto job = primary.CreateDrainJob({{source.name}, {target.name}, 1}); + ASSERT_TRUE(job); + MasterServiceTestPeer(primary).ProcessDrainJobsForTesting(); + auto tasks = primary.FetchTasks(owner, 1); + ASSERT_TRUE(tasks); + ASSERT_EQ(1u, tasks->size()); + const auto first_id = tasks->front().id; + ASSERT_TRUE( + CompleteDrainTask(primary, owner, first_id, TaskStatus::FAILED)); + MasterServiceTestPeer(primary).ProcessDrainJobsForTesting(); + tasks = primary.FetchTasks(owner, 1); + ASSERT_TRUE(tasks); + ASSERT_EQ(1u, tasks->size()); + auto active_id = tasks->front().id; + EXPECT_NE(first_id, active_id); + + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + ASSERT_TRUE(promoted.MountSegment(source, owner)); + ASSERT_TRUE(promoted.MountSegment(target, owner)); + auto progress = promoted.QueryDrainJob(*job); + ASSERT_TRUE(progress); + EXPECT_EQ(1u, progress->failed_units); + EXPECT_EQ(1u, progress->active_units); + ASSERT_TRUE(promoted.QueryTask(active_id)); + EXPECT_EQ(TaskStatus::PROCESSING, promoted.QueryTask(active_id)->status); + for (int attempt = 2; attempt <= 3; ++attempt) { + ASSERT_TRUE( + CompleteDrainTask(promoted, owner, active_id, TaskStatus::FAILED)); + MasterServiceTestPeer(promoted).ProcessDrainJobsForTesting(); + progress = promoted.QueryDrainJob(*job); + ASSERT_TRUE(progress); + EXPECT_EQ(static_cast(attempt), progress->failed_units); + tasks = promoted.FetchTasks(owner, 1); + ASSERT_TRUE(tasks); + if (attempt == 2) { + ASSERT_EQ(1u, tasks->size()); + EXPECT_NE(active_id, tasks->front().id); + active_id = tasks->front().id; + } else { + EXPECT_TRUE(tasks->empty()); + EXPECT_EQ(JobStatus::FAILED, progress->status); + EXPECT_EQ(0u, progress->active_units); + } + } + EXPECT_EQ(SegmentStatus::OK, + promoted.QuerySegmentStatus(source.name).value()); + const auto canceled = + promoted.CreateDrainJob({{source.name}, {target.name}, 1}); + ASSERT_TRUE(canceled); + ASSERT_TRUE(promoted.CancelDrainJob(*canceled)); + + MasterService promoted_again(config); + StopDrainWorkers(promoted_again); + ASSERT_TRUE( + ReplayAndPromoteDrain(promoted, promoted_again, backend, cluster)); + ASSERT_TRUE(promoted_again.MountSegment(source, owner)); + ASSERT_TRUE(promoted_again.MountSegment(target, owner)); + MasterServiceTestPeer(promoted_again).ProcessDrainJobsForTesting(); + const auto failed_job = promoted_again.QueryDrainJob(*job); + const auto canceled_job = promoted_again.QueryDrainJob(*canceled); + ASSERT_TRUE(failed_job); + ASSERT_TRUE(canceled_job); + EXPECT_EQ(JobStatus::FAILED, failed_job->status); + EXPECT_EQ(3u, failed_job->failed_units); + EXPECT_EQ(JobStatus::CANCELED, canceled_job->status); + EXPECT_EQ(0u, canceled_job->active_units); + EXPECT_EQ(SegmentStatus::OK, + promoted_again.QuerySegmentStatus(source.name).value()); + EXPECT_TRUE(promoted_again.FetchTasks(owner, 8)->empty()); +} + +TEST_F(MasterServiceHATest, + DrainRecoveryTreatsPrunedActiveTaskAsTerminalFailure) { + const std::string cluster = "drain_recovery_missing_task"; + auto config = DrainRecoveryConfig(cluster); + config.task_manager_config.max_total_finished_tasks = 0; + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + const auto source = MakeSegment("pruned-source"); + const auto target = + MakeSegment("pruned-target", kDefaultSegmentBase + kDefaultSegmentSize); + const auto owner = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(source, owner)); + ASSERT_TRUE(primary.MountSegment(target, owner)); + PutObjectOnSegment(primary, owner, "pruned-key", source.name); + const auto job = primary.CreateDrainJob({{source.name}, {target.name}, 1}); + ASSERT_TRUE(job); + MasterServiceTestPeer(primary).ProcessDrainJobsForTesting(); + const auto tasks = primary.FetchTasks(owner, 1); + ASSERT_TRUE(tasks); + ASSERT_EQ(1u, tasks->size()); + ASSERT_TRUE(CompleteDrainTask(primary, owner, tasks->front().id, + TaskStatus::FAILED)); + MasterServiceTestPeer::TaskManager(primary) + .get_write_access() + .prune_finished_tasks(); + ASSERT_FALSE(primary.QueryTask(tasks->front().id)); + + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + ASSERT_TRUE(promoted.MountSegment(source, owner)); + ASSERT_TRUE(promoted.MountSegment(target, owner)); + MasterServiceTestPeer(promoted).ProcessDrainJobsForTesting(); + MasterServiceTestPeer(promoted).ProcessDrainJobsForTesting(); + const auto failed = promoted.QueryDrainJob(*job); + ASSERT_TRUE(failed); + EXPECT_EQ(JobStatus::FAILED, failed->status); + EXPECT_EQ(1u, failed->failed_units); + EXPECT_EQ(0u, failed->active_units); + EXPECT_TRUE(promoted.FetchTasks(owner, 8)->empty()); + EXPECT_EQ(SegmentStatus::OK, + promoted.QuerySegmentStatus(source.name).value()); +} + +TEST_F(MasterServiceHATest, + DrainRecoveryReportsOrphanAndRequiresExactRemountWithEmptyName) { + const std::string cluster = "drain_recovery_orphan"; + const auto config = DrainRecoveryConfig(cluster); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + auto source = MakeSegment(""); + source.te_endpoint = "empty-name-source-endpoint"; + const auto owner = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(source, owner)); + ASSERT_TRUE(primary.SetSegmentStatus(source.name, SegmentStatus::DRAINING)); + const auto orphan_count = [] { + std::istringstream metrics( + MasterMetricManager::instance().serialize_metrics()); + for (std::string line; std::getline(metrics, line);) { + std::istringstream sample(line); + std::string name; + uint64_t value = 0; + if (sample >> name >> value && + name == "master_orphaned_draining_restore_total") + return value; + } + return uint64_t{0}; + }; + const auto before = orphan_count(); + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + EXPECT_EQ(before + 1, orphan_count()); + const auto status = promoted.QuerySegmentStatus(source.name); + ASSERT_TRUE(status); + EXPECT_EQ(SegmentStatus::DRAINING, *status); + auto wrong_owner = promoted.MountSegment(source, generate_uuid()); + ASSERT_FALSE(wrong_owner); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, wrong_owner.error()); + auto wrong_geometry = source; + wrong_geometry.base += 4096; + const auto remapped = promoted.MountSegment(wrong_geometry, owner); + ASSERT_FALSE(remapped); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, remapped.error()); + ASSERT_TRUE(promoted.MountSegment(source, owner)); + EXPECT_EQ(SegmentStatus::DRAINING, + promoted.QuerySegmentStatus(source.name).value()); + ReplicateConfig placement; + placement.replica_num = 1; + EXPECT_FALSE( + promoted.PutStart(owner, "blocked", kDefaultTenant, 1024, placement)); + ASSERT_TRUE(promoted.SetSegmentStatus(source.name, SegmentStatus::OK)); + ASSERT_TRUE( + promoted.PutStart(owner, "resumed", kDefaultTenant, 1024, placement)); + ASSERT_TRUE( + promoted.PutEnd(owner, "resumed", kDefaultTenant, ReplicaType::MEMORY)); +} + +TEST_F(MasterServiceHATest, + DrainRecoveryRestoresCopyMoveQuotaAndSourcePinsUntilCompletionOrRevoke) { + const std::string cluster = "drain_recovery_replication"; + auto config = DrainRecoveryConfig(cluster); + config.enable_multi_tenants = true; + config.tenant_quota_connector_type = "file"; + config.tenant_quota_connector_uri = WriteTenantPolicyFile( + {{kDefaultTenant.value(), kStrictTenantQuotaBytes}}); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + const auto source = MakeSegment("runtime-source"); + const auto target = MakeSegment("runtime-target", + kDefaultSegmentBase + kDefaultSegmentSize); + const auto owner = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(source, owner)); + ASSERT_TRUE(primary.MountSegment(target, owner)); + for (const auto& key : {"copy-end", "copy-revoke", "move-revoke"}) { + PutObjectOnSegment(primary, owner, key, source.name); + } + ASSERT_TRUE(primary.CopyStart(owner, "copy-end", kDefaultTenant, + source.name, {target.name})); + ASSERT_TRUE(primary.CopyStart(owner, "copy-revoke", kDefaultTenant, + source.name, {target.name})); + ASSERT_TRUE(primary.MoveStart(owner, "move-revoke", kDefaultTenant, + source.name, target.name)); + // An ordinary metadata update must not settle the in-flight operation. + for (const auto& key : {"copy-end", "copy-revoke", "move-revoke"}) { + Replica disk(owner, 1024, "runtime-local-disk", + ReplicaStatus::COMPLETE); + ASSERT_TRUE(primary.AddReplica(owner, key, kDefaultTenant, disk)); + } + ASSERT_TRUE(primary.GetTenantQuotaSnapshot(kDefaultTenant)); + EXPECT_EQ(6144u, + primary.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + ASSERT_TRUE(promoted.MountSegment(source, owner)); + ASSERT_TRUE(promoted.MountSegment(target, owner)); + ASSERT_TRUE(promoted.GetTenantQuotaSnapshot(kDefaultTenant)); + EXPECT_EQ(6144u, + promoted.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + const auto check_source_pin = [&](const std::string& key, + uint32_t expected) { + MasterServiceTestPeer::MetadataAccessorRO access(&promoted, + {kDefaultTenant, key}); + ASSERT_TRUE(access.Exists()); + uint32_t pins = 0; + for (const auto& replica : access.Get().GetAllReplicas()) { + pins += replica.get_refcnt(); + } + EXPECT_EQ(expected, pins) << key; + }; + for (const auto& key : {"copy-end", "copy-revoke", "move-revoke"}) { + check_source_pin(key, 1); + } + const auto intruder = generate_uuid(); + const auto illegal = + promoted.CopyRevoke(intruder, "copy-revoke", kDefaultTenant); + ASSERT_FALSE(illegal); + EXPECT_EQ(6144u, + promoted.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + ASSERT_TRUE(promoted.CopyEnd(owner, "copy-end", kDefaultTenant)); + ASSERT_TRUE(promoted.CopyRevoke(owner, "copy-revoke", kDefaultTenant)); + ASSERT_TRUE(promoted.MoveRevoke(owner, "move-revoke", kDefaultTenant)); + EXPECT_EQ(4096u, + promoted.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + for (const auto& key : {"copy-end", "copy-revoke", "move-revoke"}) { + check_source_pin(key, 0); + const auto replicas = promoted.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas); + EXPECT_EQ(std::string(key) == "copy-end" ? 3u : 2u, + replicas->replicas.size()); + } + const auto already_revoked = + promoted.MoveRevoke(owner, "move-revoke", kDefaultTenant); + ASSERT_FALSE(already_revoked); + EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, already_revoked.error()); + check_source_pin("move-revoke", 0); + EXPECT_EQ(4096u, + promoted.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + MasterService promoted_again(config); + StopDrainWorkers(promoted_again); + ASSERT_TRUE( + ReplayAndPromoteDrain(promoted, promoted_again, backend, cluster)); + ASSERT_TRUE(promoted_again.MountSegment(source, owner)); + ASSERT_TRUE(promoted_again.MountSegment(target, owner)); + EXPECT_EQ( + 4096u, + promoted_again.GetTenantQuotaSnapshot(kDefaultTenant)->charged_bytes); + for (const auto& key : {"copy-end", "copy-revoke", "move-revoke"}) { + MasterServiceTestPeer::MetadataAccessorRO access(&promoted_again, + {kDefaultTenant, key}); + ASSERT_TRUE(access.Exists()); + EXPECT_FALSE(access.GetState().replication_task); + for (const auto& replica : access.Get().GetAllReplicas()) { + EXPECT_EQ(0u, replica.get_refcnt()); + } + } +} + +TEST_F(MasterServiceHATest, + DrainRecoveryReservesMovedSourceDuringReleaseGrace) { + const std::string cluster = "drain_recovery_source_grace"; + auto config = DrainRecoveryConfig(cluster); + config.default_kv_lease_ttl = 30000; + auto backend = std::make_shared(); + // Real, page-aligned process buffers exercise the persisted byte addresses. + std::vector source_memory(kDefaultSegmentSize + 4096); + std::vector target_memory(kDefaultSegmentSize + 4096); + const auto aligned_base = [](const auto& memory) { + return (reinterpret_cast(memory.data()) + 4095) & + ~uintptr_t{4095}; + }; + const auto source = + MakeSegment("grace-source", aligned_base(source_memory)); + const auto target = + MakeSegment("grace-target", aligned_base(target_memory)); + const auto owner = generate_uuid(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + ASSERT_TRUE(primary.MountSegment(source, owner)); + ASSERT_TRUE(primary.MountSegment(target, owner)); + PutObjectOnSegment(primary, owner, "grace-key", source.name, + kDefaultSegmentSize); + // Keep a real pre-move reader lease across promotion. + ASSERT_TRUE(primary.GetReplicaList("grace-key", kDefaultTenant)); + const auto move = primary.MoveStart(owner, "grace-key", kDefaultTenant, + source.name, target.name); + ASSERT_TRUE(move); + ASSERT_TRUE(move->target); + const auto old_address = + move->source.get_memory_descriptor().buffer_descriptor.buffer_address_; + const auto target_address = + move->target->get_memory_descriptor().buffer_descriptor.buffer_address_; + std::memset(reinterpret_cast(old_address), 0x5a, + kDefaultSegmentSize); + std::memcpy(reinterpret_cast(target_address), + reinterpret_cast(old_address), kDefaultSegmentSize); + ASSERT_TRUE(primary.MoveEnd(owner, "grace-key", kDefaultTenant)); + + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + ASSERT_TRUE(promoted.MountSegment(source, owner)); + ASSERT_TRUE(promoted.MountSegment(target, owner)); + ReplicateConfig placement; + placement.replica_num = 1; + placement.preferred_segments = {source.name}; + // Both segments are full: the target is live and the source is retained + // for existing readers, so neither range may satisfy another full-size put. + const auto replacement = promoted.PutStart( + owner, "too-early", kDefaultTenant, kDefaultSegmentSize, placement); + EXPECT_FALSE(replacement); + EXPECT_EQ(0, std::memcmp(reinterpret_cast(old_address), + reinterpret_cast(target_address), + kDefaultSegmentSize)); + const auto replicas = promoted.GetReplicaList("grace-key", kDefaultTenant); + ASSERT_TRUE(replicas); + ASSERT_EQ(1u, replicas->replicas.size()); + EXPECT_EQ(target_address, replicas->replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_); +} + +TEST_F(MasterServiceHATest, DrainRecoveryReleasesUnremountedSegmentNamespaces) { + for (bool offboard : {false, true}) { + const std::string cluster = + offboard ? "recovered_offboarding" : "recovered_unmount"; + const auto config = DrainRecoveryConfig(cluster); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ(ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting( + backend)); + auto segment = MakeSegment("unremounted-source"); + const auto owner = generate_uuid(); + ASSERT_TRUE(primary.MountSegment(segment, owner)); + PutObjectOnSegment(primary, owner, "unremounted-key", segment.name); + ASSERT_TRUE( + primary.SetSegmentStatus(segment.name, SegmentStatus::DRAINING)); + MasterService promoted(config); + StopDrainWorkers(promoted); + ASSERT_TRUE(ReplayAndPromoteDrain(primary, promoted, backend, cluster)); + if (offboard) { + const auto liveness = ClientRecordForTesting(promoted, owner); + ASSERT_TRUE(liveness); + const auto now = ClientLivenessRecord::Clock::now(); + ASSERT_EQ(ClientLivenessTransition::BECAME_SUSPECTED, + liveness->Evaluate(now, std::chrono::seconds::zero(), + std::chrono::seconds::zero())); + MasterMetricManager::instance().client_liveness_became_suspected(); + ASSERT_EQ(ClientLivenessTransition::BECAME_OFFLINE, + liveness->Evaluate(now, std::chrono::seconds::zero(), + std::chrono::seconds::zero())); + MasterMetricManager::instance().client_liveness_became_offline(); + ClientOffboardingJob job; + job.client_id = owner; + job.liveness = liveness; + ASSERT_TRUE(ProcessClientOffboardingForTesting(promoted, job)); + EXPECT_FALSE(ClientRecordForTesting(promoted, owner)); + } else { + auto illegal = promoted.UnmountSegment(segment.id, generate_uuid()); + ASSERT_FALSE(illegal); + EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, illegal.error()); + ASSERT_TRUE(promoted.UnmountSegment(segment.id, owner)); + } + EXPECT_FALSE(promoted.QuerySegmentStatusById(segment.id)); + EXPECT_FALSE( + promoted.GetReplicaListForAdmin("unremounted-key", kDefaultTenant)); + segment.id = generate_uuid(); + segment.te_endpoint = "replacement-endpoint"; + ASSERT_TRUE(promoted.MountSegment(segment, generate_uuid())); + EXPECT_EQ(SegmentStatus::OK, + promoted.QuerySegmentStatus(segment.name).value()); + } +} + +TEST_F(MasterServiceHATest, + DrainRecoveryRejectsInconsistentAssociationsWithoutPublishingState) { + const std::string cluster = "drain_recovery_invalid_associations"; + const auto config = DrainRecoveryConfig(cluster); + auto backend = std::make_shared(); + MasterService primary(config); + StopDrainWorkers(primary); + ASSERT_EQ( + ErrorCode::OK, + MasterServiceTestPeer(primary).SetBatchOpLogBackendForTesting(backend)); + const auto source = PrepareSimpleSegment(primary, "invalid-source"); + PrepareSimpleSegment(primary, "invalid-target", + kDefaultSegmentBase + kDefaultSegmentSize); + PutObjectOnSegment(primary, source.client_id, "invalid-key", + "invalid-source"); + const auto job = + primary.CreateDrainJob({{"invalid-source"}, {"invalid-target"}, 1}); + ASSERT_TRUE(job); + MasterServiceTestPeer(primary).ProcessDrainJobsForTesting(); + const auto tasks = primary.FetchTasks(source.client_id, 1); + ASSERT_TRUE(tasks); + ASSERT_EQ(1u, tasks->size()); + primary.StopBatchOpLogWriter(); + + for (const bool invalid_source : {true, false}) { + SCOPED_TRACE(invalid_source ? "unit source outside job" + : "counter mismatch"); + StandbyMetadataStore metadata; + OpLogApplier applier(&metadata, cluster); + OpLogBatchStandbyReader reader(cluster, *backend, applier); + ASSERT_EQ(ErrorCode::OK, reader.PollOnce().error); + ASSERT_TRUE(reader.GetLastAppliedDurablePrefix()); + const auto& record = metadata.DrainRecovery().records().at( + DrainRecoveryKey(RecoveryRecordKind::Job, UuidToString(*job))); + DrainJobRecord inconsistent; + ASSERT_TRUE(DecodeRecoveryValue(record.payload, inconsistent)); + if (invalid_source) + inconsistent.sources = {"unrelated-source"}; + else + inconsistent.succeeded = 1; + // Each row is valid on its own; restore must validate their relation. + ASSERT_TRUE(ValidateDrainJob(inconsistent)); + ASSERT_TRUE(metadata.DrainRecovery().Apply( + DrainJobUpdate{inconsistent, std::nullopt, {}})); + ReplicaID max_replica_id = 0; + ASSERT_TRUE(metadata.ValidateReplicaIds(max_replica_id)); + MasterService promoted(config); + StopDrainWorkers(promoted); + const auto restored = promoted.RestoreFromBatchOpLogPromotion( + metadata, *reader.GetLastAppliedDurablePrefix(), + applier.GetSegmentRegistry().GetAllSegments(), max_replica_id, 1); + EXPECT_FALSE(restored); + EXPECT_FALSE(promoted.QueryDrainJob(*job)); + EXPECT_FALSE(promoted.QueryTask(tasks->front().id)); + EXPECT_FALSE( + promoted.GetReplicaListForAdmin("invalid-key", kDefaultTenant)); + } +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp b/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp index 6c7c0820fe..c2a1589aff 100644 --- a/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp +++ b/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp @@ -11,6 +11,7 @@ #include #include "metadata_store.h" +#include "ha/oplog/oplog_batch_codec.h" #include "mock_metadata_store.h" #include "ha/oplog/oplog_types.h" #include "ha/standby_metadata_store.h" @@ -114,6 +115,25 @@ class OpLogApplierTest : public ::testing::Test { std::string cluster_id_; }; +TEST_F(OpLogApplierTest, NativeBatchPrunesExpiredRetiredMoveSources) { + StandbyMetadataStore metadata; + OpLogApplier applier(&metadata, cluster_id_); + Replica::Descriptor source{1, MemoryDescriptor{{64, 4096, "tcp", "source"}}, + ReplicaStatus::COMPLETE}; + ASSERT_TRUE(metadata.DrainRecovery().Retire({source, 1})); + ASSERT_EQ(1u, metadata.DrainRecovery().records().size()); + OpLogEntry entry = MakeEntry(1, OpType::REMOVE, "unrelated", ""); + OpLogBatchRecord batch; + batch.batch_id = batch.first_seq = batch.last_seq = 1; + batch.entries = {entry}; + OpLogBatchRecord decoded; + ASSERT_TRUE( + DecodeOpLogBatchRecord(EncodeOpLogBatchRecord(batch), &decoded)); + ASSERT_EQ(0, decoded.entries.front().timestamp_ms); + ASSERT_TRUE(applier.ApplyOpLogEntry(decoded.entries.front())); + EXPECT_TRUE(metadata.DrainRecovery().records().empty()); +} + TEST_F(OpLogApplierTest, AppliesWeightMetadataAndRejectsStaleGeneration) { auto generation_one = MakeWeightMetadata(1); const auto key = MakeWeightRevisionMetadataKey(generation_one.identity); diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp index 2698883ed2..a2650b2be2 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp @@ -100,6 +100,22 @@ TEST(BatchOpLogSnapshotCodecTest, OldObjectChunkDefaultsHardPinnedToFalse) { EXPECT_FALSE(decoded->objects[0].metadata.session_id.has_value()); } +TEST(BatchOpLogSnapshotCodecTest, RecoveryChunkChecksEnvelopeAndTruncation) { + DrainJobRecord job{.id = {1, 2}, .sources = {"source"}, .concurrency = 1}; + const auto record = + MakeRecoveryRecord(RecoveryRecordKind::Job, UuidToString(job.id), job); + auto encoded = EncodeBatchOpLogRecoveryChunk(7, {record}); + auto decoded = DecodeBatchOpLogRecoveryChunk(encoded, 7, 1); + ASSERT_TRUE(decoded) << decoded.error(); + ASSERT_EQ(1u, decoded->records.size()); + EXPECT_EQ(record.key, decoded->records[0].key); + EXPECT_EQ(record.payload, decoded->records[0].payload); + EXPECT_FALSE(DecodeBatchOpLogRecoveryChunk(encoded, 8, 1)); + EXPECT_FALSE(DecodeBatchOpLogRecoveryChunk(encoded, 7, 2)); + encoded.pop_back(); + EXPECT_FALSE(DecodeBatchOpLogRecoveryChunk(encoded, 7, 1)); +} + TEST(BatchOpLogSnapshotCodecTest, RejectsTruncatedArtifacts) { auto segments = EncodeBatchOpLogSnapshotSegments({}); auto objects = EncodeBatchOpLogSnapshotObjectChunk( diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/gc_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/gc_test.cpp index a9d40f1719..d41cf5ba67 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/gc_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/gc_test.cpp @@ -134,13 +134,13 @@ TEST(BatchOpLogSnapshotGcTest, PointerMismatchSkipsDeletion) { EXPECT_TRUE(object_store.deleted.empty()); } -TEST(BatchOpLogSnapshotGcTest, MissingWeightArtifactSkipsDeletion) { +TEST(BatchOpLogSnapshotGcTest, MissingWeightOrRecoveryArtifactSkipsDeletion) { FakeBackend backend; RecordingObjectStore object_store; auto lease = SnapshotMaintenanceLease::MakeForTesting("c", "1"); auto descriptor = MakeDescriptor("snapshots", 2); - descriptor.schema_version = ha::kBatchOpLogWeightSnapshotSchemaVersion; - descriptor.snapshot_format = ha::kBatchOpLogWeightSnapshotFormat; + descriptor.schema_version = ha::kBatchOpLogRecoverySnapshotSchemaVersion; + descriptor.snapshot_format = ha::kBatchOpLogRecoverySnapshotFormat; ha::BatchOpLogSnapshotManifest manifest; manifest.schema_version = descriptor.schema_version; manifest.snapshot_format = descriptor.snapshot_format; @@ -155,6 +155,11 @@ TEST(BatchOpLogSnapshotGcTest, MissingWeightArtifactSkipsDeletion) { ha::BuildBatchOpLogSnapshotWeightMetadataKey("snapshots", descriptor.snapshot_id), 1, Crc32cValue("w", 1)}; + manifest.recovery_chunks.push_back( + {0, + ha::BuildBatchOpLogSnapshotRecoveryChunkKey("snapshots", + descriptor.snapshot_id, 0), + 1, 1, Crc32cValue("r", 1)}); const auto manifest_json = ha::EncodeBatchOpLogSnapshotManifest(manifest); descriptor.manifest_size = manifest_json.size(); descriptor.manifest_crc32c = @@ -165,14 +170,19 @@ TEST(BatchOpLogSnapshotGcTest, MissingWeightArtifactSkipsDeletion) { manifest_json.end()}; object_store.objects[manifest.segments.key] = {'s'}; object_store.objects[manifest.weight_metadata->key] = {'w'}; + object_store.objects[manifest.recovery_chunks.front().key] = {'r'}; const auto stale_prefix = ha::BuildBatchOpLogSnapshotArtifactPrefix("snapshots", "1-1"); object_store.objects[stale_prefix + "marker"] = {1}; - object_store.objects.erase(manifest.weight_metadata->key); BatchOpLogSnapshotGc gc(backend, object_store, "c", "snapshots"); - EXPECT_NE(ErrorCode::OK, gc.Run(*lease, pointer, std::nullopt)); - EXPECT_TRUE(object_store.deleted.empty()); - object_store.objects[manifest.weight_metadata->key] = {'w'}; + for (const auto& key : {manifest.weight_metadata->key, + manifest.recovery_chunks.front().key}) { + const auto bytes = object_store.objects.at(key); + object_store.objects.erase(key); + EXPECT_NE(ErrorCode::OK, gc.Run(*lease, pointer, std::nullopt)); + EXPECT_TRUE(object_store.deleted.empty()); + object_store.objects[key] = bytes; + } EXPECT_EQ(ErrorCode::OK, gc.Run(*lease, pointer, std::nullopt)); ASSERT_EQ(1u, object_store.deleted.size()); EXPECT_EQ(stale_prefix, object_store.deleted.front()); diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/metadata_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/metadata_test.cpp index f4a74ca161..9cc553d882 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/metadata_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/metadata_test.cpp @@ -125,6 +125,7 @@ TEST(BatchOpLogSnapshotTypesTest, ManifestRoundTripsChunksAndAllowsEmptySet) { EncodeBatchOpLogSnapshotManifest(empty)); ASSERT_TRUE(decoded_empty.has_value()) << decoded_empty.error(); EXPECT_TRUE(decoded_empty->object_chunks.empty()); + EXPECT_TRUE(decoded_empty->recovery_chunks.empty()); } TEST(BatchOpLogSnapshotTypesTest, WeightSnapshotRequiresVersionedArtifact) { @@ -155,6 +156,52 @@ TEST(BatchOpLogSnapshotTypesTest, WeightSnapshotRequiresVersionedArtifact) { EXPECT_FALSE(legacy->weight_metadata); } +TEST(BatchOpLogSnapshotTypesTest, RecoverySnapshotRequiresVersionedChunks) { + auto manifest = MakeManifest(); + manifest.schema_version = kBatchOpLogRecoverySnapshotSchemaVersion; + manifest.snapshot_format = kBatchOpLogRecoverySnapshotFormat; + manifest.recovery_chunks = {{0, + BuildBatchOpLogSnapshotRecoveryChunkKey( + "snapshots", manifest.snapshot_id, 0), + 2, 42, 7}}; + const auto encoded = EncodeBatchOpLogSnapshotManifest(manifest); + auto decoded = DecodeBatchOpLogSnapshotManifest(encoded); + ASSERT_TRUE(decoded) << decoded.error(); + ASSERT_EQ(1u, decoded->recovery_chunks.size()); + EXPECT_EQ(2u, decoded->recovery_chunks.front().object_count); + EXPECT_FALSE(decoded->weight_metadata); + auto json = ParseJson(encoded); + json.removeMember("recovery_chunks"); + ExpectManifestRejected(WriteJson(json)); + json = ParseJson(encoded); + json["recovery_chunks"][0]["chunk_index"] = Json::UInt64(1); + ExpectManifestRejected(WriteJson(json)); + json = ParseJson(encoded); + json["recovery_chunks"][0]["object_count"] = Json::UInt64(0); + ExpectManifestRejected(WriteJson(json)); + for (const auto version : {kBatchOpLogSnapshotSchemaVersion, + kBatchOpLogWeightSnapshotSchemaVersion}) { + json = ParseJson(encoded); + json["schema_version"] = version; + json["snapshot_format"] = version == kBatchOpLogSnapshotSchemaVersion + ? kBatchOpLogSnapshotFormat + : kBatchOpLogWeightSnapshotFormat; + if (version == kBatchOpLogWeightSnapshotSchemaVersion) { + json["weight_metadata"] = json["segments"]; + } + ExpectManifestRejected(WriteJson(json)); + json.removeMember("recovery_chunks"); + auto legacy = DecodeBatchOpLogSnapshotManifest(WriteJson(json)); + ASSERT_TRUE(legacy) << legacy.error(); + EXPECT_TRUE(legacy->recovery_chunks.empty()); + } + auto descriptor = MakeDescriptor(); + descriptor.schema_version = kBatchOpLogRecoverySnapshotSchemaVersion; + descriptor.snapshot_format = kBatchOpLogRecoverySnapshotFormat; + EXPECT_TRUE(DecodeBatchOpLogSnapshotDescriptor( + EncodeBatchOpLogSnapshotDescriptor(descriptor))); +} + TEST(BatchOpLogSnapshotTypesTest, JsonRoundTripPreservesEscapedKeys) { auto descriptor = MakeDescriptor(); descriptor.manifest_key = R"(snapshots/"quoted"/manifest\key.json)"; diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/promotion_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/promotion_test.cpp index 8be15cadfb..e28508a73b 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/promotion_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/promotion_test.cpp @@ -174,6 +174,56 @@ TEST(BatchOpLogPromotionTest, EmptyStoreWithoutDurablePrefixUsesZeroCursor) { EXPECT_FALSE(standby->ExportMetadataSnapshot(retained)); } +TEST(BatchOpLogPromotionTest, DetachesRecoveryStoreWithoutSnapshotProvider) { + auto backend = std::make_shared(); + DrainJobRecord job{.id = {1, 2}, .sources = {"source"}, .concurrency = 1}; + OpLogEntry update; + update.sequence_id = 1; + update.op_type = OpType::DRAIN_JOB_UPDATE; + update.object_key = "drain"; + update.payload = EncodeRecoveryValue(DrainJobUpdate{job, std::nullopt, {}}); + update.checksum = ComputeOpLogChecksum(update.payload); + OpLogEntry put; + put.sequence_id = 2; + put.op_type = OpType::PUT_END; + put.object_key = "object"; + OpLogBatchRecord batch; + batch.batch_id = 1; + batch.first_seq = 1; + batch.last_seq = 2; + batch.entries = {update, put}; + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildBatchRecordKey("promotion-test", 1), + EncodeOpLogBatchRecord(batch))); + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildDurablePrefixKey("promotion-test"), + EncodeDurablePrefix({1, 2}))); + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildProducerViewKey("promotion-test"), "7")); + HotStandbyConfig config; + config.enable_snapshot_bootstrap = false; + config.enable_oplog_following = true; + config.enable_verification = false; + config.oplog_poll_interval_ms = 1; + HotStandbyService standby(config); + standby.SetCatchUpBatchKvBackendForTesting(backend); + ASSERT_EQ(ErrorCode::OK, standby.Start("", "", "promotion-test")); + EXPECT_FALSE(standby.IsBatchOpLogSnapshotMode()); + + auto handoff = standby.PromoteAndDetachBatchOpLogStore(); + + ASSERT_TRUE(handoff) << toString(handoff.error()); + ASSERT_TRUE(handoff->metadata_store); + const DurablePrefix expected_cursor{.batch_id = 1, .last_seq = 2}; + EXPECT_EQ(expected_cursor, handoff->applied_cursor); + EXPECT_EQ(7, handoff->producer_view_version); + EXPECT_TRUE(handoff->metadata_store->Exists("object")); + const auto& records = handoff->metadata_store->DrainRecovery().records(); + ASSERT_EQ(1u, records.size()); + EXPECT_EQ(EncodeRecoveryValue(job), records.begin()->second.payload); + EXPECT_EQ(StandbyState::STOPPED, standby.GetState()); +} + TEST(BatchOpLogPromotionTest, SnapshotOnlyProviderFallsBackToLegacyExport) { auto backend = std::make_shared(); LocalFileSnapshotObjectStore object_store( diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/pruning_coordinator_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/pruning_coordinator_test.cpp index 1e860bce79..dd3ebec17d 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/pruning_coordinator_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/pruning_coordinator_test.cpp @@ -170,7 +170,8 @@ class RecordingObjectStore final : public SnapshotObjectStore { }; std::string AddSnapshot(RecordingObjectStore& store, uint64_t batch, - int64_t lease_id = 101, bool with_weights = false) { + int64_t lease_id = 101, bool with_weights = false, + bool with_recovery = false) { ha::BatchOpLogSnapshotManifest manifest; manifest.snapshot_id = ha::BuildBatchOpLogSnapshotId(batch, lease_id); manifest.last_included_batch_id = batch; @@ -195,6 +196,16 @@ std::string AddSnapshot(RecordingObjectStore& store, uint64_t batch, 1, Crc32cValue("w", 1)}; store.objects[manifest.weight_metadata->key] = {'w'}; } + if (with_recovery) { + manifest.schema_version = ha::kBatchOpLogRecoverySnapshotSchemaVersion; + manifest.snapshot_format = ha::kBatchOpLogRecoverySnapshotFormat; + manifest.recovery_chunks.push_back( + {0, + ha::BuildBatchOpLogSnapshotRecoveryChunkKey( + root, manifest.snapshot_id, 0), + 1, 1, Crc32cValue("r", 1)}); + store.objects[manifest.recovery_chunks.front().key] = {'r'}; + } const auto bytes = ha::EncodeBatchOpLogSnapshotManifest(manifest); ha::BatchOpLogSnapshotDescriptor descriptor; descriptor.schema_version = manifest.schema_version; @@ -263,6 +274,28 @@ TEST_F(PruningTest, MissingWeightArtifactDoesNotAdvanceFloorOrDeleteLog) { EXPECT_EQ(1u, backend.delete_count); } +TEST_F(PruningTest, InvalidRecoveryChunkDoesNotAdvanceFloorOrDeleteLog) { + latest = AddSnapshot(store, 20, 101, true, true); + fallback = AddSnapshot(store, 10, 101, true, true); + backend.Put(latest_key, latest); + backend.Put(fallback_key, fallback); + for (const auto& id : {"20-101", "10-101"}) { + const auto key = + ha::BuildBatchOpLogSnapshotRecoveryChunkKey("snapshots", id, 0); + const auto bytes = store.objects.at(key); + store.objects.erase(key); + EXPECT_NE(ErrorCode::OK, Run()); + store.objects[key] = {'x'}; + EXPECT_NE(ErrorCode::OK, Run()); + EXPECT_FALSE(backend.values.contains(floor_key)); + EXPECT_EQ(0u, backend.delete_count); + store.objects[key] = bytes; + } + EXPECT_EQ(ErrorCode::OK, Run()); + EXPECT_EQ("10", backend.values.at(floor_key)); + EXPECT_EQ(1u, backend.delete_count); +} + TEST_F(PruningTest, FirstSnapshotDoesNotPublishFloor) { backend.values.erase(fallback_key); EXPECT_EQ(ErrorCode::OK, Run()); diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/writer_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/writer_test.cpp index 07f70bbf6a..b80f7edcb7 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/writer_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/writer_test.cpp @@ -23,6 +23,9 @@ #include "ha/snapshot/batch_oplog/metadata.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "hot_standby_service.h" +#include "rpc_types.h" +#include "segment/status.h" +#include "task_manager.h" namespace mooncake::test { @@ -161,6 +164,9 @@ class FakeObjectStore final : public SnapshotObjectStore { } std::string GetConnectionInfo() const override { return "fake"; } + bool IsNotFoundError(const std::string& error) const override { + return error == "not found"; + } void PutExisting(const std::string& key) { objects_[key] = {1}; } bool Contains(const std::string& key) const { @@ -225,7 +231,8 @@ class BatchOpLogSnapshotWriterTest : public ::testing::Test { std::optional StartCapture( size_t object_count, - const WeightMetadataSnapshot* weight_metadata = nullptr) { + const WeightMetadataSnapshot* weight_metadata = nullptr, + std::vector recovery_entries = {}) { backend_ = std::make_shared(); auto batch = MakeObjectBatch(object_count); const auto append = [&](OpType type, const std::string& key, @@ -262,6 +269,11 @@ class BatchOpLogSnapshotWriterTest : public ::testing::Test { "weight-lease:" + std::to_string(lease.lease_id), lease); } } + for (auto& entry : recovery_entries) { + entry.sequence_id = ++batch.last_seq; + entry.checksum = ComputeOpLogChecksum(entry.payload); + batch.entries.push_back(std::move(entry)); + } if (!batch.entries.empty()) { EXPECT_EQ(ErrorCode::OK, backend_->Put(BuildBatchRecordKey(kClusterId, 1), @@ -487,6 +499,207 @@ TEST_F(BatchOpLogSnapshotWriterTest, EXPECT_TRUE(restored.SnapshotWeightMetadata().metadata.empty()); } +TEST_F(BatchOpLogSnapshotWriterTest, + RestoresRecoveryChunksAndSuffixAfterCompaction) { + DrainJobRecord job{.id = {1, 1}, + .status = static_cast(JobStatus::RUNNING), + .sources = {"source"}, + .targets = {"target"}, + .concurrency = 2, + .created_ms = 100, + .updated_ms = 101, + .message = "Drain job running", + .failed = 2}; + DrainUnitRecord unit{.job_id = job.id, + .task_id = {2, 2}, + .unit_key = DrainUnitKey("tenant", "moving", "source"), + .tenant = "tenant", + .key = "moving", + .source = "source", + .target = "target", + .bytes = 64, + .retries = 2}; + DurableTaskRecord task{ + .id = unit.task_id, + .client_id = {3, 3}, + .type = static_cast(TaskType::REPLICA_MOVE), + .status = static_cast(TaskStatus::PROCESSING), + .payload = R"({"tenant_id":"tenant","key":"moving",)" + R"("source":"source","target":"target"})", + .created_ms = 101, + .updated_ms = 102, + .max_retries = 3, + .drain = unit}; + DrainSegmentRecord segment{ + .id = {4, 4}, + .client_id = task.client_id, + .name = "source", + .endpoint = "host:1234", + .base = 4096, + .size = 4096, + .status = static_cast(SegmentStatus::DRAINING)}; + Replica::Descriptor source{10, + MemoryDescriptor{{64, 4096, "tcp", "host:1234"}}, + ReplicaStatus::COMPLETE}; + Replica::Descriptor target{11, + MemoryDescriptor{{64, 8192, "tcp", "host:5678"}}, + ReplicaStatus::PROCESSING}; + ReplicationRecoveryRecord runtime{.tenant = "tenant", + .key = "moving", + .client_id = task.client_id, + .type = 1, + .started_ms = 102, + .source_id = source.id, + .targets = {target.id}, + .pending_bytes = 64}; + MetadataPayload moving; + moving.client_id = task.client_id; + moving.size = 64; + moving.replicas = {source, target}; + auto retired_source = source; + retired_source.id = 20; + retired_source.get_memory_descriptor().buffer_descriptor.buffer_address_ += + 64; + RetiredReplicaRecord retired{ + retired_source, std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count() + + 60000}; + MetadataPayload completed; + completed.size = 64; + completed.retired_replicas = std::vector{retired}; + std::vector entries; + const auto append = [&](OpType type, const std::string& key, + const auto& payload) { + OpLogEntry entry; + entry.op_type = type; + entry.tenant_id = "tenant"; + entry.object_key = key; + entry.payload = EncodeRecoveryValue(payload); + entries.push_back(std::move(entry)); + }; + append(OpType::DRAIN_JOB_UPDATE, "drain", + DrainJobUpdate{job, std::nullopt, {segment}}); + append(OpType::TASK_UPDATE, "task", task); + append(OpType::REPLICATION_START, "moving", + ReplicationStartRecord{runtime, EncodeRecoveryValue(moving), {}}); + append(OpType::PUT_END, "completed", completed); + auto capture = StartCapture(0, nullptr, std::move(entries)); + ASSERT_TRUE(capture); + ASSERT_EQ(4u, capture->last_included_seq); + FakeObjectStore object_store; + BatchOpLogSnapshotWriter writer(object_store); + const auto latest = + writer.Write(*standby_, *capture, "snapshots", "1-42", 2, 1234); + ASSERT_TRUE(latest) << latest.error(); + auto descriptor = ha::DecodeBatchOpLogSnapshotDescriptor(*latest); + ASSERT_TRUE(descriptor); + EXPECT_EQ(ha::kBatchOpLogRecoverySnapshotSchemaVersion, + descriptor->schema_version); + EXPECT_EQ(ha::kBatchOpLogRecoverySnapshotFormat, + descriptor->snapshot_format); + std::string manifest_json; + ASSERT_TRUE( + object_store.DownloadString(descriptor->manifest_key, manifest_json)); + auto manifest = ha::DecodeBatchOpLogSnapshotManifest(manifest_json); + ASSERT_TRUE(manifest); + ASSERT_EQ(3u, manifest->recovery_chunks.size()); + for (const auto& chunk : manifest->recovery_chunks) { + EXPECT_EQ(2u, chunk.object_count); + } + auto fallback_capture = standby_->BeginBatchOpLogSnapshotCapture(); + ASSERT_TRUE(fallback_capture); + const auto fallback = writer.Write(*standby_, *fallback_capture, + "snapshots", "1-41", 2, 1233); + ASSERT_TRUE(fallback) << fallback.error(); + standby_->Stop(); + + FakeHaKvBackend compacted_backend; + ASSERT_EQ(ErrorCode::OK, + compacted_backend.Put( + ha::BuildBatchOpLogSnapshotLatestKey(kClusterId), *latest)); + ASSERT_EQ( + ErrorCode::OK, + compacted_backend.Put( + ha::BuildBatchOpLogSnapshotFallbackKey(kClusterId), *fallback)); + ASSERT_EQ( + ErrorCode::OK, + compacted_backend.Put( + ha::BuildBatchOpLogSnapshotCompactionFloorKey(kClusterId), "1")); + ASSERT_EQ(ErrorCode::OK, + compacted_backend.Put(BuildDurablePrefixKey(kClusterId), + EncodeDurablePrefix({1, 4}))); + BatchOpLogSnapshotProvider provider(kClusterId, compacted_backend, + object_store, "snapshots"); + StandbyMetadataStore restored; + StandbySegmentRegistry registry; + const auto expected = std::vector{ + MakeRecoveryRecord(RecoveryRecordKind::Job, UuidToString(job.id), job), + MakeRecoveryRecord(RecoveryRecordKind::Unit, + UuidToString(job.id) + ":" + unit.unit_key, unit), + MakeRecoveryRecord(RecoveryRecordKind::Task, UuidToString(task.id), + task), + MakeRecoveryRecord(RecoveryRecordKind::Segment, segment.name, segment), + MakeRecoveryRecord(RecoveryRecordKind::Replication, + TenantId("tenant").MakeScopedKey("moving"), runtime), + MakeRecoveryRecord(RecoveryRecordKind::RetiredReplica, "20", retired)}; + const auto expect_recovery = [&] { + const auto& records = restored.DrainRecovery().records(); + ASSERT_EQ(expected.size(), records.size()); + for (const auto& record : expected) { + ASSERT_TRUE(records.contains(record.key)) << record.key; + EXPECT_EQ(record.payload, records.at(record.key).payload) + << record.key; + } + }; + auto result = provider.RestoreBaseline(restored, registry); + ASSERT_TRUE(result) << toString(result.error()); + expect_recovery(); + EXPECT_EQ(20u, result->max_replica_id); + auto object = restored.GetMetadata("tenant", "moving"); + ASSERT_TRUE(object); + ASSERT_EQ(2u, object->replicas.size()); + EXPECT_EQ(ReplicaStatus::PROCESSING, object->replicas[1].status); + + // A later chunk fails after earlier recovery records were installed. + const auto& damaged = manifest->recovery_chunks.back(); + std::vector bytes; + ASSERT_TRUE(object_store.DownloadBuffer(damaged.key, bytes)); + auto corrupt = bytes; + corrupt.back() ^= 1; + ASSERT_TRUE(object_store.UploadBuffer(damaged.key, corrupt)); + ASSERT_TRUE(provider.RestoreBaseline(restored, registry)); + expect_recovery(); + ASSERT_TRUE(object_store.DeleteObjectsWithPrefix(damaged.key)); + ASSERT_TRUE(provider.RestoreBaseline(restored, registry)); + expect_recovery(); + ASSERT_TRUE(object_store.UploadBuffer(damaged.key, bytes)); + + task.removed = true; + OpLogEntry removal; + removal.sequence_id = 5; + removal.op_type = OpType::TASK_UPDATE; + removal.payload = EncodeRecoveryValue(task); + removal.checksum = ComputeOpLogChecksum(removal.payload); + OpLogBatchRecord suffix; + suffix.batch_id = 2; + suffix.first_seq = suffix.last_seq = 5; + suffix.entries = {removal}; + ASSERT_EQ(ErrorCode::OK, + compacted_backend.Put(BuildBatchRecordKey(kClusterId, 2), + EncodeOpLogBatchRecord(suffix))); + ASSERT_EQ(ErrorCode::OK, + compacted_backend.Put(BuildDurablePrefixKey(kClusterId), + EncodeDurablePrefix({2, 5}))); + result = provider.RestoreBaseline(restored, registry); + ASSERT_TRUE(result) << toString(result.error()); + EXPECT_EQ(2u, result->last_applied_batch_id); + EXPECT_EQ(5u, result->last_applied_seq); + EXPECT_EQ(expected.size() - 1, restored.DrainRecovery().records().size()); + EXPECT_FALSE(restored.DrainRecovery().records().contains( + DrainRecoveryKey(RecoveryRecordKind::Task, UuidToString(task.id)))); +} + TEST_F(BatchOpLogSnapshotWriterTest, RejectsExistingCandidateWithoutDeletingIt) { auto capture = StartCapture(1); diff --git a/mooncake-store/tests/master_service/master_service_test_peer.h b/mooncake-store/tests/master_service/master_service_test_peer.h index a68981c573..fb3f744244 100644 --- a/mooncake-store/tests/master_service/master_service_test_peer.h +++ b/mooncake-store/tests/master_service/master_service_test_peer.h @@ -270,6 +270,15 @@ class MasterServiceTestPeer { return service.task_manager_; } + void StopJobDispatchForTesting() { + service_.job_dispatch_running_ = false; + if (service_.job_dispatch_thread_.joinable()) { + service_.job_dispatch_thread_.join(); + } + } + + void ProcessDrainJobsForTesting() { service_.ProcessDrainJobs(); } + static auto& TenantQuotaPolicyMutex(MasterService& service) { return service.tenant_quota_manager_.policy_mutex_; } diff --git a/mooncake-store/tests/task_manager_test.cpp b/mooncake-store/tests/task_manager_test.cpp index 7c935f4f06..660de37551 100644 --- a/mooncake-store/tests/task_manager_test.cpp +++ b/mooncake-store/tests/task_manager_test.cpp @@ -605,4 +605,94 @@ INSTANTIATE_TEST_SUITE_P( TaskStatus::FAILED), ::testing::Bool())); +TEST_F(ClientTaskManagerTest, + DurableClaimsAndCompletionPublishOnlyAfterCommit) { + bool reject = false; + std::vector committed; + size_t calls = 0; + ClientTaskManager manager( + {0, 3, 2, 0, 0, 3}, [&](std::span records) { + ++calls; + if (reject) return ErrorCode::ETCD_OPERATION_ERROR; + committed.assign(records.begin(), records.end()); + return ErrorCode::OK; + }); + const auto client = generate_uuid(); + const ReplicaMovePayload payload{"default", "key", "source", "target"}; + std::vector ids; + for (int i = 0; i < 3; ++i) { + auto id = + manager.get_write_access() + .submit_task_typed(client, payload); + ASSERT_TRUE(id); + ids.push_back(*id); + } + reject = true; + EXPECT_TRUE(manager.get_write_access().pop_tasks(client, 2).empty()); + for (auto id : ids) { + EXPECT_EQ(manager.get_read_access().find_task_by_id(id)->status, + TaskStatus::PENDING); + } + reject = false; + calls = 0; + auto claimed = manager.get_write_access().pop_tasks(client, 3); + ASSERT_EQ(claimed.size(), 2); + EXPECT_EQ(calls, 1); + ASSERT_EQ(committed.size(), 2); + EXPECT_EQ(claimed[0].id, ids[0]); + EXPECT_EQ(claimed[1].id, ids[1]); + EXPECT_TRUE(manager.get_write_access().pop_tasks(client, 1).empty()); + reject = true; + EXPECT_EQ(manager.get_write_access().complete_task( + client, ids[0], TaskStatus::SUCCESS, "done"), + ErrorCode::ETCD_OPERATION_ERROR); + EXPECT_EQ(manager.get_read_access().find_task_by_id(ids[0])->status, + TaskStatus::PROCESSING); + reject = false; + ASSERT_EQ(manager.get_write_access().complete_task( + client, ids[0], TaskStatus::SUCCESS, "done"), + ErrorCode::OK); + reject = true; + manager.get_write_access().prune_finished_tasks(); + EXPECT_TRUE(manager.get_read_access().find_task_by_id(ids[0])); + reject = false; + manager.get_write_access().prune_finished_tasks(); + ASSERT_EQ(committed.size(), 1); + EXPECT_TRUE(committed.front().removed); + EXPECT_FALSE(manager.get_read_access().find_task_by_id(ids[0])); + auto tail = manager.get_write_access().pop_tasks(client, 1); + ASSERT_EQ(tail.size(), 1); + EXPECT_EQ(tail.front().id, ids[2]); +} + +TEST_F(ClientTaskManagerTest, + DurableTaskBindingRoundTripsAndRejectsDifferentWork) { + const auto client = generate_uuid(); + ClientTaskManager manager({10, 10, 10, 0, 0, 7}); + DrainUnitRecord unit{ + generate_uuid(), {}, DrainUnitKey("tenant", "key", "source"), + "tenant", "key", "source", + "target", 4096}; + auto id = + manager.get_write_access().submit_task_typed( + client, {"tenant", "key", "source", "target"}, unit); + ASSERT_TRUE(id); + auto task = manager.get_read_access().find_task_by_id(*id); + ASSERT_TRUE(task); + auto record = ToDurableTask(*task); + ASSERT_TRUE(ValidateDurableTask(record)); + DurableTaskRecord decoded; + ASSERT_TRUE(DecodeRecoveryValue(EncodeRecoveryValue(record), decoded)); + auto restored = FromDurableTask(decoded); + ASSERT_TRUE(restored.drain); + EXPECT_EQ(restored.id, *id); + EXPECT_EQ(restored.drain->task_id, *id); + EXPECT_EQ(restored.max_retry_attempts, 7); + EXPECT_EQ(restored.drain->unit_key, unit.unit_key); + decoded.drain->target = "another-target"; + EXPECT_FALSE(ValidateDurableTask(decoded)); + decoded.payload = "invalid json"; + EXPECT_FALSE(ValidateDurableTask(decoded)); +} + } // namespace mooncake