[Store] Enable Client SSD Offload And Storage Persistence

.github/workflows/ci.yml

@@ -105,6 +105,17 @@ jobs:
         MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_STORE_MEMCPY=false make test -j ARGS="-V"
       shell: bash
 
+    - name: Test (in build env with ssd)
+      run: |
+        cd build
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+        ldconfig -v || echo "always continue"
+        mkdir -p /tmp/mooncake_test_ssd
+        export MOONCAKE_STORAGE_ROOT_DIR=/tmp/mooncake_test_ssd
+        MC_METADATA_SERVER=http://127.0.0.1:8080/metadata make test -j ARGS="-V"
+        rm -rf /tmp/mooncake_test_ssd
+      shell: bash
+
     - name: Stop Mooncake Master Service
       run: |
         pkill mooncake_master || true
@@ -169,13 +180,31 @@ jobs:
         mooncake_http_metadata_server --port 8080 &
       shell: bash
 
+    - name: Run tests with ssd
+      run: |
+        source test_env/bin/activate
+        mkdir -p /tmp/mooncake_test_ssd
+        export MOONCAKE_STORAGE_ROOT_DIR=/tmp/mooncake_test_ssd
+        MC_STORE_MEMCPY=false ./scripts/run_tests.sh
+        deactivate
+        rm -rf /tmp/mooncake_test_ssd
+      shell: bash
+
+    - name: Restart metadata server
+      run: |
+        PID=$(pgrep -f "mooncake_http_metadata_server"); kill -9 $PID
+        source test_env/bin/activate
+        mooncake_http_metadata_server --port 8080 &
+      shell: bash
+
     - name: Run tests
       run: |
         source test_env/bin/activate
         MC_STORE_MEMCPY=false ./scripts/run_tests.sh
         deactivate
       shell: bash
 
+
   build-flags:
     runs-on: ubuntu-22.04
     strategy:

mooncake-integration/store/store_py.cpp

@@ -218,6 +218,21 @@ int DistributedObjectStore::initAll(const std::string &protocol_,
                  buffer_allocator_size, protocol_, device_name);
 }
 
