[Store] Recover drain jobs through batch OpLog

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

mooncake-store/include/drain_recovery.h

@@ -0,0 +1,162 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+#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<std::string> 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<DrainJobRecord> job;
+    std::optional<DrainUnitRecord> unit;
+    std::vector<DrainSegmentRecord> 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<DrainUnitRecord> 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<ReplicaID> 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<DrainSegmentRecord> 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<std::string, DrainRecoveryRecord>;
+
+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<Replica::Descriptor>& 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 <typename T>
+std::string EncodeRecoveryValue(const T& value) {
+    const auto encoded = struct_pack::serialize(value);
+    return {encoded.begin(), encoded.end()};
+}
+
+template <typename T>
+bool DecodeRecoveryValue(const std::string& encoded, T& value) {
+    return struct_pack::deserialize_to(value, encoded) == struct_pack::errc::ok;
+}
+
+template <typename T>
+DrainRecoveryRecord MakeRecoveryRecord(RecoveryRecordKind kind,
+                                       const std::string& key, const T& value) {
+    return {kind, DrainRecoveryKey(kind, key), EncodeRecoveryValue(value)};
+}
+
+}  // namespace mooncake

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

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,
 };
 

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<LeaseState> lease_state_;
 };

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<DrainRecoveryRecord> records;
+    YLT_REFL(BatchOpLogRecoveryChunk, chunk_index, records);
+};
+
+std::vector<uint8_t> EncodeBatchOpLogRecoveryChunk(
+    uint64_t index, std::vector<DrainRecoveryRecord> records);
+tl::expected<BatchOpLogRecoveryChunk, std::string>
+DecodeBatchOpLogRecoveryChunk(const std::vector<uint8_t>& encoded,
+                              uint64_t index, uint64_t count);
+
 std::vector<uint8_t> EncodeBatchOpLogSnapshotSegments(
     const std::vector<StandbySegmentInfo>& segments);
 tl::expected<std::vector<StandbySegmentInfo>, std::string>

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<BatchOpLogSnapshotChunkDescriptor> object_chunks;
     std::optional<BatchOpLogSnapshotObjectDescriptor> weight_metadata;
+    std::vector<BatchOpLogSnapshotChunkDescriptor> 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);

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<DrainRecoveryRecord>& out);
     tl::expected<BatchOpLogPromotionHandoff, ErrorCode>
     PromoteAndDetachBatchOpLogStore();
 

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