[TransferEngine]: Support specifying transport specific args when installing

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

Author: nickyc975

mooncake-transfer-engine/include/multi_transport.h

@@ -49,6 +49,10 @@ class MultiTransport {
     Transport *installTransport(const std::string &proto,
                                 std::shared_ptr<Topology> topo);
 
+    bool transportNeedArgs(const std::string &proto);
+
+    Transport *installTransport(const std::string &proto, void **args);
+
     Transport *getTransport(const std::string &proto);
 
     std::vector<Transport *> listTransports();

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

@@ -257,6 +257,9 @@ class Transport {
                         std::shared_ptr<TransferMetadata> meta,
                         std::shared_ptr<Topology> topo);
 
+    virtual int install(std::string &local_server_name,
+                        std::shared_ptr<TransferMetadata> meta, void **args);
+
     std::string local_server_name_;
     std::shared_ptr<TransferMetadata> metadata_;
 

mooncake-transfer-engine/src/multi_transport.cpp

@@ -240,6 +240,32 @@ Transport *MultiTransport::installTransport(const std::string &proto,
     return transport;
 }
 
+bool MultiTransport::transportNeedArgs(const std::string &proto) {
+    return false;
+}
+
+Transport *MultiTransport::installTransport(const std::string &proto,
+                                            void **args) {
+    std::shared_ptr<Transport> transport = nullptr;
+
+    // Add transport creation logic here.
+
+    if (!transport) {
+        LOG(ERROR) << "Unsupported transport " << proto
+                   << ", please rebuild Mooncake";
+        return nullptr;
+    }
+
+    int rc = transport->install(local_server_name_, metadata_, args);
+    if (rc != 0) {
+        LOG(ERROR) << "Failed to install transport " << proto << ", rc=" << rc;
+        return nullptr;
+    }
+
+    transport_map_[proto] = transport;
+    return transport.get();
+}
+
 Status MultiTransport::selectTransport(const TransferRequest &entry,
                                        Transport *&transport) {
     auto target_segment_desc = metadata_->getSegmentDescByID(entry.target_id);

mooncake-transfer-engine/src/transfer_engine.cpp

@@ -259,16 +259,22 @@ Transport *TransferEngine::installTransport(const std::string &proto,
         return transport;
     }
 
-    if (args != nullptr && args[0] != nullptr) {
-        const std::string nic_priority_matrix = static_cast<char *>(args[0]);
-        int ret = local_topology_->parse(nic_priority_matrix);
-        if (ret) {
-            LOG(ERROR) << "Failed to parse NIC priority matrix";
-            return nullptr;
+    if (multi_transports_->transportNeedArgs(proto)) {
+        transport = multi_transports_->installTransport(proto, args);
+    } else {
+        if (args != nullptr && args[0] != nullptr) {
+            const std::string nic_priority_matrix =
+                static_cast<char *>(args[0]);
+            int ret = local_topology_->parse(nic_priority_matrix);
+            if (ret) {
+                LOG(ERROR) << "Failed to parse NIC priority matrix";
+                return nullptr;
+            }
         }
+
+        transport = multi_transports_->installTransport(proto, local_topology_);
     }
 
-    transport = multi_transports_->installTransport(proto, local_topology_);
     if (!transport) return nullptr;
 
     // Since installTransport() is only called once during initialization

mooncake-transfer-engine/src/transport/transport.cpp

@@ -64,4 +64,11 @@ int Transport::install(std::string &local_server_name,
     metadata_ = meta;
     return 0;
 }
+
+int Transport::install(std::string &local_server_name,
+                       std::shared_ptr<TransferMetadata> meta, void **args) {
+    local_server_name_ = local_server_name;
+    metadata_ = meta;
+    return 0;
+}
 }  // namespace mooncake