+int DistributedObjectStore::allocateSlices(std::vector<Slice> &slices,
+                                           size_t length) {
+    uint64_t offset = 0;
+    while (offset < length) {
+        auto chunk_size = std::min(length - offset, kMaxSliceSize);
+        auto ptr = client_buffer_allocator_->allocate(chunk_size);
+        if (!ptr) {
+            return 1;  // SliceGuard will handle cleanup
+        }
+        slices.emplace_back(Slice{ptr, chunk_size});
+        offset += chunk_size;
+    }
+    return 0;
+}
+
 int DistributedObjectStore::allocateSlices(std::vector<Slice> &slices,
                                            const std::string &value) {
     uint64_t offset = 0;
@@ -302,15 +317,22 @@ int DistributedObjectStore::allocateSlices(
     length = 0;
     if (object_info.replica_list.empty()) return -1;
     auto &replica = object_info.replica_list[0];
-    for (auto &handle : replica.buffer_descriptors) {
-        auto chunk_size = handle.size_;
-        assert(chunk_size <= kMaxSliceSize);
-        auto ptr = client_buffer_allocator_->allocate(chunk_size);
-        if (!ptr) {
-            return 1;  // SliceGuard will handle cleanup
+    if(replica.is_memory_replica() == false) {
+        auto &disk_descriptor =replica.get_disk_descriptor();
+        length = disk_descriptor.file_size;
+        return allocateSlices(slices, length);
+    }else{
+        auto &memory_descriptors = replica.get_memory_descriptor();
+        for (auto &handle : memory_descriptors.buffer_descriptors) {
+            auto chunk_size = handle.size_;
+            assert(chunk_size <= kMaxSliceSize);
+            auto ptr = client_buffer_allocator_->allocate(chunk_size);
+            if (!ptr) {
+                return 1;  // SliceGuard will handle cleanup
+            }
+            slices.emplace_back(Slice{ptr, chunk_size});
+            length += chunk_size;
         }
-        slices.emplace_back(Slice{ptr, chunk_size});
-        length += chunk_size;
     }
     return 0;
 }
@@ -355,15 +377,26 @@ int DistributedObjectStore::allocateBatchedSlices(
         // Get first replica
         auto &replica = object_info_it->second[0];
         uint64_t length = 0;
-        for (auto &handle : replica.buffer_descriptors) {
-            auto chunk_size = handle.size_;
-            assert(chunk_size <= kMaxSliceSize);
-            auto ptr = client_buffer_allocator_->allocate(chunk_size);
-            if (!ptr) {
+
+        if(replica.is_memory_replica() == false) {
+            auto &disk_descriptor =replica.get_disk_descriptor();
+            length = disk_descriptor.file_size;
+            auto result = allocateSlices(batched_slices[key], length);
+            if(result) {
                 return 1;
             }
-            batched_slices[key].emplace_back(Slice{ptr, chunk_size});
-            length += chunk_size;
+        }else{
+            auto &memory_descriptors = replica.get_memory_descriptor();
+            for (auto &handle : memory_descriptors.buffer_descriptors) {
+                auto chunk_size = handle.size_;
+                assert(chunk_size <= kMaxSliceSize);
+                auto ptr = client_buffer_allocator_->allocate(chunk_size);
+                if (!ptr) {
+                    return 1;
+                }
+                batched_slices[key].emplace_back(Slice{ptr, chunk_size});
+                length += chunk_size;
+            }
         }
         str_length_map.emplace(key, length);
     }
@@ -712,8 +745,14 @@ int64_t DistributedObjectStore::getSize(const std::string &key) {
     int64_t total_size = 0;
     if (!object_info.replica_list.empty()) {
         auto &replica = object_info.replica_list[0];
-        for (auto &handle : replica.buffer_descriptors) {
-            total_size += handle.size_;
+        if(replica.is_memory_replica() == false) {
+            auto &disk_descriptor = replica.get_disk_descriptor();
+            total_size = disk_descriptor.file_size;
+        }else{
+            auto &memory_descriptors = replica.get_memory_descriptor();
+            for (auto &handle : memory_descriptors.buffer_descriptors) {
+                total_size += handle.size_;
+            }
         }
     } else {
         LOG(ERROR) << "Internal error: object_info.replica_list_size() is 0";
@@ -850,8 +889,13 @@ int DistributedObjectStore::get_into(const std::string &key, void *buffer,
     }
 
     auto &replica = object_info.replica_list[0];
-    for (auto &handle : replica.buffer_descriptors) {
-        total_size += handle.size_;
+    if(replica.is_memory_replica() == false) {
+        auto &disk_descriptor = replica.get_disk_descriptor();
+        total_size = disk_descriptor.file_size;
+    }else{
+        for (auto &handle : replica.get_memory_descriptor().buffer_descriptors) {
+            total_size += handle.size_;
+        }
     }
 
     // Check if user buffer is large enough
@@ -865,11 +909,19 @@ int DistributedObjectStore::get_into(const std::string &key, void *buffer,
     std::vector<mooncake::Slice> slices;
     uint64_t offset = 0;
 
-    for (auto &handle : replica.buffer_descriptors) {
-        auto chunk_size = handle.size_;
-        void *chunk_ptr = static_cast<char *>(buffer) + offset;
-        slices.emplace_back(Slice{chunk_ptr, chunk_size});
-        offset += chunk_size;
+    if(replica.is_memory_replica() == false) {
+        while(offset < total_size){
+            auto chunk_size = std::min(total_size - offset, kMaxSliceSize);
+            void *chunk_ptr = static_cast<char *>(buffer) + offset;
+            slices.emplace_back(Slice{chunk_ptr, chunk_size});
+            offset += chunk_size;
+        }
+    }else{
+        for (auto &handle : replica.get_memory_descriptor().buffer_descriptors) {
+            void *chunk_ptr = static_cast<char *>(buffer) + offset;
+            slices.emplace_back(Slice{chunk_ptr, handle.size_});
+            offset += handle.size_;
+        }
     }
 
     // Step 3: Read data directly into user buffer
@@ -984,8 +1036,13 @@ std::vector<int> DistributedObjectStore::batch_get_into(
 
         auto &replica = replica_list[0];
         uint64_t total_size = 0;
-        for (auto &handle : replica.buffer_descriptors) {
-            total_size += handle.size_;
+        if(replica.is_memory_replica() == false) {
+            auto &disk_descriptor = replica.get_disk_descriptor();
+            total_size = disk_descriptor.file_size;
+        }else{
+            for (auto &handle : replica.get_memory_descriptor().buffer_descriptors) {
+                total_size += handle.size_;
+            }
         }
 
         if (sizes[i] < total_size) {
@@ -998,10 +1055,19 @@ std::vector<int> DistributedObjectStore::batch_get_into(
 
         uint64_t offset = 0;
         std::vector<mooncake::Slice> key_slices;
-        for (auto &handle : replica.buffer_descriptors) {
-            void *chunk_ptr = static_cast<char *>(buffers[i]) + offset;
-            key_slices.emplace_back(Slice{chunk_ptr, handle.size_});
-            offset += handle.size_;
+        if(replica.is_memory_replica() == false) {
+            while(offset < total_size){
+                auto chunk_size = std::min(total_size - offset, kMaxSliceSize);
+                void *chunk_ptr = static_cast<char *>(buffers[i]) + offset;
+                key_slices.emplace_back(Slice{chunk_ptr, chunk_size});
+                offset += chunk_size;
+            }
+        }else{
+            for (auto &handle : replica.get_memory_descriptor().buffer_descriptors) {
+                void *chunk_ptr = static_cast<char *>(buffers[i]) + offset;
+                key_slices.emplace_back(Slice{chunk_ptr, handle.size_});
+                offset += handle.size_;
+            }
         }
         all_slices[key] = key_slices;
         results[i] = static_cast<int>(total_size);

mooncake-integration/store/store_py.h

@@ -217,6 +217,9 @@ class DistributedObjectStore {
     int64_t getSize(const std::string &key);
 
    private:
+    int allocateSlices(std::vector<mooncake::Slice> &slices,
+                                           size_t length);
+
     int allocateSlices(std::vector<mooncake::Slice> &slices,
                        const std::string &value);
 

mooncake-store/include/client.h

@@ -13,6 +13,8 @@
 #include "transfer_engine.h"
 #include "transfer_task.h"
 #include "types.h"
+#include "thread_pool.h"
+#include "storage_backend.h"
 
 namespace mooncake {
 
@@ -95,9 +97,9 @@ class Client {
      * @param slices Vector of slices to store the data
      * @return ErrorCode indicating success/failure
      */
-    ErrorCode Get(const std::string& object_key, const ObjectInfo& object_info,
+    ErrorCode Get(const std::string& object_key, ObjectInfo& object_info,
                   std::vector<Slice>& slices);
-
+                             
     /**
      * @brief Transfers data using pre-queried object information
      * @param object_keys Keys of the objects
@@ -204,7 +206,8 @@ class Client {
      * @brief Private constructor to enforce creation through Create() method
      */
     Client(const std::string& local_hostname,
-           const std::string& metadata_connstring);
+           const std::string& metadata_connstring,
+           const std::string& storage_root_dir);
 
     /**
      * @brief Internal helper functions for initialization and data transfer
@@ -215,26 +218,36 @@ class Client {
                                  const std::string& protocol,
                                  void** protocol_args);
     ErrorCode TransferData(
-        const std::vector<AllocatedBuffer::Descriptor>& handles,
+        const Replica::Descriptor &replica,
         std::vector<Slice>& slices, TransferRequest::OpCode op_code);
     ErrorCode TransferWrite(
-        const std::vector<AllocatedBuffer::Descriptor>& handles,
+        const Replica::Descriptor &replica,
         std::vector<Slice>& slices);
     ErrorCode TransferRead(
-        const std::vector<AllocatedBuffer::Descriptor>& handles,
+        const Replica::Descriptor &replica,
         std::vector<Slice>& slices);
 
+    /**
+     * @brief Prepare and use the storage backend for persisting data
+     */
+    void PrepareStorageBackend(const std::string& storage_root_dir, const std::string& fsdir);
+
+    ErrorCode GetFromLocalFile(const std::string& object_key,
+                             std::vector<Slice>& slices, ObjectInfo& object_info);
+
+    void PutToLocalFile(const std::string& object_key,
+                             std::vector<Slice>& slices);
+
     /**
      * @brief Find the first complete replica from a replica list
      * @param replica_list List of replicas to search through
-     * @param handles Output vector to store the buffer handles of the found
-     * replica
+     * @param replica the first complete replica (file or memory)
      * @return ErrorCode::OK if found, ErrorCode::INVALID_REPLICA if no complete
      * replica
      */
     ErrorCode FindFirstCompleteReplica(
         const std::vector<Replica::Descriptor>& replica_list,
-        std::vector<AllocatedBuffer::Descriptor>& handles);
+        Replica::Descriptor& replica);
 
     // Core components
     TransferEngine transfer_engine_;
@@ -248,6 +261,11 @@ class Client {
     // Configuration
     const std::string local_hostname_;
     const std::string metadata_connstring_;
+    const std::string storage_root_dir_; 
+
+    // Client persistent thread pool for async operations
+    ThreadPool write_thread_pool_;
+    std::shared_ptr<StorageBackend> storage_backend_;
 
     // For high availability
     MasterViewHelper master_view_helper_;

mooncake-store/include/ha_helper.h

@@ -86,7 +86,8 @@ class MasterServiceSupervisor {
         std::chrono::steady_clock::duration rpc_conn_timeout =
             std::chrono::seconds(
                 0),  // Client connection timeout. 0 = no timeout (infinite)
-        bool rpc_enable_tcp_no_delay = true);
+        bool rpc_enable_tcp_no_delay = true,
+        const std::string& cluster_id = DEFAULT_CLUSTER_ID);
     int Start();
     ~MasterServiceSupervisor();
 
@@ -115,6 +116,8 @@ class MasterServiceSupervisor {
 
     // Local hostname for leader election
     std::string local_hostname_;
+
+    std::string cluster_id_;
 };
 
 }  // namespace mooncake