[Store]: Support mounting file segments

Author: nickyc975

mooncake-store/include/allocator.h

@@ -53,12 +53,13 @@ class AllocatedBuffer {
     struct Descriptor;
 
     AllocatedBuffer(std::shared_ptr<BufferAllocatorBase> allocator,
-                    std::string segment_name, void* buffer_ptr,
-                    std::size_t size,
+                    std::string segment_name, FileBufferID file_id,
+                    void* buffer_ptr, std::size_t size,
                     std::optional<offset_allocator::OffsetAllocationHandle>&&
                         offset_handle = std::nullopt)
         : allocator_(std::move(allocator)),
           segment_name_(std::move(segment_name)),
+          file_id_(file_id),
           buffer_ptr_(buffer_ptr),
           size_(size),
           offset_handle_(std::move(offset_handle)) {}
@@ -92,17 +93,20 @@ class AllocatedBuffer {
     // Represents the serializable state
     struct Descriptor {
         std::string segment_name_;
+        FileBufferID file_id_;
         uint64_t size_;
         uintptr_t buffer_address_;
         BufStatus status_;
-        YLT_REFL(Descriptor, segment_name_, size_, buffer_address_, status_);
+        YLT_REFL(Descriptor, segment_name_, file_id_, size_, buffer_address_,
+                 status_);
     };
 
     void mark_complete() { status = BufStatus::COMPLETE; }
 
    private:
     std::weak_ptr<BufferAllocatorBase> allocator_;
     std::string segment_name_;
+    FileBufferID file_id_;
     BufStatus status{BufStatus::INIT};
     void* buffer_ptr_{nullptr};
     std::size_t size_{0};
@@ -152,7 +156,8 @@ class CachelibBufferAllocator
     : public BufferAllocatorBase,
       public std::enable_shared_from_this<CachelibBufferAllocator> {
    public:
-    CachelibBufferAllocator(std::string segment_name, size_t base, size_t size);
+    CachelibBufferAllocator(std::string segment_name, size_t base, size_t size,
+                            FileBufferID file_id = 0);
 
     ~CachelibBufferAllocator() override;
 
@@ -170,6 +175,7 @@ class CachelibBufferAllocator
     const size_t base_;
     const size_t total_size_;
     std::atomic_size_t cur_size_;
+    const FileBufferID file_id_;
 
     // metrics - removed allocated_bytes_ member
     // ylt::metric::gauge_t* allocated_bytes_{nullptr};
@@ -189,7 +195,8 @@ class OffsetBufferAllocator
     : public BufferAllocatorBase,
       public std::enable_shared_from_this<OffsetBufferAllocator> {
    public:
-    OffsetBufferAllocator(std::string segment_name, size_t base, size_t size);
+    OffsetBufferAllocator(std::string segment_name, size_t base, size_t size,
+                          FileBufferID file_id = 0);
 
     ~OffsetBufferAllocator() override;
 
@@ -207,6 +214,7 @@ class OffsetBufferAllocator
     const size_t base_;
     const size_t total_size_;
     std::atomic_size_t cur_size_;
+    const FileBufferID file_id_;
 
     // offset allocator implementation
     std::shared_ptr<offset_allocator::OffsetAllocator> offset_allocator_;

mooncake-store/include/client.h

@@ -179,6 +179,20 @@ class Client {
     tl::expected<void, ErrorCode> UnmountSegment(const void* buffer,
                                                  size_t size);
 
+    /**
+     * @brief Register a file to master for allocation
+     * @param path The file path
+     * @return ErrorCode indicating success/failure
+     */
+    tl::expected<void, ErrorCode> MountFileSegment(const std::string& path);
+
+    /**
+     * @brief Unregisters a file segment from master
+     * @param path File path to unregister
+     * @return ErrorCode indicating success/failure
+     */
+    tl::expected<void, ErrorCode> UnmountFileSegment(const std::string& path);
+
     /**
      * @brief Registers memory buffer with TransferEngine for data transfer
      * @param addr Memory address to register

mooncake-store/include/types.h

@@ -45,6 +45,7 @@ class Replica;
 using ObjectKey = std::string;
 using Version = uint64_t;
 using SegmentId = int64_t;
+using FileBufferID = uint32_t;
 using TaskID = int64_t;
 using BufHandleList = std::vector<std::shared_ptr<AllocatedBuffer>>;
 // using ReplicaList = std::vector<ReplicaInfo>;
@@ -164,20 +165,61 @@ const static uint64_t kMaxSliceSize =
     facebook::cachelib::Slab::kSize - 16;  // should be lower than limit
 
 /**
- * @brief Represents a contiguous memory region
+ * @brief Type of segments.
+ */
+enum class SegmentType {
+    UNKNOWN = -1,
+    MEMORY,
+    FILE,
+};
+
+/**
+ * @brief Stream operator for SegmentType
+ */
+inline std::ostream& operator<<(std::ostream& os,
+                                const SegmentType& type) noexcept {
+    static const std::unordered_map<SegmentType, std::string_view> type_strings{
+        {SegmentType::UNKNOWN, "UNKNOWN"},
+        {SegmentType::MEMORY, "MEMORY"},
+        {SegmentType::FILE, "FILE"}};
+
+    os << (type_strings.count(type) ? type_strings.at(type) : "UNKNOWN");
+    return os;
+}
+
+/**
+ * @brief Represents a contiguous storage region, could be memory or file.
  */
 struct Segment {
     UUID id{0, 0};
+    SegmentType type{SegmentType::UNKNOWN};
     std::string name{};  // The name of the segment, also might be the
                          // hostname of the server that owns the segment
     uintptr_t base{0};
     size_t size{0};
+    // For a file segment, this will be the path of the file.
+    std::string path{};
+    // For a file segment, this will be the id of the file buffer.
+    FileBufferID file_id{0};
     Segment() = default;
     Segment(const UUID& id, const std::string& name, uintptr_t base,
             size_t size)
-        : id(id), name(name), base(base), size(size) {}
+        : id(id),
+          type(SegmentType::MEMORY),
+          name(name),
+          base(base),
+          size(size) {}
+    Segment(const UUID& id, const std::string& name, uintptr_t base,
+            size_t size, const std::string& path, FileBufferID file_id)
+        : id(id),
+          type(SegmentType::FILE),
+          name(name),
+          base(base),
+          size(size),
+          path(path),
+          file_id(file_id) {}
 };
-YLT_REFL(Segment, id, name, base, size);
+YLT_REFL(Segment, id, type, name, base, size, path, file_id);
 
 /**
  * @brief Client status from the master's perspective

mooncake-store/src/allocator.cpp

@@ -23,29 +23,32 @@ AllocatedBuffer::~AllocatedBuffer() {
 
 // Implementation of get_descriptor
 AllocatedBuffer::Descriptor AllocatedBuffer::get_descriptor() const {
-    return {segment_name_, static_cast<uint64_t>(size()),
+    return {segment_name_, file_id_, static_cast<uint64_t>(size()),
             reinterpret_cast<uintptr_t>(buffer_ptr_), status};
 }
 
 // Define operator<< using public accessors or get_descriptor if appropriate
 std::ostream& operator<<(std::ostream& os, const AllocatedBuffer& buffer) {
     return os << "AllocatedBuffer: { "
               << "segment_name: " << buffer.segment_name_ << ", "
+              << "file_id: " << buffer.file_id_ << ", "
               << "size: " << buffer.size() << ", "
               << "status: " << buffer.status << ", "
               << "buffer_ptr: " << static_cast<void*>(buffer.data()) << " }";
 }
 
 // Removed allocated_bytes parameter and member initialization
 CachelibBufferAllocator::CachelibBufferAllocator(std::string segment_name,
-                                                 size_t base, size_t size)
+                                                 size_t base, size_t size,
+                                                 FileBufferID file_id)
     : segment_name_(segment_name),
       base_(base),
       total_size_(size),
-      cur_size_(0) {
+      cur_size_(0),
+      file_id_(file_id) {
     VLOG(1) << "initializing_buffer_allocator segment_name=" << segment_name
             << " base_address=" << reinterpret_cast<void*>(base)
-            << " size=" << size;
+            << " size=" << size << " file_id=" << file_id;
 
     // Calculate the size of the header region.
     header_region_size_ =
@@ -107,7 +110,7 @@ std::unique_ptr<AllocatedBuffer> CachelibBufferAllocator::allocate(
     cur_size_.fetch_add(size);
     MasterMetricManager::instance().inc_allocated_size(size);
     return std::make_unique<AllocatedBuffer>(shared_from_this(), segment_name_,
-                                             buffer, size);
+                                             file_id_, buffer, size);
 }
 
 void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) {
@@ -133,14 +136,16 @@ void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) {
 
 // OffsetBufferAllocator implementation
 OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name,
-                                             size_t base, size_t size)
+                                             size_t base, size_t size,
+                                             FileBufferID file_id)
     : segment_name_(segment_name),
       base_(base),
       total_size_(size),
-      cur_size_(0) {
+      cur_size_(0),
+      file_id_(file_id) {
     VLOG(1) << "initializing_offset_buffer_allocator segment_name="
             << segment_name << " base_address=" << reinterpret_cast<void*>(base)
-            << " size=" << size;
+            << " size=" << size << " file_id=" << file_id;
 
     try {
         // 1k <= init_capacity <= 64k
@@ -196,7 +201,7 @@ std::unique_ptr<AllocatedBuffer> OffsetBufferAllocator::allocate(size_t size) {
         // Create a custom AllocatedBuffer that manages the
         // OffsetAllocationHandle
         allocated_buffer = std::make_unique<AllocatedBuffer>(
-            shared_from_this(), segment_name_, buffer_ptr, size,
+            shared_from_this(), segment_name_, file_id_, buffer_ptr, size,
             std::move(allocation_handle));
         VLOG(1) << "allocation_succeeded size=" << size
                 << " segment=" << segment_name_ << " address=" << buffer_ptr;