[Store] L2->L1 promotion-on-hit

Author: yzhan1

mooncake-integration/store/store_py.cpp

@@ -1616,6 +1616,9 @@ PYBIND11_MODULE(store, m) {
         .def("is_disk_replica",
              static_cast<bool (Replica::Descriptor::*)() const noexcept>(
                  &Replica::Descriptor::is_disk_replica))
+        .def("is_local_disk_replica",
+             static_cast<bool (Replica::Descriptor::*)() const noexcept>(
+                 &Replica::Descriptor::is_local_disk_replica))
         .def(
             "get_memory_descriptor",
             static_cast<const MemoryDescriptor &(Replica::Descriptor::*)()

mooncake-store/include/client_service.h

@@ -60,7 +60,7 @@ class QueryResult {
  */
 class Client {
    public:
-    ~Client();
+    virtual ~Client();
 
     const UUID& getClientId() const { return client_id_; }
 
@@ -405,6 +405,42 @@ class Client {
     tl::expected<void, ErrorCode> ReportSsdCapacity(
         int64_t ssd_total_capacity_bytes);
 
+    /**
+     * @brief Heartbeat-driven pull of pending L2->L1 promotion work for this
+     * client. Mirror of OffloadObjectHeartbeat. Returns key->size pairs the
+     * caller (FileStorage) must read from local SSD and stage as MEMORY
+     * replicas via PromotionAllocStart + NotifyPromotionSuccess.
+     */
+    // Virtual to enable subclassing in unit tests.
+    virtual tl::expected<void, ErrorCode> PromotionObjectHeartbeat(
+        std::unordered_map<std::string, int64_t>& promotion_objects);
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key during
+     * L2->L1 promotion. Returns the new replica's descriptor that the caller
+     * writes via Transfer Engine before calling NotifyPromotionSuccess.
+     */
+    virtual tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const std::string& key, uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; called after the
+     * client has written the bytes via Transfer Engine.
+     */
+    virtual tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const std::string& key);
+
+    /**
+     * @brief Write `slices` into the memory replica described by
+     * `memory_descriptor` via Transfer Engine. Used by FileStorage to fill a
+     * PROCESSING memory replica staged by PromotionAllocStart before calling
+     * NotifyPromotionSuccess.
+     */
+    virtual ErrorCode PromotionWrite(
+        const Replica::Descriptor& memory_descriptor,
+        std::vector<Slice>& slices);
+
     /**
      * @brief Performs a batched read of multiple objects using a
      * high-throughput Transfer Engine.
@@ -570,14 +606,16 @@ class Client {
 
     bool IsReplicaOnLocalMemory(const Replica::Descriptor& replica);
 
-   private:
+   protected:
     /**
-     * @brief Private constructor to enforce creation through Create() method
+     * @brief Constructor exposed to subclasses for testing only; production
+     * code must go through Create().
      */
     Client(const std::string& local_hostname,
            const std::string& metadata_connstring, const std::string& protocol,
            const std::map<std::string, std::string>& labels = {});
 
+   private:
     /**
      * @brief Internal helper functions for initialization and data transfer
      */

mooncake-store/include/file_storage.h

@@ -51,6 +51,7 @@ class FileStorage {
 
    private:
     friend class FileStorageTest;
+    friend class FileStoragePromotionTest;
     struct AllocatedBatch {
         uint64_t batch_id;
         std::vector<BufferHandle> handles;
@@ -82,10 +83,24 @@ class FileStorage {
      * client.
      * 2. Receives feedback on which objects should be offloaded.
      * 3. Triggers asynchronous offloading of pending objects.
+     * 4. Pulls and processes any pending L2->L1 promotion tasks queued by the
+     *    master (mirror of step 1+2 in the reverse direction).
      * @return tl::expected<void, ErrorCode> indicating operation status.
      */
     tl::expected<void, ErrorCode> Heartbeat();
 
+    /**
+     * @brief Drives the L2->L1 promotion pipeline for one heartbeat tick.
+     * Pulls promotion work from the master, stages a MEMORY replica for each
+     * key, copies the bytes from local SSD into that replica, and notifies the
+     * master on success. A failure on any single key is logged and skipped;
+     * the master-side reaper decrements the source replica's refcnt and
+     * erases the task entry on TTL expiry, and any orphaned PROCESSING
+     * MEMORY replica is reaped via the standard discarded-replicas path.
+     * @return tl::expected<void, ErrorCode> indicating operation status.
+     */
+    tl::expected<void, ErrorCode> ProcessPromotionTasks();
+
     tl::expected<bool, ErrorCode> IsEnableOffloading();
 
     tl::expected<void, ErrorCode> BatchLoad(
@@ -123,4 +138,4 @@ class FileStorage {
     std::thread client_buffer_gc_thread_;
 };
 
-}  // namespace mooncake
+}  // namespace mooncake

mooncake-store/include/master_client.h

@@ -363,6 +363,32 @@ class MasterClient {
         const UUID& client_id, const std::vector<std::string>& keys,
         const std::vector<StorageObjectMetadata>& metadatas);
 
+    /**
+     * @brief Heartbeat-driven pull of pending L2->L1 promotion work for a
+     * client. Returns key->size pairs the caller should read from local
+     * SSD and stage as MEMORY replicas via PromotionAllocStart +
+     * NotifyPromotionSuccess.
+     */
+    [[nodiscard]] tl::expected<std::unordered_map<std::string, int64_t>,
+                               ErrorCode>
+    PromotionObjectHeartbeat(const UUID& client_id);
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key during
+     * promotion. Returns the new replica's descriptor that the caller writes
+     * via Transfer Engine.
+     */
+    [[nodiscard]] tl::expected<PromotionAllocStartResponse, ErrorCode>
+    PromotionAllocStart(const std::string& key, uint64_t size,
+                        const std::vector<std::string>& preferred_segments);
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; called after the
+     * client has written the bytes via Transfer Engine.
+     */
+    [[nodiscard]] tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const UUID& client_id, const std::string& key);
+
     /**
      * @brief Start a copy operation
      * @param key Object key

mooncake-store/include/master_config.h

@@ -85,6 +85,12 @@ struct MasterConfig {
     // Offload-on-evict: defer LOCAL_DISK offload to eviction time
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+
+    // Promotion-on-hit: when Get observes a LOCAL_DISK-only key, queue an
+    // async copy back to MEMORY so the next Get is fast.
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
 };
 
 class MasterServiceSupervisorConfig {
@@ -146,6 +152,9 @@ class MasterServiceSupervisorConfig {
     bool enable_cxl = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
     MasterServiceSupervisorConfig() = default;
 
     // From MasterConfig
@@ -163,6 +172,9 @@ class MasterServiceSupervisorConfig {
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
         rpc_port = static_cast<int>(config.rpc_port);
         rpc_thread_num = static_cast<size_t>(config.rpc_thread_num);
 
@@ -277,6 +289,9 @@ class WrappedMasterServiceConfig {
     bool enable_offload = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -334,6 +349,9 @@ class WrappedMasterServiceConfig {
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         if (ha_backend_connstring.empty()) {
@@ -414,6 +432,9 @@ class WrappedMasterServiceConfig {
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         if (ha_backend_connstring.empty()) {
@@ -752,6 +773,9 @@ class MasterServiceConfig {
     bool enable_offload = false;
     bool offload_on_evict = false;
     bool offload_force_evict = false;
+    bool promotion_on_hit = false;
+    uint32_t promotion_admission_threshold = 2;
+    uint32_t promotion_queue_limit = 50000;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -805,6 +829,9 @@ class MasterServiceConfig {
         enable_offload = config.enable_offload;
         offload_on_evict = config.offload_on_evict;
         offload_force_evict = config.offload_force_evict;
+        promotion_on_hit = config.promotion_on_hit;
+        promotion_admission_threshold = config.promotion_admission_threshold;
+        promotion_queue_limit = config.promotion_queue_limit;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         cluster_id = config.cluster_id;

mooncake-store/include/master_service.h

@@ -19,6 +19,7 @@
 #include <ylt/util/tl/expected.hpp>
 
 #include "allocation_strategy.h"
+#include "count_min_sketch.h"
 #include "master_metric_manager.h"
 #include "mutex.h"
 #include "segment.h"
@@ -456,6 +457,34 @@ class MasterService {
         const std::vector<StorageObjectMetadata>& metadatas)
         -> tl::expected<void, ErrorCode>;
 
+    /**
+     * @brief Heartbeat-driven pull of pending promotion work for a client.
+     * Returns the per-client promotion_objects map (key -> object size) and
+     * clears it. The per-shard promotion_tasks map remains populated as the
+     * source of truth until NotifyPromotionSuccess commits the new MEMORY
+     * replica.
+     */
+    auto PromotionObjectHeartbeat(const UUID& client_id)
+        -> tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode>;
+
+    /**
+     * @brief Stage a PROCESSING MEMORY replica for an existing key. Allocates
+     * DRAM via the existing AllocationStrategy, optionally biased toward the
+     * caller's local memory segment via preferred_segments. The new replica is
+     * invisible to readers until NotifyPromotionSuccess flips it to COMPLETE.
+     */
+    auto PromotionAllocStart(const std::string& key, uint64_t size,
+                             const std::vector<std::string>& preferred_segments)
+        -> tl::expected<PromotionAllocStartResponse, ErrorCode>;
+
+    /**
+     * @brief Commit a staged MEMORY replica to COMPLETE; decrement source
+     * refcnt; erase per-shard and per-client task entries. Mirror of
+     * NotifyOffloadSuccess.
+     */
+    auto NotifyPromotionSuccess(const UUID& client_id, const std::string& key)
+        -> tl::expected<void, ErrorCode>;
+
     /**
      * @brief Create a copy task to copy an object's replicas to target segments
      * @return Copy task ID on success, ErrorCode on failure
@@ -848,6 +877,19 @@ class MasterService {
         std::chrono::system_clock::time_point start_time;
     };
 
+    // Tracks an in-flight LOCAL_DISK -> MEMORY copy. The source LOCAL_DISK
+    // replica is refcnt-pinned for the duration of the task so it cannot be
+    // evicted. alloc_id pins down which staged replica
+    // NotifyPromotionSuccess should commit, so a concurrent Put creating
+    // another PROCESSING MEMORY replica cannot be confused with ours.
+    // alloc_id is 0 until PromotionAllocStart records the new replica.
+    struct PromotionTask {
+        ReplicaID source_id;    // the LOCAL_DISK replica being promoted
+        ReplicaID alloc_id{0};  // the new MEMORY replica staged by AllocStart
+        uint64_t object_size;
+        std::chrono::system_clock::time_point start_time;
+    };
+
     static constexpr size_t kNumShards = 1024;  // Number of metadata shards
 
     // Sharded metadata maps and their mutexes
@@ -860,6 +902,8 @@ class MasterService {
             GUARDED_BY(mutex);
         std::unordered_map<std::string, const OffloadingTask> offloading_tasks
             GUARDED_BY(mutex);
+        std::unordered_map<std::string, const PromotionTask> promotion_tasks
+            GUARDED_BY(mutex);
     };
     std::array<MetadataShard, kNumShards> metadata_shards_;
 
@@ -935,6 +979,25 @@ class MasterService {
     tl::expected<void, ErrorCode> PushOffloadingQueue(const std::string& key,
                                                       Replica& replica);
 
+    /**
+     * @brief Mirror of PushOffloadingQueue for promotion-on-hit. Inserts an
+     * entry into the holder client's LocalDiskSegment::promotion_objects map.
+     * Caller is responsible for refcnt-pinning the source replica and
+     * recording the task in the shard's promotion_tasks map.
+     */
+    tl::expected<void, ErrorCode> PushPromotionQueue(const std::string& key,
+                                                     Replica& source_replica);
+
+    /**
+     * @brief Helper invoked from GetReplicaList when an only-LOCAL_DISK key is
+     * observed. Applies the gating chain (frequency / watermark / dedup /
+     * cap), refcnt-pins the source LOCAL_DISK replica, records a
+     * PromotionTask, and pushes onto the holder client's promotion_objects
+     * map. Acquires its own RW shard accessor; safe to call after
+     * GetReplicaList's RO accessor has been released.
+     */
+    void TryPushPromotionQueue(const std::string& key);
+
     // Lease related members
     const uint64_t default_kv_lease_ttl_;     // in milliseconds
     const uint64_t default_kv_soft_pin_ttl_;  // in milliseconds
@@ -1180,6 +1243,16 @@ class MasterService {
     // offload_on_evict_=true)
     bool offload_force_evict_{false};
 
+    // Promotion-on-hit: opt-in flag enabling LOCAL_DISK -> MEMORY promotion
+    // when a Get observes a key with only LOCAL_DISK replicas.
+    bool promotion_on_hit_{false};
+    uint32_t promotion_admission_threshold_{2};
+    uint32_t promotion_queue_limit_{50000};
+    // Master-side frequency sketch. Constructed only when promotion_on_hit_ is
+    // true. CountMinSketch is mutex-protected internally so we can call into it
+    // from any GetReplicaList caller without additional locking.
+    std::unique_ptr<CountMinSketch> promotion_sketch_;
+
     const std::string ha_backend_type_;
 
     const std::string ha_backend_connstring_;

mooncake-store/include/rpc_service.h

@@ -152,6 +152,17 @@ class WrappedMasterService {
         const UUID& client_id, const std::vector<std::string>& keys,
         const std::vector<StorageObjectMetadata>& metadatas);
 
+    // Promotion-on-hit RPCs.
+    tl::expected<std::unordered_map<std::string, int64_t>, ErrorCode>
+    PromotionObjectHeartbeat(const UUID& client_id);
+
+    tl::expected<PromotionAllocStartResponse, ErrorCode> PromotionAllocStart(
+        const std::string& key, uint64_t size,
+        const std::vector<std::string>& preferred_segments);
+
+    tl::expected<void, ErrorCode> NotifyPromotionSuccess(
+        const UUID& client_id, const std::string& key);
+
     tl::expected<UUID, ErrorCode> CreateDrainJob(
         const CreateDrainJobRequest& request);
 

mooncake-store/include/rpc_types.h

@@ -67,6 +67,15 @@ struct CopyStartResponse {
 };
 YLT_REFL(CopyStartResponse, source, targets);
 
+/**
+ * @brief Response structure for PromotionAllocStart (L2->L1 promotion-on-hit).
+ * Carries the staged PROCESSING MEMORY replica descriptor.
+ */
+struct PromotionAllocStartResponse {
+    Replica::Descriptor memory_descriptor;
+};
+YLT_REFL(PromotionAllocStartResponse, memory_descriptor);
+
 /**
  * @brief Response structure for MoveStart operation
  */

mooncake-store/include/segment.h

@@ -53,6 +53,13 @@ struct LocalDiskSegment {
     int64_t ssd_total_capacity_bytes = 0;  // last reported by client heartbeat
     std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
         offloading_objects;
+    // Promotion-on-hit pending work for this client. Populated by master's
+    // TryPushPromotionQueue when a Get hits a LOCAL_DISK-only key on this
+    // client. Drained by PromotionObjectHeartbeat. Same locking as
+    // offloading_objects (offloading_mutex_) to keep call-site discipline
+    // identical to #1899's offload path.
+    std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
+        promotion_objects;
     explicit LocalDiskSegment(bool enable_offloading)
         : enable_offloading(enable_offloading) {}