[TENT] Integrate receiver-credit pull control with adaptive dispatch protection

mooncake-transfer-engine/benchmark/tent_backend.cpp

@@ -53,6 +53,53 @@ std::shared_ptr<Config> loadConfig() {
     config->set("rpc_server_port", XferBenchConfig::rpc_server_port);
     config->set("transports/rdma/deadline_bw_arbitration",
                 XferBenchConfig::deadline_bw_arbitration);
+    config->set("enable_runtime_queue", XferBenchConfig::enable_runtime_queue);
+    if (XferBenchConfig::enable_runtime_queue) {
+        config->set("runtime_queue/max_dispatch_owners",
+                    XferBenchConfig::runtime_queue_max_dispatch_owners);
+    }
+    if (XferBenchConfig::receiver_credit_mode != "disabled") {
+        config->set("receiver_credit/mode",
+                    XferBenchConfig::receiver_credit_mode);
+        config->set("receiver_credit/capacity/data_bytes",
+                    XferBenchConfig::receiver_credit_capacity_bytes);
+        config->set("receiver_credit/capacity/request_slots",
+                    XferBenchConfig::receiver_credit_capacity_slots);
+        config->set("receiver_credit/grant_batch/data_bytes",
+                    XferBenchConfig::receiver_credit_grant_bytes);
+        config->set("receiver_credit/grant_batch/request_slots",
+                    XferBenchConfig::receiver_credit_grant_slots);
+        config->set("receiver_credit/control/freshness_ttl_ms", 60000U);
+        config->set("receiver_credit/control/retry_after_us", 100U);
+        config->set("receiver_credit/control/poll_interval_us", 100U);
+        config->set("receiver_credit/control/adaptive_dispatch/min_owners",
+                    XferBenchConfig::receiver_credit_adaptive_min_owners);
+        config->set("receiver_credit/control/adaptive_dispatch/initial_owners",
+                    XferBenchConfig::receiver_credit_adaptive_initial_owners);
+        config->set("receiver_credit/control/adaptive_dispatch/max_owners",
+                    XferBenchConfig::receiver_credit_adaptive_max_owners);
+        config->set("receiver_credit/control/adaptive_dispatch/slow_rtt_us",
+                    XferBenchConfig::receiver_credit_adaptive_slow_rtt_us);
+        config->set(
+            "receiver_credit/control/adaptive_dispatch/healthy_pulls_per_"
+            "increase",
+            XferBenchConfig::receiver_credit_adaptive_healthy_pulls);
+        config->set("receiver_credit/limits/max_peers", size_t{4096});
+    }
+    if (!XferBenchConfig::tent_rdma_devices.empty()) {
+        std::vector<std::string> devices;
+        std::stringstream input(XferBenchConfig::tent_rdma_devices);
+        for (std::string device; std::getline(input, device, ',');) {
+            if (!device.empty()) devices.push_back(std::move(device));
+        }
+        if (!devices.empty()) {
+            json rule = {{"name", "tebench_rdma_devices"},
+                         {"segment_type", "memory"},
+                         {"devices", std::move(devices)},
+                         {"transports", json::array({"rdma"})}};
+            config->set("policy", json::array({std::move(rule)}));
+        }
+    }
 
     // Configure transport types based on xport_type parameter
     if (!XferBenchConfig::xport_type.empty()) {
@@ -126,6 +173,15 @@ int TENTBenchRunner::allocateBuffers() {
     if (seg_type == "DRAM") {
         device_prefix = "cpu";
         num_buffers = numa_num_configured_nodes();
+        if (XferBenchConfig::tent_dram_numa_node >= 0) {
+            start_idx = XferBenchConfig::tent_dram_numa_node;
+            if (start_idx >= num_buffers) {
+                LOG(ERROR) << "tent_dram_numa_node " << start_idx
+                           << " out of range [0, " << num_buffers << ")";
+                return -1;
+            }
+            num_buffers = 1;
+        }
 #if defined(USE_CUDA) || defined(USE_SUNRISE)
     } else if (seg_type == "VRAM") {
         device_prefix = "cuda";

mooncake-transfer-engine/benchmark/utils.cpp

@@ -72,6 +72,35 @@ DEFINE_string(
     tent_transport_hint, "unspec",
     "tent only: per-request transport_hint. "
     "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|sunrise_link");
+DEFINE_string(tent_rdma_devices, "",
+              "tent only: comma-separated RDMA devices allowed by tebench");
+DEFINE_int32(tent_dram_numa_node, -1,
+             "tent only: allocate one DRAM buffer on this NUMA node");
+DEFINE_bool(enable_runtime_queue, false,
+            "tent only: route submissions through the runtime queue");
+DEFINE_uint64(runtime_queue_max_dispatch_owners, 64,
+              "tent only: maximum concurrently dispatched queue owners");
+DEFINE_string(
+    receiver_credit_mode, "disabled",
+    "tent only: receiver credit rollout mode: disabled|optional|required");
+DEFINE_uint64(receiver_credit_capacity_bytes, 1UL << 30,
+              "tent only: receiver-wide outstanding byte capacity");
+DEFINE_uint64(receiver_credit_capacity_slots, 1024,
+              "tent only: receiver-wide outstanding request capacity");
+DEFINE_uint64(receiver_credit_grant_bytes, 64UL << 20,
+              "tent only: maximum byte grant per pull");
+DEFINE_uint64(receiver_credit_grant_slots, 64,
+              "tent only: maximum request-slot grant per pull");
+DEFINE_uint64(receiver_credit_adaptive_min_owners, 1,
+              "tent only: adaptive dispatch lower bound");
+DEFINE_uint64(receiver_credit_adaptive_initial_owners, 2,
+              "tent only: adaptive dispatch startup window");
+DEFINE_uint64(receiver_credit_adaptive_max_owners, 2,
+              "tent only: adaptive dispatch exploration ceiling");
+DEFINE_uint32(receiver_credit_adaptive_slow_rtt_us, 20000,
+              "tent only: credit RPC RTT that triggers backoff");
+DEFINE_uint32(receiver_credit_adaptive_healthy_pulls, 512,
+              "tent only: healthy pulls required for additive recovery");
 DEFINE_string(tent_intent_type, "unspec",
               "tent only: intent_type attached to every benchmark request. "
               "unspec|foreground_get|background_prefetch|migration|checkpoint|"
@@ -108,6 +137,20 @@ std::string XferBenchConfig::xport_type;
 std::string XferBenchConfig::backend;
 bool XferBenchConfig::notifi = false;
 std::string XferBenchConfig::tent_transport_hint;
+std::string XferBenchConfig::tent_rdma_devices;
+int XferBenchConfig::tent_dram_numa_node = -1;
+bool XferBenchConfig::enable_runtime_queue = false;
+size_t XferBenchConfig::runtime_queue_max_dispatch_owners = 0;
+std::string XferBenchConfig::receiver_credit_mode;
+size_t XferBenchConfig::receiver_credit_capacity_bytes = 0;
+size_t XferBenchConfig::receiver_credit_capacity_slots = 0;
+size_t XferBenchConfig::receiver_credit_grant_bytes = 0;
+size_t XferBenchConfig::receiver_credit_grant_slots = 0;
+size_t XferBenchConfig::receiver_credit_adaptive_min_owners = 0;
+size_t XferBenchConfig::receiver_credit_adaptive_initial_owners = 0;
+size_t XferBenchConfig::receiver_credit_adaptive_max_owners = 0;
+uint32_t XferBenchConfig::receiver_credit_adaptive_slow_rtt_us = 0;
+uint32_t XferBenchConfig::receiver_credit_adaptive_healthy_pulls = 0;
 std::string XferBenchConfig::tent_intent_type;
 
 int XferBenchConfig::local_gpu_id = 0;
@@ -144,6 +187,25 @@ void XferBenchConfig::loadFromFlags() {
     backend = FLAGS_backend;
     notifi = FLAGS_notifi;
     tent_transport_hint = FLAGS_tent_transport_hint;
+    tent_rdma_devices = FLAGS_tent_rdma_devices;
+    tent_dram_numa_node = FLAGS_tent_dram_numa_node;
+    enable_runtime_queue = FLAGS_enable_runtime_queue;
+    runtime_queue_max_dispatch_owners = FLAGS_runtime_queue_max_dispatch_owners;
+    receiver_credit_mode = FLAGS_receiver_credit_mode;
+    receiver_credit_capacity_bytes = FLAGS_receiver_credit_capacity_bytes;
+    receiver_credit_capacity_slots = FLAGS_receiver_credit_capacity_slots;
+    receiver_credit_grant_bytes = FLAGS_receiver_credit_grant_bytes;
+    receiver_credit_grant_slots = FLAGS_receiver_credit_grant_slots;
+    receiver_credit_adaptive_min_owners =
+        FLAGS_receiver_credit_adaptive_min_owners;
+    receiver_credit_adaptive_initial_owners =
+        FLAGS_receiver_credit_adaptive_initial_owners;
+    receiver_credit_adaptive_max_owners =
+        FLAGS_receiver_credit_adaptive_max_owners;
+    receiver_credit_adaptive_slow_rtt_us =
+        FLAGS_receiver_credit_adaptive_slow_rtt_us;
+    receiver_credit_adaptive_healthy_pulls =
+        FLAGS_receiver_credit_adaptive_healthy_pulls;
     tent_intent_type = FLAGS_tent_intent_type;
 
     local_gpu_id = FLAGS_local_gpu_id;

mooncake-transfer-engine/benchmark/utils.h

@@ -83,6 +83,20 @@ struct XferBenchConfig {
     static std::string backend;
     static bool notifi;
     static std::string tent_transport_hint;
+    static std::string tent_rdma_devices;
+    static int tent_dram_numa_node;
+    static bool enable_runtime_queue;
+    static size_t runtime_queue_max_dispatch_owners;
+    static std::string receiver_credit_mode;
+    static size_t receiver_credit_capacity_bytes;
+    static size_t receiver_credit_capacity_slots;
+    static size_t receiver_credit_grant_bytes;
+    static size_t receiver_credit_grant_slots;
+    static size_t receiver_credit_adaptive_min_owners;
+    static size_t receiver_credit_adaptive_initial_owners;
+    static size_t receiver_credit_adaptive_max_owners;
+    static uint32_t receiver_credit_adaptive_slow_rtt_us;
+    static uint32_t receiver_credit_adaptive_healthy_pulls;
     static std::string tent_intent_type;
 
     static int local_gpu_id;

mooncake-transfer-engine/tent/include/tent/rpc/rpc.h

@@ -47,6 +47,9 @@ enum RpcFuncID {
     Unpin,
     SubscribeSegmentUpdate,
     NotifySegmentUpdated,
+    // Sender-pull control message for receiver-advertised credits. Keep this
+    // ID append-only so mixed-version peers never reinterpret an existing RPC.
+    PullReceiverCredit = 12,
 };
 
 class ClientPool;
@@ -57,35 +60,35 @@ class CoroRpcAgent {
 
     virtual ~CoroRpcAgent();
 
-    CoroRpcAgent(const CoroRpcAgent &) = delete;
-    CoroRpcAgent &operator=(const CoroRpcAgent &) = delete;
+    CoroRpcAgent(const CoroRpcAgent&) = delete;
+    CoroRpcAgent& operator=(const CoroRpcAgent&) = delete;
 
    public:
-    using Function = std::function<void(const std::string_view & /* request */,
-                                        std::string & /* response */)>;
-    Status registerFunction(int func_id, const Function &func);
+    using Function = std::function<void(const std::string_view& /* request */,
+                                        std::string& /* response */)>;
+    Status registerFunction(int func_id, const Function& func);
 
-    Status start(uint16_t &port, bool ipv6 = false);
+    Status start(uint16_t& port, bool ipv6 = false);
 
     Status stop();
 
-    Status call(const std::string &server_addr, int func_id,
-                const std::string_view &request, std::string &response);
+    Status call(const std::string& server_addr, int func_id,
+                const std::string_view& request, std::string& response);
 
     using AsyncCallback = std::function<void(Status, std::string)>;
-    void callAsync(const std::string &server_addr, int func_id,
-                   const std::string &request, AsyncCallback callback);
+    void callAsync(const std::string& server_addr, int func_id,
+                   const std::string& request, AsyncCallback callback);
 
     async_simple::coro::Lazy<std::pair<Status, std::string>> callCoroutine(
         std::string server_addr, int func_id, std::string request);
 
    private:
     void process(int func_id);
 
-    std::shared_ptr<ClientPool> getOrCreatePool(const std::string &server_addr);
+    std::shared_ptr<ClientPool> getOrCreatePool(const std::string& server_addr);
 
    private:
-    coro_rpc::coro_rpc_server *server_ = nullptr;
+    coro_rpc::coro_rpc_server* server_ = nullptr;
 
     std::mutex pools_mutex_;
     std::unordered_map<std::string, std::shared_ptr<ClientPool>> pools_;

mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h

@@ -131,6 +131,8 @@ class LocalTransferAdmissionQueue {
         size_t max_owners, size_t max_bytes,
         std::vector<QueueOwnerId>* dropped_owner_ids = nullptr);
 
+    Status deferDispatch(QueueOwnerId owner_id);
+
     // Install the step-3 degradation policy inputs. Optional; without it the
     // queue never drops (default behavior). now defaults to steady_clock.
     void setDegradationPolicy(BandwidthProvider bandwidth_provider,

mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h

@@ -28,6 +28,7 @@
 #include <variant>
 
 #include "tent/runtime/metastore.h"
+#include "tent/runtime/receiver_credit_allocator.h"
 #include "tent/runtime/segment.h"
 #include "tent/runtime/segment_manager.h"
 #include "tent/rpc/rpc.h"
@@ -107,6 +108,20 @@ class ControlClient {
     static void notifySegmentUpdatedAsync(
         const std::string& server_addr, const std::string& segment_name,
         const onNotifySegmentUpdateFailure& on_failure);
+
+    static Status pullReceiverCredit(const std::string& server_addr,
+                                     const ReceiverCreditPullRequestV1& request,
+                                     ReceiverCreditPullResponseV1& response);
+
+    using OnReceiverCreditPull =
+        std::function<void(Status, ReceiverCreditPullResponseV1)>;
+    // The caller owns `agent`; retaining it in the callback owner makes the
+    // in-flight coroutine independent from thread-local client teardown.
+    static void pullReceiverCreditAsync(
+        const std::shared_ptr<CoroRpcAgent>& agent,
+        const std::string& server_addr,
+        const ReceiverCreditPullRequestV1& request,
+        OnReceiverCreditPull callback);
 };
 
 class ControlService {
@@ -129,6 +144,9 @@ class ControlService {
         notify_callback_ = callback;
     }
 
+    void setReceiverCreditAllocator(
+        std::shared_ptr<ReceiverCreditAllocator> allocator);
+
     Status start(uint16_t& port, bool ipv6_ = false);
 
    private:
@@ -160,13 +178,18 @@ class ControlService {
     void onSegmentUpdated(const std::string_view& request,
                           std::string& response);
 
+    void onPullReceiverCredit(const std::string_view& request,
+                              std::string& response);
+
    private:
     std::unique_ptr<SegmentManager> manager_;
     std::shared_ptr<CoroRpcAgent> rpc_server_;
 
     OnReceiveBootstrap bootstrap_callback_;
     OnNotify notify_callback_;
     TransferEngineImpl* impl_;
+    std::mutex receiver_credit_mutex_;
+    std::shared_ptr<ReceiverCreditAllocator> receiver_credit_allocator_;
 };
 
 }  // namespace tent

mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h

@@ -61,6 +61,17 @@ enum class CreditUpdateDisposition : uint8_t {
     SequenceGap
 };
 
+// A consistent sender-side view used to report cumulative usage to the
+// receiver. Completion is deliberately separate from consumption: completing
+// work does not mint sender credit locally; only a later receiver grant can do
+// that.
+struct CreditLedgerSnapshot {
+    uint64_t epoch{0}, last_sequence{0};
+    bool has_update{false};
+    std::array<uint64_t, kCreditResourceCount> grants{}, consumed{},
+        completed{};
+};
+
 // Private, sender-side state model. It has no network or Admission integration.
 class SenderCreditLedger {
    public:
@@ -74,17 +85,28 @@ class SenderCreditLedger {
     Status applyUpdate(const CreditKey&, const ReceiverCreditUpdateV1&,
                        CreditUpdateDisposition&);
     Status tryReserve(const CreditKey&, const CreditCharge&);
+    Status tryReserve(const CreditKey&, uint64_t expected_epoch,
+                      const CreditCharge&);
     // Only for work not yet handed to a transport; completions need a new
     // grant.
-    Status rollbackReservation(const CreditKey&, const CreditCharge&);
+    Status rollbackReservation(const CreditKey&, uint64_t expected_epoch,
+                               const CreditCharge&);
+    // Records transport-owned work reaching a terminal state. This advances a
+    // cumulative completion counter without changing consumed or available
+    // credit. The receiver must explicitly issue a new cumulative grant.
+    Status recordCompletion(const CreditKey&, uint64_t expected_epoch,
+                            const CreditCharge&);
+    Status snapshot(const CreditKey&, uint64_t expected_epoch,
+                    CreditLedgerSnapshot&) const;
     Status available(const CreditKey&, CreditResource, uint64_t&) const;
     Status consumed(const CreditKey&, CreditResource, uint64_t&) const;
 
    private:
     struct Entry {
         uint64_t epoch{0}, last_sequence{0};
         bool has_update{false};
-        std::array<uint64_t, kCreditResourceCount> grants{}, consumed{};
+        std::array<uint64_t, kCreditResourceCount> grants{}, consumed{},
+            completed{};
     };
     static Status resourceIndex(CreditResource, size_t&);
     static Status normalize(const CreditCharge&,