[TransferEngine]: Support registering files as shared buffers

Signed-off-by: Jinlong Chen <[email protected]>

Author: nickyc975

mooncake-transfer-engine/include/transfer_engine.h

@@ -46,12 +46,14 @@ using SegmentHandle = Transport::SegmentHandle;
 using SegmentID = Transport::SegmentID;
 using BatchID = Transport::BatchID;
 using BufferEntry = Transport::BufferEntry;
+using FileBufferID = TransferMetadata::FileBufferID;
 
 class TransferEngine {
    public:
     TransferEngine(bool auto_discover = false)
         : metadata_(nullptr),
           local_topology_(std::make_shared<Topology>()),
+          next_file_id_(1),
           auto_discover_(auto_discover) {
 #ifdef WITH_METRICS
         InitializeMetricsConfig();
@@ -62,6 +64,7 @@ class TransferEngine {
     TransferEngine(bool auto_discover, const std::vector<std::string> &filter)
         : metadata_(nullptr),
           local_topology_(std::make_shared<Topology>()),
+          next_file_id_(1),
           auto_discover_(auto_discover),
           filter_(filter) {
 #ifdef WITH_METRICS
@@ -111,6 +114,25 @@ class TransferEngine {
 
     int unregisterLocalMemoryBatch(const std::vector<void *> &addr_list);
 
+    bool supportFileBuffer();
+
+    /**
+     * @brief Register a local file as a shared buffer.
+     * @param[in] path Local path of the file.
+     * @param[in] size Available size of the file.
+     * @param[out] id The id of the registered file buffer.
+     * @return 0 on success, or error number on failure.
+     */
+    int registerLocalFile(const std::string &path, size_t size,
+                          FileBufferID &id);
+
+    /**
+     * @brief Unregister a previously registered file.
+     * @param[in] path The path of the registered file buffer.
+     * @return 0 on success, or error number on failure.
+     */
+    int unregisterLocalFile(const std::string &path);
+
     BatchID allocateBatchID(size_t batch_size) {
         return multi_transports_->allocateBatchID(batch_size);
     }
@@ -224,13 +246,22 @@ class TransferEngine {
         bool remote_accessible;
     };
 
+    struct LocalFile {
+        FileBufferID id;
+        std::string path;
+        std::size_t size;
+    };
+
     std::shared_ptr<TransferMetadata> metadata_;
     std::string local_server_name_;
     std::shared_ptr<MultiTransport> multi_transports_;
     std::shared_mutex mutex_;
     std::vector<MemoryRegion> local_memory_regions_;
     std::shared_ptr<Topology> local_topology_;
 
+    std::atomic<FileBufferID> next_file_id_;
+    std::unordered_map<std::string, LocalFile> local_files_;
+
     RWSpinlock send_notifies_lock_;
     std::unordered_map<BatchID,
                        std::pair<SegmentID, TransferMetadata::NotifyDesc>>

mooncake-transfer-engine/include/transfer_engine_c.h

@@ -24,6 +24,7 @@ extern "C" {
 
 #define segment_handle_t int32_t
 #define segment_id_t int32_t
+#define file_id_t uint32_t
 #define batch_id_t uint64_t
 #define LOCAL_SEGMENT (0)
 #define INVALID_BATCH UINT64_MAX
@@ -35,6 +36,7 @@ struct transfer_request {
     int opcode;
     void *source;
     segment_id_t target_id;
+    file_id_t file_id;
     uint64_t target_offset;
     uint64_t length;
 };
@@ -135,6 +137,13 @@ int registerLocalMemoryBatch(transfer_engine_t engine,
 int unregisterLocalMemoryBatch(transfer_engine_t engine, void **addr_list,
                                size_t addr_len);
 
+bool supportFileBuffer(transfer_engine_t engine);
+
+int registerLocalFile(transfer_engine_t engine, const char *path, size_t size,
+                      file_id_t *id);
+
+int unregisterLocalFile(transfer_engine_t engine, const char *path);
+
 batch_id_t allocateBatchID(transfer_engine_t engine, size_t batch_size);
 
 int submitTransfer(transfer_engine_t engine, batch_id_t batch_id,

mooncake-transfer-engine/include/transfer_metadata.h

@@ -64,6 +64,16 @@ class TransferMetadata {
         std::unordered_map<std::string, std::string> local_path_map;
     };
 
+    // Identify a single file in a segment's file buffers.
+    using FileBufferID = uint32_t;
+
+    struct FileBufferDesc {
+        FileBufferID id;
+        std::string path;
+        std::size_t size;
+        std::size_t align;  // For future usage.
+    };
+
     struct RankInfoDesc {
         uint64_t rankId = 0xFFFFFFFF;  // rank id, user rank
         std::string hostIp;
@@ -87,6 +97,8 @@ class TransferMetadata {
         std::vector<BufferDesc> buffers;
         // this is for nvmeof.
         std::vector<NVMeoFBufferDesc> nvmeof_buffers;
+        // Generic file buffers.
+        std::vector<FileBufferDesc> file_buffers;
         // this is for cxl.
         std::string cxl_name;
         uint64_t cxl_base_addr;
@@ -148,6 +160,10 @@ class TransferMetadata {
 
     int removeLocalMemoryBuffer(void *addr, bool update_metadata);
 
+    int addFileBuffer(const FileBufferDesc &buffer_desc, bool update_metadata);
+
+    int removeFileBuffer(FileBufferID id, bool update_metadata);
+
     int addLocalSegment(SegmentID segment_id, const std::string &segment_name,
                         std::shared_ptr<SegmentDesc> &&desc);
 

mooncake-transfer-engine/include/transport/transport.h

@@ -43,6 +43,8 @@ class Transport {
     using SegmentID = uint64_t;
     using SegmentHandle = SegmentID;
 
+    using FileBufferID = TransferMetadata::FileBufferID;
+
     using BatchID = uint64_t;
     const static BatchID INVALID_BATCH_ID = UINT64_MAX;
 
@@ -60,6 +62,7 @@ class Transport {
         uint64_t target_offset;
         size_t length;
         int advise_retry_cnt = 0;
+        FileBufferID file_id;
     };
 
     enum TransferStatusEnum {
@@ -92,6 +95,7 @@ class Transport {
         SliceStatus status;
         TransferTask *task;
         bool from_cache;
+        FileBufferID file_id;
 
         union {
             struct {
@@ -284,6 +288,17 @@ class Transport {
     virtual int unregisterLocalMemoryBatch(
         const std::vector<void *> &addr_list) = 0;
 
+    virtual bool supportFileBuffer() { return false; };
+
+    virtual int registerLocalFile(FileBufferID id, const std::string &path,
+                                  size_t size) {
+        return ERR_NOT_IMPLEMENTED;
+    }
+
+    virtual int unregisterLocalFile(FileBufferID id) {
+        return ERR_NOT_IMPLEMENTED;
+    }
+
     virtual const char *getName() const = 0;
 };
 }  // namespace mooncake

mooncake-transfer-engine/src/transfer_engine.cpp

@@ -287,6 +287,15 @@ Transport *TransferEngine::installTransport(const std::string &proto,
             entry.addr, entry.length, entry.location, entry.remote_accessible);
         if (ret < 0) return nullptr;
     }
+
+    if (transport->supportFileBuffer()) {
+        for (auto &file : local_files_) {
+            int ret = transport->registerLocalFile(
+                file.second.id, file.second.path, file.second.size);
+            if (ret < 0) return nullptr;
+        }
+    }
+
     return transport;
 }
 
@@ -436,6 +445,77 @@ int TransferEngine::unregisterLocalMemoryBatch(
     return 0;
 }
 
+bool TransferEngine::supportFileBuffer() {
+    bool supported = false;
+    for (auto &transport : multi_transports_->listTransports()) {
+        supported = supported || transport->supportFileBuffer();
+    }
+    return supported;
+}
+
+int TransferEngine::registerLocalFile(const std::string &path, size_t size,
+                                      FileBufferID &id) {
+    if (!supportFileBuffer()) {
+        LOG(ERROR) << "File buffers not suppotred";
+        return ERR_NOT_IMPLEMENTED;
+    }
+
+    std::unique_lock<std::shared_mutex> lock(mutex_);
+    if (local_files_.count(path) > 0) {
+        LOG(ERROR) << "Registering an already registered file: " << path;
+        return ERR_ADDRESS_OVERLAPPED;
+    }
+
+    const auto id_ = next_file_id_.fetch_add(1);
+
+    for (auto &transport : multi_transports_->listTransports()) {
+        if (!transport->supportFileBuffer()) {
+            continue;
+        }
+
+        int ret = transport->registerLocalFile(id_, path, size);
+        if (ret != 0) {
+            LOG(ERROR) << "Failed to register file " << path << " to transport "
+                       << transport->getName() << ", ret=" << ret;
+            return ret;
+        }
+    }
+
+    local_files_[path] = {id_, path, size};
+    id = id_;
+    return 0;
+}
+
+int TransferEngine::unregisterLocalFile(const std::string &path) {
+    if (!supportFileBuffer()) {
+        LOG(ERROR) << "File buffers not suppotred";
+        return ERR_NOT_IMPLEMENTED;
+    }
+
+    std::unique_lock<std::shared_mutex> lock(mutex_);
+    auto it = local_files_.find(path);
+    if (it == local_files_.end()) {
+        return ERR_ADDRESS_NOT_REGISTERED;
+    }
+
+    for (auto &transport : multi_transports_->listTransports()) {
+        if (!transport->supportFileBuffer()) {
+            continue;
+        }
+
+        int ret = transport->unregisterLocalFile(it->second.id);
+        if (ret != 0 && ret != ERR_ADDRESS_NOT_REGISTERED) {
+            LOG(ERROR) << "Failed to unregister file " << path
+                       << " from transport " << transport->getName()
+                       << ", ret=" << ret;
+            return ret;
+        }
+    }
+
+    local_files_.erase(it);
+    return 0;
+}
+
 #ifdef WITH_METRICS
 // Helper function to convert string to lowercase for case-insensitive
 // comparison

mooncake-transfer-engine/src/transfer_engine_c.cpp

@@ -117,6 +117,22 @@ int unregisterLocalMemoryBatch(transfer_engine_t engine, void **addr_list,
     return native->unregisterLocalMemoryBatch(native_addr_list);
 }
 
+bool supportFileBuffer(transfer_engine_t engine) {
+    TransferEngine *native = (TransferEngine *)engine;
+    return native->supportFileBuffer();
+}
+
+int registerLocalFile(transfer_engine_t engine, const char *path, size_t size,
+                      file_id_t *id) {
+    TransferEngine *native = (TransferEngine *)engine;
+    return native->registerLocalFile(path, size, *id);
+}
+
+int unregisterLocalFile(transfer_engine_t engine, const char *path) {
+    TransferEngine *native = (TransferEngine *)engine;
+    return native->unregisterLocalFile(path);
+}
+
 batch_id_t allocateBatchID(transfer_engine_t engine, size_t batch_size) {
     TransferEngine *native = (TransferEngine *)engine;
     return (batch_id_t)native->allocateBatchID(batch_size);
@@ -132,6 +148,7 @@ int submitTransfer(transfer_engine_t engine, batch_id_t batch_id,
             (Transport::TransferRequest::OpCode)entries[index].opcode;
         native_entries[index].source = entries[index].source;
         native_entries[index].target_id = entries[index].target_id;
+        native_entries[index].file_id = entries[index].file_id;
         native_entries[index].target_offset = entries[index].target_offset;
         native_entries[index].length = entries[index].length;
     }

mooncake-transfer-engine/src/transfer_metadata.cpp

@@ -231,6 +231,18 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc,
                    << desc.name << " protocol " << desc.protocol;
         return ERR_METADATA;
     }
+
+    Json::Value fileBuffersJson(Json::arrayValue);
+    for (const auto &fileBuffer : desc.file_buffers) {
+        Json::Value bufferJSON;
+        bufferJSON["id"] = fileBuffer.id;
+        bufferJSON["path"] = fileBuffer.path;
+        bufferJSON["size"] = fileBuffer.size;
+        bufferJSON["align"] = fileBuffer.align;
+        fileBuffersJson.append(bufferJSON);
+    }
+    segmentJSON["file_buffers"] = fileBuffersJson;
+
     return 0;
 }
 
@@ -415,6 +427,16 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON,
                    << " protocol " << desc->protocol;
         return nullptr;
     }
+
+    for (const auto &bufferJSON : segmentJSON["file_buffers"]) {
+        FileBufferDesc buffer;
+        buffer.id = bufferJSON["id"].asUInt();
+        buffer.path = bufferJSON["path"].asString();
+        buffer.size = bufferJSON["size"].asUInt64();
+        buffer.align = bufferJSON["align"].asUInt64();
+        desc->file_buffers.push_back(buffer);
+    }
+
     return desc;
 }
 
@@ -605,6 +627,38 @@ int TransferMetadata::removeLocalMemoryBuffer(void *addr,
     return ERR_ADDRESS_NOT_REGISTERED;
 }
 
+int TransferMetadata::addFileBuffer(const FileBufferDesc &buffer_desc,
+                                    bool update_metadata) {
+    {
+        RWSpinlock::WriteGuard guard(segment_lock_);
+        auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID];
+        segment_desc->file_buffers.push_back(buffer_desc);
+    }
+    if (update_metadata) return updateLocalSegmentDesc();
+    return 0;
+}
+
+int TransferMetadata::removeFileBuffer(FileBufferID id, bool update_metadata) {
+    bool buffer_exist = false;
+    {
+        RWSpinlock::WriteGuard guard(segment_lock_);
+        auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID];
+        for (auto iter = segment_desc->file_buffers.begin();
+             iter != segment_desc->file_buffers.end(); ++iter) {
+            if (iter->id == id) {
+                segment_desc->file_buffers.erase(iter);
+                buffer_exist = true;
+                break;
+            }
+        }
+    }
+    if (buffer_exist) {
+        if (update_metadata) return updateLocalSegmentDesc();
+        return 0;
+    }
+    return ERR_ADDRESS_NOT_REGISTERED;
+}
+
 int TransferMetadata::addRpcMetaEntry(const std::string &server_name,
                                       RpcMetaDesc &desc) {
     local_rpc_meta_ = desc;