diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 3ca9156306..abab30f7d1 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.cpp +++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp @@ -53,6 +53,53 @@ std::shared_ptr 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 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"; diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp index ca69e87020..b93e99b8c7 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/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; diff --git a/mooncake-transfer-engine/benchmark/utils.h b/mooncake-transfer-engine/benchmark/utils.h index ddb4e73b08..c56026118c 100644 --- a/mooncake-transfer-engine/benchmark/utils.h +++ b/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; diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 3105c6fdaf..3ad4d8f60b 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/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,24 +60,24 @@ 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; - Status registerFunction(int func_id, const Function &func); + using Function = std::function; + 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 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> callCoroutine( std::string server_addr, int func_id, std::string request); @@ -82,10 +85,10 @@ class CoroRpcAgent { private: void process(int func_id); - std::shared_ptr getOrCreatePool(const std::string &server_addr); + std::shared_ptr 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> pools_; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index 4253adb191..b6476f4736 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/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* 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, diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h index a4739f27aa..61ec30e4b6 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -28,6 +28,7 @@ #include #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; + // 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& 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 allocator); + Status start(uint16_t& port, bool ipv6_ = false); private: @@ -160,6 +178,9 @@ 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 manager_; std::shared_ptr rpc_server_; @@ -167,6 +188,8 @@ class ControlService { OnReceiveBootstrap bootstrap_callback_; OnNotify notify_callback_; TransferEngineImpl* impl_; + std::mutex receiver_credit_mutex_; + std::shared_ptr receiver_credit_allocator_; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h index 09b243ab13..c88751356b 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h +++ b/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 grants{}, consumed{}, + completed{}; +}; + // Private, sender-side state model. It has no network or Admission integration. class SenderCreditLedger { public: @@ -74,9 +85,19 @@ 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; @@ -84,7 +105,8 @@ class SenderCreditLedger { struct Entry { uint64_t epoch{0}, last_sequence{0}; bool has_update{false}; - std::array grants{}, consumed{}; + std::array grants{}, consumed{}, + completed{}; }; static Status resourceIndex(CreditResource, size_t&); static Status normalize(const CreditCharge&, diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_allocator.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_allocator.h new file mode 100644 index 0000000000..b99e88f994 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_allocator.h @@ -0,0 +1,102 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_ALLOCATOR_H +#define TENT_RUNTIME_RECEIVER_CREDIT_ALLOCATOR_H + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/receiver_credit_protocol.h" + +namespace mooncake::tent { + +struct ReceiverCreditAllocatorConfig { + std::array capacity{}; + std::array max_grant_per_pull{}; + size_t max_entries{0}; + uint32_t ttl_ms{0}; + uint32_t retry_after_us{0}; + ReceiverSessionId receiver_session_id; + uint64_t epoch{0}; +}; + +struct ReceiverCreditAllocatorSnapshot { + std::array capacity{}; + std::array committed{}; + std::array free{}; + size_t entries{0}; + uint64_t pull_requests{0}; +}; + +// Receiver-wide allocator for a configured ingress budget. Credits represent +// bounded bytes/request slots that the receiver is willing to have outstanding; +// this class deliberately does not allocate memory extents or expose rkeys. +class ReceiverCreditAllocator { + public: + static Status create(const ReceiverCreditAllocatorConfig& config, + std::unique_ptr& allocator); + + // A syntactically valid request always produces a protocol response. Retry, + // restart fencing and sequence rejection are represented by + // response.status. The output is unchanged only when the request itself is + // malformed or an internal invariant is violated. + Status pull(const ReceiverCreditPullRequestV1& request, + ReceiverCreditPullResponseV1& response); + + // Recomputes all accounting under the lock and returns an error instead of + // publishing a snapshot if any invariant has been violated. + Status snapshot(ReceiverCreditAllocatorSnapshot& snapshot) const; + + private: + struct EntryKey { + uint64_t sender_peer{0}; + uint32_t qos_class{0}; + bool operator==(const EntryKey& other) const { + return sender_peer == other.sender_peer && + qos_class == other.qos_class; + } + }; + + struct EntryKeyHash { + size_t operator()(const EntryKey& key) const noexcept; + }; + + struct Entry { + std::array granted{}; + std::array consumed{}; + std::array completed{}; + uint64_t update_sequence{0}; + bool has_request{false}; + ReceiverCreditPullRequestV1 last_request; + ReceiverCreditPullResponseV1 last_response; + }; + + explicit ReceiverCreditAllocator(ReceiverCreditAllocatorConfig config) + : config_(std::move(config)) {} + + static Status validateConfig(const ReceiverCreditAllocatorConfig& config); + ReceiverCreditPullResponseV1 makeResponse( + ReceiverCreditPullStatus status, uint32_t retry_after_us, + uint32_t qos_class, + const std::array& grants, + uint64_t update_sequence) const; + Status validateInvariantsLocked( + ReceiverCreditAllocatorSnapshot* snapshot) const; + + const ReceiverCreditAllocatorConfig config_; + std::atomic pull_requests_{0}; + mutable std::mutex mutex_; + std::array committed_{}; + std::unordered_map entries_; +}; + +} // namespace mooncake::tent + +#endif // TENT_RUNTIME_RECEIVER_CREDIT_ALLOCATOR_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_config.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_config.h new file mode 100644 index 0000000000..f0786a778f --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_config.h @@ -0,0 +1,56 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_CONFIG_H +#define TENT_RUNTIME_RECEIVER_CREDIT_CONFIG_H + +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/runtime/receiver_credit_control.h" + +namespace mooncake::tent { + +// Strict, default-off runtime configuration for receiver-credit control. +// Resource arrays use the CreditResource wire order (enum value minus one). +struct ReceiverCreditRuntimeConfig { + static constexpr size_t kDefaultMaxPeers = 1024; + static constexpr uint32_t kDefaultFreshnessTtlMs = 1000; + static constexpr uint32_t kDefaultRetryAfterUs = 1000; + static constexpr uint32_t kDefaultProgressIntervalUs = 1000; + static constexpr size_t kDefaultAdaptiveMinOwners = 1; + static constexpr size_t kDefaultAdaptiveInitialOwners = 2; + static constexpr size_t kDefaultAdaptiveMaxOwners = 2; + // The observed failure mode is a TCP minimum retransmission-timeout event + // (~200 ms). Keep ample distance from healthy scheduler jitter, which can + // reach ~10 ms. + static constexpr uint32_t kDefaultAdaptiveSlowRttUs = 20000; + static constexpr uint32_t kDefaultAdaptiveHealthyPulls = 512; + + CreditRolloutMode mode{CreditRolloutMode::Disabled}; + std::array capacity{}; + std::array max_grant_per_pull{}; + size_t max_peers{kDefaultMaxPeers}; + uint32_t freshness_ttl_ms{kDefaultFreshnessTtlMs}; + uint32_t retry_after_us{kDefaultRetryAfterUs}; + uint32_t progress_interval_us{kDefaultProgressIntervalUs}; + uint32_t default_qos_class{0}; + bool adaptive_dispatch_enabled{true}; + size_t adaptive_dispatch_min_owners{kDefaultAdaptiveMinOwners}; + size_t adaptive_dispatch_initial_owners{kDefaultAdaptiveInitialOwners}; + size_t adaptive_dispatch_max_owners{kDefaultAdaptiveMaxOwners}; + uint32_t adaptive_dispatch_slow_rtt_us{kDefaultAdaptiveSlowRttUs}; + uint32_t adaptive_dispatch_healthy_pulls{kDefaultAdaptiveHealthyPulls}; +}; + +// Loads and validates the receiver_credit section. The output is changed only +// after the complete section has passed validation. +Status loadReceiverCreditConfig(const Config& config, + bool runtime_queue_enabled, + ReceiverCreditRuntimeConfig& output); + +} // namespace mooncake::tent + +#endif // TENT_RUNTIME_RECEIVER_CREDIT_CONFIG_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_control.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_control.h new file mode 100644 index 0000000000..cd73fb28c9 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_control.h @@ -0,0 +1,177 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_CONTROL_H +#define TENT_RUNTIME_RECEIVER_CREDIT_CONTROL_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/receiver_credit.h" + +namespace mooncake::tent { + +enum class CreditRolloutMode : uint8_t { Disabled, Optional, Required }; +enum class CreditPeerState : uint8_t { + Disabled, + Negotiating, + Legacy, + Active, + Stale, + Failed +}; + +class CreditCapabilityState { + public: + explicit CreditCapabilityState(CreditRolloutMode mode); + Status beginNegotiation(); + Status completeNegotiation(const std::vector& peer_versions); + Status markStale(); + Status refresh(uint16_t version); + CreditPeerState state() const { return state_; } + uint16_t version() const { return version_; } + + private: + CreditRolloutMode mode_; + CreditPeerState state_; + uint16_t version_{0}; +}; + +class CreditCapabilityCodecV1 { + public: + static constexpr size_t kHeaderBytes = 8; + static constexpr size_t kMaxVersions = 8; + static constexpr size_t kMaxWireBytes = + kHeaderBytes + kMaxVersions * sizeof(uint16_t); + + static Status encode(const std::vector& versions, + std::string& wire); + static Status decode(std::string_view wire, + std::vector& versions); +}; + +struct CreditActivationV1 { + uint16_t schema_version{1}; + uint16_t chosen_version{1}; + ReceiverSessionId receiver_session_id; + uint64_t epoch{0}; + uint32_t freshness_ttl_ms{0}; +}; + +class CreditActivationCodecV1 { + public: + static constexpr size_t kWireBytes = 40; + static Status encode(const CreditActivationV1& activation, + std::string& wire); + static Status decode(std::string_view wire, CreditActivationV1& activation); +}; + +struct CreditPeerContextSnapshot { + CreditKey key; + uint64_t epoch{0}; + uint32_t freshness_ttl_ms{0}; + std::chrono::steady_clock::time_point refreshed_at{}; + + bool isFresh(std::chrono::steady_clock::time_point now = + std::chrono::steady_clock::now()) const { + if (freshness_ttl_ms == 0) return true; + return now - refreshed_at <= + std::chrono::milliseconds(freshness_ttl_ms); + } +}; + +// Bounded sender-side mapping established by capability activation. Runtime +// admission snapshots this context so a later receiver restart cannot retag +// already queued work. +class CreditPeerContextTable { + public: + explicit CreditPeerContextTable(size_t max_entries = 1024) + : max_entries_(max_entries) {} + + Status activate(uint64_t target_id, uint64_t sender_peer, + uint32_t qos_class, const CreditActivationV1& activation); + Status lookup(uint64_t target_id, uint32_t qos_class, + CreditPeerContextSnapshot& snapshot) const; + Status lookupFresh(uint64_t target_id, uint32_t qos_class, + CreditPeerContextSnapshot& snapshot) const; + Status deactivate(uint64_t target_id, uint32_t qos_class, + const ReceiverSessionId& receiver_session, + uint64_t epoch); + size_t size() const; + + private: + struct LookupKey { + uint64_t target_id{0}; + uint32_t qos_class{0}; + bool operator==(const LookupKey& other) const { + return target_id == other.target_id && qos_class == other.qos_class; + } + }; + struct LookupKeyHash { + size_t operator()(const LookupKey& key) const noexcept; + }; + + mutable std::mutex mutex_; + const size_t max_entries_; + std::unordered_map + contexts_; +}; + +struct CreditControlEnvelope { + CreditKey key; + ReceiverCreditUpdateV1 update; +}; + +// Bounded, nonblocking publisher queue. Control callbacks only validate and +// enqueue; the runtime owner drains and mutates the ledger. +class BoundedCreditUpdateInbox { + public: + explicit BoundedCreditUpdateInbox(size_t capacity) : capacity_(capacity) {} + Status tryPublish(CreditControlEnvelope envelope); + size_t drain(std::vector& output, + size_t max_updates); + size_t size() const; + + private: + const size_t capacity_; + mutable std::mutex mutex_; + std::deque queue_; +}; + +class ReceiverCreditCodecV1 { + public: + static constexpr size_t kHeaderBytes = 52; + static constexpr size_t kGrantBytes = 12; + static constexpr size_t kMaxWireBytes = + kHeaderBytes + kCreditResourceCount * kGrantBytes; + + static Status encode(const ReceiverCreditUpdateV1& update, + std::string& wire); + static Status decode(std::string_view wire, ReceiverCreditUpdateV1& update); +}; + +// Transport-neutral ingress used by a control-plane callback. It performs no +// ledger mutation and never blocks waiting for queue capacity. +class ReceiverCreditIngress { + public: + ReceiverCreditIngress(BoundedCreditUpdateInbox& inbox, CreditKey key, + uint64_t epoch) + : inbox_(inbox), key_(key), epoch_(epoch) {} + + Status tryAccept(std::string_view wire); + + private: + BoundedCreditUpdateInbox& inbox_; + CreditKey key_; + uint64_t epoch_; +}; + +} // namespace mooncake::tent +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_controller.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_controller.h new file mode 100644 index 0000000000..8623f47a8c --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_controller.h @@ -0,0 +1,163 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_CONTROLLER_H +#define TENT_RUNTIME_RECEIVER_CREDIT_CONTROLLER_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/rpc/rpc.h" +#include "tent/runtime/receiver_credit_config.h" +#include "tent/runtime/receiver_credit_dispatch.h" +#include "tent/runtime/receiver_credit_protocol.h" + +namespace mooncake::tent { + +struct AdaptiveDispatchSnapshot { + size_t current_owners{0}; + size_t learned_ceiling{0}; + size_t suspect_level{0}; + uint64_t slow_or_failed_pulls{0}; + uint64_t reductions{0}; + uint64_t increases{0}; +}; + +// Process-lifetime safety controller for the shared RDMA/control path. A slow +// or failed pull multiplicatively reduces the dispatch window and records the +// failed level as unsafe. Healthy pulls additively recover, but never probe a +// level already shown to starve control traffic during this process lifetime. +class AdaptiveCreditDispatchLimiter { + public: + explicit AdaptiveCreditDispatchLimiter( + const ReceiverCreditRuntimeConfig& config, + size_t hard_max_owners = std::numeric_limits::max()); + + void observe(std::chrono::nanoseconds elapsed, bool rpc_ok, + size_t owners_at_start); + size_t ownerLimit() const; + AdaptiveDispatchSnapshot snapshot() const; + + private: + const bool enabled_; + const size_t min_owners_; + const uint64_t slow_rtt_ns_; + const uint32_t healthy_pulls_per_increase_; + + std::atomic current_owners_; + mutable std::mutex mutex_; + size_t learned_ceiling_; + size_t suspect_level_{0}; + uint32_t healthy_pulls_{0}; + uint64_t slow_or_failed_pulls_{0}; + uint64_t reductions_{0}; + uint64_t increases_{0}; +}; + +// Sender-side, nonblocking pull coordinator. At most one RPC is outstanding +// for each (target, QoS class); callers only publish bounded demand and return. +// RPC callbacks mutate the separately synchronized context/ledger models and +// never call back into TransferEngineImpl, so engine teardown cannot race a +// raw `this` callback. Runtime-queue fallback progress observes the update. +class ReceiverCreditPullController + : public std::enable_shared_from_this { + public: + static Status create( + const ReceiverCreditRuntimeConfig& config, uint64_t sender_peer, + std::shared_ptr contexts, + std::shared_ptr ledger, + std::shared_ptr& controller, + size_t dispatch_owner_ceiling = std::numeric_limits::max()); + + ~ReceiverCreditPullController(); + + ReceiverCreditPullController(const ReceiverCreditPullController&) = delete; + ReceiverCreditPullController& operator=( + const ReceiverCreditPullController&) = delete; + + // Publishes the minimum complete charge needed by blocked work. The + // desired pull window is the larger of this charge and configured batch. + Status request(uint64_t target_id, const std::string& server_addr, + uint32_t qos_class, const CreditCharge& minimum_charge); + + CreditPeerState peerState(uint64_t target_id, uint32_t qos_class) const; + size_t peerCount() const; + size_t dispatchOwnerLimit() const; + AdaptiveDispatchSnapshot adaptiveDispatchSnapshot() const; + + // Prevents new pulls. Outstanding callbacks retain this object until they + // finish but become no-ops, avoiding a callback-to-engine lifetime edge. + void stop(); + + private: + struct PeerKey { + uint64_t target_id{0}; + uint32_t qos_class{0}; + bool operator==(const PeerKey& other) const { + return target_id == other.target_id && qos_class == other.qos_class; + } + }; + + struct PeerKeyHash { + size_t operator()(const PeerKey& key) const noexcept; + }; + + struct Peer { + std::string server_addr; + std::array minimum{}; + uint64_t next_request_sequence{1}; + bool in_flight{false}; + bool dirty{false}; + CreditPeerState state{CreditPeerState::Negotiating}; + std::chrono::steady_clock::time_point pull_started_at{}; + size_t pull_dispatch_owners{0}; + }; + + ReceiverCreditPullController( + ReceiverCreditRuntimeConfig config, uint64_t sender_peer, + std::shared_ptr contexts, + std::shared_ptr ledger, + size_t dispatch_owner_ceiling); + + static Status normalize( + const CreditCharge& charge, + std::array& normalized); + void launch(PeerKey key); + void finish(PeerKey key, Status rpc_status, + ReceiverCreditPullResponseV1 response); + Status applyResponse(const PeerKey& key, + const ReceiverCreditPullResponseV1& response, + bool& request_again); + bool oldGenerationDrained(const CreditPeerContextSnapshot& context) const; + + const ReceiverCreditRuntimeConfig config_; + const uint64_t sender_peer_; + const std::shared_ptr contexts_; + const std::shared_ptr ledger_; + // Engine-independent client lifetime: an in-flight callback retains the + // controller, which in turn retains this agent until the coroutine ends. + const std::shared_ptr rpc_agent_; + AdaptiveCreditDispatchLimiter dispatch_limiter_; + + mutable std::mutex mutex_; + bool stopped_{false}; + bool summary_logged_{false}; + uint64_t pulls_started_{0}; + uint64_t pulls_completed_{0}; + uint64_t pull_latency_ns_sum_{0}; + uint64_t pull_latency_ns_max_{0}; + std::array pull_latency_buckets_{}; + std::unordered_map peers_; +}; + +} // namespace mooncake::tent + +#endif // TENT_RUNTIME_RECEIVER_CREDIT_CONTROLLER_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_dispatch.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_dispatch.h new file mode 100644 index 0000000000..3f7a90ef7b --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_dispatch.h @@ -0,0 +1,60 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_DISPATCH_H +#define TENT_RUNTIME_RECEIVER_CREDIT_DISPATCH_H + +#include + +#include "tent/runtime/receiver_credit.h" +#include "tent/runtime/receiver_credit_control.h" + +namespace mooncake::tent { + +struct CreditDispatchSnapshot { + uint64_t target_id{0}; + uint32_t qos_class{0}; + CreditKey key; + uint64_t epoch{0}; + CreditCharge charge; +}; + +enum class CreditReservationState : uint8_t { + Empty, + Reserved, + Committed, + RolledBack, + Released +}; + +struct CreditDispatchReservation { + CreditDispatchSnapshot snapshot; + CreditReservationState state{CreditReservationState::Empty}; +}; + +// Runtime-owner helper for the dispatch boundary. The owner must serialize +// peer activation/update processing with reserve/commit/rollback calls. +class ReceiverCreditDispatchGate { + public: + ReceiverCreditDispatchGate(const CreditPeerContextTable& contexts, + SenderCreditLedger& ledger) + : contexts_(contexts), ledger_(ledger) {} + + Status snapshot(uint64_t target_id, uint32_t qos_class, CreditCharge charge, + CreditDispatchSnapshot& output) const; + Status tryReserve(const CreditDispatchSnapshot& snapshot, + CreditDispatchReservation& reservation); + Status commit(CreditDispatchReservation& reservation); + Status rollback(CreditDispatchReservation& reservation); + // Releases transport-owned capacity at terminal completion. It does not + // restore local availability; the receiver must issue the next grant. + Status release(CreditDispatchReservation& reservation); + + private: + const CreditPeerContextTable& contexts_; + SenderCreditLedger& ledger_; +}; + +} // namespace mooncake::tent + +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_protocol.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_protocol.h new file mode 100644 index 0000000000..d79806d326 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit_protocol.h @@ -0,0 +1,115 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_PROTOCOL_H +#define TENT_RUNTIME_RECEIVER_CREDIT_PROTOCOL_H + +#include +#include +#include +#include +#include + +#include "tent/runtime/receiver_credit_control.h" + +namespace mooncake::tent { + +// A sender reports cumulative usage and asks the receiver to preserve at least +// minimum_available, and preferably desired_available, after this pull. +struct ReceiverCreditResourceUsageV1 { + CreditResource resource{CreditResource::DataBytes}; + uint64_t consumed_total{0}; + uint64_t completed_total{0}; + uint64_t minimum_available{0}; + uint64_t desired_available{0}; + + bool operator==(const ReceiverCreditResourceUsageV1& other) const { + return resource == other.resource && + consumed_total == other.consumed_total && + completed_total == other.completed_total && + minimum_available == other.minimum_available && + desired_available == other.desired_available; + } +}; + +struct ReceiverCreditPullRequestV1 { + uint16_t schema_version{1}; + uint16_t flags{0}; + uint64_t sender_peer{0}; + uint32_t qos_class{0}; + // The first pull uses an all-zero session and epoch. Once activated, all + // three values must exactly match the receiver response. + ReceiverSessionId expected_receiver_session_id; + uint64_t expected_epoch{0}; + uint64_t request_sequence{0}; + uint64_t last_update_sequence{0}; + std::vector resources; + + bool operator==(const ReceiverCreditPullRequestV1& other) const { + return schema_version == other.schema_version && flags == other.flags && + sender_peer == other.sender_peer && + qos_class == other.qos_class && + expected_receiver_session_id == + other.expected_receiver_session_id && + expected_epoch == other.expected_epoch && + request_sequence == other.request_sequence && + last_update_sequence == other.last_update_sequence && + resources == other.resources; + } +}; + +enum class ReceiverCreditPullStatus : uint16_t { + Granted = 1, + Retry = 2, + SessionChanged = 3, + Unsupported = 4, + Rejected = 5, +}; + +struct ReceiverCreditPullResponseV1 { + uint16_t schema_version{1}; + uint16_t flags{0}; + ReceiverCreditPullStatus status{ReceiverCreditPullStatus::Rejected}; + uint32_t retry_after_us{0}; + CreditActivationV1 activation; + // Allocator responses carry a full, cumulative four-resource update. + // Unsupported/Rejected may instead use an all-zero activation/update so a + // ControlService without an allocator can still return a typed response. + ReceiverCreditUpdateV1 update; +}; + +class ReceiverCreditPullRequestCodecV1 { + public: + static constexpr size_t kHeaderBytes = 64; + static constexpr size_t kResourceBytes = 36; + static constexpr size_t kMaxWireBytes = + kHeaderBytes + kCreditResourceCount * kResourceBytes; + + static Status validate(const ReceiverCreditPullRequestV1& request); + static Status encode(const ReceiverCreditPullRequestV1& request, + std::string& wire); + static Status decode(std::string_view wire, + ReceiverCreditPullRequestV1& request); +}; + +class ReceiverCreditPullResponseCodecV1 { + public: + static constexpr size_t kHeaderBytes = 100; + static constexpr size_t kGrantBytes = 12; + // The frame stays fixed-size for a zero payload; its grant slots become + // validated zero padding. + static constexpr size_t kWireBytes = + kHeaderBytes + kCreditResourceCount * kGrantBytes; + + static Status makeZeroPayload(ReceiverCreditPullStatus status, + ReceiverCreditPullResponseV1& response); + static Status validate(const ReceiverCreditPullResponseV1& response); + static Status encode(const ReceiverCreditPullResponseV1& response, + std::string& wire); + static Status decode(std::string_view wire, + ReceiverCreditPullResponseV1& response); +}; + +} // namespace mooncake::tent + +#endif // TENT_RUNTIME_RECEIVER_CREDIT_PROTOCOL_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h index e913abe9d2..1814be139a 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h @@ -178,6 +178,15 @@ struct SegmentDesc { // of `name`. std::string rdma_server_name; + // Process-lifetime identity and optional control-plane capabilities. + // These fields are deliberately independent from RDMA bootstrap: a TENT + // process can serve receiver credits for non-RDMA plans as well. Older + // peers ignore the extra JSON fields, while newer peers treat an absent + // capability list as unsupported rather than as unlimited capacity. + uint64_t peer_session_high{0}; + uint64_t peer_session_low{0}; + std::vector receiver_credit_versions; + public: BufferDesc* findBuffer(uint64_t base, uint64_t length); DeviceDesc* findDevice(const std::string& name); @@ -201,6 +210,13 @@ inline void to_json(json& j, const SegmentDesc& s) { if (!s.rdma_server_name.empty()) { j["rdma_server_name"] = s.rdma_server_name; } + if (s.peer_session_high != 0 || s.peer_session_low != 0) { + j["peer_session_high"] = s.peer_session_high; + j["peer_session_low"] = s.peer_session_low; + } + if (!s.receiver_credit_versions.empty()) { + j["receiver_credit_versions"] = s.receiver_credit_versions; + } if (s.type == SegmentType::Memory) { j["detail"] = std::get(s.detail); } else { @@ -216,6 +232,15 @@ inline void from_json(const json& j, SegmentDesc& s) { if (j.contains("rdma_server_name")) { j.at("rdma_server_name").get_to(s.rdma_server_name); } + if (j.contains("peer_session_high")) { + j.at("peer_session_high").get_to(s.peer_session_high); + } + if (j.contains("peer_session_low")) { + j.at("peer_session_low").get_to(s.peer_session_low); + } + if (j.contains("receiver_credit_versions")) { + j.at("receiver_credit_versions").get_to(s.receiver_credit_versions); + } if (s.type == SegmentType::Memory) { s.detail = j.at("detail").get(); } else { @@ -295,4 +320,4 @@ inline std::vector coalesceRegions( } // namespace tent } // namespace mooncake -#endif // SEGMENT_H \ No newline at end of file +#endif // SEGMENT_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 5ae33785bc..e7d9126ac2 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -20,8 +20,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -31,6 +33,10 @@ #include "tent/common/status.h" #include "tent/common/types.h" #include "tent/runtime/admission_queue.h" +#include "tent/runtime/receiver_credit_dispatch.h" +#include "tent/runtime/receiver_credit_allocator.h" +#include "tent/runtime/receiver_credit_config.h" +#include "tent/runtime/receiver_credit_controller.h" #include "tent/runtime/transport.h" #include "tent/runtime/transport_selector.h" @@ -181,6 +187,16 @@ class TransferEngineImpl { } } + using CreditQosProvider = std::function; + // Installs the control-plane state used by the opt-in receiver-credit + // dispatch gate. The engine retains shared ownership so queued snapshots + // and reservations cannot outlive their context or ledger. Must be called + // before submitting transfers when receiver_credit/enabled is true. + Status installReceiverCreditDispatch( + std::shared_ptr contexts, + std::shared_ptr ledger, + CreditQosProvider qos_provider = nullptr); + // Wake the optional event-driven progress worker for `batch_id`. No-op if // enable_progress_worker is false. Transport completion paths use this as // an idempotent "maybe ready" signal. @@ -301,6 +317,8 @@ class TransferEngineImpl { struct RuntimeQueueConfig { bool enabled{false}; + bool receiver_credit_enabled{false}; + uint32_t receiver_credit_default_qos_class{0}; QueueLimits limits{}; size_t max_dispatch_owners{0}; size_t max_dispatch_bytes{0}; @@ -313,6 +331,11 @@ class TransferEngineImpl { std::vector public_task_ids; size_t byte_charge{0}; bool in_dispatch_window{false}; + std::optional credit_snapshot; + std::optional credit_reservation; + std::optional credit_charge; + std::string credit_server_addr; + uint32_t credit_qos_class{0}; }; private: @@ -342,8 +365,19 @@ class TransferEngineImpl { bool enable_auto_failover_on_poll_{true}; bool enable_progress_worker_{false}; RuntimeQueueConfig runtime_queue_config_; + ReceiverCreditRuntimeConfig receiver_credit_config_; + bool receiver_credit_production_enabled_{false}; + ReceiverSessionId local_peer_session_; + uint64_t local_sender_peer_{0}; std::unique_ptr runtime_queue_; std::unordered_map queued_owners_; + std::shared_ptr receiver_credit_contexts_; + std::shared_ptr receiver_credit_ledger_; + std::shared_ptr receiver_credit_dispatch_gate_; + std::shared_ptr receiver_credit_allocator_; + std::shared_ptr + receiver_credit_pull_controller_; + CreditQosProvider receiver_credit_qos_provider_; size_t dispatch_inflight_owners_{0}; size_t dispatch_inflight_bytes_{0}; uint64_t next_batch_token_{1}; diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 6c552dea07..edc2fa67ea 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -340,6 +340,18 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( return picked; } +Status LocalTransferAdmissionQueue::deferDispatch(QueueOwnerId owner_id) { + auto owner_it = owners_.find(owner_id); + if (owner_it == owners_.end()) + return Status::InvalidEntry("queue owner not found" LOC_MARK); + if (owner_it->second.state != QueueState::Dispatching) + return Status::InvalidEntry( + "only dispatching owner can be deferred" LOC_MARK); + owner_it->second.state = QueueState::Queued; + fifo_.push_front(owner_id); + return Status::OK(); +} + Status LocalTransferAdmissionQueue::complete( QueueOwnerId owner_id, TransferStatusEnum terminal_status) { if (owner_id == 0) { diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index ca5facb0e6..f70dcc32fe 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -97,6 +97,55 @@ Status ControlClient::probe(const std::string& server_addr) { return tl_rpc_agent.call(server_addr, Probe, request, response); } +Status ControlClient::pullReceiverCredit( + const std::string& server_addr, const ReceiverCreditPullRequestV1& request, + ReceiverCreditPullResponseV1& response) { + std::string request_wire, response_wire; + CHECK_STATUS( + ReceiverCreditPullRequestCodecV1::encode(request, request_wire)); + CHECK_STATUS(tl_rpc_agent.call(server_addr, PullReceiverCredit, + request_wire, response_wire)); + if (response_wire.empty()) + return Status::NotImplemented( + "peer does not implement receiver credit pull" LOC_MARK); + return ReceiverCreditPullResponseCodecV1::decode(response_wire, response); +} + +void ControlClient::pullReceiverCreditAsync( + const std::shared_ptr& agent, const std::string& server_addr, + const ReceiverCreditPullRequestV1& request, OnReceiverCreditPull callback) { + if (!agent) { + callback(Status::InvalidArgument( + "receiver credit pull RPC agent is null" LOC_MARK), + {}); + return; + } + std::string request_wire; + auto status = + ReceiverCreditPullRequestCodecV1::encode(request, request_wire); + if (!status.ok()) { + callback(std::move(status), {}); + return; + } + agent->callAsync(server_addr, PullReceiverCredit, request_wire, + [agent, callback = std::move(callback)]( + Status rpc_status, std::string response_wire) mutable { + ReceiverCreditPullResponseV1 response; + if (rpc_status.ok()) { + if (response_wire.empty()) { + rpc_status = Status::NotImplemented( + "peer does not implement receiver credit " + "pull" LOC_MARK); + } else { + rpc_status = + ReceiverCreditPullResponseCodecV1::decode( + response_wire, response); + } + } + callback(std::move(rpc_status), std::move(response)); + }); +} + inline void to_json(json& j, const Request& r) { j = json{{"opcode", r.opcode == Request::READ ? "READ" : "WRITE"}, {"source", reinterpret_cast(r.source)}, @@ -216,14 +265,39 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onSegmentUpdated(request, response); }); + rpc_server_->registerFunction( + PullReceiverCredit, + [this](const std::string_view& request, std::string& response) { + onPullReceiverCredit(request, response); + }); } -ControlService::~ControlService() {} +ControlService::~ControlService() { + std::shared_ptr allocator; + { + std::lock_guard lock(receiver_credit_mutex_); + allocator = receiver_credit_allocator_; + } + if (allocator) { + ReceiverCreditAllocatorSnapshot snapshot; + if (allocator->snapshot(snapshot).ok()) { + LOG(INFO) << "Receiver credit control summary: pulls=" + << snapshot.pull_requests + << " peers=" << snapshot.entries; + } + } +} Status ControlService::start(uint16_t& port, bool ipv6_) { return rpc_server_->start(port, ipv6_); } +void ControlService::setReceiverCreditAllocator( + std::shared_ptr allocator) { + std::lock_guard lock(receiver_credit_mutex_); + receiver_credit_allocator_ = std::move(allocator); +} + void ControlService::onGetSegmentDesc(const std::string_view& request, std::string& response) { // Re-use the cached dump shared across concurrent peer fetches. @@ -344,6 +418,41 @@ void ControlService::onSegmentUpdated(const std::string_view& request, << " due to remote update notification"; } +void ControlService::onPullReceiverCredit(const std::string_view& request, + std::string& response) { + ReceiverCreditPullResponseV1 reply; + ReceiverCreditPullRequestV1 decoded; + auto status = ReceiverCreditPullRequestCodecV1::decode(request, decoded); + if (!status.ok()) { + ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus::Rejected, reply); + } else { + std::shared_ptr allocator; + { + std::lock_guard lock(receiver_credit_mutex_); + allocator = receiver_credit_allocator_; + } + if (!allocator) { + ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus::Unsupported, reply); + } else { + status = allocator->pull(decoded, reply); + if (!status.ok()) + ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus::Rejected, reply); + } + } + auto encode_status = + ReceiverCreditPullResponseCodecV1::encode(reply, response); + if (!encode_status.ok()) { + // An empty response is interpreted as a protocol/RPC failure by new + // clients. Never expose decoder text or allocate an unbounded error. + response.clear(); + LOG(ERROR) << "Failed to encode receiver credit pull response: " + << encode_status.ToString(); + } +} + void ControlClient::subscribeSegmentUpdateAsync( const std::string& server_addr, const std::string& subscriber_addr) { json j = subscriber_addr; diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp index 5ae6c14a2e..b40bedec01 100644 --- a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp @@ -121,11 +121,28 @@ Status SenderCreditLedger::applyUpdate(const CreditKey& k, Status SenderCreditLedger::tryReserve(const CreditKey& k, const CreditCharge& c) { + uint64_t epoch = 0; + { + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end()) + return Status::InvalidEntry("credit unavailable" LOC_MARK); + epoch = it->second.epoch; + } + return tryReserve(k, epoch, c); +} + +Status SenderCreditLedger::tryReserve(const CreditKey& k, + uint64_t expected_epoch, + const CreditCharge& c) { + if (!expected_epoch) + return Status::InvalidArgument("zero expected credit epoch" LOC_MARK); std::array n; CHECK_STATUS(normalize(c, n)); std::lock_guard lock(mutex_); auto it = entries_.find(k); - if (it == entries_.end() || !it->second.has_update) + if (it == entries_.end() || it->second.epoch != expected_epoch || + !it->second.has_update) return Status::InvalidEntry("credit unavailable" LOC_MARK); auto& e = it->second; for (size_t i = 0; i < kCreditResourceCount; ++i) @@ -136,15 +153,20 @@ Status SenderCreditLedger::tryReserve(const CreditKey& k, } Status SenderCreditLedger::rollbackReservation(const CreditKey& k, + uint64_t expected_epoch, const CreditCharge& c) { + if (!expected_epoch) + return Status::InvalidArgument("zero expected credit epoch" LOC_MARK); std::array n; CHECK_STATUS(normalize(c, n)); std::lock_guard lock(mutex_); auto it = entries_.find(k); - if (it == entries_.end()) - return Status::InvalidEntry("credit session inactive" LOC_MARK); + if (it == entries_.end() || it->second.epoch != expected_epoch) + return Status::InvalidEntry( + "inactive or stale credit rollback epoch" LOC_MARK); for (size_t i = 0; i < kCreditResourceCount; ++i) - if (n[i] > it->second.consumed[i]) + if (it->second.completed[i] > it->second.consumed[i] || + n[i] > it->second.consumed[i] - it->second.completed[i]) return Status::InvalidArgument( "credit rollback underflow" LOC_MARK); for (size_t i = 0; i < kCreditResourceCount; ++i) @@ -152,6 +174,49 @@ Status SenderCreditLedger::rollbackReservation(const CreditKey& k, return Status::OK(); } +Status SenderCreditLedger::recordCompletion(const CreditKey& k, + uint64_t expected_epoch, + const CreditCharge& c) { + if (!expected_epoch) + return Status::InvalidArgument("zero expected credit epoch" LOC_MARK); + std::array n; + CHECK_STATUS(normalize(c, n)); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end() || it->second.epoch != expected_epoch) + return Status::InvalidEntry( + "inactive or stale credit completion epoch" LOC_MARK); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (it->second.completed[i] > it->second.consumed[i] || + n[i] > it->second.consumed[i] - it->second.completed[i]) + return Status::InvalidArgument( + "credit completion exceeds consumption" LOC_MARK); + } + for (size_t i = 0; i < kCreditResourceCount; ++i) + it->second.completed[i] += n[i]; + return Status::OK(); +} + +Status SenderCreditLedger::snapshot(const CreditKey& k, uint64_t expected_epoch, + CreditLedgerSnapshot& output) const { + if (!expected_epoch) + return Status::InvalidArgument("zero expected credit epoch" LOC_MARK); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end() || it->second.epoch != expected_epoch) + return Status::InvalidEntry( + "inactive or stale credit snapshot epoch" LOC_MARK); + CreditLedgerSnapshot next; + next.epoch = it->second.epoch; + next.last_sequence = it->second.last_sequence; + next.has_update = it->second.has_update; + next.grants = it->second.grants; + next.consumed = it->second.consumed; + next.completed = it->second.completed; + output = next; + return Status::OK(); +} + Status SenderCreditLedger::available(const CreditKey& k, CreditResource r, uint64_t& v) const { size_t i = 0; diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_allocator.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_allocator.cpp new file mode 100644 index 0000000000..604b92a16f --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_allocator.cpp @@ -0,0 +1,324 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_allocator.h" + +#include +#include +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t resourceIndex(CreditResource resource) { + return static_cast(resource) - 1; +} + +bool isZeroSession(const ReceiverSessionId& session) { + return session.high == 0 && session.low == 0; +} + +} // namespace + +size_t ReceiverCreditAllocator::EntryKeyHash::operator()( + const EntryKey& key) const noexcept { + size_t hash = std::hash{}(key.sender_peer); + hash ^= std::hash{}(key.qos_class) + 0x9e3779b97f4a7c15ULL + + (hash << 6) + (hash >> 2); + return hash; +} + +Status ReceiverCreditAllocator::validateConfig( + const ReceiverCreditAllocatorConfig& config) { + if (isZeroSession(config.receiver_session_id) || config.epoch == 0 || + config.max_entries == 0 || + config.max_entries > std::numeric_limits::max() || + config.ttl_ms == 0 || config.retry_after_us == 0) + return Status::InvalidArgument( + "invalid receiver credit allocator identity or bounds" LOC_MARK); + + bool has_ingress_budget = false; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const bool supported = i == resourceIndex(CreditResource::DataBytes) || + i == resourceIndex(CreditResource::RequestSlots); + const uint64_t capacity = config.capacity[i]; + const uint64_t maximum = config.max_grant_per_pull[i]; + if ((!supported && (capacity != 0 || maximum != 0)) || + ((capacity == 0) != (maximum == 0)) || maximum > capacity) + return Status::InvalidArgument( + "invalid receiver credit allocator resource bounds" LOC_MARK); + has_ingress_budget = has_ingress_budget || capacity != 0; + } + if (!has_ingress_budget) + return Status::InvalidArgument( + "receiver credit allocator has no ingress budget" LOC_MARK); + return Status::OK(); +} + +Status ReceiverCreditAllocator::create( + const ReceiverCreditAllocatorConfig& config, + std::unique_ptr& allocator) { + CHECK_STATUS(validateConfig(config)); + std::unique_ptr created( + new ReceiverCreditAllocator(config)); + allocator = std::move(created); + return Status::OK(); +} + +ReceiverCreditPullResponseV1 ReceiverCreditAllocator::makeResponse( + ReceiverCreditPullStatus status, uint32_t retry_after_us, + uint32_t qos_class, + const std::array& grants, + uint64_t update_sequence) const { + ReceiverCreditPullResponseV1 response; + response.status = status; + response.retry_after_us = retry_after_us; + response.activation.receiver_session_id = config_.receiver_session_id; + response.activation.epoch = config_.epoch; + response.activation.freshness_ttl_ms = config_.ttl_ms; + response.update.qos_class = qos_class; + response.update.receiver_session_id = config_.receiver_session_id; + response.update.epoch = config_.epoch; + response.update.sequence = update_sequence; + response.update.freshness_ttl_ms = config_.ttl_ms; + response.update.grants.reserve(kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + response.update.grants.push_back( + {static_cast(i + 1), grants[i]}); + } + return response; +} + +Status ReceiverCreditAllocator::pull(const ReceiverCreditPullRequestV1& request, + ReceiverCreditPullResponseV1& response) { + CHECK_STATUS(ReceiverCreditPullRequestCodecV1::validate(request)); + pull_requests_.fetch_add(1, std::memory_order_relaxed); + + std::lock_guard lock(mutex_); + CHECK_STATUS(validateInvariantsLocked(nullptr)); + const std::array zero_grants{}; + + if (request.qos_class != 0) { + response = makeResponse(ReceiverCreditPullStatus::Unsupported, 0, + request.qos_class, zero_grants, 1); + return Status::OK(); + } + + const bool initial_session = + isZeroSession(request.expected_receiver_session_id); + const EntryKey key{request.sender_peer, request.qos_class}; + auto existing = entries_.find(key); + if (existing != entries_.end() && existing->second.has_request) { + Entry& entry = existing->second; + if (request.request_sequence == entry.last_request.request_sequence) { + if (request == entry.last_request) { + response = entry.last_response; + } else { + response = makeResponse(ReceiverCreditPullStatus::Rejected, 0, + request.qos_class, entry.granted, + entry.update_sequence); + } + return Status::OK(); + } + if (!initial_session && + request.request_sequence < entry.last_request.request_sequence) { + response = makeResponse(ReceiverCreditPullStatus::Rejected, 0, + request.qos_class, entry.granted, + entry.update_sequence); + return Status::OK(); + } + } + + if (initial_session && existing != entries_.end()) { + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const auto& entry = existing->second; + if (entry.completed[i] > entry.granted[i]) + return Status::InternalError( + "receiver credit restart entry is inconsistent" LOC_MARK); + const uint64_t outstanding = entry.granted[i] - entry.completed[i]; + if (outstanding > committed_[i]) + return Status::InternalError( + "receiver credit restart reclaim underflow" LOC_MARK); + committed_[i] -= outstanding; + } + entries_.erase(existing); + existing = entries_.end(); + } + + const bool matching_session = + request.expected_receiver_session_id == config_.receiver_session_id && + request.expected_epoch == config_.epoch; + if (!initial_session && !matching_session) { + const auto& grants = + existing == entries_.end() ? zero_grants : existing->second.granted; + const uint64_t update_sequence = + existing == entries_.end() + ? 1 + : std::max(1, existing->second.update_sequence); + response = makeResponse(ReceiverCreditPullStatus::SessionChanged, 0, + request.qos_class, grants, update_sequence); + return Status::OK(); + } + + Entry candidate; + if (existing != entries_.end()) { + candidate = existing->second; + } else { + if (request.last_update_sequence != 0) { + response = makeResponse(ReceiverCreditPullStatus::Rejected, 0, + request.qos_class, zero_grants, 1); + return Status::OK(); + } + if (entries_.size() >= config_.max_entries) { + response = makeResponse(ReceiverCreditPullStatus::Retry, + config_.retry_after_us, request.qos_class, + zero_grants, 1); + return Status::OK(); + } + } + + if (request.last_update_sequence > candidate.update_sequence || + candidate.update_sequence == std::numeric_limits::max()) { + response = + makeResponse(ReceiverCreditPullStatus::Rejected, 0, + request.qos_class, candidate.granted, + std::max(1, candidate.update_sequence)); + return Status::OK(); + } + + auto proposed_consumed = candidate.consumed; + auto proposed_completed = candidate.completed; + std::array present{}; + std::array minimum{}; + std::array desired{}; + for (const auto& usage : request.resources) { + const size_t index = resourceIndex(usage.resource); + present[index] = true; + minimum[index] = usage.minimum_available; + desired[index] = usage.desired_available; + if (usage.consumed_total < candidate.consumed[index] || + usage.completed_total < candidate.completed[index] || + usage.completed_total > usage.consumed_total || + usage.consumed_total > candidate.granted[index]) { + response = + makeResponse(ReceiverCreditPullStatus::Rejected, 0, + request.qos_class, candidate.granted, + std::max(1, candidate.update_sequence)); + return Status::OK(); + } + proposed_consumed[index] = usage.consumed_total; + proposed_completed[index] = usage.completed_total; + } + + // Completion is applied before evaluating a new pull. This makes released + // ingress budget immediately available while keeping an all-or-nothing + // minimum across resources. + auto proposed_committed = committed_; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const uint64_t completion_delta = + proposed_completed[i] - candidate.completed[i]; + if (completion_delta > proposed_committed[i]) + return Status::InternalError( + "receiver credit completion exceeds committed budget" LOC_MARK); + proposed_committed[i] -= completion_delta; + } + + std::array additional{}; + bool all_minimums_fit = true; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!present[i]) continue; + const uint64_t existing_available = + candidate.granted[i] - proposed_consumed[i]; + const uint64_t minimum_additional = + minimum[i] > existing_available ? minimum[i] - existing_available + : 0; + const uint64_t desired_additional = + desired[i] > existing_available ? desired[i] - existing_available + : 0; + const uint64_t free = config_.capacity[i] - proposed_committed[i]; + additional[i] = + std::min({desired_additional, free, config_.max_grant_per_pull[i]}); + if (additional[i] < minimum_additional) all_minimums_fit = false; + } + if (!all_minimums_fit) additional.fill(0); + + candidate.consumed = proposed_consumed; + candidate.completed = proposed_completed; + committed_ = proposed_committed; + if (all_minimums_fit) { + for (size_t i = 0; i < kCreditResourceCount; ++i) { + candidate.granted[i] += additional[i]; + committed_[i] += additional[i]; + } + } + ++candidate.update_sequence; + candidate.has_request = true; + candidate.last_request = request; + candidate.last_response = makeResponse( + all_minimums_fit ? ReceiverCreditPullStatus::Granted + : ReceiverCreditPullStatus::Retry, + all_minimums_fit ? 0 : config_.retry_after_us, request.qos_class, + candidate.granted, candidate.update_sequence); + + if (existing == entries_.end()) { + entries_.emplace(key, candidate); + } else { + existing->second = candidate; + } + CHECK_STATUS(validateInvariantsLocked(nullptr)); + response = candidate.last_response; + return Status::OK(); +} + +Status ReceiverCreditAllocator::validateInvariantsLocked( + ReceiverCreditAllocatorSnapshot* output) const { + if (entries_.size() > config_.max_entries) + return Status::InternalError( + "receiver credit entry bound violated" LOC_MARK); + + std::array recomputed{}; + for (const auto& [key, entry] : entries_) { + if (key.sender_peer == 0 || key.qos_class != 0 || !entry.has_request || + entry.update_sequence == 0) + return Status::InternalError( + "invalid receiver credit allocator entry" LOC_MARK); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (entry.completed[i] > entry.consumed[i] || + entry.consumed[i] > entry.granted[i]) + return Status::InternalError( + "receiver credit entry totals are inconsistent" LOC_MARK); + const uint64_t outstanding = entry.granted[i] - entry.completed[i]; + if (outstanding > + std::numeric_limits::max() - recomputed[i]) + return Status::InternalError( + "receiver credit committed total overflow" LOC_MARK); + recomputed[i] += outstanding; + } + } + + ReceiverCreditAllocatorSnapshot snapshot; + snapshot.capacity = config_.capacity; + snapshot.entries = entries_.size(); + snapshot.pull_requests = pull_requests_.load(std::memory_order_relaxed); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (recomputed[i] != committed_[i] || + committed_[i] > config_.capacity[i]) + return Status::InternalError( + "receiver credit global budget invariant violated" LOC_MARK); + snapshot.committed[i] = committed_[i]; + snapshot.free[i] = config_.capacity[i] - committed_[i]; + } + if (output != nullptr) *output = snapshot; + return Status::OK(); +} + +Status ReceiverCreditAllocator::snapshot( + ReceiverCreditAllocatorSnapshot& snapshot) const { + std::lock_guard lock(mutex_); + ReceiverCreditAllocatorSnapshot checked; + CHECK_STATUS(validateInvariantsLocked(&checked)); + snapshot = checked; + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_config.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_config.cpp new file mode 100644 index 0000000000..45932de7d9 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_config.cpp @@ -0,0 +1,340 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_config.h" + +#include +#include +#include +#include +#include +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t kMaxPeers = 65536; +constexpr size_t kMaxDispatchOwners = 65536; +constexpr size_t kDataBytesIndex = + static_cast(CreditResource::DataBytes) - 1; +constexpr size_t kRequestSlotsIndex = + static_cast(CreditResource::RequestSlots) - 1; + +Status invalidConfig(const std::string& path, std::string_view reason) { + return Status::InvalidArgument("invalid receiver-credit config at " + path + + ": " + std::string(reason) + LOC_MARK); +} + +bool isKnownKey(std::string_view key, + std::initializer_list known) { + for (auto candidate : known) { + if (key == candidate) return true; + } + return false; +} + +Status rejectUnknownKeys(const json& object, const std::string& path, + std::initializer_list known) { + for (auto it = object.begin(); it != object.end(); ++it) { + if (!isKnownKey(it.key(), known)) + return invalidConfig(path + "." + it.key(), "unknown key"); + } + return Status::OK(); +} + +Status requireObject(const json& value, const std::string& path) { + if (!value.is_object()) return invalidConfig(path, "expected object"); + return Status::OK(); +} + +Status readPositiveUint64(const json& object, const char* key, + const std::string& path, uint64_t& output) { + auto it = object.find(key); + if (it == object.end()) + return invalidConfig(path + "." + key, "missing required key"); + if (!it->is_number_unsigned()) + return invalidConfig(path + "." + key, "expected a positive integer"); + uint64_t value = it->get(); + if (value == 0) + return invalidConfig(path + "." + key, "must be greater than zero"); + output = value; + return Status::OK(); +} + +Status readPositiveUint32(const json& object, const char* key, + const std::string& path, uint32_t& output) { + auto it = object.find(key); + if (it == object.end()) return Status::OK(); + if (!it->is_number_unsigned()) + return invalidConfig(path + "." + key, "expected a positive integer"); + uint64_t value = it->get(); + if (value == 0 || value > std::numeric_limits::max()) + return invalidConfig(path + "." + key, "out of range"); + output = static_cast(value); + return Status::OK(); +} + +Status readPositiveSize(const json& object, const char* key, + const std::string& path, size_t maximum, + size_t& output) { + auto it = object.find(key); + if (it == object.end()) return Status::OK(); + if (!it->is_number_unsigned()) + return invalidConfig(path + "." + key, "expected a positive integer"); + uint64_t value = it->get(); + if (value == 0 || value > maximum) + return invalidConfig(path + "." + key, "out of range"); + output = static_cast(value); + return Status::OK(); +} + +Status readMode(const json& value, CreditRolloutMode& output) { + if (!value.is_string()) + return invalidConfig("receiver_credit.mode", "expected string"); + const auto& mode = value.get_ref(); + if (mode == "disabled") { + output = CreditRolloutMode::Disabled; + } else if (mode == "optional") { + output = CreditRolloutMode::Optional; + } else if (mode == "required") { + output = CreditRolloutMode::Required; + } else { + return invalidConfig("receiver_credit.mode", + "expected disabled, optional, or required"); + } + return Status::OK(); +} + +CreditRolloutMode legacyMode(bool enabled) { + return enabled ? CreditRolloutMode::Required : CreditRolloutMode::Disabled; +} + +Status mergeLegacyMode(std::optional& result, + CreditRolloutMode candidate, + const std::string& source_path) { + if (result && *result != candidate) + return invalidConfig(source_path, + "conflicts with another rollout setting"); + result = candidate; + return Status::OK(); +} + +Status parseResourcePair(const json& section, const char* key, + const std::string& path, + std::array& output, + bool required, bool& present) { + auto it = section.find(key); + present = it != section.end(); + if (!present) { + if (required) return invalidConfig(path, "missing required object"); + return Status::OK(); + } + CHECK_STATUS(requireObject(*it, path)); + CHECK_STATUS(rejectUnknownKeys(*it, path, {"data_bytes", "request_slots"})); + CHECK_STATUS( + readPositiveUint64(*it, "data_bytes", path, output[kDataBytesIndex])); + CHECK_STATUS(readPositiveUint64(*it, "request_slots", path, + output[kRequestSlotsIndex])); + return Status::OK(); +} + +} // namespace + +Status loadReceiverCreditConfig(const Config& config, + bool runtime_queue_enabled, + ReceiverCreditRuntimeConfig& output) { + ReceiverCreditRuntimeConfig parsed; + json root; + try { + root = json::parse(config.dump()); + } catch (const std::exception& e) { + return invalidConfig("", e.what()); + } + + // A default-constructed Config is null and is equivalent to an empty + // object. Other non-object roots are configuration errors. + if (root.is_null()) root = json::object(); + if (!root.is_object()) return invalidConfig("", "expected object"); + + const json* section = nullptr; + auto section_it = root.find("receiver_credit"); + if (section_it != root.end()) { + CHECK_STATUS(requireObject(*section_it, "receiver_credit")); + section = &*section_it; + CHECK_STATUS(rejectUnknownKeys( + *section, "receiver_credit", + {"mode", "enabled", "default_qos_class", "capacity", "grant_batch", + "control", "limits"})); + } + + std::optional selected_mode; + if (section) { + auto mode_it = section->find("mode"); + if (mode_it != section->end()) { + CreditRolloutMode explicit_mode; + CHECK_STATUS(readMode(*mode_it, explicit_mode)); + selected_mode = explicit_mode; + } + + auto enabled_it = section->find("enabled"); + if (enabled_it != section->end()) { + if (!enabled_it->is_boolean()) + return invalidConfig("receiver_credit.enabled", + "expected boolean"); + CHECK_STATUS(mergeLegacyMode(selected_mode, + legacyMode(enabled_it->get()), + "receiver_credit.enabled")); + } + } + + // Config supports historical top-level keys containing '/'. Validate the + // JSON type before using Config::get so a malformed alias cannot silently + // fall back to the supplied default. + auto flat_enabled_it = root.find("receiver_credit/enabled"); + if (flat_enabled_it != root.end()) { + if (!flat_enabled_it->is_boolean()) + return invalidConfig("receiver_credit/enabled", "expected boolean"); + bool enabled = flat_enabled_it->get(); + if (!section || section->find("enabled") == section->end()) { + if (!config.contains("receiver_credit/enabled")) + return invalidConfig("receiver_credit/enabled", + "alias lookup failed"); + enabled = config.get("receiver_credit/enabled", !enabled); + } + CHECK_STATUS(mergeLegacyMode(selected_mode, legacyMode(enabled), + "receiver_credit/enabled")); + } + + parsed.mode = selected_mode.value_or(CreditRolloutMode::Disabled); + + if (parsed.mode != CreditRolloutMode::Disabled && !section) + return invalidConfig("receiver_credit.capacity", + "missing required object"); + + if (section) { + auto qos_it = section->find("default_qos_class"); + if (qos_it != section->end()) { + if (!qos_it->is_number_unsigned()) + return invalidConfig("receiver_credit.default_qos_class", + "expected integer zero"); + auto qos_class = qos_it->get(); + if (qos_class != 0) + return invalidConfig("receiver_credit.default_qos_class", + "only QoS class 0 is supported"); + parsed.default_qos_class = 0; + } + + const bool active = parsed.mode != CreditRolloutMode::Disabled; + bool capacity_present = false; + bool grant_present = false; + CHECK_STATUS( + parseResourcePair(*section, "capacity", "receiver_credit.capacity", + parsed.capacity, active, capacity_present)); + CHECK_STATUS(parseResourcePair( + *section, "grant_batch", "receiver_credit.grant_batch", + parsed.max_grant_per_pull, active, grant_present)); + if (capacity_present != grant_present) + return invalidConfig( + "receiver_credit", + "capacity and grant_batch must be specified together"); + if (capacity_present) { + for (size_t i : {kDataBytesIndex, kRequestSlotsIndex}) { + if (parsed.max_grant_per_pull[i] > parsed.capacity[i]) + return invalidConfig( + "receiver_credit.grant_batch", + "resource grant must not exceed capacity"); + } + } + + auto control_it = section->find("control"); + if (control_it != section->end()) { + CHECK_STATUS(requireObject(*control_it, "receiver_credit.control")); + CHECK_STATUS( + rejectUnknownKeys(*control_it, "receiver_credit.control", + {"freshness_ttl_ms", "retry_after_us", + "poll_interval_us", "adaptive_dispatch"})); + CHECK_STATUS(readPositiveUint32(*control_it, "freshness_ttl_ms", + "receiver_credit.control", + parsed.freshness_ttl_ms)); + CHECK_STATUS(readPositiveUint32(*control_it, "retry_after_us", + "receiver_credit.control", + parsed.retry_after_us)); + CHECK_STATUS(readPositiveUint32(*control_it, "poll_interval_us", + "receiver_credit.control", + parsed.progress_interval_us)); + + auto adaptive_it = control_it->find("adaptive_dispatch"); + if (adaptive_it != control_it->end()) { + const std::string path = + "receiver_credit.control.adaptive_dispatch"; + CHECK_STATUS(requireObject(*adaptive_it, path)); + CHECK_STATUS(rejectUnknownKeys( + *adaptive_it, path, + {"enabled", "min_owners", "initial_owners", "max_owners", + "slow_rtt_us", "healthy_pulls_per_increase"})); + auto enabled_it = adaptive_it->find("enabled"); + if (enabled_it != adaptive_it->end()) { + if (!enabled_it->is_boolean()) + return invalidConfig(path + ".enabled", + "expected boolean"); + parsed.adaptive_dispatch_enabled = enabled_it->get(); + } + CHECK_STATUS(readPositiveSize( + *adaptive_it, "min_owners", path, kMaxDispatchOwners, + parsed.adaptive_dispatch_min_owners)); + CHECK_STATUS(readPositiveSize( + *adaptive_it, "initial_owners", path, kMaxDispatchOwners, + parsed.adaptive_dispatch_initial_owners)); + CHECK_STATUS(readPositiveSize( + *adaptive_it, "max_owners", path, kMaxDispatchOwners, + parsed.adaptive_dispatch_max_owners)); + CHECK_STATUS( + readPositiveUint32(*adaptive_it, "slow_rtt_us", path, + parsed.adaptive_dispatch_slow_rtt_us)); + CHECK_STATUS(readPositiveUint32( + *adaptive_it, "healthy_pulls_per_increase", path, + parsed.adaptive_dispatch_healthy_pulls)); + if (parsed.adaptive_dispatch_min_owners > + parsed.adaptive_dispatch_initial_owners || + parsed.adaptive_dispatch_initial_owners > + parsed.adaptive_dispatch_max_owners) + return invalidConfig( + path, + "expected min_owners <= initial_owners <= max_owners"); + } + } + + auto limits_it = section->find("limits"); + if (limits_it != section->end()) { + CHECK_STATUS(requireObject(*limits_it, "receiver_credit.limits")); + CHECK_STATUS(rejectUnknownKeys(*limits_it, "receiver_credit.limits", + {"max_peers"})); + auto max_peers_it = limits_it->find("max_peers"); + if (max_peers_it != limits_it->end()) { + if (!max_peers_it->is_number_unsigned()) + return invalidConfig("receiver_credit.limits.max_peers", + "expected a positive integer"); + uint64_t max_peers = max_peers_it->get(); + if (max_peers == 0 || max_peers > kMaxPeers) + return invalidConfig("receiver_credit.limits.max_peers", + "must be in range [1, 65536]"); + parsed.max_peers = static_cast(max_peers); + } + } + } + + if (parsed.mode != CreditRolloutMode::Disabled && !runtime_queue_enabled) + return invalidConfig("receiver_credit.mode", + "runtime queue must be enabled"); + + const uint64_t ttl_us = + static_cast(parsed.freshness_ttl_ms) * 1000; + if (ttl_us <= parsed.progress_interval_us) + return invalidConfig("receiver_credit.control.freshness_ttl_ms", + "must be longer than control.poll_interval_us"); + + output = parsed; + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_control.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_control.cpp new file mode 100644 index 0000000000..eb527bdf4c --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_control.cpp @@ -0,0 +1,402 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_control.h" + +#include +#include + +namespace mooncake::tent { + +CreditCapabilityState::CreditCapabilityState(CreditRolloutMode mode) + : mode_(mode), + state_(mode == CreditRolloutMode::Disabled + ? CreditPeerState::Disabled + : CreditPeerState::Negotiating) {} + +Status CreditCapabilityState::beginNegotiation() { + if (mode_ == CreditRolloutMode::Disabled) return Status::OK(); + state_ = CreditPeerState::Negotiating; + version_ = 0; + return Status::OK(); +} + +Status CreditCapabilityState::completeNegotiation( + const std::vector& peer_versions) { + if (mode_ == CreditRolloutMode::Disabled) return Status::OK(); + if (state_ != CreditPeerState::Negotiating) + return Status::InvalidEntry("credit peer is not negotiating" LOC_MARK); + auto supported = std::find(peer_versions.begin(), peer_versions.end(), 1); + if (supported != peer_versions.end()) { + version_ = 1; + state_ = CreditPeerState::Active; + return Status::OK(); + } + if (mode_ == CreditRolloutMode::Optional) { + state_ = CreditPeerState::Legacy; + return Status::OK(); + } + state_ = CreditPeerState::Failed; + return Status::NotImplemented( + "receiver credit is required but unsupported" LOC_MARK); +} + +Status CreditCapabilityState::markStale() { + if (state_ != CreditPeerState::Active) + return Status::InvalidEntry( + "only active credit can become stale" LOC_MARK); + state_ = CreditPeerState::Stale; + return Status::OK(); +} + +Status CreditCapabilityState::refresh(uint16_t version) { + if (state_ != CreditPeerState::Stale || version != version_) + return Status::InvalidArgument("invalid stale credit refresh" LOC_MARK); + state_ = CreditPeerState::Active; + return Status::OK(); +} + +Status BoundedCreditUpdateInbox::tryPublish(CreditControlEnvelope envelope) { + if (capacity_ == 0) + return Status::TooManyRequests("credit update inbox disabled" LOC_MARK); + std::lock_guard lock(mutex_); + if (queue_.size() >= capacity_) + return Status::TooManyRequests("credit update inbox full" LOC_MARK); + queue_.push_back(std::move(envelope)); + return Status::OK(); +} + +size_t BoundedCreditUpdateInbox::drain( + std::vector& output, size_t max_updates) { + std::lock_guard lock(mutex_); + size_t count = std::min(max_updates, queue_.size()); + output.reserve(output.size() + count); + for (size_t i = 0; i < count; ++i) { + output.push_back(std::move(queue_.front())); + queue_.pop_front(); + } + return count; +} + +size_t BoundedCreditUpdateInbox::size() const { + std::lock_guard lock(mutex_); + return queue_.size(); +} + +namespace { +constexpr uint32_t kCreditMagic = 0x54435231; // "TCR1" +constexpr uint32_t kCapabilityMagic = 0x54434331; // "TCC1" +constexpr uint32_t kActivationMagic = 0x54434131; // "TCA1" +void append16(std::string& out, uint16_t v) { + out.push_back(static_cast(v >> 8)); + out.push_back(static_cast(v)); +} +void append32(std::string& out, uint32_t v) { + for (int shift = 24; shift >= 0; shift -= 8) + out.push_back(static_cast(v >> shift)); +} +void append64(std::string& out, uint64_t v) { + for (int shift = 56; shift >= 0; shift -= 8) + out.push_back(static_cast(v >> shift)); +} +uint16_t read16(std::string_view in, size_t& p) { + uint16_t v = 0; + for (int i = 0; i < 2; ++i) v = (v << 8) | uint8_t(in[p++]); + return v; +} +uint32_t read32(std::string_view in, size_t& p) { + uint32_t v = 0; + for (int i = 0; i < 4; ++i) v = (v << 8) | uint8_t(in[p++]); + return v; +} +uint64_t read64(std::string_view in, size_t& p) { + uint64_t v = 0; + for (int i = 0; i < 8; ++i) v = (v << 8) | uint8_t(in[p++]); + return v; +} +} // namespace + +Status CreditCapabilityCodecV1::encode(const std::vector& versions, + std::string& wire) { + if (versions.size() > kMaxVersions) + return Status::InvalidArgument( + "too many receiver credit versions" LOC_MARK); + std::vector encoded_versions; + encoded_versions.reserve(versions.size()); + std::string encoded; + encoded.reserve(kHeaderBytes + versions.size() * sizeof(uint16_t)); + append32(encoded, kCapabilityMagic); + append16(encoded, static_cast(versions.size())); + append16(encoded, 0); + for (uint16_t version : versions) { + if (version == 0 || + std::find(encoded_versions.begin(), encoded_versions.end(), + version) != encoded_versions.end()) + return Status::InvalidArgument( + "invalid or duplicate receiver credit version" LOC_MARK); + encoded_versions.push_back(version); + append16(encoded, version); + } + wire.swap(encoded); + return Status::OK(); +} + +Status CreditCapabilityCodecV1::decode(std::string_view wire, + std::vector& versions) { + if (wire.size() < kHeaderBytes || wire.size() > kMaxWireBytes) + return Status::InvalidArgument( + "invalid receiver credit capability length" LOC_MARK); + size_t p = 0; + if (read32(wire, p) != kCapabilityMagic) + return Status::InvalidArgument( + "invalid receiver credit capability magic" LOC_MARK); + uint16_t count = read16(wire, p); + uint16_t reserved = read16(wire, p); + if (reserved != 0 || count > kMaxVersions || + wire.size() != kHeaderBytes + count * sizeof(uint16_t)) + return Status::InvalidArgument( + "invalid receiver credit capability header" LOC_MARK); + std::vector decoded; + decoded.reserve(count); + for (uint16_t i = 0; i < count; ++i) { + uint16_t version = read16(wire, p); + if (version == 0 || + std::find(decoded.begin(), decoded.end(), version) != decoded.end()) + return Status::InvalidArgument( + "invalid receiver credit capability version" LOC_MARK); + decoded.push_back(version); + } + versions = std::move(decoded); + return Status::OK(); +} + +Status CreditActivationCodecV1::encode(const CreditActivationV1& activation, + std::string& wire) { + if (activation.schema_version != 1 || activation.chosen_version != 1 || + (!activation.receiver_session_id.high && + !activation.receiver_session_id.low) || + !activation.epoch) + return Status::InvalidArgument( + "invalid receiver credit activation" LOC_MARK); + std::string encoded; + encoded.reserve(kWireBytes); + append32(encoded, kActivationMagic); + append16(encoded, activation.schema_version); + append16(encoded, activation.chosen_version); + append64(encoded, activation.receiver_session_id.high); + append64(encoded, activation.receiver_session_id.low); + append64(encoded, activation.epoch); + append32(encoded, activation.freshness_ttl_ms); + append32(encoded, 0); + wire.swap(encoded); + return Status::OK(); +} + +Status CreditActivationCodecV1::decode(std::string_view wire, + CreditActivationV1& activation) { + if (wire.size() != kWireBytes) + return Status::InvalidArgument( + "invalid receiver credit activation length" LOC_MARK); + size_t p = 0; + if (read32(wire, p) != kActivationMagic) + return Status::InvalidArgument( + "invalid receiver credit activation magic" LOC_MARK); + CreditActivationV1 decoded; + decoded.schema_version = read16(wire, p); + decoded.chosen_version = read16(wire, p); + decoded.receiver_session_id.high = read64(wire, p); + decoded.receiver_session_id.low = read64(wire, p); + decoded.epoch = read64(wire, p); + decoded.freshness_ttl_ms = read32(wire, p); + uint32_t reserved = read32(wire, p); + if (decoded.schema_version != 1 || decoded.chosen_version != 1 || + (!decoded.receiver_session_id.high && + !decoded.receiver_session_id.low) || + !decoded.epoch || reserved != 0) + return Status::InvalidArgument( + "invalid receiver credit activation header" LOC_MARK); + activation = decoded; + return Status::OK(); +} + +size_t CreditPeerContextTable::LookupKeyHash::operator()( + const LookupKey& key) const noexcept { + size_t h = std::hash{}(key.target_id); + h ^= std::hash{}(key.qos_class) + 0x9e3779b97f4a7c15ULL + + (h << 6) + (h >> 2); + return h; +} + +Status CreditPeerContextTable::activate(uint64_t target_id, + uint64_t sender_peer, + uint32_t qos_class, + const CreditActivationV1& activation) { + if (!sender_peer || activation.schema_version != 1 || + activation.chosen_version != 1 || + (!activation.receiver_session_id.high && + !activation.receiver_session_id.low) || + !activation.epoch) + return Status::InvalidArgument( + "invalid receiver credit peer context" LOC_MARK); + LookupKey lookup_key{target_id, qos_class}; + CreditPeerContextSnapshot next{ + {activation.receiver_session_id, sender_peer, qos_class}, + activation.epoch, + activation.freshness_ttl_ms, + std::chrono::steady_clock::now()}; + std::lock_guard lock(mutex_); + auto it = contexts_.find(lookup_key); + if (it == contexts_.end()) { + if (contexts_.size() >= max_entries_) + return Status::TooManyRequests( + "receiver credit peer context table full" LOC_MARK); + contexts_.emplace(lookup_key, next); + return Status::OK(); + } + const auto& current = it->second; + if (current.key.sender_peer != sender_peer) + return Status::InvalidEntry( + "receiver credit sender identity changed" LOC_MARK); + if (current.key.receiver_session == activation.receiver_session_id && + activation.epoch < current.epoch) + return Status::InvalidEntry( + "stale receiver credit peer activation" LOC_MARK); + if (current.key.receiver_session == activation.receiver_session_id && + activation.epoch == current.epoch) { + // A valid pull response is also a freshness heartbeat. Refreshing an + // existing generation must not reset ledger consumption, so only the + // context timestamp and advertised TTL are updated here. + it->second.freshness_ttl_ms = activation.freshness_ttl_ms; + it->second.refreshed_at = next.refreshed_at; + return Status::OK(); + } + it->second = next; + return Status::OK(); +} + +Status CreditPeerContextTable::lookup( + uint64_t target_id, uint32_t qos_class, + CreditPeerContextSnapshot& snapshot) const { + std::lock_guard lock(mutex_); + auto it = contexts_.find({target_id, qos_class}); + if (it == contexts_.end()) + return Status::InvalidEntry( + "receiver credit peer context not found" LOC_MARK); + snapshot = it->second; + return Status::OK(); +} + +Status CreditPeerContextTable::lookupFresh( + uint64_t target_id, uint32_t qos_class, + CreditPeerContextSnapshot& snapshot) const { + CHECK_STATUS(lookup(target_id, qos_class, snapshot)); + if (!snapshot.isFresh()) + return Status::InvalidEntry( + "receiver credit peer context is stale" LOC_MARK); + return Status::OK(); +} + +Status CreditPeerContextTable::deactivate( + uint64_t target_id, uint32_t qos_class, + const ReceiverSessionId& receiver_session, uint64_t epoch) { + std::lock_guard lock(mutex_); + auto it = contexts_.find({target_id, qos_class}); + if (it == contexts_.end()) return Status::OK(); + if (!(it->second.key.receiver_session == receiver_session) || + it->second.epoch != epoch) + return Status::InvalidEntry( + "receiver credit peer cleanup is stale" LOC_MARK); + contexts_.erase(it); + return Status::OK(); +} + +size_t CreditPeerContextTable::size() const { + std::lock_guard lock(mutex_); + return contexts_.size(); +} + +Status ReceiverCreditCodecV1::encode(const ReceiverCreditUpdateV1& u, + std::string& wire) { + if (u.schema_version != 1 || !u.epoch || !u.sequence || + u.grants.size() > kCreditResourceCount) + return Status::InvalidArgument("invalid credit wire update" LOC_MARK); + std::array present{}; + for (const auto& grant : u.grants) { + auto raw = static_cast(grant.resource); + if (raw < 1 || raw > kCreditResourceCount || present[raw - 1]) + return Status::InvalidArgument( + "invalid or duplicate wire resource" LOC_MARK); + present[raw - 1] = true; + } + std::string encoded; + encoded.reserve(kHeaderBytes + u.grants.size() * kGrantBytes); + append32(encoded, kCreditMagic); + append16(encoded, u.schema_version); + append16(encoded, u.flags); + append32(encoded, u.qos_class); + append64(encoded, u.receiver_session_id.high); + append64(encoded, u.receiver_session_id.low); + append64(encoded, u.epoch); + append64(encoded, u.sequence); + append32(encoded, u.freshness_ttl_ms); + append16(encoded, static_cast(u.grants.size())); + append16(encoded, 0); + for (const auto& grant : u.grants) { + append16(encoded, static_cast(grant.resource)); + append16(encoded, 0); + append64(encoded, grant.grant_total); + } + wire.swap(encoded); // do not mutate output on validation failure + return Status::OK(); +} + +Status ReceiverCreditCodecV1::decode(std::string_view wire, + ReceiverCreditUpdateV1& update) { + if (wire.size() < kHeaderBytes || wire.size() > kMaxWireBytes) + return Status::InvalidArgument("invalid credit wire length" LOC_MARK); + size_t p = 0; + if (read32(wire, p) != kCreditMagic) + return Status::InvalidArgument("invalid credit wire magic" LOC_MARK); + ReceiverCreditUpdateV1 decoded; + decoded.schema_version = read16(wire, p); + decoded.flags = read16(wire, p); + decoded.qos_class = read32(wire, p); + decoded.receiver_session_id.high = read64(wire, p); + decoded.receiver_session_id.low = read64(wire, p); + decoded.epoch = read64(wire, p); + decoded.sequence = read64(wire, p); + decoded.freshness_ttl_ms = read32(wire, p); + uint16_t count = read16(wire, p); + uint16_t reserved = read16(wire, p); + if (decoded.schema_version != 1 || !decoded.epoch || !decoded.sequence || + reserved != 0 || count > kCreditResourceCount || + wire.size() != kHeaderBytes + count * kGrantBytes) + return Status::InvalidArgument("invalid credit wire header" LOC_MARK); + std::array present{}; + decoded.grants.reserve(count); + for (uint16_t i = 0; i < count; ++i) { + uint16_t raw = read16(wire, p); + uint16_t grant_reserved = read16(wire, p); + uint64_t total = read64(wire, p); + if (raw < 1 || raw > kCreditResourceCount || present[raw - 1] || + grant_reserved != 0) + return Status::InvalidArgument( + "invalid credit wire grant" LOC_MARK); + present[raw - 1] = true; + decoded.grants.push_back({static_cast(raw), total}); + } + update = std::move(decoded); // atomic decode result + return Status::OK(); +} + +Status ReceiverCreditIngress::tryAccept(std::string_view wire) { + ReceiverCreditUpdateV1 decoded; + CHECK_STATUS(ReceiverCreditCodecV1::decode(wire, decoded)); + if (!(decoded.receiver_session_id == key_.receiver_session) || + decoded.qos_class != key_.qos_class || decoded.epoch != epoch_) + return Status::InvalidEntry( + "credit update does not match active ingress" LOC_MARK); + return inbox_.tryPublish({key_, std::move(decoded)}); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_controller.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_controller.cpp new file mode 100644 index 0000000000..d31f24cef3 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_controller.cpp @@ -0,0 +1,495 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_controller.h" + +#include +#include +#include +#include + +#include "tent/runtime/control_plane.h" + +namespace mooncake::tent { +namespace { + +constexpr size_t resourceIndex(CreditResource resource) { + return static_cast(resource) - 1; +} + +constexpr std::array kProductionResources{ + CreditResource::DataBytes, CreditResource::RequestSlots}; + +constexpr std::array kPullLatencyBucketUpperNs{ + 100'000, 250'000, 500'000, 1'000'000, 10'000'000, 100'000'000}; + +bool sameGeneration(const CreditPeerContextSnapshot& context, + const CreditActivationV1& activation) { + return context.key.receiver_session == activation.receiver_session_id && + context.epoch == activation.epoch; +} + +} // namespace + +AdaptiveCreditDispatchLimiter::AdaptiveCreditDispatchLimiter( + const ReceiverCreditRuntimeConfig& config, size_t hard_max_owners) + : enabled_(config.adaptive_dispatch_enabled), + min_owners_( + std::min(config.adaptive_dispatch_min_owners, hard_max_owners)), + slow_rtt_ns_(static_cast(config.adaptive_dispatch_slow_rtt_us) * + 1000), + healthy_pulls_per_increase_(config.adaptive_dispatch_healthy_pulls), + current_owners_(enabled_ + ? std::min(config.adaptive_dispatch_initial_owners, + hard_max_owners) + : std::numeric_limits::max()), + learned_ceiling_(enabled_ ? std::min(config.adaptive_dispatch_max_owners, + hard_max_owners) + : std::numeric_limits::max()) {} + +void AdaptiveCreditDispatchLimiter::observe(std::chrono::nanoseconds elapsed, + bool rpc_ok, + size_t owners_at_start) { + if (!enabled_) return; + const bool unhealthy = !rpc_ok || static_cast(std::max( + 0, elapsed.count())) >= slow_rtt_ns_; + + std::lock_guard lock(mutex_); + size_t current = current_owners_.load(std::memory_order_relaxed); + if (unhealthy) { + ++slow_or_failed_pulls_; + healthy_pulls_ = 0; + const size_t failed_level = std::min(owners_at_start, learned_ceiling_); + if (failed_level > min_owners_) { + if (suspect_level_ == failed_level) { + learned_ceiling_ = failed_level - 1; + suspect_level_ = 0; + } else { + suspect_level_ = failed_level; + } + } + if (current > min_owners_) { + const size_t reduced = std::max(min_owners_, current / 2); + current_owners_.store(reduced, std::memory_order_release); + ++reductions_; + } + return; + } + + if (owners_at_start != current || current >= learned_ceiling_) { + healthy_pulls_ = 0; + return; + } + if (++healthy_pulls_ < healthy_pulls_per_increase_) return; + healthy_pulls_ = 0; + current_owners_.store(current + 1, std::memory_order_release); + ++increases_; +} + +size_t AdaptiveCreditDispatchLimiter::ownerLimit() const { + return current_owners_.load(std::memory_order_acquire); +} + +AdaptiveDispatchSnapshot AdaptiveCreditDispatchLimiter::snapshot() const { + std::lock_guard lock(mutex_); + return {.current_owners = current_owners_.load(std::memory_order_relaxed), + .learned_ceiling = learned_ceiling_, + .suspect_level = suspect_level_, + .slow_or_failed_pulls = slow_or_failed_pulls_, + .reductions = reductions_, + .increases = increases_}; +} + +size_t ReceiverCreditPullController::PeerKeyHash::operator()( + const PeerKey& key) const noexcept { + size_t h = std::hash{}(key.target_id); + h ^= std::hash{}(key.qos_class) + 0x9e3779b97f4a7c15ULL + + (h << 6) + (h >> 2); + return h; +} + +ReceiverCreditPullController::ReceiverCreditPullController( + ReceiverCreditRuntimeConfig config, uint64_t sender_peer, + std::shared_ptr contexts, + std::shared_ptr ledger, size_t dispatch_owner_ceiling) + : config_(std::move(config)), + sender_peer_(sender_peer), + contexts_(std::move(contexts)), + ledger_(std::move(ledger)), + rpc_agent_(std::make_shared()), + dispatch_limiter_(config_, dispatch_owner_ceiling) {} + +ReceiverCreditPullController::~ReceiverCreditPullController() { stop(); } + +Status ReceiverCreditPullController::create( + const ReceiverCreditRuntimeConfig& config, uint64_t sender_peer, + std::shared_ptr contexts, + std::shared_ptr ledger, + std::shared_ptr& controller, + size_t dispatch_owner_ceiling) { + if (config.mode == CreditRolloutMode::Disabled || sender_peer == 0 || + !contexts || !ledger || config.max_peers == 0 || + dispatch_owner_ceiling == 0) + return Status::InvalidArgument( + "invalid receiver credit pull controller configuration" LOC_MARK); + auto next = std::shared_ptr( + new ReceiverCreditPullController(config, sender_peer, + std::move(contexts), std::move(ledger), + dispatch_owner_ceiling)); + controller = std::move(next); + return Status::OK(); +} + +Status ReceiverCreditPullController::normalize( + const CreditCharge& charge, + std::array& normalized) { + normalized.fill(0); + if (charge.resources.empty()) + return Status::InvalidArgument( + "empty receiver credit pull demand" LOC_MARK); + for (const auto& [resource, amount] : charge.resources) { + const auto raw = static_cast(resource); + if (raw < static_cast(CreditResource::DataBytes) || + raw > static_cast(CreditResource::ConsumerSlots) || + amount == 0) + return Status::InvalidArgument( + "invalid receiver credit pull demand" LOC_MARK); + if (resource != CreditResource::DataBytes && + resource != CreditResource::RequestSlots) + return Status::NotImplemented( + "production receiver credit currently models only direct " + "WRITE bytes and request slots" LOC_MARK); + auto index = resourceIndex(resource); + if (normalized[index] != 0) + return Status::InvalidArgument( + "duplicate receiver credit pull resource" LOC_MARK); + normalized[index] = amount; + } + return Status::OK(); +} + +Status ReceiverCreditPullController::request( + uint64_t target_id, const std::string& server_addr, uint32_t qos_class, + const CreditCharge& minimum_charge) { + if (target_id == 0 || server_addr.empty()) + return Status::InvalidArgument( + "receiver credit pull requires a remote RPC address" LOC_MARK); + if (qos_class != 0) + return Status::InvalidArgument( + "receiver credit MVP authorizes only QoS class 0" LOC_MARK); + + std::array minimum; + CHECK_STATUS(normalize(minimum_charge, minimum)); + + PeerKey key{target_id, qos_class}; + bool should_launch = false; + { + std::lock_guard lock(mutex_); + if (stopped_) + return Status::InvalidEntry( + "receiver credit pull controller stopped" LOC_MARK); + auto it = peers_.find(key); + if (it == peers_.end()) { + if (peers_.size() >= config_.max_peers) + return Status::TooManyRequests( + "receiver credit peer limit reached" LOC_MARK); + it = peers_.emplace(key, Peer{}).first; + } + auto& peer = it->second; + bool demand_changed = peer.server_addr != server_addr; + peer.server_addr = server_addr; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (minimum[i] > peer.minimum[i]) { + peer.minimum[i] = minimum[i]; + demand_changed = true; + } + } + if (!peer.in_flight) { + peer.in_flight = true; + should_launch = true; + } else if (demand_changed) { + // Repeated progress attempts for the same blocked owner must not + // create an RPC train behind the outstanding pull. A fresh attempt + // after completion can still launch if the refill was insufficient. + peer.dirty = true; + } + } + if (should_launch) launch(key); + return Status::OK(); +} + +void ReceiverCreditPullController::launch(PeerKey key) { + std::string server_addr; + std::array minimum{}; + uint64_t request_sequence = 0; + { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (stopped_ || it == peers_.end()) return; + auto& peer = it->second; + server_addr = peer.server_addr; + minimum = peer.minimum; + request_sequence = peer.next_request_sequence++; + peer.dirty = false; + } + + ReceiverCreditPullRequestV1 request; + request.sender_peer = sender_peer_; + request.qos_class = key.qos_class; + request.request_sequence = request_sequence; + + CreditPeerContextSnapshot context; + CreditLedgerSnapshot ledger_snapshot; + const auto context_status = + contexts_->lookup(key.target_id, key.qos_class, context); + if (context_status.ok()) { + request.expected_receiver_session_id = context.key.receiver_session; + request.expected_epoch = context.epoch; + auto ledger_status = + ledger_->snapshot(context.key, context.epoch, ledger_snapshot); + if (!ledger_status.ok()) { + finish(key, ledger_status, {}); + return; + } + request.last_update_sequence = ledger_snapshot.last_sequence; + } else if (!context_status.IsInvalidEntry()) { + finish(key, context_status, {}); + return; + } + + request.resources.reserve(kProductionResources.size()); + for (auto resource : kProductionResources) { + const size_t index = resourceIndex(resource); + ReceiverCreditResourceUsageV1 usage; + usage.resource = resource; + if (context_status.ok()) { + usage.consumed_total = ledger_snapshot.consumed[index]; + usage.completed_total = ledger_snapshot.completed[index]; + } + usage.minimum_available = minimum[index]; + usage.desired_available = + std::max(minimum[index], config_.max_grant_per_pull[index]); + request.resources.push_back(usage); + } + + { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (stopped_ || it == peers_.end()) return; + it->second.pull_started_at = std::chrono::steady_clock::now(); + it->second.pull_dispatch_owners = dispatch_limiter_.ownerLimit(); + ++pulls_started_; + } + + auto self = shared_from_this(); + ControlClient::pullReceiverCreditAsync( + rpc_agent_, server_addr, request, + [self = std::move(self), key](Status status, + ReceiverCreditPullResponseV1 response) { + self->finish(key, std::move(status), std::move(response)); + }); +} + +bool ReceiverCreditPullController::oldGenerationDrained( + const CreditPeerContextSnapshot& context) const { + CreditLedgerSnapshot snapshot; + if (!ledger_->snapshot(context.key, context.epoch, snapshot).ok()) + return false; + for (size_t i = 0; i < kCreditResourceCount; ++i) + if (snapshot.consumed[i] != snapshot.completed[i]) return false; + return true; +} + +Status ReceiverCreditPullController::applyResponse( + const PeerKey& key, const ReceiverCreditPullResponseV1& response, + bool& request_again) { + request_again = false; + if (response.status == ReceiverCreditPullStatus::Unsupported) { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end()) + it->second.state = config_.mode == CreditRolloutMode::Optional + ? CreditPeerState::Legacy + : CreditPeerState::Failed; + return Status::OK(); + } + if (response.status == ReceiverCreditPullStatus::Rejected) { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end()) it->second.state = CreditPeerState::Failed; + return Status::InvalidArgument( + "receiver rejected credit pull" LOC_MARK); + } + + CreditPeerContextSnapshot old_context; + const auto old_status = + contexts_->lookup(key.target_id, key.qos_class, old_context); + + if (response.status == ReceiverCreditPullStatus::SessionChanged || + (old_status.ok() && + !sameGeneration(old_context, response.activation))) { + if (old_status.ok() && !oldGenerationDrained(old_context)) { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end()) it->second.state = CreditPeerState::Stale; + return Status::InvalidEntry( + "receiver generation changed with work still in " + "flight" LOC_MARK); + } + if (old_status.ok()) { + CHECK_STATUS(contexts_->deactivate(key.target_id, key.qos_class, + old_context.key.receiver_session, + old_context.epoch)); + CHECK_STATUS( + ledger_->deactivate(old_context.key, old_context.epoch)); + } + { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end()) { + it->second.state = CreditPeerState::Negotiating; + it->second.dirty = true; + } + } + request_again = true; + return Status::OK(); + } + + if (response.status != ReceiverCreditPullStatus::Granted && + response.status != ReceiverCreditPullStatus::Retry) + return Status::InvalidArgument( + "unknown receiver credit pull disposition" LOC_MARK); + + CreditKey credit_key{response.activation.receiver_session_id, sender_peer_, + key.qos_class}; + CHECK_STATUS(ledger_->activate(credit_key, response.activation.epoch)); + CreditUpdateDisposition disposition; + CHECK_STATUS( + ledger_->applyUpdate(credit_key, response.update, disposition)); + // Publish the context last. A dispatcher that can see this generation is + // therefore guaranteed to find an activated, updated ledger entry. + CHECK_STATUS(contexts_->activate(key.target_id, sender_peer_, key.qos_class, + response.activation)); + { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end()) it->second.state = CreditPeerState::Active; + } + return Status::OK(); +} + +void ReceiverCreditPullController::finish( + PeerKey key, Status rpc_status, ReceiverCreditPullResponseV1 response) { + std::chrono::nanoseconds pull_elapsed{0}; + bool pull_measured = false; + size_t pull_dispatch_owners = 0; + { + const auto now = std::chrono::steady_clock::now(); + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (it != peers_.end() && + it->second.pull_started_at.time_since_epoch().count() != 0) { + const auto elapsed = static_cast( + std::chrono::duration_cast( + now - it->second.pull_started_at) + .count()); + pull_elapsed = std::chrono::nanoseconds(elapsed); + pull_measured = true; + pull_dispatch_owners = it->second.pull_dispatch_owners; + it->second.pull_started_at = {}; + ++pulls_completed_; + pull_latency_ns_sum_ += elapsed; + pull_latency_ns_max_ = std::max(pull_latency_ns_max_, elapsed); + size_t bucket = 0; + while (bucket < kPullLatencyBucketUpperNs.size() && + elapsed > kPullLatencyBucketUpperNs[bucket]) + ++bucket; + ++pull_latency_buckets_[bucket]; + } + } + + if (pull_measured) + dispatch_limiter_.observe(pull_elapsed, rpc_status.ok(), + pull_dispatch_owners); + + bool request_again = false; + Status result = std::move(rpc_status); + if (result.ok()) result = applyResponse(key, response, request_again); + + bool launch_again = false; + { + std::lock_guard lock(mutex_); + auto it = peers_.find(key); + if (stopped_ || it == peers_.end()) return; + auto& peer = it->second; + peer.in_flight = false; + if (!result.ok() && peer.state == CreditPeerState::Active) + peer.state = CreditPeerState::Stale; + if (!result.ok() && peer.state == CreditPeerState::Negotiating && + config_.mode == CreditRolloutMode::Required && + result.IsInvalidArgument()) + peer.state = CreditPeerState::Failed; + peer.dirty = peer.dirty || request_again; + if (peer.dirty) { + peer.in_flight = true; + launch_again = true; + } + } + if (launch_again) launch(key); +} + +CreditPeerState ReceiverCreditPullController::peerState( + uint64_t target_id, uint32_t qos_class) const { + std::lock_guard lock(mutex_); + auto it = peers_.find({target_id, qos_class}); + return it == peers_.end() ? CreditPeerState::Negotiating : it->second.state; +} + +size_t ReceiverCreditPullController::peerCount() const { + std::lock_guard lock(mutex_); + return peers_.size(); +} + +size_t ReceiverCreditPullController::dispatchOwnerLimit() const { + return dispatch_limiter_.ownerLimit(); +} + +AdaptiveDispatchSnapshot +ReceiverCreditPullController::adaptiveDispatchSnapshot() const { + return dispatch_limiter_.snapshot(); +} + +void ReceiverCreditPullController::stop() { + std::lock_guard lock(mutex_); + stopped_ = true; + if (!summary_logged_) { + const auto adaptive = dispatch_limiter_.snapshot(); + const double average_us = + pulls_completed_ == 0 + ? 0.0 + : static_cast(pull_latency_ns_sum_) / + static_cast(pulls_completed_) / 1000.0; + LOG(INFO) << "Receiver credit pull summary: started=" << pulls_started_ + << " completed=" << pulls_completed_ + << " avg_us=" << average_us << " max_us=" + << static_cast(pull_latency_ns_max_) / 1000.0 + << " buckets_le_100us=" << pull_latency_buckets_[0] + << " le_250us=" << pull_latency_buckets_[1] + << " le_500us=" << pull_latency_buckets_[2] + << " le_1ms=" << pull_latency_buckets_[3] + << " le_10ms=" << pull_latency_buckets_[4] + << " le_100ms=" << pull_latency_buckets_[5] + << " gt_100ms=" << pull_latency_buckets_[6]; + LOG(INFO) << "Receiver credit adaptive dispatch summary: owners=" + << adaptive.current_owners + << " learned_ceiling=" << adaptive.learned_ceiling + << " suspect_level=" << adaptive.suspect_level + << " slow_or_failed_pulls=" << adaptive.slow_or_failed_pulls + << " reductions=" << adaptive.reductions + << " increases=" << adaptive.increases; + summary_logged_ = true; + } + peers_.clear(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_dispatch.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_dispatch.cpp new file mode 100644 index 0000000000..8737c90d26 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_dispatch.cpp @@ -0,0 +1,78 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_dispatch.h" + +namespace mooncake::tent { + +Status ReceiverCreditDispatchGate::snapshot( + uint64_t target_id, uint32_t qos_class, CreditCharge charge, + CreditDispatchSnapshot& output) const { + if (charge.resources.empty()) + return Status::InvalidArgument( + "invalid receiver credit dispatch snapshot" LOC_MARK); + CreditPeerContextSnapshot peer; + CHECK_STATUS(contexts_.lookupFresh(target_id, qos_class, peer)); + CreditDispatchSnapshot next; + next.target_id = target_id; + next.qos_class = qos_class; + next.key = peer.key; + next.epoch = peer.epoch; + next.charge = std::move(charge); + output = std::move(next); + return Status::OK(); +} + +Status ReceiverCreditDispatchGate::tryReserve( + const CreditDispatchSnapshot& snapshot, + CreditDispatchReservation& reservation) { + if (reservation.state != CreditReservationState::Empty) + return Status::InvalidEntry( + "receiver credit reservation token already used" LOC_MARK); + CreditPeerContextSnapshot current; + CHECK_STATUS( + contexts_.lookupFresh(snapshot.target_id, snapshot.qos_class, current)); + if (!(current.key == snapshot.key) || current.epoch != snapshot.epoch) + return Status::InvalidEntry( + "receiver credit dispatch snapshot is stale" LOC_MARK); + CHECK_STATUS( + ledger_.tryReserve(snapshot.key, snapshot.epoch, snapshot.charge)); + reservation.snapshot = snapshot; + reservation.state = CreditReservationState::Reserved; + return Status::OK(); +} + +Status ReceiverCreditDispatchGate::commit( + CreditDispatchReservation& reservation) { + if (reservation.state != CreditReservationState::Reserved) + return Status::InvalidEntry( + "receiver credit reservation is not pending" LOC_MARK); + reservation.state = CreditReservationState::Committed; + return Status::OK(); +} + +Status ReceiverCreditDispatchGate::rollback( + CreditDispatchReservation& reservation) { + if (reservation.state != CreditReservationState::Reserved) + return Status::InvalidEntry( + "receiver credit reservation is not pending" LOC_MARK); + CHECK_STATUS(ledger_.rollbackReservation(reservation.snapshot.key, + reservation.snapshot.epoch, + reservation.snapshot.charge)); + reservation.state = CreditReservationState::RolledBack; + return Status::OK(); +} + +Status ReceiverCreditDispatchGate::release( + CreditDispatchReservation& reservation) { + if (reservation.state != CreditReservationState::Committed) + return Status::InvalidEntry( + "receiver credit reservation is not committed" LOC_MARK); + CHECK_STATUS(ledger_.recordCompletion(reservation.snapshot.key, + reservation.snapshot.epoch, + reservation.snapshot.charge)); + reservation.state = CreditReservationState::Released; + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit_protocol.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_protocol.cpp new file mode 100644 index 0000000000..ed9552b60d --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit_protocol.cpp @@ -0,0 +1,357 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_protocol.h" + +#include +#include +#include + +namespace mooncake::tent { +namespace { + +constexpr uint32_t kPullRequestMagic = 0x54435251; // "TCRQ" +constexpr uint32_t kPullResponseMagic = 0x54435250; // "TCRP" + +void append16(std::string& output, uint16_t value) { + output.push_back(static_cast(value >> 8)); + output.push_back(static_cast(value)); +} + +void append32(std::string& output, uint32_t value) { + for (int shift = 24; shift >= 0; shift -= 8) + output.push_back(static_cast(value >> shift)); +} + +void append64(std::string& output, uint64_t value) { + for (int shift = 56; shift >= 0; shift -= 8) + output.push_back(static_cast(value >> shift)); +} + +uint16_t read16(std::string_view input, size_t& position) { + uint16_t value = 0; + for (int i = 0; i < 2; ++i) + value = static_cast((value << 8) | + static_cast(input[position++])); + return value; +} + +uint32_t read32(std::string_view input, size_t& position) { + uint32_t value = 0; + for (int i = 0; i < 4; ++i) + value = (value << 8) | static_cast(input[position++]); + return value; +} + +uint64_t read64(std::string_view input, size_t& position) { + uint64_t value = 0; + for (int i = 0; i < 8; ++i) + value = (value << 8) | static_cast(input[position++]); + return value; +} + +bool isKnownResource(CreditResource resource, size_t& index) { + const uint16_t raw = static_cast(resource); + if (raw < 1 || raw > kCreditResourceCount) return false; + index = raw - 1; + return true; +} + +bool isKnownStatus(ReceiverCreditPullStatus status) { + switch (status) { + case ReceiverCreditPullStatus::Granted: + case ReceiverCreditPullStatus::Retry: + case ReceiverCreditPullStatus::SessionChanged: + case ReceiverCreditPullStatus::Unsupported: + case ReceiverCreditPullStatus::Rejected: + return true; + } + return false; +} + +bool isZeroPayload(const ReceiverCreditPullResponseV1& response) { + const auto& activation = response.activation; + const auto& update = response.update; + return activation.schema_version == 0 && activation.chosen_version == 0 && + activation.receiver_session_id.high == 0 && + activation.receiver_session_id.low == 0 && activation.epoch == 0 && + activation.freshness_ttl_ms == 0 && update.schema_version == 0 && + update.flags == 0 && update.qos_class == 0 && + update.receiver_session_id.high == 0 && + update.receiver_session_id.low == 0 && update.epoch == 0 && + update.sequence == 0 && update.freshness_ttl_ms == 0 && + update.grants.empty(); +} + +} // namespace + +Status ReceiverCreditPullRequestCodecV1::validate( + const ReceiverCreditPullRequestV1& request) { + if (request.schema_version != 1 || request.flags != 0 || + request.sender_peer == 0 || request.request_sequence == 0 || + request.resources.empty() || + request.resources.size() > kCreditResourceCount) + return Status::InvalidArgument( + "invalid receiver credit pull request header" LOC_MARK); + + const bool zero_session = request.expected_receiver_session_id.high == 0 && + request.expected_receiver_session_id.low == 0; + if (zero_session != (request.expected_epoch == 0)) + return Status::InvalidArgument( + "partial receiver credit session expectation" LOC_MARK); + + std::array present{}; + for (const auto& usage : request.resources) { + size_t index = 0; + if (!isKnownResource(usage.resource, index) || present[index]) + return Status::InvalidArgument( + "invalid or duplicate receiver credit pull resource" LOC_MARK); + present[index] = true; + if (usage.completed_total > usage.consumed_total || + usage.minimum_available > usage.desired_available || + usage.desired_available > + std::numeric_limits::max() - usage.consumed_total) + return Status::InvalidArgument( + "invalid or overflowing receiver credit pull usage" LOC_MARK); + } + return Status::OK(); +} + +Status ReceiverCreditPullRequestCodecV1::encode( + const ReceiverCreditPullRequestV1& request, std::string& wire) { + CHECK_STATUS(validate(request)); + std::string encoded; + encoded.reserve(kHeaderBytes + request.resources.size() * kResourceBytes); + append32(encoded, kPullRequestMagic); + append16(encoded, request.schema_version); + append16(encoded, request.flags); + append64(encoded, request.sender_peer); + append32(encoded, request.qos_class); + append64(encoded, request.expected_receiver_session_id.high); + append64(encoded, request.expected_receiver_session_id.low); + append64(encoded, request.expected_epoch); + append64(encoded, request.request_sequence); + append64(encoded, request.last_update_sequence); + append16(encoded, static_cast(request.resources.size())); + append16(encoded, 0); + for (const auto& usage : request.resources) { + append16(encoded, static_cast(usage.resource)); + append16(encoded, 0); + append64(encoded, usage.consumed_total); + append64(encoded, usage.completed_total); + append64(encoded, usage.minimum_available); + append64(encoded, usage.desired_available); + } + wire.swap(encoded); + return Status::OK(); +} + +Status ReceiverCreditPullRequestCodecV1::decode( + std::string_view wire, ReceiverCreditPullRequestV1& request) { + if (wire.size() < kHeaderBytes || wire.size() > kMaxWireBytes) + return Status::InvalidArgument( + "invalid receiver credit pull request length" LOC_MARK); + size_t position = 0; + if (read32(wire, position) != kPullRequestMagic) + return Status::InvalidArgument( + "invalid receiver credit pull request magic" LOC_MARK); + + ReceiverCreditPullRequestV1 decoded; + decoded.schema_version = read16(wire, position); + decoded.flags = read16(wire, position); + decoded.sender_peer = read64(wire, position); + decoded.qos_class = read32(wire, position); + decoded.expected_receiver_session_id.high = read64(wire, position); + decoded.expected_receiver_session_id.low = read64(wire, position); + decoded.expected_epoch = read64(wire, position); + decoded.request_sequence = read64(wire, position); + decoded.last_update_sequence = read64(wire, position); + const uint16_t count = read16(wire, position); + const uint16_t reserved = read16(wire, position); + if (reserved != 0 || count == 0 || count > kCreditResourceCount || + wire.size() != kHeaderBytes + count * kResourceBytes) + return Status::InvalidArgument( + "invalid receiver credit pull request framing" LOC_MARK); + + decoded.resources.reserve(count); + for (uint16_t i = 0; i < count; ++i) { + ReceiverCreditResourceUsageV1 usage; + usage.resource = static_cast(read16(wire, position)); + const uint16_t resource_reserved = read16(wire, position); + usage.consumed_total = read64(wire, position); + usage.completed_total = read64(wire, position); + usage.minimum_available = read64(wire, position); + usage.desired_available = read64(wire, position); + if (resource_reserved != 0) + return Status::InvalidArgument( + "nonzero receiver credit pull resource reserved " + "field" LOC_MARK); + decoded.resources.push_back(usage); + } + CHECK_STATUS(validate(decoded)); + request = std::move(decoded); + return Status::OK(); +} + +Status ReceiverCreditPullResponseCodecV1::validate( + const ReceiverCreditPullResponseV1& response) { + if (response.schema_version != 1 || response.flags != 0 || + !isKnownStatus(response.status)) + return Status::InvalidArgument( + "invalid receiver credit pull response header" LOC_MARK); + if ((response.status == ReceiverCreditPullStatus::Retry) != + (response.retry_after_us != 0)) + return Status::InvalidArgument( + "invalid receiver credit pull retry policy" LOC_MARK); + + const bool zero_payload_allowed = + response.status == ReceiverCreditPullStatus::Unsupported || + response.status == ReceiverCreditPullStatus::Rejected; + if (zero_payload_allowed && isZeroPayload(response)) return Status::OK(); + + const auto& activation = response.activation; + const auto& update = response.update; + if (activation.schema_version != 1 || activation.chosen_version != 1 || + (activation.receiver_session_id.high == 0 && + activation.receiver_session_id.low == 0) || + activation.epoch == 0 || activation.freshness_ttl_ms == 0 || + update.schema_version != 1 || update.flags != 0 || + !(update.receiver_session_id == activation.receiver_session_id) || + update.epoch != activation.epoch || update.sequence == 0 || + update.freshness_ttl_ms != activation.freshness_ttl_ms || + update.grants.size() != kCreditResourceCount) + return Status::InvalidArgument( + "invalid receiver credit pull activation or update" LOC_MARK); + + std::array present{}; + for (const auto& grant : update.grants) { + size_t index = 0; + if (!isKnownResource(grant.resource, index) || present[index]) + return Status::InvalidArgument( + "invalid receiver credit pull response grant" LOC_MARK); + present[index] = true; + } + return Status::OK(); +} + +Status ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus status, ReceiverCreditPullResponseV1& response) { + if (status != ReceiverCreditPullStatus::Unsupported && + status != ReceiverCreditPullStatus::Rejected) + return Status::InvalidArgument( + "zero receiver credit payload is not allowed for status" LOC_MARK); + ReceiverCreditPullResponseV1 empty; + empty.status = status; + empty.activation.schema_version = 0; + empty.activation.chosen_version = 0; + empty.update.schema_version = 0; + response = std::move(empty); + return Status::OK(); +} + +Status ReceiverCreditPullResponseCodecV1::encode( + const ReceiverCreditPullResponseV1& response, std::string& wire) { + CHECK_STATUS(validate(response)); + std::string encoded; + encoded.reserve(kWireBytes); + append32(encoded, kPullResponseMagic); + append16(encoded, response.schema_version); + append16(encoded, response.flags); + append16(encoded, static_cast(response.status)); + append16(encoded, 0); + append32(encoded, response.retry_after_us); + + append16(encoded, response.activation.schema_version); + append16(encoded, response.activation.chosen_version); + append64(encoded, response.activation.receiver_session_id.high); + append64(encoded, response.activation.receiver_session_id.low); + append64(encoded, response.activation.epoch); + append32(encoded, response.activation.freshness_ttl_ms); + append32(encoded, 0); + + append16(encoded, response.update.schema_version); + append16(encoded, response.update.flags); + append32(encoded, response.update.qos_class); + append64(encoded, response.update.receiver_session_id.high); + append64(encoded, response.update.receiver_session_id.low); + append64(encoded, response.update.epoch); + append64(encoded, response.update.sequence); + append32(encoded, response.update.freshness_ttl_ms); + append16(encoded, static_cast(response.update.grants.size())); + append16(encoded, 0); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (i < response.update.grants.size()) { + const auto& grant = response.update.grants[i]; + append16(encoded, static_cast(grant.resource)); + append16(encoded, 0); + append64(encoded, grant.grant_total); + } else { + append16(encoded, 0); + append16(encoded, 0); + append64(encoded, 0); + } + } + wire.swap(encoded); + return Status::OK(); +} + +Status ReceiverCreditPullResponseCodecV1::decode( + std::string_view wire, ReceiverCreditPullResponseV1& response) { + if (wire.size() != kWireBytes) + return Status::InvalidArgument( + "invalid receiver credit pull response length" LOC_MARK); + size_t position = 0; + if (read32(wire, position) != kPullResponseMagic) + return Status::InvalidArgument( + "invalid receiver credit pull response magic" LOC_MARK); + + ReceiverCreditPullResponseV1 decoded; + decoded.schema_version = read16(wire, position); + decoded.flags = read16(wire, position); + decoded.status = + static_cast(read16(wire, position)); + const uint16_t response_reserved = read16(wire, position); + decoded.retry_after_us = read32(wire, position); + + decoded.activation.schema_version = read16(wire, position); + decoded.activation.chosen_version = read16(wire, position); + decoded.activation.receiver_session_id.high = read64(wire, position); + decoded.activation.receiver_session_id.low = read64(wire, position); + decoded.activation.epoch = read64(wire, position); + decoded.activation.freshness_ttl_ms = read32(wire, position); + const uint32_t activation_reserved = read32(wire, position); + + decoded.update.schema_version = read16(wire, position); + decoded.update.flags = read16(wire, position); + decoded.update.qos_class = read32(wire, position); + decoded.update.receiver_session_id.high = read64(wire, position); + decoded.update.receiver_session_id.low = read64(wire, position); + decoded.update.epoch = read64(wire, position); + decoded.update.sequence = read64(wire, position); + decoded.update.freshness_ttl_ms = read32(wire, position); + const uint16_t count = read16(wire, position); + const uint16_t update_reserved = read16(wire, position); + if (response_reserved != 0 || activation_reserved != 0 || + update_reserved != 0 || (count != 0 && count != kCreditResourceCount)) + return Status::InvalidArgument( + "invalid receiver credit pull response framing" LOC_MARK); + + decoded.update.grants.reserve(count); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + CreditAmount grant; + grant.resource = static_cast(read16(wire, position)); + const uint16_t grant_reserved = read16(wire, position); + grant.grant_total = read64(wire, position); + if (grant_reserved != 0 || + (count == 0 && (static_cast(grant.resource) != 0 || + grant.grant_total != 0))) + return Status::InvalidArgument( + "invalid receiver credit pull grant or padding" LOC_MARK); + if (count != 0) decoded.update.grants.push_back(grant); + } + CHECK_STATUS(validate(decoded)); + response = std::move(decoded); + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 641a8a62bd..a6de4cf40a 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -46,6 +46,24 @@ namespace tent { namespace { constexpr uint8_t kRedisMaxDbIndex = 255; constexpr uint8_t kRedisDefaultDbIndex = 0; + +bool hasProductionReceiverCreditConfig(const Config& config) { + return config.contains("receiver_credit/mode") || + config.contains("receiver_credit/capacity") || + config.contains("receiver_credit/grant_batch") || + config.contains("receiver_credit/control") || + config.contains("receiver_credit/limits"); +} + +uint64_t randomNonzero64() { + static std::atomic sequence{1}; + std::random_device random; + uint64_t value = (static_cast(random()) << 32) ^ random(); + value ^= static_cast(getCurrentTimeInNano()); + value ^= sequence.fetch_add(1, std::memory_order_relaxed) * + 0x9e3779b97f4a7c15ULL; + return value == 0 ? 1 : value; +} } // namespace struct Batch { @@ -270,6 +288,11 @@ Status TransferEngineImpl::setupLocalSegment() { segment.type = SegmentType::Memory; segment.machine_id = getMachineID(); segment.rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); + if (receiver_credit_production_enabled_) { + segment.peer_session_high = local_peer_session_.high; + segment.peer_session_low = local_peer_session_.low; + segment.receiver_credit_versions = {1}; + } auto& detail = std::get(segment.detail); detail.topology = *(topology_.get()); return Status::OK(); @@ -292,6 +315,31 @@ Status TransferEngineImpl::construct() { conf_->get("enable_auto_failover_on_poll", true); enable_progress_worker_ = conf_->get("enable_progress_worker", false); runtime_queue_config_.enabled = conf_->get("enable_runtime_queue", false); + receiver_credit_production_enabled_ = + hasProductionReceiverCreditConfig(*conf_); + if (receiver_credit_production_enabled_) { + CHECK_STATUS(loadReceiverCreditConfig( + *conf_, runtime_queue_config_.enabled, receiver_credit_config_)); + receiver_credit_production_enabled_ = + receiver_credit_config_.mode != CreditRolloutMode::Disabled; + runtime_queue_config_.receiver_credit_enabled = + receiver_credit_production_enabled_; + runtime_queue_config_.receiver_credit_default_qos_class = + receiver_credit_config_.default_qos_class; + } else { + // Compatibility for the stacked dispatch-gate prototype. Production + // automatic wiring requires the strict mode/capacity schema above; + // the legacy boolean still supports the explicit install test API. + runtime_queue_config_.receiver_credit_enabled = + conf_->get("receiver_credit/enabled", false); + runtime_queue_config_.receiver_credit_default_qos_class = + conf_->get("receiver_credit/default_qos_class", uint32_t{0}); + } + if (runtime_queue_config_.receiver_credit_enabled && + !runtime_queue_config_.enabled) { + return Status::InvalidArgument( + "receiver credit dispatch requires the runtime queue" LOC_MARK); + } if (runtime_queue_config_.enabled) enable_progress_worker_ = true; runtime_queue_config_.limits.max_outstanding_owners = conf_->get("runtime_queue/max_outstanding_owners", 1024UL); @@ -314,6 +362,11 @@ Status TransferEngineImpl::construct() { runtime_queue_config_.progress_fallback_interval = std::chrono::microseconds( conf_->get("runtime_queue/progress_fallback_interval_us", 50000UL)); + if (receiver_credit_production_enabled_) { + runtime_queue_config_.progress_fallback_interval = + std::chrono::microseconds( + receiver_credit_config_.progress_interval_us); + } if (runtime_queue_config_.enabled && (runtime_queue_config_.max_dispatch_owners == 0 || runtime_queue_config_.max_dispatch_bytes == 0)) { @@ -322,6 +375,41 @@ Status TransferEngineImpl::construct() { } runtime_queue_ = std::make_unique( runtime_queue_config_.limits); + if (receiver_credit_production_enabled_) { + local_peer_session_ = {randomNonzero64(), randomNonzero64()}; + local_sender_peer_ = local_peer_session_.low; + + ReceiverCreditAllocatorConfig allocator_config; + allocator_config.capacity = receiver_credit_config_.capacity; + allocator_config.max_grant_per_pull = + receiver_credit_config_.max_grant_per_pull; + allocator_config.max_entries = receiver_credit_config_.max_peers; + allocator_config.ttl_ms = receiver_credit_config_.freshness_ttl_ms; + allocator_config.retry_after_us = + receiver_credit_config_.retry_after_us; + allocator_config.receiver_session_id = local_peer_session_; + allocator_config.epoch = 1; + std::unique_ptr allocator; + CHECK_STATUS( + ReceiverCreditAllocator::create(allocator_config, allocator)); + receiver_credit_allocator_ = std::move(allocator); + + receiver_credit_contexts_ = std::make_shared( + receiver_credit_config_.max_peers); + receiver_credit_ledger_ = std::make_shared( + receiver_credit_config_.max_peers); + receiver_credit_dispatch_gate_ = + std::make_shared( + *receiver_credit_contexts_, *receiver_credit_ledger_); + receiver_credit_qos_provider_ = [](const Request&) -> uint32_t { + return 0; + }; + CHECK_STATUS(ReceiverCreditPullController::create( + receiver_credit_config_, local_sender_peer_, + receiver_credit_contexts_, receiver_credit_ledger_, + receiver_credit_pull_controller_, + runtime_queue_config_.max_dispatch_owners)); + } if (!hostname_.empty()) CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_)); else @@ -334,6 +422,9 @@ Status TransferEngineImpl::construct() { metadata_ = std::make_shared(metadata_type, metadata_servers, this); + if (receiver_credit_allocator_) + metadata_->setReceiverCreditAllocator(receiver_credit_allocator_); + CHECK_STATUS(metadata_->start(port_, ipv6_)); if (metadata_type == "p2p") @@ -445,9 +536,52 @@ Status TransferEngineImpl::construct() { return Status::OK(); } +Status TransferEngineImpl::installReceiverCreditDispatch( + std::shared_ptr contexts, + std::shared_ptr ledger, + CreditQosProvider qos_provider) { + std::lock_guard lk(progress_mutex_); + if (receiver_credit_production_enabled_) { + return Status::InvalidEntry( + "production receiver credit control is installed " + "automatically" LOC_MARK); + } + if (!runtime_queue_config_.receiver_credit_enabled) { + return Status::InvalidArgument( + "receiver credit dispatch is not enabled" LOC_MARK); + } + if (!contexts || !ledger) { + return Status::InvalidArgument( + "receiver credit context and ledger are required" LOC_MARK); + } + if (!queued_owners_.empty()) { + return Status::InvalidEntry( + "receiver credit dispatch must be installed before " + "submit" LOC_MARK); + } + receiver_credit_contexts_ = std::move(contexts); + receiver_credit_ledger_ = std::move(ledger); + receiver_credit_dispatch_gate_ = + std::make_shared(*receiver_credit_contexts_, + *receiver_credit_ledger_); + if (qos_provider) { + receiver_credit_qos_provider_ = std::move(qos_provider); + } else { + const auto qos_class = + runtime_queue_config_.receiver_credit_default_qos_class; + receiver_credit_qos_provider_ = [qos_class](const Request&) { + return qos_class; + }; + } + return Status::OK(); +} + Status TransferEngineImpl::deconstruct() { // Metrics cleanup is handled automatically by TentMetrics destructor + if (receiver_credit_pull_controller_) + receiver_credit_pull_controller_->stop(); + // Stop the progress worker first so it cannot race with batch teardown // below (it dereferences BatchID into Batch* via progressBatch). Keep the // object alive until transports are destroyed: completion paths may still @@ -1603,6 +1737,12 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, submit.batch_token = batch_token; submit.batch_slots_left = batch->max_size - batch->task_list.size(); submit.owners.reserve(prepared.owners.size()); + std::vector> credit_snapshots( + prepared.owners.size()); + std::vector> credit_charges( + prepared.owners.size()); + std::vector credit_server_addrs(prepared.owners.size()); + std::vector credit_qos_classes(prepared.owners.size(), 0); for (const auto& owner : prepared.owners) { if (owner.request.length > runtime_queue_config_.max_dispatch_bytes) { return Status::TooManyRequests( @@ -1617,6 +1757,95 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, owner.route.transport == RDMA && !owner.staging; submit.owners.push_back(std::move(input)); } + if (runtime_queue_config_.receiver_credit_enabled && + !receiver_credit_dispatch_gate_) { + return Status::InvalidEntry( + "receiver credit dispatch is enabled but not installed" LOC_MARK); + } + if (receiver_credit_production_enabled_) { + if (!receiver_credit_pull_controller_ || + !receiver_credit_dispatch_gate_) + return Status::InvalidEntry( + "production receiver credit control is not installed" LOC_MARK); + for (size_t i = 0; i < prepared.owners.size(); ++i) { + const auto& owner = prepared.owners[i]; + const auto& request = owner.request; + const bool consumes_remote_receiver = + request.target_id != LOCAL_SEGMENT_ID && + request.opcode == Request::WRITE; + if (!consumes_remote_receiver) continue; + + const bool supported_direct_write = + owner_kind == QueueOwnerKind::User && !owner.staging && + owner.route.transport == RDMA; + if (!supported_direct_write) { + if (receiver_credit_config_.mode == CreditRolloutMode::Required) + return Status::NotImplemented( + "receiver credit required mode currently supports " + "only direct remote RDMA WRITE" LOC_MARK); + continue; + } + + SegmentDescRef remote; + CHECK_STATUS(metadata_->segmentManager().getRemoteCached( + remote, request.target_id)); + if (!remote) + return Status::InvalidEntry( + "remote segment descriptor is null" LOC_MARK); + const bool supports_v1 = + std::find(remote->receiver_credit_versions.begin(), + remote->receiver_credit_versions.end(), + uint16_t{1}) != + remote->receiver_credit_versions.end(); + if (!supports_v1) { + if (receiver_credit_config_.mode == CreditRolloutMode::Required) + return Status::NotImplemented( + "remote peer does not advertise receiver credit " + "v1" LOC_MARK); + continue; + } + if (remote->rpc_server_addr.empty()) + return Status::InvalidEntry( + "receiver credit peer has no RPC address" LOC_MARK); + + CreditCharge charge{{{CreditResource::DataBytes, request.length}, + {CreditResource::RequestSlots, 1}}}; + const uint32_t qos_class = 0; + CreditDispatchSnapshot snapshot; + auto snapshot_status = receiver_credit_dispatch_gate_->snapshot( + request.target_id, qos_class, charge, snapshot); + if (snapshot_status.ok()) { + credit_snapshots[i] = std::move(snapshot); + } else if (snapshot_status.IsInvalidEntry()) { + // Pull only while the peer has no fresh generation. Available + // credit is reserved later by the queue owner; an exhausted + // ledger triggers a coalesced refill from that path. Pulling + // here unconditionally turns every submit into a control RPC. + CHECK_STATUS(receiver_credit_pull_controller_->request( + request.target_id, remote->rpc_server_addr, qos_class, + charge)); + } else { + return snapshot_status; + } + credit_charges[i] = std::move(charge); + credit_server_addrs[i] = remote->rpc_server_addr; + credit_qos_classes[i] = qos_class; + } + } else if (receiver_credit_dispatch_gate_) { + if (!receiver_credit_qos_provider_) + return Status::InvalidEntry( + "receiver credit QoS provider is missing" LOC_MARK); + for (size_t i = 0; i < prepared.owners.size(); ++i) { + const auto& request = prepared.owners[i].request; + CreditCharge charge{{{CreditResource::DataBytes, request.length}, + {CreditResource::RequestSlots, 1}}}; + CreditDispatchSnapshot snapshot; + CHECK_STATUS(receiver_credit_dispatch_gate_->snapshot( + request.target_id, receiver_credit_qos_provider_(request), + std::move(charge), snapshot)); + credit_snapshots[i] = std::move(snapshot); + } + } std::vector admitted_owner_ids; CHECK_STATUS(runtime_queue_->tryAdmit(submit, admitted_owner_ids)); @@ -1650,6 +1879,10 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, queued.public_task_ids.end(), prepared.owners[i].derived_task_ids.begin(), prepared.owners[i].derived_task_ids.end()); + queued.credit_snapshot = std::move(credit_snapshots[i]); + queued.credit_charge = std::move(credit_charges[i]); + queued.credit_server_addr = std::move(credit_server_addrs[i]); + queued.credit_qos_class = credit_qos_classes[i]; queued_owners_.emplace(admitted_owner_ids[i], queued); } return Status::OK(); @@ -1662,6 +1895,19 @@ Status TransferEngineImpl::finishQueuedOwner( return Status::InvalidEntry("queued owner not found" LOC_MARK); } auto& queued = queued_it->second; + if (queued.credit_reservation && + queued.credit_reservation->state == CreditReservationState::Reserved) { + CHECK_STATUS(receiver_credit_dispatch_gate_->rollback( + *queued.credit_reservation)); + } + if (queued.credit_reservation && + queued.credit_reservation->state == CreditReservationState::Committed) { + CHECK_STATUS(receiver_credit_dispatch_gate_->release( + *queued.credit_reservation)); + // Completion is cumulative local state, not a push notification. The + // next on-demand pull reports it while asking for a refill; sending an + // RPC here makes control traffic proportional to transfer count. + } if (queued.in_dispatch_window) { if (dispatch_inflight_owners_ == 0 || dispatch_inflight_bytes_ < queued.byte_charge) { @@ -1728,7 +1974,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (queued_it == queued_owners_.end()) { return Status::InternalError("queued owner metadata missing" LOC_MARK); } - const auto queued = queued_it->second; + auto& queued = queued_it->second; auto* batch = queued.batch; auto& task = batch->task_list[queued.owner_task_id]; task.dispatch_time = std::chrono::steady_clock::now(); @@ -1740,6 +1986,72 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { return finishQueuedOwner(owner_id, FAILED); } + if (receiver_credit_production_enabled_ && queued.credit_charge && + task.type != RDMA) { + if (receiver_credit_config_.mode == CreditRolloutMode::Required) + return finishQueuedOwner(owner_id, FAILED); + queued.credit_charge.reset(); + queued.credit_snapshot.reset(); + } + + if (receiver_credit_production_enabled_ && queued.credit_charge && + !queued.credit_snapshot) { + CreditDispatchSnapshot snapshot; + auto snapshot_status = receiver_credit_dispatch_gate_->snapshot( + task.request.target_id, queued.credit_qos_class, + *queued.credit_charge, snapshot); + if (snapshot_status.ok()) { + queued.credit_snapshot = std::move(snapshot); + } else if (snapshot_status.IsInvalidEntry()) { + auto peer_state = receiver_credit_pull_controller_->peerState( + task.request.target_id, queued.credit_qos_class); + if (peer_state == CreditPeerState::Failed) + return finishQueuedOwner(owner_id, FAILED); + if (peer_state == CreditPeerState::Legacy && + receiver_credit_config_.mode == CreditRolloutMode::Optional) { + queued.credit_charge.reset(); + } else { + CHECK_STATUS(receiver_credit_pull_controller_->request( + task.request.target_id, queued.credit_server_addr, + queued.credit_qos_class, *queued.credit_charge)); + CHECK_STATUS(runtime_queue_->deferDispatch(owner_id)); + return Status::TooManyRequests( + "waiting for receiver credit" LOC_MARK); + } + } else { + return finishQueuedOwner(owner_id, FAILED); + } + } + + if (queued.credit_snapshot) { + queued.credit_reservation.emplace(); + auto reserve_status = receiver_credit_dispatch_gate_->tryReserve( + *queued.credit_snapshot, *queued.credit_reservation); + if (reserve_status.IsTooManyRequests()) { + queued.credit_reservation.reset(); + if (receiver_credit_production_enabled_ && queued.credit_charge) + CHECK_STATUS(receiver_credit_pull_controller_->request( + task.request.target_id, queued.credit_server_addr, + queued.credit_qos_class, *queued.credit_charge)); + CHECK_STATUS(runtime_queue_->deferDispatch(owner_id)); + return reserve_status; + } + if (!reserve_status.ok()) { + queued.credit_reservation.reset(); + if (receiver_credit_production_enabled_ && queued.credit_charge && + reserve_status.IsInvalidEntry()) { + queued.credit_snapshot.reset(); + CHECK_STATUS(receiver_credit_pull_controller_->request( + task.request.target_id, queued.credit_server_addr, + queued.credit_qos_class, *queued.credit_charge)); + CHECK_STATUS(runtime_queue_->deferDispatch(owner_id)); + return Status::TooManyRequests( + "receiver credit generation is refreshing" LOC_MARK); + } + return finishQueuedOwner(owner_id, FAILED); + } + } + if (task.type == TCP) { std::vector staging_params; findStagingPolicy(task.request, staging_params); @@ -1749,6 +2061,9 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { staging_proxy_->submit(&task, (BatchID)batch, staging_params); if (!status.ok()) return finishQueuedOwner(owner_id, FAILED); task.post_time = std::chrono::steady_clock::now(); + if (queued.credit_reservation) + CHECK_STATUS(receiver_credit_dispatch_gate_->commit( + *queued.credit_reservation)); return markQueuedOwnerSubmitted(owner_id); } } @@ -1776,25 +2091,39 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { return finishQueuedOwner(owner_id, FAILED); } task.post_time = std::chrono::steady_clock::now(); + if (queued.credit_reservation) + CHECK_STATUS( + receiver_credit_dispatch_gate_->commit(*queued.credit_reservation)); return markQueuedOwnerSubmitted(owner_id); } Status TransferEngineImpl::refillDispatchWindow() { std::lock_guard lk(progress_mutex_); if (!runtime_queue_config_.enabled) return Status::OK(); - if (dispatch_inflight_owners_ >= - runtime_queue_config_.max_dispatch_owners || + size_t max_dispatch_owners = runtime_queue_config_.max_dispatch_owners; + if (receiver_credit_production_enabled_ && receiver_credit_pull_controller_) + max_dispatch_owners = + std::min(max_dispatch_owners, + receiver_credit_pull_controller_->dispatchOwnerLimit()); + if (dispatch_inflight_owners_ >= max_dispatch_owners || dispatch_inflight_bytes_ >= runtime_queue_config_.max_dispatch_bytes) { return Status::OK(); } - const size_t owner_budget = - runtime_queue_config_.max_dispatch_owners - dispatch_inflight_owners_; + const size_t owner_budget = max_dispatch_owners - dispatch_inflight_owners_; const size_t byte_budget = runtime_queue_config_.max_dispatch_bytes - dispatch_inflight_bytes_; auto picked = runtime_queue_->pickForDispatch(owner_budget, byte_budget); - for (const auto owner_id : picked) { - CHECK_STATUS(dispatchQueuedOwner(owner_id)); + for (size_t i = 0; i < picked.size(); ++i) { + auto dispatch_status = dispatchQueuedOwner(picked[i]); + if (dispatch_status.ok()) continue; + // pickForDispatch transitions the complete returned set to + // Dispatching. If dispatch stops early, every unprocessed owner must + // be returned to Queued or it can never be selected again. + for (size_t j = i + 1; j < picked.size(); ++j) + CHECK_STATUS(runtime_queue_->deferDispatch(picked[j])); + if (dispatch_status.IsTooManyRequests()) return Status::OK(); + return dispatch_status; } return Status::OK(); } @@ -1869,6 +2198,21 @@ Status TransferEngineImpl::submitTransfer( PreparedSubmit prepared; CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, prepared)); + if (receiver_credit_production_enabled_ && + receiver_credit_config_.mode == CreditRolloutMode::Required) { + for (const auto& owner : prepared.owners) { + const bool remote_write = + owner.request.target_id != LOCAL_SEGMENT_ID && + owner.request.opcode == Request::WRITE; + if (remote_write && + (owner_kind != QueueOwnerKind::User || owner.staging || + owner.route.transport != RDMA)) + return Status::NotImplemented( + "receiver credit required mode currently supports only " + "direct remote RDMA WRITE" LOC_MARK); + } + } + if (shouldQueueSubmit(prepared, owner_kind)) { CHECK_STATUS( enqueuePreparedSubmit(batch_ref.get(), prepared, owner_kind)); diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 6778937d3f..c853f58cc5 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -34,14 +34,107 @@ target_include_directories(receiver_credit_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME receiver_credit_test COMMAND receiver_credit_test) +add_executable(receiver_credit_control_test receiver_credit_control_test.cpp + ../src/runtime/receiver_credit.cpp + ../src/runtime/receiver_credit_control.cpp) +target_link_libraries(receiver_credit_control_test + PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_control_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_control_test COMMAND receiver_credit_control_test) + +add_executable(receiver_credit_protocol_test + receiver_credit_protocol_test.cpp + ../src/runtime/receiver_credit_protocol.cpp) +target_link_libraries(receiver_credit_protocol_test + PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_protocol_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_protocol_test + COMMAND receiver_credit_protocol_test) + +add_executable(receiver_credit_allocator_test + receiver_credit_allocator_test.cpp + ../src/runtime/receiver_credit_protocol.cpp + ../src/runtime/receiver_credit_allocator.cpp) +target_link_libraries(receiver_credit_allocator_test + PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_allocator_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_allocator_test + COMMAND receiver_credit_allocator_test) + +add_executable(receiver_credit_config_test + receiver_credit_config_test.cpp + ../src/runtime/receiver_credit_config.cpp) +target_link_libraries(receiver_credit_config_test + PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_config_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_config_test COMMAND receiver_credit_config_test) + +add_executable(receiver_credit_control_service_test + receiver_credit_control_service_test.cpp) +target_link_libraries(receiver_credit_control_service_test + PRIVATE gtest gtest_main tent_link_group) +if(TARGET asio_shared) + target_link_libraries(receiver_credit_control_service_test PRIVATE asio_shared) +endif() +target_include_directories(receiver_credit_control_service_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_control_service_test + COMMAND receiver_credit_control_service_test) + +add_executable(receiver_credit_controller_test + receiver_credit_controller_test.cpp) +target_link_libraries(receiver_credit_controller_test + PRIVATE gtest gtest_main tent_link_group) +if(TARGET asio_shared) + target_link_libraries(receiver_credit_controller_test PRIVATE asio_shared) +endif() +target_include_directories(receiver_credit_controller_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_controller_test + COMMAND receiver_credit_controller_test) + +# Manual cross-host benchmark for the production receiver-credit ControlService +# RPC. It is intentionally not registered with ctest. +add_executable(receiver_credit_rpc_bench receiver_credit_rpc_bench.cpp) +target_link_libraries(receiver_credit_rpc_bench PRIVATE tent_link_group) +if(TARGET asio_shared) + target_link_libraries(receiver_credit_rpc_bench PRIVATE asio_shared) +endif() +target_include_directories(receiver_credit_rpc_bench + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + +add_executable(receiver_credit_dispatch_test receiver_credit_dispatch_test.cpp + ../src/runtime/receiver_credit.cpp + ../src/runtime/receiver_credit_control.cpp + ../src/runtime/receiver_credit_dispatch.cpp) +target_link_libraries(receiver_credit_dispatch_test + PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_dispatch_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_dispatch_test + COMMAND receiver_credit_dispatch_test) + +# Manual cross-host harness. It is intentionally not registered with ctest. +add_executable(receiver_credit_rpc_harness receiver_credit_rpc_harness.cpp + ../src/runtime/receiver_credit.cpp + ../src/runtime/receiver_credit_control.cpp) +target_link_libraries(receiver_credit_rpc_harness + PRIVATE tent_rpc tent_common asio_shared glog) +target_include_directories(receiver_credit_rpc_harness + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_executable(bw_arbitration_test bw_arbitration_test.cpp) target_link_libraries(bw_arbitration_test PRIVATE tent_common gtest gtest_main) target_include_directories(bw_arbitration_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME bw_arbitration_test COMMAND bw_arbitration_test) -# Reproducible hot-path microbenchmark; intentionally not registered with -# ctest. Run manually when changing deadline promotion partitioning. +# Reproducible hot-path microbenchmark; intentionally not registered with ctest. +# Run manually when changing deadline promotion partitioning. add_executable(deadline_promotion_bench deadline_promotion_bench.cpp ../src/runtime/admission_queue.cpp) target_link_libraries(deadline_promotion_bench PRIVATE tent_common) @@ -67,12 +160,6 @@ target_include_directories(thread_local_storage_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME thread_local_storage_test COMMAND thread_local_storage_test) -add_executable(tent_rw_spinlock_test rw_spinlock_test.cpp) -target_link_libraries(tent_rw_spinlock_test PRIVATE gtest gtest_main) -target_include_directories(tent_rw_spinlock_test - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_rw_spinlock_test COMMAND tent_rw_spinlock_test) - add_executable(tent_ip_utils_test ip_utils_test.cpp) target_link_libraries(tent_ip_utils_test PRIVATE tent_common gtest gtest_main) target_include_directories(tent_ip_utils_test @@ -130,13 +217,6 @@ target_include_directories(tent_tcp_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_tcp_transport_test COMMAND tent_tcp_transport_test) -add_executable(tent_shm_transport_test shm_transport_test.cpp) -target_link_libraries(tent_shm_transport_test PRIVATE gtest gtest_main - tent_link_group) -target_include_directories(tent_shm_transport_test - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_shm_transport_test COMMAND tent_shm_transport_test) - add_executable(tent_failover_test failover_test.cpp) target_link_libraries(tent_failover_test PRIVATE gtest gtest_main tent_link_group) @@ -267,8 +347,7 @@ add_test(NAME tent_runtime_queue_dispatch_test # Causal chain stage decomposition: validates that dispatch_time and post_time # timestamps are populated on both queue and direct-commit paths. add_executable(causal_chain_test causal_chain_test.cpp) -target_link_libraries(causal_chain_test PRIVATE gtest gtest_main - tent_link_group) +target_link_libraries(causal_chain_test PRIVATE gtest gtest_main tent_link_group) if(TARGET asio_shared) target_link_libraries(causal_chain_test PRIVATE asio_shared) endif() diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index 27f9681a2a..960fc71c24 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -171,6 +171,24 @@ TEST(AdmissionQueueTest, RejectsExistingPublicTaskConflictWithoutMutation) { EXPECT_EQ(admitted_ids[0], 2u); } +TEST(AdmissionQueueTest, DeferredDispatchReturnsOwnerToQueue) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 16); + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(queue.deferDispatch(picked[0]).ok()); + EXPECT_TRUE(queue.deferDispatch(picked[0]).IsInvalidEntry()); + + auto retried = queue.pickForDispatch(1, 16); + ASSERT_EQ(retried.size(), 1u); + EXPECT_EQ(retried[0], picked[0]); + EXPECT_TRUE(queue.complete(retried[0], TransferStatusEnum::COMPLETED).ok()); +} + TEST(AdmissionQueueTest, AccountsPublicSlotsSeparatelyFromQueueOwners) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_allocator_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_allocator_test.cpp new file mode 100644 index 0000000000..7ebbfc6142 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_allocator_test.cpp @@ -0,0 +1,334 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_allocator.h" + +#include +#include +#include +#include + +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t index(CreditResource resource) { + return static_cast(resource) - 1; +} + +ReceiverCreditAllocatorConfig config(uint64_t bytes = 1000, + uint64_t requests = 10) { + ReceiverCreditAllocatorConfig config; + config.capacity[index(CreditResource::DataBytes)] = bytes; + config.capacity[index(CreditResource::RequestSlots)] = requests; + config.max_grant_per_pull[index(CreditResource::DataBytes)] = bytes; + config.max_grant_per_pull[index(CreditResource::RequestSlots)] = requests; + config.max_entries = 64; + config.ttl_ms = 1000; + config.retry_after_us = 50; + config.receiver_session_id = {11, 22}; + config.epoch = 7; + return config; +} + +std::unique_ptr allocator( + const ReceiverCreditAllocatorConfig& allocator_config = config()) { + std::unique_ptr allocator; + EXPECT_TRUE( + ReceiverCreditAllocator::create(allocator_config, allocator).ok()); + return allocator; +} + +ReceiverCreditPullRequestV1 pullRequest(uint64_t peer, uint64_t sequence, + uint64_t data_minimum, + uint64_t data_desired, + uint64_t slot_minimum, + uint64_t slot_desired) { + ReceiverCreditPullRequestV1 request; + request.sender_peer = peer; + request.request_sequence = sequence; + if (sequence > 1) { + request.expected_receiver_session_id = {11, 22}; + request.expected_epoch = 7; + request.last_update_sequence = sequence - 1; + } + request.resources = { + {CreditResource::DataBytes, 0, 0, data_minimum, data_desired}, + {CreditResource::RequestSlots, 0, 0, slot_minimum, slot_desired}, + }; + return request; +} + +uint64_t grant(const ReceiverCreditPullResponseV1& response, + CreditResource resource) { + for (const auto& amount : response.update.grants) + if (amount.resource == resource) return amount.grant_total; + ADD_FAILURE() << "missing full credit resource"; + return 0; +} + +TEST(ReceiverCreditAllocator, RejectsInvalidConfigurationAtomically) { + auto valid = config(); + std::unique_ptr output; + ASSERT_TRUE(ReceiverCreditAllocator::create(valid, output).ok()); + auto* original = output.get(); + + auto invalid = valid; + invalid.receiver_session_id = {}; + EXPECT_TRUE( + ReceiverCreditAllocator::create(invalid, output).IsInvalidArgument()); + EXPECT_EQ(output.get(), original); + + invalid = valid; + invalid.capacity[index(CreditResource::StagingSlots)] = 1; + invalid.max_grant_per_pull[index(CreditResource::StagingSlots)] = 1; + EXPECT_TRUE( + ReceiverCreditAllocator::create(invalid, output).IsInvalidArgument()); + EXPECT_EQ(output.get(), original); + + invalid = valid; + invalid.max_grant_per_pull[index(CreditResource::DataBytes)] = 1001; + EXPECT_TRUE( + ReceiverCreditAllocator::create(invalid, output).IsInvalidArgument()); + EXPECT_EQ(output.get(), original); +} + +TEST(ReceiverCreditAllocator, FirstPullActivatesAndReturnsFullGrantVector) { + auto receiver = allocator(); + auto request = pullRequest(101, 1, 100, 500, 1, 4); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE(receiver->pull(request, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(response.activation.receiver_session_id, + (ReceiverSessionId{11, 22})); + EXPECT_EQ(response.activation.epoch, 7); + EXPECT_EQ(response.update.grants.size(), kCreditResourceCount); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 500); + EXPECT_EQ(grant(response, CreditResource::RequestSlots), 4); + EXPECT_EQ(grant(response, CreditResource::StagingSlots), 0); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 500); + EXPECT_EQ(snapshot.free[index(CreditResource::DataBytes)], 500); + EXPECT_EQ(snapshot.entries, 1); +} + +TEST(ReceiverCreditAllocator, DuplicateIsIdempotentAndChangedBodyIsRejected) { + auto receiver = allocator(); + auto request = pullRequest(101, 1, 100, 500, 1, 4); + ReceiverCreditPullResponseV1 first; + ReceiverCreditPullResponseV1 duplicate; + ASSERT_TRUE(receiver->pull(request, first).ok()); + ASSERT_TRUE(receiver->pull(request, duplicate).ok()); + EXPECT_EQ(duplicate.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(duplicate.update.sequence, first.update.sequence); + EXPECT_EQ(grant(duplicate, CreditResource::DataBytes), 500); + + request.resources[0].desired_available = 600; + ReceiverCreditPullResponseV1 changed; + ASSERT_TRUE(receiver->pull(request, changed).ok()); + EXPECT_EQ(changed.status, ReceiverCreditPullStatus::Rejected); + EXPECT_EQ(grant(changed, CreditResource::DataBytes), 500); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 500); +} + +TEST(ReceiverCreditAllocator, OldSequenceAndFutureUpdateAreRejected) { + auto receiver = allocator(); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE( + receiver->pull(pullRequest(101, 1, 1, 100, 1, 1), response).ok()); + auto second = pullRequest(101, 2, 1, 100, 1, 1); + second.last_update_sequence = 99; + ASSERT_TRUE(receiver->pull(second, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Rejected); + + auto old = pullRequest(101, 1, 1, 100, 1, 1); + old.expected_receiver_session_id = {11, 22}; + old.expected_epoch = 7; + ASSERT_TRUE(receiver->pull(old, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Rejected); +} + +TEST(ReceiverCreditAllocator, MinimumIsAtomicAcrossResources) { + auto receiver = allocator(); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE( + receiver->pull(pullRequest(101, 1, 900, 900, 1, 1), response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + + ASSERT_TRUE( + receiver->pull(pullRequest(102, 1, 200, 200, 2, 2), response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Retry); + EXPECT_EQ(response.retry_after_us, 50); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 0); + EXPECT_EQ(grant(response, CreditResource::RequestSlots), 0); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 900); + EXPECT_EQ(snapshot.committed[index(CreditResource::RequestSlots)], 1); +} + +TEST(ReceiverCreditAllocator, CompletionIsAppliedBeforeARejectedMinimum) { + auto receiver = allocator(); + ReceiverCreditPullResponseV1 response; + auto first = pullRequest(101, 1, 1000, 1000, 0, 0); + ASSERT_TRUE(receiver->pull(first, response).ok()); + ASSERT_EQ(grant(response, CreditResource::DataBytes), 1000); + + ReceiverCreditPullRequestV1 next; + next.sender_peer = 101; + next.expected_receiver_session_id = {11, 22}; + next.expected_epoch = 7; + next.request_sequence = 2; + next.last_update_sequence = 1; + next.resources = { + {CreditResource::DataBytes, 1000, 500, 600, 600}, + }; + ASSERT_TRUE(receiver->pull(next, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Retry); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 1000); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 500); + EXPECT_EQ(snapshot.free[index(CreditResource::DataBytes)], 500); +} + +TEST(ReceiverCreditAllocator, CompletionAllowsAnotherSenderToAcquireBudget) { + auto receiver = allocator(); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE( + receiver->pull(pullRequest(101, 1, 1000, 1000, 10, 10), response).ok()); + ASSERT_TRUE( + receiver->pull(pullRequest(102, 1, 1, 500, 1, 5), response).ok()); + ASSERT_EQ(response.status, ReceiverCreditPullStatus::Retry); + + ReceiverCreditPullRequestV1 release; + release.sender_peer = 101; + release.expected_receiver_session_id = {11, 22}; + release.expected_epoch = 7; + release.request_sequence = 2; + release.last_update_sequence = 1; + release.resources = { + {CreditResource::DataBytes, 1000, 1000, 0, 0}, + {CreditResource::RequestSlots, 10, 10, 0, 0}, + }; + ASSERT_TRUE(receiver->pull(release, response).ok()); + ASSERT_EQ(response.status, ReceiverCreditPullStatus::Granted); + + auto retry = pullRequest(102, 2, 1, 500, 1, 5); + ASSERT_TRUE(receiver->pull(retry, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 500); + EXPECT_EQ(grant(response, CreditResource::RequestSlots), 5); +} + +TEST(ReceiverCreditAllocator, MaxGrantCapsDesiredButPreservesMinimum) { + auto limited = config(); + limited.max_grant_per_pull[index(CreditResource::DataBytes)] = 200; + limited.max_grant_per_pull[index(CreditResource::RequestSlots)] = 2; + auto receiver = allocator(limited); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE( + receiver->pull(pullRequest(101, 1, 100, 900, 1, 9), response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 200); + EXPECT_EQ(grant(response, CreditResource::RequestSlots), 2); +} + +TEST(ReceiverCreditAllocator, RestartAndUnsupportedQosNeverAllocate) { + auto receiver = allocator(); + auto stale = pullRequest(101, 1, 100, 500, 1, 4); + stale.expected_receiver_session_id = {99, 100}; + stale.expected_epoch = 1; + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE(receiver->pull(stale, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::SessionChanged); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 0); + + auto unsupported = pullRequest(102, 1, 100, 500, 1, 4); + unsupported.qos_class = 1; + ASSERT_TRUE(receiver->pull(unsupported, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Unsupported); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.entries, 0); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 0); +} + +TEST(ReceiverCreditAllocator, InitialSessionRestartReclaimsOutstandingGrant) { + auto receiver = allocator(); + ReceiverCreditPullResponseV1 response; + + ASSERT_TRUE( + receiver->pull(pullRequest(101, 1, 100, 500, 1, 4), response).ok()); + ASSERT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 500); + + auto restart = pullRequest(101, 2, 100, 600, 1, 5); + restart.expected_receiver_session_id = {}; + restart.expected_epoch = 0; + restart.last_update_sequence = 0; + ASSERT_TRUE(receiver->pull(restart, response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(grant(response, CreditResource::DataBytes), 600); + EXPECT_EQ(grant(response, CreditResource::RequestSlots), 5); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.entries, 1); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 600); + EXPECT_EQ(snapshot.committed[index(CreditResource::RequestSlots)], 5); + EXPECT_EQ(snapshot.free[index(CreditResource::DataBytes)], 400); +} + +TEST(ReceiverCreditAllocator, EntryTableIsBounded) { + auto bounded = config(); + bounded.max_entries = 1; + auto receiver = allocator(bounded); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE(receiver->pull(pullRequest(101, 1, 1, 1, 1, 1), response).ok()); + ASSERT_TRUE(receiver->pull(pullRequest(102, 1, 1, 1, 1, 1), response).ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Retry); + EXPECT_EQ(response.retry_after_us, 50); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(snapshot.entries, 1); +} + +TEST(ReceiverCreditAllocator, ConcurrentPullsNeverExceedGlobalBudget) { + auto concurrent = config(1000, 1000); + concurrent.max_entries = 128; + auto receiver = allocator(concurrent); + std::atomic granted{0}; + std::vector threads; + for (uint64_t peer = 1; peer <= 64; ++peer) { + threads.emplace_back([&, peer] { + ReceiverCreditPullResponseV1 response; + auto request = pullRequest(peer, 1, 100, 100, 1, 1); + if (receiver->pull(request, response).ok() && + response.status == ReceiverCreditPullStatus::Granted) + ++granted; + }); + } + for (auto& thread : threads) thread.join(); + + ReceiverCreditAllocatorSnapshot snapshot; + ASSERT_TRUE(receiver->snapshot(snapshot).ok()); + EXPECT_EQ(granted, 10); + EXPECT_EQ(snapshot.committed[index(CreditResource::DataBytes)], 1000); + EXPECT_EQ(snapshot.free[index(CreditResource::DataBytes)], 0); + EXPECT_EQ(snapshot.committed[index(CreditResource::RequestSlots)], 10); + EXPECT_EQ(snapshot.entries, 64); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_config_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_config_test.cpp new file mode 100644 index 0000000000..376961dd49 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_config_test.cpp @@ -0,0 +1,344 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_config.h" + +#include + +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t resourceIndex(CreditResource resource) { + return static_cast(resource) - 1; +} + +Status load(const std::string& contents, bool runtime_queue_enabled, + ReceiverCreditRuntimeConfig& output) { + Config config; + auto status = config.load(contents); + if (!status.ok()) return status; + return loadReceiverCreditConfig(config, runtime_queue_enabled, output); +} + +TEST(ReceiverCreditConfig, MissingSectionIsStrictlyDefaultOff) { + Config config; + ReceiverCreditRuntimeConfig output; + output.mode = CreditRolloutMode::Required; + + ASSERT_TRUE(loadReceiverCreditConfig(config, false, output).ok()); + EXPECT_EQ(output.mode, CreditRolloutMode::Disabled); + EXPECT_EQ(output.capacity, (std::array{})); + EXPECT_EQ(output.max_grant_per_pull, + (std::array{})); + EXPECT_EQ(output.max_peers, ReceiverCreditRuntimeConfig::kDefaultMaxPeers); + EXPECT_TRUE(output.adaptive_dispatch_enabled); + EXPECT_EQ(output.adaptive_dispatch_min_owners, 1); + EXPECT_EQ(output.adaptive_dispatch_initial_owners, 2); + EXPECT_EQ(output.adaptive_dispatch_max_owners, 2); + EXPECT_EQ(output.adaptive_dispatch_slow_rtt_us, 20000); +} + +TEST(ReceiverCreditConfig, LoadsCompleteNestedRequiredConfiguration) { + ReceiverCreditRuntimeConfig output; + ASSERT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "default_qos_class": 0, + "capacity": {"data_bytes": 1048576, "request_slots": 64}, + "grant_batch": {"data_bytes": 262144, "request_slots": 8}, + "control": { + "freshness_ttl_ms": 500, + "retry_after_us": 200, + "poll_interval_us": 50, + "adaptive_dispatch": { + "enabled": true, + "min_owners": 1, + "initial_owners": 2, + "max_owners": 8, + "slow_rtt_us": 1500, + "healthy_pulls_per_increase": 1024 + } + }, + "limits": {"max_peers": 4096} + } + })", + true, output) + .ok()); + + EXPECT_EQ(output.mode, CreditRolloutMode::Required); + EXPECT_EQ(output.capacity[resourceIndex(CreditResource::DataBytes)], + 1048576); + EXPECT_EQ(output.capacity[resourceIndex(CreditResource::RequestSlots)], 64); + EXPECT_EQ( + output.max_grant_per_pull[resourceIndex(CreditResource::DataBytes)], + 262144); + EXPECT_EQ( + output.max_grant_per_pull[resourceIndex(CreditResource::RequestSlots)], + 8); + EXPECT_EQ(output.capacity[resourceIndex(CreditResource::StagingSlots)], 0); + EXPECT_EQ(output.max_peers, 4096); + EXPECT_EQ(output.freshness_ttl_ms, 500); + EXPECT_EQ(output.retry_after_us, 200); + EXPECT_EQ(output.progress_interval_us, 50); + EXPECT_TRUE(output.adaptive_dispatch_enabled); + EXPECT_EQ(output.adaptive_dispatch_min_owners, 1); + EXPECT_EQ(output.adaptive_dispatch_initial_owners, 2); + EXPECT_EQ(output.adaptive_dispatch_max_owners, 8); + EXPECT_EQ(output.adaptive_dispatch_slow_rtt_us, 1500); + EXPECT_EQ(output.adaptive_dispatch_healthy_pulls, 1024); +} + +TEST(ReceiverCreditConfig, SupportsOptionalModeAndLegacyAliases) { + ReceiverCreditRuntimeConfig output; + ASSERT_TRUE(load(R"({ + "receiver_credit": { + "mode": "optional", + "capacity": {"data_bytes": 1024, "request_slots": 4}, + "grant_batch": {"data_bytes": 512, "request_slots": 2} + } + })", + true, output) + .ok()); + EXPECT_EQ(output.mode, CreditRolloutMode::Optional); + + ASSERT_TRUE( + load(R"({"receiver_credit": {"enabled": false}})", false, output).ok()); + EXPECT_EQ(output.mode, CreditRolloutMode::Disabled); + + ASSERT_TRUE(load(R"({ + "receiver_credit": { + "capacity": {"data_bytes": 1024, "request_slots": 4}, + "grant_batch": {"data_bytes": 512, "request_slots": 2} + }, + "receiver_credit/enabled": true + })", + true, output) + .ok()); + EXPECT_EQ(output.mode, CreditRolloutMode::Required); +} + +TEST(ReceiverCreditConfig, RejectsConflictingRolloutSettings) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE(load(R"({ + "receiver_credit": {"mode": "optional", "enabled": true} + })", + true, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"enabled": false}, + "receiver_credit/enabled": true + })", + true, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, AllowsConsistentModeAndLegacyAlias) { + ReceiverCreditRuntimeConfig output; + ASSERT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "enabled": true, + "capacity": {"data_bytes": 1024, "request_slots": 4}, + "grant_batch": {"data_bytes": 512, "request_slots": 2} + } + })", + true, output) + .ok()); + EXPECT_EQ(output.mode, CreditRolloutMode::Required); +} + +TEST(ReceiverCreditConfig, RejectsUnknownKeysAtEverySchemaLevel) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE(load(R"({"receiver_credit": {"surprise": 1}})", false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "capacity": {"data_bytes": 1, "request_slots": 1, "gpu": 1}, + "grant_batch": {"data_bytes": 1, "request_slots": 1} + } + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit": {"control": {"period_us": 1}}})", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit": {"limits": {"entries": 1}}})", + false, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, RejectsWrongJsonTypesWithoutDefaulting) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE( + load(R"({"receiver_credit": []})", false, output).IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit": {"mode": true}})", false, output) + .IsInvalidArgument()); + EXPECT_TRUE( + load(R"({"receiver_credit": {"enabled": "true"}})", false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit/enabled": 1})", false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit": {"default_qos_class": 0.0}})", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit": {"limits": {"max_peers": -1}}})", + false, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, ActiveModeRequiresRuntimeQueueAndBothResources) { + ReceiverCreditRuntimeConfig output; + const std::string valid = R"({ + "receiver_credit": { + "mode": "required", + "capacity": {"data_bytes": 1024, "request_slots": 4}, + "grant_batch": {"data_bytes": 512, "request_slots": 2} + } + })"; + EXPECT_TRUE(load(valid, false, output).IsInvalidArgument()); + EXPECT_TRUE(load(R"({"receiver_credit/enabled": true})", true, output) + .IsInvalidArgument()); + EXPECT_TRUE( + load(R"({"receiver_credit": {"mode": "optional"}})", true, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "capacity": {"data_bytes": 1024, "request_slots": 4} + } + })", + true, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, RejectsZeroMissingOrOversizedGrants) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "capacity": {"data_bytes": 1024, "request_slots": 0}, + "grant_batch": {"data_bytes": 512, "request_slots": 1} + } + })", + true, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "capacity": {"data_bytes": 1024, "request_slots": 4}, + "grant_batch": {"data_bytes": 2048, "request_slots": 1} + } + })", + true, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "mode": "required", + "capacity": {"data_bytes": 1024}, + "grant_batch": {"data_bytes": 512, "request_slots": 1} + } + })", + true, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, EnforcesControlAndPeerBounds) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE(load(R"({ + "receiver_credit": { + "control": {"freshness_ttl_ms": 1, "poll_interval_us": 1000} + } + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"control": {"retry_after_us": 0}} + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"limits": {"max_peers": 65537}} + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"limits": {"max_peers": 0}} + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"control": {"adaptive_dispatch": { + "min_owners": 4, "initial_owners": 2, "max_owners": 8 + }}} + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"control": {"adaptive_dispatch": { + "min_owners": 1, "initial_owners": 4, "max_owners": 2 + }}} + })", + false, output) + .IsInvalidArgument()); + EXPECT_TRUE(load(R"({ + "receiver_credit": {"control": {"adaptive_dispatch": { + "enabled": "true" + }}} + })", + false, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, MvpOnlyAcceptsQosClassZero) { + ReceiverCreditRuntimeConfig output; + EXPECT_TRUE( + load(R"({"receiver_credit": {"default_qos_class": 1}})", false, output) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditConfig, FailureLeavesOutputUnchanged) { + ReceiverCreditRuntimeConfig output; + output.mode = CreditRolloutMode::Optional; + output.capacity.fill(17); + output.max_grant_per_pull.fill(19); + output.max_peers = 23; + output.freshness_ttl_ms = 29; + output.retry_after_us = 31; + output.progress_interval_us = 37; + output.default_qos_class = 41; + output.adaptive_dispatch_enabled = false; + output.adaptive_dispatch_min_owners = 43; + output.adaptive_dispatch_initial_owners = 47; + output.adaptive_dispatch_max_owners = 53; + output.adaptive_dispatch_slow_rtt_us = 59; + output.adaptive_dispatch_healthy_pulls = 61; + const auto before = output; + + EXPECT_TRUE(load(R"({"receiver_credit": {"mode": "broken"}})", true, output) + .IsInvalidArgument()); + EXPECT_EQ(output.mode, before.mode); + EXPECT_EQ(output.capacity, before.capacity); + EXPECT_EQ(output.max_grant_per_pull, before.max_grant_per_pull); + EXPECT_EQ(output.max_peers, before.max_peers); + EXPECT_EQ(output.freshness_ttl_ms, before.freshness_ttl_ms); + EXPECT_EQ(output.retry_after_us, before.retry_after_us); + EXPECT_EQ(output.progress_interval_us, before.progress_interval_us); + EXPECT_EQ(output.default_qos_class, before.default_qos_class); + EXPECT_EQ(output.adaptive_dispatch_enabled, + before.adaptive_dispatch_enabled); + EXPECT_EQ(output.adaptive_dispatch_min_owners, + before.adaptive_dispatch_min_owners); + EXPECT_EQ(output.adaptive_dispatch_initial_owners, + before.adaptive_dispatch_initial_owners); + EXPECT_EQ(output.adaptive_dispatch_max_owners, + before.adaptive_dispatch_max_owners); + EXPECT_EQ(output.adaptive_dispatch_slow_rtt_us, + before.adaptive_dispatch_slow_rtt_us); + EXPECT_EQ(output.adaptive_dispatch_healthy_pulls, + before.adaptive_dispatch_healthy_pulls); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_control_service_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_control_service_test.cpp new file mode 100644 index 0000000000..bff82bfa39 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_control_service_test.cpp @@ -0,0 +1,98 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/control_plane.h" + +#include +#include + +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t index(CreditResource resource) { + return static_cast(resource) - 1; +} + +ReceiverCreditPullRequestV1 request() { + ReceiverCreditPullRequestV1 value; + value.sender_peer = 77; + value.request_sequence = 1; + value.resources = { + {CreditResource::DataBytes, 0, 0, 4096, 8192}, + {CreditResource::RequestSlots, 0, 0, 1, 4}, + }; + return value; +} + +std::shared_ptr allocator() { + ReceiverCreditAllocatorConfig config; + config.capacity[index(CreditResource::DataBytes)] = 1ULL << 20; + config.capacity[index(CreditResource::RequestSlots)] = 64; + config.max_grant_per_pull[index(CreditResource::DataBytes)] = 64ULL << 10; + config.max_grant_per_pull[index(CreditResource::RequestSlots)] = 8; + config.max_entries = 16; + config.ttl_ms = 1000; + config.retry_after_us = 100; + config.receiver_session_id = {11, 22}; + config.epoch = 1; + std::unique_ptr result; + EXPECT_TRUE(ReceiverCreditAllocator::create(config, result).ok()); + return std::shared_ptr(std::move(result)); +} + +TEST(ReceiverCreditControlService, TypedPullReachesReceiverAllocator) { + ControlService service("p2p", "", nullptr); + service.setReceiverCreditAllocator(allocator()); + uint16_t port = 0; + ASSERT_TRUE(service.start(port).ok()); + + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE(ControlClient::pullReceiverCredit( + "127.0.0.1:" + std::to_string(port), request(), response) + .ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Granted); + EXPECT_EQ(response.activation.receiver_session_id, + (ReceiverSessionId{11, 22})); + ASSERT_EQ(response.update.grants.size(), kCreditResourceCount); + EXPECT_EQ(response.update.grants[0].grant_total, 8192); + EXPECT_EQ(response.update.grants[1].grant_total, 4); +} + +TEST(ReceiverCreditControlService, DisabledReceiverRepliesUnsupported) { + ControlService service("p2p", "", nullptr); + uint16_t port = 0; + ASSERT_TRUE(service.start(port).ok()); + + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE(ControlClient::pullReceiverCredit( + "127.0.0.1:" + std::to_string(port), request(), response) + .ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Unsupported); + EXPECT_EQ(response.activation.schema_version, 0); + EXPECT_TRUE(response.update.grants.empty()); +} + +TEST(ReceiverCreditControlService, MalformedPullGetsBoundedRejectedReply) { + ControlService service("p2p", "", nullptr); + uint16_t port = 0; + ASSERT_TRUE(service.start(port).ok()); + + CoroRpcAgent client; + std::string response_wire; + ASSERT_TRUE(client + .call("127.0.0.1:" + std::to_string(port), + PullReceiverCredit, "malformed", response_wire) + .ok()); + EXPECT_EQ(response_wire.size(), + ReceiverCreditPullResponseCodecV1::kWireBytes); + ReceiverCreditPullResponseV1 response; + ASSERT_TRUE( + ReceiverCreditPullResponseCodecV1::decode(response_wire, response) + .ok()); + EXPECT_EQ(response.status, ReceiverCreditPullStatus::Rejected); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_control_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_control_test.cpp new file mode 100644 index 0000000000..5575165607 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_control_test.cpp @@ -0,0 +1,538 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_control.h" + +#include +#include +#include +#include +#include + +#include + +namespace mooncake::tent { +namespace { +CreditKey key() { return {{1, 2}, 3, 4}; } +CreditControlEnvelope envelope(uint64_t sequence, uint64_t total) { + CreditControlEnvelope e; + e.key = key(); + e.update.receiver_session_id = e.key.receiver_session; + e.update.qos_class = e.key.qos_class; + e.update.epoch = 7; + e.update.sequence = sequence; + e.update.grants = {{CreditResource::DataBytes, total}}; + return e; +} + +TEST(ReceiverCreditControl, NegotiatesOptionalAndRequiredModes) { + CreditCapabilityState optional(CreditRolloutMode::Optional); + ASSERT_TRUE(optional.completeNegotiation({}).ok()); + EXPECT_EQ(optional.state(), CreditPeerState::Legacy); + + CreditCapabilityState required(CreditRolloutMode::Required); + EXPECT_TRUE(required.completeNegotiation({2}).IsNotImplemented()); + EXPECT_EQ(required.state(), CreditPeerState::Failed); + + ASSERT_TRUE(required.beginNegotiation().ok()); + ASSERT_TRUE(required.completeNegotiation({2, 1}).ok()); + EXPECT_EQ(required.state(), CreditPeerState::Active); + EXPECT_EQ(required.version(), 1); + ASSERT_TRUE(required.markStale().ok()); + EXPECT_TRUE(required.refresh(2).IsInvalidArgument()); + ASSERT_TRUE(required.refresh(1).ok()); +} + +TEST(ReceiverCreditControl, CapabilityWireDrivesMixedVersionNegotiation) { + std::string wire; + ASSERT_TRUE(CreditCapabilityCodecV1::encode({2, 1}, wire).ok()); + std::vector versions; + ASSERT_TRUE(CreditCapabilityCodecV1::decode(wire, versions).ok()); + CreditCapabilityState supported(CreditRolloutMode::Required); + ASSERT_TRUE(supported.completeNegotiation(versions).ok()); + EXPECT_EQ(supported.state(), CreditPeerState::Active); + EXPECT_EQ(supported.version(), 1); + + ASSERT_TRUE(CreditCapabilityCodecV1::encode({2}, wire).ok()); + ASSERT_TRUE(CreditCapabilityCodecV1::decode(wire, versions).ok()); + CreditCapabilityState optional(CreditRolloutMode::Optional); + ASSERT_TRUE(optional.completeNegotiation(versions).ok()); + EXPECT_EQ(optional.state(), CreditPeerState::Legacy); + CreditCapabilityState required(CreditRolloutMode::Required); + EXPECT_TRUE(required.completeNegotiation(versions).IsNotImplemented()); + EXPECT_EQ(required.state(), CreditPeerState::Failed); +} + +TEST(ReceiverCreditControl, CapabilityWireRejectsMalformedOffersAtomically) { + std::string wire = "sentinel"; + EXPECT_TRUE( + CreditCapabilityCodecV1::encode({1, 1}, wire).IsInvalidArgument()); + EXPECT_EQ(wire, "sentinel"); + EXPECT_TRUE(CreditCapabilityCodecV1::encode({0}, wire).IsInvalidArgument()); + EXPECT_EQ(wire, "sentinel"); + EXPECT_TRUE( + CreditCapabilityCodecV1::encode( + std::vector(CreditCapabilityCodecV1::kMaxVersions + 1, 1), + wire) + .IsInvalidArgument()); + EXPECT_EQ(wire, "sentinel"); + + ASSERT_TRUE(CreditCapabilityCodecV1::encode({1, 2}, wire).ok()); + std::vector output{99}; + for (size_t length = 0; length < wire.size(); ++length) { + EXPECT_TRUE(CreditCapabilityCodecV1::decode( + std::string_view(wire.data(), length), output) + .IsInvalidArgument()); + EXPECT_EQ(output, std::vector({99})); + } + std::string duplicate = wire; + duplicate[10] = 0; + duplicate[11] = 1; + EXPECT_TRUE( + CreditCapabilityCodecV1::decode(duplicate, output).IsInvalidArgument()); + EXPECT_EQ(output, std::vector({99})); +} + +TEST(ReceiverCreditControl, ActivationWireCreatesEpochFencedLedgerContext) { + CreditActivationV1 original; + original.receiver_session_id = {11, 22}; + original.epoch = 7; + original.freshness_ttl_ms = 900; + std::string wire; + ASSERT_TRUE(CreditActivationCodecV1::encode(original, wire).ok()); + EXPECT_EQ(wire.size(), CreditActivationCodecV1::kWireBytes); + CreditActivationV1 decoded; + ASSERT_TRUE(CreditActivationCodecV1::decode(wire, decoded).ok()); + EXPECT_EQ(decoded.receiver_session_id, original.receiver_session_id); + EXPECT_EQ(decoded.epoch, 7); + + CreditKey activation_key{decoded.receiver_session_id, 33, 4}; + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(activation_key, decoded.epoch).ok()); + ASSERT_TRUE(ledger.activate(activation_key, decoded.epoch + 1).ok()); + EXPECT_TRUE( + ledger.deactivate(activation_key, decoded.epoch).IsInvalidEntry()); +} + +TEST(ReceiverCreditControl, ActivationWireRejectsMalformedInputAtomically) { + CreditActivationV1 original; + original.receiver_session_id = {11, 22}; + original.epoch = 7; + std::string wire; + ASSERT_TRUE(CreditActivationCodecV1::encode(original, wire).ok()); + CreditActivationV1 output; + output.epoch = 99; + for (size_t length = 0; length < wire.size(); ++length) { + EXPECT_TRUE(CreditActivationCodecV1::decode( + std::string_view(wire.data(), length), output) + .IsInvalidArgument()); + EXPECT_EQ(output.epoch, 99); + } + std::string reserved = wire; + reserved.back() = 1; + EXPECT_TRUE( + CreditActivationCodecV1::decode(reserved, output).IsInvalidArgument()); + EXPECT_EQ(output.epoch, 99); + + original.receiver_session_id = {}; + std::string unchanged = "sentinel"; + EXPECT_TRUE(CreditActivationCodecV1::encode(original, unchanged) + .IsInvalidArgument()); + EXPECT_EQ(unchanged, "sentinel"); +} + +TEST(ReceiverCreditControl, PeerContextSnapshotsActivationByTargetAndQos) { + CreditPeerContextTable contexts; + CreditActivationV1 activation; + activation.receiver_session_id = {11, 22}; + activation.epoch = 7; + activation.freshness_ttl_ms = 900; + ASSERT_TRUE(contexts.activate(100, 200, 3, activation).ok()); + CreditPeerContextSnapshot snapshot; + ASSERT_TRUE(contexts.lookup(100, 3, snapshot).ok()); + EXPECT_EQ(snapshot.key.receiver_session, activation.receiver_session_id); + EXPECT_EQ(snapshot.key.sender_peer, 200); + EXPECT_EQ(snapshot.key.qos_class, 3); + EXPECT_EQ(snapshot.epoch, 7); + EXPECT_EQ(snapshot.freshness_ttl_ms, 900); + EXPECT_TRUE(contexts.lookup(100, 4, snapshot).IsInvalidEntry()); +} + +TEST(ReceiverCreditControl, PeerContextRestartFencesOldCleanup) { + CreditPeerContextTable contexts; + CreditActivationV1 first; + first.receiver_session_id = {11, 22}; + first.epoch = 7; + ASSERT_TRUE(contexts.activate(100, 200, 3, first).ok()); + EXPECT_TRUE( + contexts.activate(100, 200, 3, CreditActivationV1{1, 1, {11, 22}, 6, 0}) + .IsInvalidEntry()); + + CreditActivationV1 restarted; + restarted.receiver_session_id = {33, 44}; + restarted.epoch = 1; + ASSERT_TRUE(contexts.activate(100, 200, 3, restarted).ok()); + EXPECT_TRUE( + contexts.deactivate(100, 3, first.receiver_session_id, first.epoch) + .IsInvalidEntry()); + CreditPeerContextSnapshot snapshot; + ASSERT_TRUE(contexts.lookup(100, 3, snapshot).ok()); + EXPECT_EQ(snapshot.key.receiver_session, restarted.receiver_session_id); + EXPECT_EQ(snapshot.epoch, 1); +} + +TEST(ReceiverCreditControl, PeerContextRejectsSenderIdentityChange) { + CreditPeerContextTable contexts; + CreditActivationV1 activation; + activation.receiver_session_id = {11, 22}; + activation.epoch = 7; + ASSERT_TRUE(contexts.activate(100, 200, 3, activation).ok()); + + activation.epoch = 8; + EXPECT_TRUE(contexts.activate(100, 201, 3, activation).IsInvalidEntry()); + activation.receiver_session_id = {33, 44}; + activation.epoch = 1; + EXPECT_TRUE(contexts.activate(100, 201, 3, activation).IsInvalidEntry()); + + CreditPeerContextSnapshot snapshot; + ASSERT_TRUE(contexts.lookup(100, 3, snapshot).ok()); + EXPECT_EQ(snapshot.key.receiver_session, (ReceiverSessionId{11, 22})); + EXPECT_EQ(snapshot.key.sender_peer, 200); + EXPECT_EQ(snapshot.epoch, 7); +} + +TEST(ReceiverCreditControl, PeerContextCapacityRecoversAfterExactCleanup) { + CreditPeerContextTable contexts(1); + CreditActivationV1 activation; + activation.receiver_session_id = {11, 22}; + activation.epoch = 7; + ASSERT_TRUE(contexts.activate(100, 200, 3, activation).ok()); + EXPECT_TRUE(contexts.activate(101, 200, 3, activation).IsTooManyRequests()); + ASSERT_TRUE(contexts + .deactivate(100, 3, activation.receiver_session_id, + activation.epoch) + .ok()); + ASSERT_TRUE(contexts.activate(101, 200, 3, activation).ok()); + EXPECT_EQ(contexts.size(), 1); +} + +TEST(ReceiverCreditControl, PeerContextFreshnessExpiresAndRefreshesInPlace) { + CreditPeerContextTable contexts; + CreditActivationV1 activation; + activation.receiver_session_id = {11, 22}; + activation.epoch = 7; + activation.freshness_ttl_ms = 2; + ASSERT_TRUE(contexts.activate(9, 3, 4, activation).ok()); + + CreditPeerContextSnapshot snapshot; + ASSERT_TRUE(contexts.lookupFresh(9, 4, snapshot).ok()); + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + EXPECT_TRUE(contexts.lookupFresh(9, 4, snapshot).IsInvalidEntry()); + + // A heartbeat for the same generation refreshes only freshness metadata; + // it does not create a new key or epoch. + ASSERT_TRUE(contexts.activate(9, 3, 4, activation).ok()); + ASSERT_TRUE(contexts.lookupFresh(9, 4, snapshot).ok()); + EXPECT_EQ(snapshot.epoch, 7); + EXPECT_EQ(contexts.size(), 1); +} + +TEST(ReceiverCreditControl, ConcurrentPeerRestartLookupNeverTearsSnapshot) { + CreditPeerContextTable contexts; + CreditActivationV1 initial; + initial.receiver_session_id = {1, 2}; + initial.epoch = 1; + ASSERT_TRUE(contexts.activate(100, 200, 3, initial).ok()); + + std::atomic writer_done{false}; + std::atomic invalid_snapshots{0}; + std::thread writer([&] { + for (uint64_t epoch = 2; epoch <= 10000; ++epoch) { + CreditActivationV1 activation; + activation.receiver_session_id = {epoch, epoch + 1}; + activation.epoch = epoch; + if (!contexts.activate(100, 200, 3, activation).ok()) + ++invalid_snapshots; + } + writer_done = true; + }); + std::vector readers; + for (int reader = 0; reader < 16; ++reader) { + readers.emplace_back([&] { + do { + CreditPeerContextSnapshot snapshot; + if (!contexts.lookup(100, 3, snapshot).ok() || + snapshot.key.receiver_session.high != snapshot.epoch || + snapshot.key.receiver_session.low != snapshot.epoch + 1 || + snapshot.key.sender_peer != 200 || + snapshot.key.qos_class != 3) + ++invalid_snapshots; + } while (!writer_done.load()); + }); + } + writer.join(); + for (auto& reader : readers) reader.join(); + EXPECT_EQ(invalid_snapshots, 0); + CreditPeerContextSnapshot final; + ASSERT_TRUE(contexts.lookup(100, 3, final).ok()); + EXPECT_EQ(final.epoch, 10000); +} + +TEST(ReceiverCreditControl, InboxIsBoundedAndDrainIsLimited) { + BoundedCreditUpdateInbox inbox(2); + ASSERT_TRUE(inbox.tryPublish(envelope(1, 10)).ok()); + ASSERT_TRUE(inbox.tryPublish(envelope(2, 20)).ok()); + EXPECT_TRUE(inbox.tryPublish(envelope(3, 30)).IsTooManyRequests()); + std::vector drained; + EXPECT_EQ(inbox.drain(drained, 1), 1); + EXPECT_EQ(drained[0].update.sequence, 1); + EXPECT_EQ(inbox.size(), 1); +} + +TEST(ReceiverCreditControl, ConcurrentPublishNeverExceedsBound) { + BoundedCreditUpdateInbox inbox(64); + std::atomic accepted{0}; + std::vector threads; + for (int i = 0; i < 8; ++i) + threads.emplace_back([&, i] { + for (int j = 0; j < 32; ++j) + if (inbox.tryPublish(envelope(i * 32 + j + 1, 1)).ok()) + ++accepted; + }); + for (auto& thread : threads) thread.join(); + EXPECT_EQ(accepted, 64); + EXPECT_EQ(inbox.size(), 64); +} + +TEST(ReceiverCreditControl, LossDuplicateAndReorderCannotMintCredit) { + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(key(), 7).ok()); + BoundedCreditUpdateInbox inbox(8); + // Sequence 2 is lost. Sequence 4 arrives before 3, then 4 is duplicated. + ASSERT_TRUE(inbox.tryPublish(envelope(1, 10)).ok()); + ASSERT_TRUE(inbox.tryPublish(envelope(4, 40)).ok()); + ASSERT_TRUE(inbox.tryPublish(envelope(3, 30)).ok()); + ASSERT_TRUE(inbox.tryPublish(envelope(4, 400)).ok()); + std::vector drained; + inbox.drain(drained, 8); + std::vector dispositions; + for (auto& e : drained) { + CreditUpdateDisposition d; + ASSERT_TRUE(ledger.applyUpdate(e.key, e.update, d).ok()); + dispositions.push_back(d); + } + EXPECT_EQ(dispositions[1], CreditUpdateDisposition::SequenceGap); + EXPECT_EQ(dispositions[2], CreditUpdateDisposition::DuplicateOrOld); + EXPECT_EQ(dispositions[3], CreditUpdateDisposition::DuplicateOrOld); + uint64_t available; + ASSERT_TRUE( + ledger.available(key(), CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, 40); +} + +TEST(ReceiverCreditControl, DisconnectMakesActivePeerStale) { + CreditCapabilityState peer(CreditRolloutMode::Required); + ASSERT_TRUE(peer.completeNegotiation({1}).ok()); + ASSERT_TRUE(peer.markStale().ok()); + EXPECT_EQ(peer.state(), CreditPeerState::Stale); + EXPECT_TRUE(peer.completeNegotiation({1}).IsInvalidEntry()); +} + +TEST(ReceiverCreditControl, ZeroCapacityInboxAlwaysFailsFast) { + BoundedCreditUpdateInbox inbox(0); + EXPECT_TRUE(inbox.tryPublish(envelope(1, 1)).IsTooManyRequests()); + std::vector drained; + EXPECT_EQ(inbox.drain(drained, 100), 0); +} + +TEST(ReceiverCreditControl, SustainedConcurrentPublishAndDrainLosesNothing) { + constexpr int kProducers = 16; + constexpr int kPerProducer = 2000; + constexpr int kTotal = kProducers * kPerProducer; + BoundedCreditUpdateInbox inbox(127); + std::atomic producers_done{0}; + std::vector producers; + for (int producer = 0; producer < kProducers; ++producer) { + producers.emplace_back([&, producer] { + for (int i = 0; i < kPerProducer; ++i) { + uint64_t sequence = producer * kPerProducer + i + 1; + while (!inbox.tryPublish(envelope(sequence, sequence)).ok()) + std::this_thread::yield(); + } + ++producers_done; + }); + } + + std::vector received; + received.reserve(kTotal); + while (producers_done != kProducers || inbox.size() != 0) { + inbox.drain(received, 31); + std::this_thread::yield(); + } + for (auto& producer : producers) producer.join(); + EXPECT_EQ(received.size(), static_cast(kTotal)); + std::unordered_set unique; + for (const auto& item : received) unique.insert(item.update.sequence); + EXPECT_EQ(unique.size(), static_cast(kTotal)); +} + +TEST(ReceiverCreditControl, DeterministicReorderDuplicateFuzzNeverMints) { + constexpr uint64_t kUpdates = 10000; + std::vector traffic; + traffic.reserve(kUpdates * 2); + for (uint64_t sequence = 1; sequence <= kUpdates; ++sequence) { + traffic.push_back(envelope(sequence, sequence * 8)); + if (sequence % 3 == 0) + traffic.push_back(envelope(sequence, sequence * 8)); + } + std::mt19937_64 random(0x2849); + std::shuffle(traffic.begin(), traffic.end(), random); + + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(key(), 7).ok()); + for (const auto& item : traffic) { + CreditUpdateDisposition disposition; + ASSERT_TRUE( + ledger.applyUpdate(item.key, item.update, disposition).ok()); + } + uint64_t available; + ASSERT_TRUE( + ledger.available(key(), CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, kUpdates * 8); +} + +TEST(ReceiverCreditControl, ReconnectRequiresNegotiationAndNewLedgerEpoch) { + CreditCapabilityState peer(CreditRolloutMode::Required); + ASSERT_TRUE(peer.completeNegotiation({1}).ok()); + ASSERT_TRUE(peer.markStale().ok()); + ASSERT_TRUE(peer.beginNegotiation().ok()); + ASSERT_TRUE(peer.completeNegotiation({1}).ok()); + + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(key(), 7).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE( + ledger.applyUpdate(key(), envelope(1, 100).update, disposition).ok()); + ASSERT_TRUE(ledger.activate(key(), 8).ok()); + EXPECT_TRUE(ledger.applyUpdate(key(), envelope(2, 200).update, disposition) + .IsInvalidEntry()); + auto fresh = envelope(1, 10); + fresh.update.epoch = 8; + ASSERT_TRUE(ledger.applyUpdate(key(), fresh.update, disposition).ok()); +} + +TEST(ReceiverCreditControl, WireCodecRoundTripsAndUsesNetworkByteOrder) { + auto original = + envelope(0x0102030405060708ULL, 0x1122334455667788ULL).update; + original.flags = 3; + original.freshness_ttl_ms = 900; + original.grants.push_back({CreditResource::RequestSlots, 17}); + std::string wire; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(original, wire).ok()); + EXPECT_EQ(wire.size(), ReceiverCreditCodecV1::kHeaderBytes + + 2 * ReceiverCreditCodecV1::kGrantBytes); + EXPECT_EQ(static_cast(wire[0]), 0x54); + EXPECT_EQ(static_cast(wire[1]), 0x43); + ReceiverCreditUpdateV1 decoded; + ASSERT_TRUE(ReceiverCreditCodecV1::decode(wire, decoded).ok()); + EXPECT_EQ(decoded.sequence, original.sequence); + EXPECT_EQ(decoded.grants.size(), 2); + EXPECT_EQ(decoded.grants[0].grant_total, 0x1122334455667788ULL); + EXPECT_EQ(decoded.grants[1].resource, CreditResource::RequestSlots); +} + +TEST(ReceiverCreditControl, EveryTruncationFailsWithoutMutatingOutput) { + auto original = envelope(7, 100).update; + std::string wire; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(original, wire).ok()); + for (size_t length = 0; length < wire.size(); ++length) { + ReceiverCreditUpdateV1 output; + output.sequence = 999; + EXPECT_TRUE(ReceiverCreditCodecV1::decode( + std::string_view(wire.data(), length), output) + .IsInvalidArgument()); + EXPECT_EQ(output.sequence, 999); + } +} + +TEST(ReceiverCreditControl, OversizedDuplicateAndUnknownWireFieldsFail) { + auto original = envelope(7, 100).update; + std::string wire; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(original, wire).ok()); + ReceiverCreditUpdateV1 output; + std::string oversized(ReceiverCreditCodecV1::kMaxWireBytes + 1, 'x'); + EXPECT_TRUE( + ReceiverCreditCodecV1::decode(oversized, output).IsInvalidArgument()); + + original.grants.push_back({CreditResource::DataBytes, 200}); + std::string unchanged = "sentinel"; + EXPECT_TRUE( + ReceiverCreditCodecV1::encode(original, unchanged).IsInvalidArgument()); + EXPECT_EQ(unchanged, "sentinel"); +} + +TEST(ReceiverCreditControl, DeterministicMalformedWireFuzzIsMemorySafe) { + std::mt19937_64 random(0x2860); + ReceiverCreditUpdateV1 output; + for (int iteration = 0; iteration < 100000; ++iteration) { + size_t length = random() % (ReceiverCreditCodecV1::kMaxWireBytes + 33); + std::string wire(length, '\0'); + for (char& byte : wire) byte = static_cast(random()); + auto status = ReceiverCreditCodecV1::decode(wire, output); + if (status.ok()) { + EXPECT_EQ(output.schema_version, 1); + EXPECT_LE(output.grants.size(), kCreditResourceCount); + EXPECT_NE(output.epoch, 0); + EXPECT_NE(output.sequence, 0); + } else { + EXPECT_TRUE(status.IsInvalidArgument()); + } + } +} + +TEST(ReceiverCreditControl, IngressValidatesBeforePublishing) { + BoundedCreditUpdateInbox inbox(2); + ReceiverCreditIngress ingress(inbox, key(), 7); + std::string wire; + ASSERT_TRUE( + ReceiverCreditCodecV1::encode(envelope(1, 10).update, wire).ok()); + ASSERT_TRUE(ingress.tryAccept(wire).ok()); + + auto wrong_session = envelope(2, 20).update; + ++wrong_session.receiver_session_id.low; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(wrong_session, wire).ok()); + EXPECT_TRUE(ingress.tryAccept(wire).IsInvalidEntry()); + + auto wrong_qos = envelope(2, 20).update; + ++wrong_qos.qos_class; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(wrong_qos, wire).ok()); + EXPECT_TRUE(ingress.tryAccept(wire).IsInvalidEntry()); + + auto wrong_epoch = envelope(2, 20).update; + ++wrong_epoch.epoch; + ASSERT_TRUE(ReceiverCreditCodecV1::encode(wrong_epoch, wire).ok()); + EXPECT_TRUE(ingress.tryAccept(wire).IsInvalidEntry()); + EXPECT_EQ(inbox.size(), 1); +} + +TEST(ReceiverCreditControl, IngressQueueFullIsExplicitAndRetryable) { + BoundedCreditUpdateInbox inbox(1); + ReceiverCreditIngress ingress(inbox, key(), 7); + std::string first_wire, second_wire; + ASSERT_TRUE( + ReceiverCreditCodecV1::encode(envelope(1, 10).update, first_wire).ok()); + ASSERT_TRUE( + ReceiverCreditCodecV1::encode(envelope(2, 20).update, second_wire) + .ok()); + ASSERT_TRUE(ingress.tryAccept(first_wire).ok()); + EXPECT_TRUE(ingress.tryAccept(second_wire).IsTooManyRequests()); + + std::vector drained; + ASSERT_EQ(inbox.drain(drained, 1), 1); + ASSERT_TRUE(ingress.tryAccept(second_wire).ok()); + drained.clear(); + ASSERT_EQ(inbox.drain(drained, 1), 1); + EXPECT_EQ(drained.front().update.sequence, 2); +} +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_controller_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_controller_test.cpp new file mode 100644 index 0000000000..2b43a2cd3e --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_controller_test.cpp @@ -0,0 +1,204 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/control_plane.h" +#include "tent/runtime/receiver_credit_controller.h" + +#include +#include +#include +#include +#include + +#include + +namespace mooncake::tent { +namespace { + +constexpr size_t index(CreditResource resource) { + return static_cast(resource) - 1; +} + +ReceiverCreditRuntimeConfig controllerConfig(CreditRolloutMode mode) { + ReceiverCreditRuntimeConfig config; + config.mode = mode; + config.max_peers = 16; + config.max_grant_per_pull[index(CreditResource::DataBytes)] = 8192; + config.max_grant_per_pull[index(CreditResource::RequestSlots)] = 4; + return config; +} + +std::shared_ptr allocator() { + ReceiverCreditAllocatorConfig config; + config.capacity[index(CreditResource::DataBytes)] = 1ULL << 20; + config.capacity[index(CreditResource::RequestSlots)] = 64; + config.max_grant_per_pull[index(CreditResource::DataBytes)] = 8192; + config.max_grant_per_pull[index(CreditResource::RequestSlots)] = 4; + config.max_entries = 16; + config.ttl_ms = 1000; + config.retry_after_us = 100; + config.receiver_session_id = {11, 22}; + config.epoch = 1; + std::unique_ptr result; + EXPECT_TRUE(ReceiverCreditAllocator::create(config, result).ok()); + return std::shared_ptr(std::move(result)); +} + +CreditCharge charge() { + return { + {{CreditResource::DataBytes, 4096}, {CreditResource::RequestSlots, 1}}}; +} + +template +bool waitUntil(Predicate predicate) { + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + return predicate(); +} + +TEST(ReceiverCreditController, PullActivatesLedgerAndReportsCompletion) { + ControlService service("p2p", "", nullptr); + service.setReceiverCreditAllocator(allocator()); + uint16_t port = 0; + ASSERT_TRUE(service.start(port).ok()); + const std::string address = "127.0.0.1:" + std::to_string(port); + + auto contexts = std::make_shared(16); + auto ledger = std::make_shared(16); + std::shared_ptr controller; + ASSERT_TRUE(ReceiverCreditPullController::create( + controllerConfig(CreditRolloutMode::Required), 77, contexts, + ledger, controller) + .ok()); + ASSERT_TRUE(controller->request(9, address, 0, charge()).ok()); + + CreditPeerContextSnapshot peer; + ASSERT_TRUE( + waitUntil([&] { return contexts->lookupFresh(9, 0, peer).ok(); })); + uint64_t available = 0; + ASSERT_TRUE( + ledger->available(peer.key, CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, 8192); + + ReceiverCreditDispatchGate gate(*contexts, *ledger); + CreditDispatchSnapshot snapshot; + ASSERT_TRUE(gate.snapshot(9, 0, charge(), snapshot).ok()); + CreditDispatchReservation reservation; + ASSERT_TRUE(gate.tryReserve(snapshot, reservation).ok()); + ASSERT_TRUE(gate.commit(reservation).ok()); + ASSERT_TRUE(gate.release(reservation).ok()); + ASSERT_TRUE(controller->request(9, address, 0, charge()).ok()); + + ASSERT_TRUE(waitUntil([&] { + CreditLedgerSnapshot usage; + return ledger->snapshot(peer.key, peer.epoch, usage).ok() && + usage.grants[index(CreditResource::DataBytes)] == 12288; + })); + controller->stop(); +} + +TEST(ReceiverCreditController, ExplicitUnsupportedObeysRolloutMode) { + ControlService service("p2p", "", nullptr); + uint16_t port = 0; + ASSERT_TRUE(service.start(port).ok()); + const std::string address = "127.0.0.1:" + std::to_string(port); + + for (auto [mode, expected] : + {std::pair{CreditRolloutMode::Optional, CreditPeerState::Legacy}, + std::pair{CreditRolloutMode::Required, CreditPeerState::Failed}}) { + auto contexts = std::make_shared(16); + auto ledger = std::make_shared(16); + std::shared_ptr controller; + ASSERT_TRUE(ReceiverCreditPullController::create(controllerConfig(mode), + 77, contexts, ledger, + controller) + .ok()); + ASSERT_TRUE(controller->request(9, address, 0, charge()).ok()); + ASSERT_TRUE( + waitUntil([&] { return controller->peerState(9, 0) == expected; })); + controller->stop(); + } +} + +TEST(AdaptiveCreditDispatchLimiter, + SlowPullReducesAndNeverReprobesLearnedUnsafeLevel) { + auto config = controllerConfig(CreditRolloutMode::Required); + config.adaptive_dispatch_min_owners = 1; + config.adaptive_dispatch_initial_owners = 4; + config.adaptive_dispatch_max_owners = 8; + config.adaptive_dispatch_slow_rtt_us = 1000; + config.adaptive_dispatch_healthy_pulls = 2; + AdaptiveCreditDispatchLimiter limiter(config); + + limiter.observe(std::chrono::microseconds(100), true, 4); + limiter.observe(std::chrono::microseconds(100), true, 4); + ASSERT_EQ(limiter.ownerLimit(), 5); + + limiter.observe(std::chrono::milliseconds(200), true, 5); + EXPECT_EQ(limiter.ownerLimit(), 2); + auto reduced = limiter.snapshot(); + EXPECT_EQ(reduced.learned_ceiling, 8); + EXPECT_EQ(reduced.suspect_level, 5); + EXPECT_EQ(reduced.slow_or_failed_pulls, 1); + EXPECT_EQ(reduced.reductions, 1); + + for (size_t i = 0; i < 6; ++i) + limiter.observe(std::chrono::microseconds(100), true, + limiter.ownerLimit()); + ASSERT_EQ(limiter.ownerLimit(), 5); + + limiter.observe(std::chrono::milliseconds(200), true, 5); + EXPECT_EQ(limiter.ownerLimit(), 2); + reduced = limiter.snapshot(); + EXPECT_EQ(reduced.learned_ceiling, 4); + EXPECT_EQ(reduced.suspect_level, 0); + + for (size_t i = 0; i < 20; ++i) + limiter.observe(std::chrono::microseconds(100), true, + limiter.ownerLimit()); + auto recovered = limiter.snapshot(); + EXPECT_EQ(recovered.current_owners, 4); + EXPECT_EQ(recovered.learned_ceiling, 4); + EXPECT_EQ(recovered.increases, 6); +} + +TEST(AdaptiveCreditDispatchLimiter, RpcFailureReducesAtAnyLatency) { + auto config = controllerConfig(CreditRolloutMode::Required); + config.adaptive_dispatch_min_owners = 1; + config.adaptive_dispatch_initial_owners = 2; + config.adaptive_dispatch_max_owners = 4; + AdaptiveCreditDispatchLimiter limiter(config); + + limiter.observe(std::chrono::microseconds(1), false, 2); + EXPECT_EQ(limiter.ownerLimit(), 1); + EXPECT_EQ(limiter.snapshot().learned_ceiling, 4); + EXPECT_EQ(limiter.snapshot().suspect_level, 2); +} + +TEST(AdaptiveCreditDispatchLimiter, DisabledLeavesStaticWindowUnbounded) { + auto config = controllerConfig(CreditRolloutMode::Required); + config.adaptive_dispatch_enabled = false; + AdaptiveCreditDispatchLimiter limiter(config); + + limiter.observe(std::chrono::seconds(1), false, 2); + EXPECT_EQ(limiter.ownerLimit(), std::numeric_limits::max()); + EXPECT_EQ(limiter.snapshot().slow_or_failed_pulls, 0); +} + +TEST(AdaptiveCreditDispatchLimiter, HonorsSmallerRuntimeQueueCeiling) { + auto config = controllerConfig(CreditRolloutMode::Required); + config.adaptive_dispatch_min_owners = 1; + config.adaptive_dispatch_initial_owners = 4; + config.adaptive_dispatch_max_owners = 8; + AdaptiveCreditDispatchLimiter limiter(config, 2); + + EXPECT_EQ(limiter.ownerLimit(), 2); + EXPECT_EQ(limiter.snapshot().learned_ceiling, 2); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_dispatch_test.cpp new file mode 100644 index 0000000000..b9415c9255 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_dispatch_test.cpp @@ -0,0 +1,185 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_dispatch.h" + +#include + +namespace mooncake::tent { +namespace { + +constexpr uint64_t kTarget = 100; +constexpr uint64_t kSender = 200; +constexpr uint32_t kQos = 3; + +CreditActivationV1 activation(ReceiverSessionId session, uint64_t epoch) { + CreditActivationV1 value; + value.receiver_session_id = session; + value.epoch = epoch; + return value; +} + +CreditCharge charge(uint64_t bytes, uint64_t request_slots) { + return {{{CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, request_slots}}}; +} + +void grant(SenderCreditLedger& ledger, const CreditKey& key, uint64_t epoch, + uint64_t bytes, uint64_t request_slots) { + ReceiverCreditUpdateV1 update; + update.receiver_session_id = key.receiver_session; + update.qos_class = key.qos_class; + update.epoch = epoch; + update.sequence = 1; + update.grants = {{CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, request_slots}}; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger.applyUpdate(key, update, disposition).ok()); +} + +TEST(ReceiverCreditDispatch, SubmitFailureRollsBackExactMultiResourceCharge) { + CreditPeerContextTable contexts; + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({11, 22}, 7)) + .ok()); + CreditPeerContextSnapshot peer; + ASSERT_TRUE(contexts.lookup(kTarget, kQos, peer).ok()); + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(peer.key, peer.epoch).ok()); + grant(ledger, peer.key, peer.epoch, 100, 4); + ReceiverCreditDispatchGate gate(contexts, ledger); + CreditDispatchSnapshot snapshot; + ASSERT_TRUE(gate.snapshot(kTarget, kQos, charge(60, 2), snapshot).ok()); + CreditDispatchReservation reservation; + ASSERT_TRUE(gate.tryReserve(snapshot, reservation).ok()); + ASSERT_TRUE(gate.rollback(reservation).ok()); + EXPECT_EQ(reservation.state, CreditReservationState::RolledBack); + uint64_t bytes = 0, slots = 0; + ASSERT_TRUE( + ledger.available(peer.key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger.available(peer.key, CreditResource::RequestSlots, slots).ok()); + EXPECT_EQ(bytes, 100); + EXPECT_EQ(slots, 4); +} + +TEST(ReceiverCreditDispatch, SubmitSuccessReleasesWithoutMintingCredit) { + CreditPeerContextTable contexts; + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({11, 22}, 7)) + .ok()); + CreditPeerContextSnapshot peer; + ASSERT_TRUE(contexts.lookup(kTarget, kQos, peer).ok()); + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(peer.key, peer.epoch).ok()); + grant(ledger, peer.key, peer.epoch, 100, 4); + ReceiverCreditDispatchGate gate(contexts, ledger); + CreditDispatchSnapshot snapshot; + ASSERT_TRUE(gate.snapshot(kTarget, kQos, charge(60, 2), snapshot).ok()); + CreditDispatchReservation reservation; + ASSERT_TRUE(gate.tryReserve(snapshot, reservation).ok()); + ASSERT_TRUE(gate.commit(reservation).ok()); + EXPECT_TRUE(gate.rollback(reservation).IsInvalidEntry()); + ASSERT_TRUE(gate.release(reservation).ok()); + EXPECT_EQ(reservation.state, CreditReservationState::Released); + EXPECT_TRUE(gate.release(reservation).IsInvalidEntry()); + uint64_t bytes = 0, slots = 0; + ASSERT_TRUE( + ledger.available(peer.key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger.available(peer.key, CreditResource::RequestSlots, slots).ok()); + EXPECT_EQ(bytes, 40); + EXPECT_EQ(slots, 2); + CreditLedgerSnapshot ledger_snapshot; + ASSERT_TRUE(ledger.snapshot(peer.key, peer.epoch, ledger_snapshot).ok()); + EXPECT_EQ(ledger_snapshot.consumed[0], 60); + EXPECT_EQ(ledger_snapshot.consumed[1], 2); + EXPECT_EQ(ledger_snapshot.completed[0], 60); + EXPECT_EQ(ledger_snapshot.completed[1], 2); +} + +TEST(ReceiverCreditDispatch, OldEpochRollbackAndReleaseLeaveNewEpochUntouched) { + CreditPeerContextTable contexts; + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({11, 22}, 7)) + .ok()); + CreditPeerContextSnapshot peer; + ASSERT_TRUE(contexts.lookup(kTarget, kQos, peer).ok()); + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(peer.key, peer.epoch).ok()); + grant(ledger, peer.key, peer.epoch, 100, 4); + ReceiverCreditDispatchGate gate(contexts, ledger); + CreditDispatchSnapshot snapshot; + ASSERT_TRUE(gate.snapshot(kTarget, kQos, charge(20, 1), snapshot).ok()); + CreditDispatchReservation pending; + CreditDispatchReservation committed; + ASSERT_TRUE(gate.tryReserve(snapshot, pending).ok()); + ASSERT_TRUE(gate.tryReserve(snapshot, committed).ok()); + ASSERT_TRUE(gate.commit(committed).ok()); + + ASSERT_TRUE(ledger.activate(peer.key, 8).ok()); + grant(ledger, peer.key, 8, 50, 2); + ASSERT_TRUE(ledger.tryReserve(peer.key, charge(10, 1)).ok()); + + EXPECT_TRUE(gate.rollback(pending).IsInvalidEntry()); + EXPECT_TRUE(gate.release(committed).IsInvalidEntry()); + EXPECT_EQ(pending.state, CreditReservationState::Reserved); + EXPECT_EQ(committed.state, CreditReservationState::Committed); + CreditLedgerSnapshot current; + ASSERT_TRUE(ledger.snapshot(peer.key, 8, current).ok()); + EXPECT_EQ(current.consumed[0], 10); + EXPECT_EQ(current.consumed[1], 1); + EXPECT_EQ(current.completed[0], 0); + EXPECT_EQ(current.completed[1], 0); +} + +TEST(ReceiverCreditDispatch, ReceiverRestartRejectsQueuedOldSnapshot) { + CreditPeerContextTable contexts; + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({11, 22}, 7)) + .ok()); + CreditPeerContextSnapshot old_peer; + ASSERT_TRUE(contexts.lookup(kTarget, kQos, old_peer).ok()); + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(old_peer.key, old_peer.epoch).ok()); + grant(ledger, old_peer.key, old_peer.epoch, 100, 4); + ReceiverCreditDispatchGate gate(contexts, ledger); + CreditDispatchSnapshot old_snapshot; + ASSERT_TRUE(gate.snapshot(kTarget, kQos, charge(60, 2), old_snapshot).ok()); + + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({33, 44}, 1)) + .ok()); + CreditDispatchReservation reservation; + EXPECT_TRUE(gate.tryReserve(old_snapshot, reservation).IsInvalidEntry()); + uint64_t available = 0; + ASSERT_TRUE( + ledger.available(old_peer.key, CreditResource::DataBytes, available) + .ok()); + EXPECT_EQ(available, 100); +} + +TEST(ReceiverCreditDispatch, ReserveFailureLeavesTokenAndLedgerUntouched) { + CreditPeerContextTable contexts; + ASSERT_TRUE( + contexts.activate(kTarget, kSender, kQos, activation({11, 22}, 7)) + .ok()); + CreditPeerContextSnapshot peer; + ASSERT_TRUE(contexts.lookup(kTarget, kQos, peer).ok()); + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(peer.key, peer.epoch).ok()); + grant(ledger, peer.key, peer.epoch, 50, 1); + ReceiverCreditDispatchGate gate(contexts, ledger); + CreditDispatchSnapshot snapshot; + ASSERT_TRUE(gate.snapshot(kTarget, kQos, charge(60, 2), snapshot).ok()); + CreditDispatchReservation reservation; + EXPECT_TRUE(gate.tryReserve(snapshot, reservation).IsTooManyRequests()); + EXPECT_EQ(reservation.state, CreditReservationState::Empty); + uint64_t available = 0; + ASSERT_TRUE( + ledger.available(peer.key, CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, 50); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_protocol_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_protocol_test.cpp new file mode 100644 index 0000000000..0c7d480695 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_protocol_test.cpp @@ -0,0 +1,233 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit_protocol.h" + +#include +#include +#include + +#include + +namespace mooncake::tent { +namespace { + +ReceiverCreditPullRequestV1 request() { + ReceiverCreditPullRequestV1 request; + request.sender_peer = 0x0102030405060708ULL; + request.qos_class = 0; + request.expected_receiver_session_id = {11, 22}; + request.expected_epoch = 7; + request.request_sequence = 9; + request.last_update_sequence = 8; + request.resources = { + {CreditResource::DataBytes, 100, 90, 10, 200}, + {CreditResource::RequestSlots, 4, 3, 1, 8}, + }; + return request; +} + +ReceiverCreditPullResponseV1 response() { + ReceiverCreditPullResponseV1 response; + response.status = ReceiverCreditPullStatus::Granted; + response.activation.receiver_session_id = {11, 22}; + response.activation.epoch = 7; + response.activation.freshness_ttl_ms = 1000; + response.update.qos_class = 0; + response.update.receiver_session_id = {11, 22}; + response.update.epoch = 7; + response.update.sequence = 8; + response.update.freshness_ttl_ms = 1000; + response.update.grants = { + {CreditResource::DataBytes, 1000}, + {CreditResource::RequestSlots, 16}, + {CreditResource::StagingSlots, 0}, + {CreditResource::ConsumerSlots, 0}, + }; + return response; +} + +TEST(ReceiverCreditProtocol, RequestRoundTripsInNetworkByteOrder) { + const auto original = request(); + std::string wire; + ASSERT_TRUE(ReceiverCreditPullRequestCodecV1::encode(original, wire).ok()); + ASSERT_EQ(wire.size(), + ReceiverCreditPullRequestCodecV1::kHeaderBytes + + 2 * ReceiverCreditPullRequestCodecV1::kResourceBytes); + EXPECT_EQ(static_cast(wire[0]), 0x54); + EXPECT_EQ(static_cast(wire[1]), 0x43); + EXPECT_EQ(static_cast(wire[2]), 0x52); + EXPECT_EQ(static_cast(wire[3]), 0x51); + for (size_t i = 0; i < sizeof(uint64_t); ++i) + EXPECT_EQ(static_cast(wire[8 + i]), i + 1); + + ReceiverCreditPullRequestV1 decoded; + ASSERT_TRUE(ReceiverCreditPullRequestCodecV1::decode(wire, decoded).ok()); + EXPECT_TRUE(decoded == original); +} + +TEST(ReceiverCreditProtocol, RequestRejectsMalformedInputAtomically) { + std::string wire; + ASSERT_TRUE(ReceiverCreditPullRequestCodecV1::encode(request(), wire).ok()); + ReceiverCreditPullRequestV1 output; + output.sender_peer = 99; + for (size_t length = 0; length < wire.size(); ++length) { + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::decode( + std::string_view(wire.data(), length), output) + .IsInvalidArgument()); + EXPECT_EQ(output.sender_peer, 99); + } + + std::string reserved = wire; + reserved[63] = 1; + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::decode(reserved, output) + .IsInvalidArgument()); + EXPECT_EQ(output.sender_peer, 99); + + std::string resource_reserved = wire; + resource_reserved[67] = 1; + EXPECT_TRUE( + ReceiverCreditPullRequestCodecV1::decode(resource_reserved, output) + .IsInvalidArgument()); + EXPECT_EQ(output.sender_peer, 99); + + auto duplicate = request(); + duplicate.resources[1].resource = CreditResource::DataBytes; + std::string unchanged = "sentinel"; + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::encode(duplicate, unchanged) + .IsInvalidArgument()); + EXPECT_EQ(unchanged, "sentinel"); +} + +TEST(ReceiverCreditProtocol, RequestRejectsPartialSessionAndOverflow) { + auto invalid = request(); + invalid.expected_receiver_session_id = {}; + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::validate(invalid) + .IsInvalidArgument()); + + invalid = request(); + invalid.resources[0].completed_total = + invalid.resources[0].consumed_total + 1; + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::validate(invalid) + .IsInvalidArgument()); + + invalid = request(); + invalid.resources[0].consumed_total = + std::numeric_limits::max() - 10; + invalid.resources[0].completed_total = 0; + invalid.resources[0].desired_available = 11; + EXPECT_TRUE(ReceiverCreditPullRequestCodecV1::validate(invalid) + .IsInvalidArgument()); +} + +TEST(ReceiverCreditProtocol, InitialRequestUsesAllZeroSession) { + auto initial = request(); + initial.expected_receiver_session_id = {}; + initial.expected_epoch = 0; + initial.last_update_sequence = 0; + std::string wire; + ASSERT_TRUE(ReceiverCreditPullRequestCodecV1::encode(initial, wire).ok()); + ReceiverCreditPullRequestV1 decoded; + ASSERT_TRUE(ReceiverCreditPullRequestCodecV1::decode(wire, decoded).ok()); + EXPECT_EQ(decoded.expected_receiver_session_id, (ReceiverSessionId{})); + EXPECT_EQ(decoded.expected_epoch, 0); +} + +TEST(ReceiverCreditProtocol, FullResponseRoundTrips) { + const auto original = response(); + std::string wire; + ASSERT_TRUE(ReceiverCreditPullResponseCodecV1::encode(original, wire).ok()); + EXPECT_EQ(wire.size(), ReceiverCreditPullResponseCodecV1::kWireBytes); + EXPECT_EQ(static_cast(wire[0]), 0x54); + EXPECT_EQ(static_cast(wire[1]), 0x43); + EXPECT_EQ(static_cast(wire[2]), 0x52); + EXPECT_EQ(static_cast(wire[3]), 0x50); + + ReceiverCreditPullResponseV1 decoded; + ASSERT_TRUE(ReceiverCreditPullResponseCodecV1::decode(wire, decoded).ok()); + EXPECT_EQ(decoded.status, original.status); + EXPECT_EQ(decoded.retry_after_us, 0); + EXPECT_EQ(decoded.activation.receiver_session_id, + original.activation.receiver_session_id); + EXPECT_EQ(decoded.update.sequence, original.update.sequence); + ASSERT_EQ(decoded.update.grants.size(), kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + EXPECT_EQ(decoded.update.grants[i].resource, + original.update.grants[i].resource); + EXPECT_EQ(decoded.update.grants[i].grant_total, + original.update.grants[i].grant_total); + } +} + +TEST(ReceiverCreditProtocol, ResponseRejectsInvalidRetryAndPartialUpdate) { + auto invalid = response(); + invalid.status = ReceiverCreditPullStatus::Retry; + std::string unchanged = "sentinel"; + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::encode(invalid, unchanged) + .IsInvalidArgument()); + EXPECT_EQ(unchanged, "sentinel"); + + invalid.retry_after_us = 10; + ASSERT_TRUE( + ReceiverCreditPullResponseCodecV1::encode(invalid, unchanged).ok()); + invalid.update.grants.pop_back(); + unchanged = "sentinel"; + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::encode(invalid, unchanged) + .IsInvalidArgument()); + EXPECT_EQ(unchanged, "sentinel"); +} + +TEST(ReceiverCreditProtocol, UnsupportedCanUseFixedSizeZeroPayload) { + ReceiverCreditPullResponseV1 unsupported; + ASSERT_TRUE(ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus::Unsupported, unsupported) + .ok()); + + std::string wire; + ASSERT_TRUE( + ReceiverCreditPullResponseCodecV1::encode(unsupported, wire).ok()); + EXPECT_EQ(wire.size(), ReceiverCreditPullResponseCodecV1::kWireBytes); + ReceiverCreditPullResponseV1 decoded; + ASSERT_TRUE(ReceiverCreditPullResponseCodecV1::decode(wire, decoded).ok()); + EXPECT_EQ(decoded.status, ReceiverCreditPullStatus::Unsupported); + EXPECT_EQ(decoded.activation.schema_version, 0); + EXPECT_EQ(decoded.update.schema_version, 0); + EXPECT_TRUE(decoded.update.grants.empty()); + + const auto prior = decoded.status; + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::makeZeroPayload( + ReceiverCreditPullStatus::Granted, decoded) + .IsInvalidArgument()); + EXPECT_EQ(decoded.status, prior); +} + +TEST(ReceiverCreditProtocol, ResponseDecodeIsAtomicForReservedAndDuplicate) { + std::string wire; + ASSERT_TRUE( + ReceiverCreditPullResponseCodecV1::encode(response(), wire).ok()); + ReceiverCreditPullResponseV1 output; + output.update.sequence = 99; + for (size_t length = 0; length < wire.size(); ++length) { + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::decode( + std::string_view(wire.data(), length), output) + .IsInvalidArgument()); + EXPECT_EQ(output.update.sequence, 99); + } + + std::string reserved = wire; + reserved[11] = 1; + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::decode(reserved, output) + .IsInvalidArgument()); + EXPECT_EQ(output.update.sequence, 99); + + std::string duplicate = wire; + // The second grant starts at 112; overwrite its resource with DataBytes. + duplicate[112] = 0; + duplicate[113] = 1; + EXPECT_TRUE(ReceiverCreditPullResponseCodecV1::decode(duplicate, output) + .IsInvalidArgument()); + EXPECT_EQ(output.update.sequence, 99); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_bench.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_bench.cpp new file mode 100644 index 0000000000..046cfa001c --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_bench.cpp @@ -0,0 +1,190 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/control_plane.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake::tent { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr size_t resourceIndex(CreditResource resource) { + return static_cast(resource) - 1; +} + +double processCpuSeconds() { + rusage usage{}; + if (getrusage(RUSAGE_SELF, &usage) != 0) return 0; + return usage.ru_utime.tv_sec + usage.ru_utime.tv_usec / 1e6 + + usage.ru_stime.tv_sec + usage.ru_stime.tv_usec / 1e6; +} + +uint64_t percentile(const std::vector& sorted, double fraction) { + if (sorted.empty()) return 0; + const auto index = + static_cast(fraction * static_cast(sorted.size() - 1)); + return sorted[index]; +} + +std::shared_ptr makeAllocator() { + ReceiverCreditAllocatorConfig config; + config.capacity[resourceIndex(CreditResource::DataBytes)] = 1ULL << 50; + config.capacity[resourceIndex(CreditResource::RequestSlots)] = 1ULL << 30; + config.max_grant_per_pull[resourceIndex(CreditResource::DataBytes)] = 64ULL + << 20; + config.max_grant_per_pull[resourceIndex(CreditResource::RequestSlots)] = 64; + config.max_entries = 4096; + config.ttl_ms = 60'000; + config.retry_after_us = 100; + config.receiver_session_id = {0x1234, 0x5678}; + config.epoch = 1; + + std::unique_ptr allocator; + auto status = ReceiverCreditAllocator::create(config, allocator); + if (!status.ok()) return nullptr; + return std::shared_ptr(std::move(allocator)); +} + +ReceiverCreditPullRequestV1 makeRequest(uint64_t sender_peer) { + ReceiverCreditPullRequestV1 request; + request.sender_peer = sender_peer; + request.qos_class = 0; + request.request_sequence = 1; + request.resources = { + {CreditResource::DataBytes, 0, 0, 4ULL << 20, 64ULL << 20}, + {CreditResource::RequestSlots, 0, 0, 1, 64}, + }; + return request; +} + +int runServer(uint16_t port, uint64_t duration_seconds) { + auto allocator = makeAllocator(); + if (!allocator) return 2; + ControlService service("p2p", "", nullptr); + service.setReceiverCreditAllocator(allocator); + auto status = service.start(port); + if (!status.ok()) { + std::cerr << status.ToString() << '\n'; + return 3; + } + + const double cpu_start = processCpuSeconds(); + const auto wall_start = Clock::now(); + std::cout << "READY port=" << port << std::endl; + std::this_thread::sleep_for(std::chrono::seconds(duration_seconds)); + const double wall_seconds = + std::chrono::duration(Clock::now() - wall_start).count(); + const double cpu_seconds = processCpuSeconds() - cpu_start; + + ReceiverCreditAllocatorSnapshot snapshot; + status = allocator->snapshot(snapshot); + if (!status.ok()) return 4; + std::cout << "{\"role\":\"server\",\"wall_s\":" << wall_seconds + << ",\"cpu_s\":" << cpu_seconds + << ",\"cpu_cores\":" << cpu_seconds / wall_seconds + << ",\"peer_entries\":" << snapshot.entries << "}" << std::endl; + return 0; +} + +int runClient(const std::string& address, uint64_t sender_peer, + uint64_t measured_calls, uint64_t warmup_calls) { + if (sender_peer == 0 || measured_calls == 0) return 2; + auto request = makeRequest(sender_peer); + const uint64_t total_calls = measured_calls + warmup_calls; + std::vector latency_ns; + latency_ns.reserve(measured_calls); + uint64_t granted = 0; + uint64_t retries = 0; + size_t request_wire_bytes = 0; + + const double cpu_start = processCpuSeconds(); + auto wall_start = Clock::now(); + for (uint64_t call = 0; call < total_calls; ++call) { + if (call == warmup_calls) wall_start = Clock::now(); + const auto start = Clock::now(); + ReceiverCreditPullResponseV1 response; + auto status = + ControlClient::pullReceiverCredit(address, request, response); + const auto elapsed = + std::chrono::duration_cast(Clock::now() - + start) + .count(); + if (!status.ok()) { + std::cerr << status.ToString() << '\n'; + return 3; + } + if (response.status == ReceiverCreditPullStatus::Retry) { + ++retries; + } else if (response.status != ReceiverCreditPullStatus::Granted) { + std::cerr << "unexpected response status " + << static_cast(response.status) << '\n'; + return 4; + } + if (call >= warmup_calls) latency_ns.push_back(elapsed); + + request.expected_receiver_session_id = + response.activation.receiver_session_id; + request.expected_epoch = response.activation.epoch; + request.last_update_sequence = response.update.sequence; + ++request.request_sequence; + for (const auto& grant : response.update.grants) + if (grant.resource == CreditResource::DataBytes) + granted = grant.grant_total; + } + const double wall_seconds = + std::chrono::duration(Clock::now() - wall_start).count(); + const double cpu_seconds = processCpuSeconds() - cpu_start; + + std::string request_wire; + if (!ReceiverCreditPullRequestCodecV1::encode(request, request_wire).ok()) + return 5; + request_wire_bytes = request_wire.size(); + std::sort(latency_ns.begin(), latency_ns.end()); + const uint64_t response_wire_bytes = + ReceiverCreditPullResponseCodecV1::kWireBytes; + const uint64_t wire_payload_bytes = + measured_calls * (request_wire_bytes + response_wire_bytes); + std::cout << "{\"role\":\"client\",\"sender_peer\":" << sender_peer + << ",\"calls\":" << measured_calls + << ",\"warmup_calls\":" << warmup_calls + << ",\"request_bytes\":" << request_wire_bytes + << ",\"response_bytes\":" << response_wire_bytes + << ",\"wire_payload_bytes\":" << wire_payload_bytes + << ",\"wall_s\":" << wall_seconds << ",\"cpu_s\":" << cpu_seconds + << ",\"qps\":" << measured_calls / wall_seconds + << ",\"rtt_p50_us\":" << percentile(latency_ns, 0.50) / 1000.0 + << ",\"rtt_p95_us\":" << percentile(latency_ns, 0.95) / 1000.0 + << ",\"rtt_p99_us\":" << percentile(latency_ns, 0.99) / 1000.0 + << ",\"rtt_max_us\":" << latency_ns.back() / 1000.0 + << ",\"retries\":" << retries << ",\"granted_bytes\":" << granted + << "}" << std::endl; + return 0; +} + +} // namespace +} // namespace mooncake::tent + +int main(int argc, char** argv) { + using namespace mooncake::tent; + if (argc == 4 && std::string(argv[1]) == "server") + return runServer(static_cast(std::stoul(argv[2])), + std::stoull(argv[3])); + if (argc == 6 && std::string(argv[1]) == "client") + return runClient(argv[2], std::stoull(argv[3]), std::stoull(argv[4]), + std::stoull(argv[5])); + std::cerr << "server | client " + " \n"; + return 1; +} diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_harness.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_harness.cpp new file mode 100644 index 0000000000..1dc50dc1e6 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_rpc_harness.cpp @@ -0,0 +1,204 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/rpc/rpc.h" +#include "tent/runtime/receiver_credit_control.h" + +#include +#include +#include +#include +#include + +namespace mooncake::tent { +namespace { +constexpr int kCreditTestRpc = 1001; +constexpr int kCapabilityTestRpc = 1002; +CreditKey key() { return {{1, 2}, 3, 4}; } +ReceiverCreditUpdateV1 update(uint64_t sequence, uint64_t epoch) { + ReceiverCreditUpdateV1 u; + u.receiver_session_id = key().receiver_session; + u.qos_class = key().qos_class; + u.epoch = epoch; + u.sequence = sequence; + u.grants = {{CreditResource::DataBytes, sequence}}; + return u; +} + +int runServer(uint16_t port, uint64_t expected_calls, uint64_t drain_delay_us, + uint64_t epoch) { + BoundedCreditUpdateInbox inbox(64); + ReceiverCreditIngress ingress(inbox, key(), epoch); + SenderCreditLedger ledger; + if (!ledger.activate(key(), epoch).ok()) return 2; + std::atomic accepted{0}; + std::atomic queue_full{0}; + std::atomic invalid{0}; + CoroRpcAgent server; + server.registerFunction(kCreditTestRpc, + [&](std::string_view wire, std::string& response) { + auto status = ingress.tryAccept(wire); + if (status.IsTooManyRequests()) { + ++queue_full; + response = "FULL"; + return; + } + if (!status.ok()) { + ++invalid; + response = "INVALID"; + return; + } + ++accepted; + response = "OK"; + }); + auto status = server.start(port); + if (!status.ok()) return 3; + uint64_t applied = 0, duplicate = 0, gaps = 0; + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(60); + while ((accepted < expected_calls || inbox.size() != 0) && + std::chrono::steady_clock::now() < deadline) { + std::vector batch; + inbox.drain(batch, 31); + for (auto& envelope : batch) { + CreditUpdateDisposition disposition; + status = + ledger.applyUpdate(envelope.key, envelope.update, disposition); + if (!status.ok()) return 4; + if (disposition == CreditUpdateDisposition::DuplicateOrOld) + ++duplicate; + else if (disposition == CreditUpdateDisposition::SequenceGap) + ++gaps; + else + ++applied; + } + if (drain_delay_us) + std::this_thread::sleep_for( + std::chrono::microseconds(drain_delay_us)); + std::this_thread::yield(); + } + uint64_t available = 0; + status = ledger.available(key(), CreditResource::DataBytes, available); + std::cout << "accepted=" << accepted << " applied=" << applied + << " duplicate=" << duplicate << " gaps=" << gaps + << " queue_full=" << queue_full << " invalid=" << invalid + << " epoch=" << epoch << " available=" << available << std::endl; + return !status.ok() || accepted != expected_calls ? 5 : 0; +} + +int runClient(const std::string& address, uint64_t first, uint64_t count, + uint64_t epoch, bool expect_invalid) { + CoroRpcAgent client; + for (uint64_t sequence = first; sequence < first + count; ++sequence) { + std::string wire, response; + auto status = + ReceiverCreditCodecV1::encode(update(sequence, epoch), wire); + if (!status.ok()) return 2; + do { + status = client.call(address, kCreditTestRpc, wire, response); + if (!status.ok()) return 3; + if (response == "FULL") std::this_thread::yield(); + } while (response == "FULL"); + if (expect_invalid) { + if (response != "INVALID") return 4; + continue; + } + if (response != "OK") return 4; + } + return 0; +} + +std::vector parseVersions(const std::string& text) { + std::vector versions; + std::stringstream stream(text); + std::string item; + while (std::getline(stream, item, ',')) + if (!item.empty()) versions.push_back(std::stoul(item)); + return versions; +} + +int runCapabilityServer(uint16_t port, uint64_t expected_calls, + const std::vector& advertised_versions) { + std::string advertised_wire; + if (!CreditCapabilityCodecV1::encode(advertised_versions, advertised_wire) + .ok()) + return 2; + std::atomic calls{0}; + std::atomic invalid{0}; + CoroRpcAgent server; + server.registerFunction(kCapabilityTestRpc, [&](std::string_view request, + std::string& response) { + std::vector offered; + if (!CreditCapabilityCodecV1::decode(request, offered).ok()) { + ++invalid; + response = "INVALID"; + return; + } + ++calls; + response = advertised_wire; + }); + auto status = server.start(port); + if (!status.ok()) return 3; + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (calls < expected_calls && + std::chrono::steady_clock::now() < deadline) + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + std::cout << "calls=" << calls << " invalid=" << invalid << std::endl; + return calls == expected_calls && invalid == 0 ? 0 : 4; +} + +int runCapabilityClient(const std::string& address, CreditRolloutMode mode, + CreditPeerState expected) { + std::string request, response; + if (!CreditCapabilityCodecV1::encode({1}, request).ok()) return 2; + CoroRpcAgent client; + auto status = client.call(address, kCapabilityTestRpc, request, response); + if (!status.ok()) return 3; + std::vector peer_versions; + if (!CreditCapabilityCodecV1::decode(response, peer_versions).ok()) + return 4; + CreditCapabilityState state(mode); + status = state.completeNegotiation(peer_versions); + if (expected == CreditPeerState::Failed) { + if (!status.IsNotImplemented()) return 5; + } else if (!status.ok()) { + return 5; + } + std::cout << "state=" << static_cast(state.state()) + << " version=" << state.version() << std::endl; + return state.state() == expected ? 0 : 6; +} +} // namespace +} // namespace mooncake::tent + +int main(int argc, char** argv) { + using namespace mooncake::tent; + if (argc >= 4 && argc <= 6 && std::string(argv[1]) == "server") + return runServer(static_cast(std::stoul(argv[2])), + std::stoull(argv[3]), + argc >= 5 ? std::stoull(argv[4]) : 0, + argc == 6 ? std::stoull(argv[5]) : 7); + if (argc >= 5 && argc <= 7 && std::string(argv[1]) == "client") + return runClient(argv[2], std::stoull(argv[3]), std::stoull(argv[4]), + argc >= 6 ? std::stoull(argv[5]) : 7, + argc == 7 && std::string(argv[6]) == "expect-invalid"); + if (argc == 5 && std::string(argv[1]) == "cap-server") + return runCapabilityServer(static_cast(std::stoul(argv[2])), + std::stoull(argv[3]), + parseVersions(argv[4])); + if (argc == 5 && std::string(argv[1]) == "cap-client") { + CreditRolloutMode mode = std::string(argv[3]) == "required" + ? CreditRolloutMode::Required + : CreditRolloutMode::Optional; + CreditPeerState expected = CreditPeerState::Active; + if (std::string(argv[4]) == "legacy") + expected = CreditPeerState::Legacy; + else if (std::string(argv[4]) == "failed") + expected = CreditPeerState::Failed; + return runCapabilityClient(argv[2], mode, expected); + } + std::cerr << "server [delay_us] [epoch] | client " + " [epoch] [expect-invalid] | cap-server " + " | cap-client " + " \n"; + return 1; +} diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp index 33faa04a2c..320a913c7d 100644 --- a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp @@ -75,24 +75,6 @@ TEST(ReceiverCredit, SequenceGapIsVisibleAndSafe) { EXPECT_EQ(d, CreditUpdateDisposition::SequenceGap); } -TEST(ReceiverCredit, PartialGrantUpdateRetainsOmittedResources) { - SenderCreditLedger l; - ASSERT_TRUE(l.activate(key(), 7).ok()); - grant(l, 1, 100, 5); - - CreditUpdateDisposition d; - ASSERT_TRUE(l.applyUpdate( - key(), update(7, 2, {{CreditResource::DataBytes, 160}}), d) - .ok()); - EXPECT_EQ(d, CreditUpdateDisposition::Applied); - - uint64_t v; - ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); - EXPECT_EQ(v, 160); - ASSERT_TRUE(l.available(key(), CreditResource::RequestSlots, v).ok()); - EXPECT_EQ(v, 5); -} - TEST(ReceiverCredit, StaleEpochFailsAndActivationFencesOldState) { SenderCreditLedger l; ASSERT_TRUE(l.activate(key(), 7).ok()); @@ -196,11 +178,72 @@ TEST(ReceiverCredit, RollbackChecksUnderflowAtomically) { grant(l, 1); ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); EXPECT_TRUE( - l.rollbackReservation(key(), charge(61, 1)).IsInvalidArgument()); + l.rollbackReservation(key(), 7, charge(61, 1)).IsInvalidArgument()); uint64_t v; ASSERT_TRUE(l.consumed(key(), CreditResource::RequestSlots, v).ok()); EXPECT_EQ(v, 1); - ASSERT_TRUE(l.rollbackReservation(key(), charge(60, 1)).ok()); + ASSERT_TRUE(l.rollbackReservation(key(), 7, charge(60, 1)).ok()); +} + +TEST(ReceiverCredit, CompletionIsCumulativeAndDoesNotMintLocalCredit) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); + + CreditLedgerSnapshot before; + ASSERT_TRUE(l.snapshot(key(), 7, before).ok()); + EXPECT_EQ(before.epoch, 7); + EXPECT_EQ(before.last_sequence, 1); + EXPECT_TRUE(before.has_update); + EXPECT_EQ(before.grants[0], 100); + EXPECT_EQ(before.consumed[0], 60); + EXPECT_EQ(before.completed[0], 0); + + ASSERT_TRUE(l.recordCompletion(key(), 7, charge(60, 1)).ok()); + uint64_t bytes = 0, slots = 0; + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE(l.available(key(), CreditResource::RequestSlots, slots).ok()); + EXPECT_EQ(bytes, 40); + EXPECT_EQ(slots, 1); + + CreditLedgerSnapshot after; + ASSERT_TRUE(l.snapshot(key(), 7, after).ok()); + EXPECT_EQ(after.consumed[0], 60); + EXPECT_EQ(after.consumed[1], 1); + EXPECT_EQ(after.completed[0], 60); + EXPECT_EQ(after.completed[1], 1); + EXPECT_TRUE(l.recordCompletion(key(), 7, charge(1, 1)).IsInvalidArgument()); + EXPECT_TRUE( + l.rollbackReservation(key(), 7, charge(1, 1)).IsInvalidArgument()); +} + +TEST(ReceiverCredit, OldEpochRollbackAndCompletionCannotMutateNewEpoch) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); + + ASSERT_TRUE(l.activate(key(), 8).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE(l.applyUpdate(key(), + update(8, 1, + {{CreditResource::DataBytes, 50}, + {CreditResource::RequestSlots, 2}}), + disposition) + .ok()); + ASSERT_TRUE(l.tryReserve(key(), charge(10, 1)).ok()); + + EXPECT_TRUE( + l.rollbackReservation(key(), 7, charge(60, 1)).IsInvalidEntry()); + EXPECT_TRUE(l.recordCompletion(key(), 7, charge(60, 1)).IsInvalidEntry()); + CreditLedgerSnapshot current; + ASSERT_TRUE(l.snapshot(key(), 8, current).ok()); + EXPECT_EQ(current.consumed[0], 10); + EXPECT_EQ(current.consumed[1], 1); + EXPECT_EQ(current.completed[0], 0); + EXPECT_EQ(current.completed[1], 0); + EXPECT_TRUE(l.snapshot(key(), 7, current).IsInvalidEntry()); } TEST(ReceiverCredit, GrantCannotDecreaseOrFallBelowConsumption) { @@ -232,5 +275,29 @@ TEST(ReceiverCredit, ConcurrentReservationsNeverExceedGrant) { ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, consumed).ok()); EXPECT_EQ(consumed, 100); } + +TEST(ReceiverCredit, ConcurrentCompletionsNeverExceedConsumption) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1, 100, 100); + ASSERT_TRUE(l.tryReserve(key(), charge(100, 100)).ok()); + std::atomic completed{0}; + std::vector threads; + for (int i = 0; i < 16; ++i) { + threads.emplace_back([&] { + for (int j = 0; j < 20; ++j) + if (l.recordCompletion(key(), 7, charge(1, 1)).ok()) + ++completed; + }); + } + for (auto& thread : threads) thread.join(); + EXPECT_EQ(completed, 100); + CreditLedgerSnapshot current; + ASSERT_TRUE(l.snapshot(key(), 7, current).ok()); + EXPECT_EQ(current.consumed[0], 100); + EXPECT_EQ(current.consumed[1], 100); + EXPECT_EQ(current.completed[0], 100); + EXPECT_EQ(current.completed[1], 100); +} } // namespace } // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp index b164b642a3..e68f62a99e 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -51,10 +51,12 @@ class FakeTransport : public Transport { explicit FakeTransport(TransportType self_type, PollStatusFactory poll_status_factory = {}, - bool notify_on_submit = false) + bool notify_on_submit = false, + bool fail_submit = false) : self_type_(self_type), poll_status_factory_(std::move(poll_status_factory)), - notify_on_submit_(notify_on_submit) { + notify_on_submit_(notify_on_submit), + fail_submit_(fail_submit) { caps.dram_to_dram = true; } @@ -83,6 +85,8 @@ class FakeTransport : public Transport { Status submitTransferTasks(SubBatchRef batch, const std::vector& requests) override { ++submit_calls; + if (fail_submit_) + return Status::InternalError("injected submit failure" LOC_MARK); auto* fake = static_cast(batch); for (const auto& request : requests) { fake->requests.push_back(request); @@ -165,6 +169,7 @@ class FakeTransport : public Transport { TransportType self_type_; PollStatusFactory poll_status_factory_; bool notify_on_submit_; + bool fail_submit_; }; std::shared_ptr makeRuntimeQueueConfig(size_t max_dispatch_owners, @@ -197,6 +202,20 @@ std::shared_ptr makeRuntimeQueueConfig(size_t max_dispatch_owners, return cfg; } +void enableProductionReceiverCredit(const std::shared_ptr& cfg, + uint64_t data_bytes = 1ULL << 20, + uint64_t request_slots = 16) { + cfg->set("receiver_credit/mode", "required"); + cfg->set("receiver_credit/capacity/data_bytes", data_bytes); + cfg->set("receiver_credit/capacity/request_slots", request_slots); + cfg->set("receiver_credit/grant_batch/data_bytes", data_bytes); + cfg->set("receiver_credit/grant_batch/request_slots", request_slots); + cfg->set("receiver_credit/control/freshness_ttl_ms", uint64_t{1000}); + cfg->set("receiver_credit/control/retry_after_us", uint64_t{100}); + cfg->set("receiver_credit/control/poll_interval_us", uint64_t{1000}); + cfg->set("receiver_credit/limits/max_peers", uint64_t{16}); +} + void installFakeRdma(TransferEngineImpl& engine, const std::shared_ptr& fake_rdma) { std::string seg_name = engine.getSegmentName(); @@ -215,6 +234,479 @@ Request makeLocalWrite(uint8_t* ptr, size_t length) { return request; } +CreditKey installCreditGate( + TransferEngineImpl& engine, + const std::shared_ptr& contexts, + const std::shared_ptr& ledger, uint64_t bytes, + uint64_t slots) { + CreditActivationV1 activation; + activation.receiver_session_id = {11, 22}; + activation.epoch = 7; + EXPECT_TRUE(contexts->activate(LOCAL_SEGMENT_ID, 200, 3, activation).ok()); + CreditPeerContextSnapshot peer; + EXPECT_TRUE(contexts->lookup(LOCAL_SEGMENT_ID, 3, peer).ok()); + EXPECT_TRUE(ledger->activate(peer.key, peer.epoch).ok()); + ReceiverCreditUpdateV1 update; + update.receiver_session_id = peer.key.receiver_session; + update.qos_class = peer.key.qos_class; + update.epoch = peer.epoch; + update.sequence = 1; + update.grants = {{CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, slots}}; + CreditUpdateDisposition disposition; + EXPECT_TRUE(ledger->applyUpdate(peer.key, update, disposition).ok()); + EXPECT_TRUE(engine + .installReceiverCreditDispatch( + contexts, ledger, [](const Request&) { return 3; }) + .ok()); + return peer.key; +} + +TEST(RuntimeQueueDispatch, RejectsReceiverCreditWithoutRuntimeQueue) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("enable_runtime_queue", false); + cfg->set("receiver_credit/enabled", true); + TransferEngineImpl engine(cfg); + EXPECT_FALSE(engine.available()); +} + +TEST(RuntimeQueueDispatch, ProductionCreditPullGatesDirectRemoteRdmaWrite) { + auto receiver_config = makeRuntimeQueueConfig(4, 1UL << 20); + auto sender_config = makeRuntimeQueueConfig(4, 1UL << 20); + enableProductionReceiverCredit(receiver_config, 4096, 1); + enableProductionReceiverCredit(sender_config, 4096, 1); + + TransferEngineImpl receiver(receiver_config); + TransferEngineImpl sender(sender_config); + ASSERT_TRUE(receiver.available()); + ASSERT_TRUE(sender.available()); + auto receiver_rdma = std::make_shared(RDMA); + auto sender_rdma = std::make_shared(RDMA); + installFakeRdma(receiver, receiver_rdma); + installFakeRdma(sender, sender_rdma); + + constexpr size_t kReqLen = 4096; + std::vector source(kReqLen, 0x42); + std::vector target(kReqLen, 0); + ASSERT_TRUE(sender.registerLocalMemory(source.data(), source.size()).ok()); + ASSERT_TRUE( + receiver.registerLocalMemory(target.data(), target.size()).ok()); + + SegmentID target_id = 0; + ASSERT_TRUE(sender.openSegment(target_id, receiver.getSegmentName()).ok()); + ASSERT_NE(target_id, LOCAL_SEGMENT_ID); + Request request; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = target_id; + request.target_offset = reinterpret_cast(target.data()); + request.length = kReqLen; + request.transport_hint = RDMA; + + BatchID batch = sender.allocateBatch(1); + ASSERT_TRUE(sender.submitTransfer(batch, {request}).ok()); + TransferStatus status{}; + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (std::chrono::steady_clock::now() < deadline) { + ASSERT_TRUE(sender.getTransferStatus(batch, 0, status).ok()); + if (status.s != TransferStatusEnum::PENDING) break; + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(sender_rdma->submit_calls.load(), 1); + + EXPECT_TRUE(sender.freeBatch(batch).ok()); + EXPECT_TRUE(sender.closeSegment(target_id).ok()); + EXPECT_TRUE( + sender.unregisterLocalMemory(source.data(), source.size()).ok()); + EXPECT_TRUE( + receiver.unregisterLocalMemory(target.data(), target.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ProductionCreditBypassesLocalWrite) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + enableProductionReceiverCredit(cfg, 4096, 1); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + std::vector buffer(4096, 0x43); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), 4096)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsCreditInstallWhenOptInIsDisabled) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto status = engine.installReceiverCreditDispatch( + std::make_shared(), + std::make_shared()); + EXPECT_TRUE(status.IsInvalidArgument()) << status.ToString(); +} + +TEST(RuntimeQueueDispatch, CreditOptInFailsClosedBeforeInstall) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x90); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + auto status = + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}); + EXPECT_TRUE(status.IsInvalidEntry()) << status.ToString(); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditOptInUsesConfiguredDefaultQos) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + cfg->set("receiver_credit/default_qos_class", uint32_t{5}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + CreditActivationV1 activation; + activation.receiver_session_id = {55, 66}; + activation.epoch = 9; + ASSERT_TRUE(contexts->activate(LOCAL_SEGMENT_ID, 200, 5, activation).ok()); + CreditPeerContextSnapshot peer; + ASSERT_TRUE(contexts->lookup(LOCAL_SEGMENT_ID, 5, peer).ok()); + ASSERT_TRUE(ledger->activate(peer.key, peer.epoch).ok()); + constexpr size_t kReqLen = 4096; + ReceiverCreditUpdateV1 update; + update.receiver_session_id = peer.key.receiver_session; + update.qos_class = peer.key.qos_class; + update.epoch = peer.epoch; + update.sequence = 1; + update.grants = {{CreditResource::DataBytes, kReqLen}, + {CreditResource::RequestSlots, 1}}; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger->applyUpdate(peer.key, update, disposition).ok()); + ASSERT_TRUE(engine.installReceiverCreditDispatch(contexts, ledger).ok()); + + std::vector buffer(kReqLen, 0x96); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + uint64_t bytes = 1; + ASSERT_TRUE( + ledger->available(peer.key, CreditResource::DataBytes, bytes).ok()); + EXPECT_EQ(bytes, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateCommitsSuccessfulTransportSubmit) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = installCreditGate(engine, contexts, ledger, kReqLen, 1); + std::vector buffer(kReqLen, 0x91); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + uint64_t bytes = 1, slots = 1; + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, 0); + EXPECT_EQ(slots, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateRollsBackFailedTransportSubmit) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared( + RDMA, FakeTransport::PollStatusFactory{}, false, true); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = installCreditGate(engine, contexts, ledger, kReqLen, 1); + std::vector buffer(kReqLen, 0x92); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + uint64_t bytes = 0, slots = 0; + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, kReqLen); + EXPECT_EQ(slots, 1); + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateDefersUntilCreditArrives) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = installCreditGate(engine, contexts, ledger, 0, 0); + std::vector buffer(kReqLen, 0x93); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + + ReceiverCreditUpdateV1 update; + update.receiver_session_id = credit_key.receiver_session; + update.qos_class = credit_key.qos_class; + update.epoch = 7; + update.sequence = 2; + update.grants = {{CreditResource::DataBytes, kReqLen}, + {CreditResource::RequestSlots, 1}}; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger->applyUpdate(credit_key, update, disposition).ok()); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + uint64_t bytes = 1, slots = 1; + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, 0); + EXPECT_EQ(slots, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateRejectsQueuedSnapshotAfterRestart) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto old_key = installCreditGate(engine, contexts, ledger, 0, 0); + std::vector buffer(kReqLen, 0x94); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + + CreditActivationV1 restarted; + restarted.receiver_session_id = {33, 44}; + restarted.epoch = 1; + ASSERT_TRUE(contexts->activate(LOCAL_SEGMENT_ID, 200, 3, restarted).ok()); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + uint64_t bytes = 1; + ASSERT_TRUE( + ledger->available(old_key, CreditResource::DataBytes, bytes).ok()); + EXPECT_EQ(bytes, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateDoesNotOverdrawAcrossQueuedOwners) { + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = installCreditGate(engine, contexts, ledger, kReqLen, 1); + std::vector buffer(kReqLen * 2, 0x95); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + uint64_t bytes = 1, slots = 1; + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, 0); + EXPECT_EQ(slots, 0); + + ReceiverCreditUpdateV1 update; + update.receiver_session_id = credit_key.receiver_session; + update.qos_class = credit_key.qos_class; + update.epoch = 7; + update.sequence = 2; + update.grants = {{CreditResource::DataBytes, kReqLen * 2}, + {CreditResource::RequestSlots, 2}}; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger->applyUpdate(credit_key, update, disposition).ok()); + + TransferStatus overall{}; + for (int i = 0; i < 4; ++i) { + ASSERT_TRUE(engine.getTransferStatus(batch, overall).ok()); + if (overall.s == TransferStatusEnum::COMPLETED) break; + } + EXPECT_EQ(overall.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, 0); + EXPECT_EQ(slots, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateFillsDispatchWindowWhenCreditAllows) { + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = + installCreditGate(engine, contexts, ledger, kReqLen * 2, 2); + std::vector buffer(kReqLen * 2, 0x97); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + uint64_t bytes = 1, slots = 1; + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::DataBytes, bytes).ok()); + ASSERT_TRUE( + ledger->available(credit_key, CreditResource::RequestSlots, slots) + .ok()); + EXPECT_EQ(bytes, 0); + EXPECT_EQ(slots, 0); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CreditGateDefersEveryUnprocessedPickedOwner) { + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enabled", true); + auto contexts = std::make_shared(); + auto ledger = std::make_shared(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + constexpr size_t kReqLen = 4096; + auto credit_key = installCreditGate(engine, contexts, ledger, 0, 0); + std::vector buffer(kReqLen * 2, 0x98); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + + ReceiverCreditUpdateV1 update; + update.receiver_session_id = credit_key.receiver_session; + update.qos_class = credit_key.qos_class; + update.epoch = 7; + update.sequence = 2; + update.grants = {{CreditResource::DataBytes, kReqLen * 2}, + {CreditResource::RequestSlots, 2}}; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger->applyUpdate(credit_key, update, disposition).ok()); + + TransferStatus overall{}; + for (int i = 0; i < 4; ++i) { + ASSERT_TRUE(engine.getTransferStatus(batch, overall).ok()); + if (overall.s == TransferStatusEnum::COMPLETED) break; + } + EXPECT_EQ(overall.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + TEST(RuntimeQueueDispatch, RejectsOverfullBatchBeforePublishingTasks) { auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); TransferEngineImpl engine(cfg); @@ -508,7 +1000,9 @@ TEST(RuntimeQueueDispatch, ProgressWorkerRefillsWindowFromTransportNotify) { {makeLocalWrite(buffer.data(), kReqLen), makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) .ok()); - EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + // The submit notification may let the progress worker dispatch the second + // owner before this thread observes the first one. + EXPECT_GE(fake_rdma->submit_calls.load(), 1); const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(1000); diff --git a/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp index f0741c82bf..16c6df32a4 100644 --- a/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp +++ b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp @@ -37,20 +37,39 @@ void waitChildWithTimeout(pid_t pid, int* status) { FAIL() << "child did not exit before timeout"; } +void notifyParentReady(int fd) { + const char ready = '1'; + if (write(fd, &ready, 1) != 1) _exit(111); + close(fd); +} + +void waitForChildReady(int fd) { + char ready = 0; + ASSERT_EQ(read(fd, &ready, 1), 1) << "child did not report readiness"; + EXPECT_EQ(ready, '1'); + close(fd); +} + } // namespace TEST(GracefulShutdownTest, SigtermTriggersCleanExit) { + int ready_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0) << "pipe() failed"; + pid_t pid = fork(); ASSERT_NE(pid, -1) << "fork() failed"; if (pid == 0) { + close(ready_pipe[0]); auto engine = std::make_unique(false); engine->enableGracefulShutdown(); + notifyParentReady(ready_pipe[1]); pause(); _exit(99); } - usleep(100000); + close(ready_pipe[1]); + waitForChildReady(ready_pipe[0]); kill(pid, SIGTERM); int status; @@ -62,17 +81,23 @@ TEST(GracefulShutdownTest, SigtermTriggersCleanExit) { } TEST(GracefulShutdownTest, SigintTriggersCleanExit) { + int ready_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0) << "pipe() failed"; + pid_t pid = fork(); ASSERT_NE(pid, -1) << "fork() failed"; if (pid == 0) { + close(ready_pipe[0]); auto engine = std::make_unique(false); engine->enableGracefulShutdown(); + notifyParentReady(ready_pipe[1]); pause(); _exit(99); } - usleep(100000); + close(ready_pipe[1]); + waitForChildReady(ready_pipe[0]); kill(pid, SIGINT); int status; @@ -89,19 +114,25 @@ TEST(GracefulShutdownTest, IdempotentEnable) { } TEST(GracefulShutdownTest, EngineDestroyedBeforeSignal) { + int ready_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0) << "pipe() failed"; + pid_t pid = fork(); ASSERT_NE(pid, -1) << "fork() failed"; if (pid == 0) { + close(ready_pipe[0]); { auto engine = std::make_unique(false); engine->enableGracefulShutdown(); } + notifyParentReady(ready_pipe[1]); pause(); _exit(99); } - usleep(100000); + close(ready_pipe[1]); + waitForChildReady(ready_pipe[0]); kill(pid, SIGTERM); int status; @@ -114,15 +145,21 @@ TEST(GracefulShutdownTest, ForkAfterInstallDoesNotHangChildSignal) { auto engine = std::make_unique(false); engine->enableGracefulShutdown(); + int ready_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0) << "pipe() failed"; + pid_t pid = fork(); ASSERT_NE(pid, -1) << "fork() failed"; if (pid == 0) { + close(ready_pipe[0]); + notifyParentReady(ready_pipe[1]); pause(); _exit(99); } - usleep(100000); + close(ready_pipe[1]); + waitForChildReady(ready_pipe[0]); kill(pid, SIGTERM); int status; @@ -134,19 +171,25 @@ TEST(GracefulShutdownTest, ForkAfterInstallDoesNotHangChildSignal) { } TEST(GracefulShutdownTest, MultipleEngines) { + int ready_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0) << "pipe() failed"; + pid_t pid = fork(); ASSERT_NE(pid, -1) << "fork() failed"; if (pid == 0) { + close(ready_pipe[0]); auto engine1 = std::make_unique(false); auto engine2 = std::make_unique(false); engine1->enableGracefulShutdown(); engine2->enableGracefulShutdown(); + notifyParentReady(ready_pipe[1]); pause(); _exit(99); } - usleep(100000); + close(ready_pipe[1]); + waitForChildReady(ready_pipe[0]); kill(pid, SIGTERM); int status;