From 137303a38ca293ac4407a88980c037b6a32fa455 Mon Sep 17 00:00:00 2001 From: Le1zyCatt <2490162471@qq.com> Date: Tue, 28 Jul 2026 11:56:51 +0800 Subject: [PATCH 1/8] [TENT] Complete native transport integration with receiver credit (cherry picked from commit d78e20dd3f9568c8223fbdbb3e98e00f6ed34c67) --- docs/source/design/tent/overview.md | 1 + .../design/tent/ub-native-validation.md | 322 ++++++++ docs/source/performance/mooncake/tebench.md | 3 +- .../tent/config/transfer-engine.json | 10 + .../tent/include/tent/common/types.h | 6 + .../tent/include/tent/rpc/rpc.h | 1 + .../include/tent/runtime/admission_queue.h | 12 + .../tent/include/tent/runtime/control_plane.h | 97 ++- .../include/tent/runtime/receiver_credit.h | 299 ++++++- .../tent/include/tent/runtime/segment.h | 13 + .../tent/runtime/transfer_engine_impl.h | 81 +- .../tent/include/tent/runtime/transport.h | 7 + .../tent/include/tent/transport/rdma/slice.h | 3 + .../tent/transport/tcp/tcp_transport.h | 11 +- .../tent/include/tent/transport/ub/endpoint.h | 6 + .../tent/include/tent/transport/ub/slice.h | 3 + .../tent/include/tent/transport/ub/workers.h | 2 + .../tent/src/runtime/admission_queue.cpp | 104 ++- .../tent/src/runtime/control_plane.cpp | 434 ++++++++-- .../tent/src/runtime/proxy_manager.cpp | 15 + .../tent/src/runtime/receiver_credit.cpp | 642 ++++++++++++--- .../tent/src/runtime/transfer_engine_impl.cpp | 740 +++++++++++++++++- .../src/transport/rdma/rdma_transport.cpp | 5 + .../tent/src/transport/rdma/workers.cpp | 18 + .../tent/src/transport/tcp/tcp_transport.cpp | 50 +- .../tent/src/transport/ub/endpoint.cpp | 24 +- .../tent/src/transport/ub/ub_transport.cpp | 108 ++- .../tent/src/transport/ub/workers.cpp | 121 ++- .../tent/tests/admission_queue_test.cpp | 210 +++++ .../tent/tests/receiver_credit_test.cpp | 358 ++++++++- .../tent/tests/request_merge_test.cpp | 1 + .../tests/runtime_queue_dispatch_test.cpp | 452 ++++++++++- .../tent/tests/transport_selector_test.cpp | 12 + .../tent/tests/ub_native_data_path_test.cpp | 55 +- 34 files changed, 3918 insertions(+), 308 deletions(-) create mode 100644 docs/source/design/tent/ub-native-validation.md diff --git a/docs/source/design/tent/overview.md b/docs/source/design/tent/overview.md index 6b37df87e7..30c5e1fdf9 100644 --- a/docs/source/design/tent/overview.md +++ b/docs/source/design/tent/overview.md @@ -139,4 +139,5 @@ failover :maxdepth: 1 testing +ub-native-validation ::: diff --git a/docs/source/design/tent/ub-native-validation.md b/docs/source/design/tent/ub-native-validation.md new file mode 100644 index 0000000000..c7294823b8 --- /dev/null +++ b/docs/source/design/tent/ub-native-validation.md @@ -0,0 +1,322 @@ +# TENT Native UB Hardware Validation Guide + +This guide validates the TENT-native UB transport on machines with a real +UMDK/URMA provider. It covers build linkage, two-node READ/WRITE correctness, +multi-rail scheduling, fault recovery, retryable shutdown, and a performance +baseline. The repository's CI uses an injected adapter; the hardware cases in +this guide are therefore required before declaring native UB production-ready. + +## 1. Record the Test Environment + +Run every case on the same Mooncake commit and record the following on both +hosts: + +```bash +git rev-parse HEAD +uname -a +gcc --version | head -1 +cmake --version | head -1 +urma_admin -l +ldconfig -p | grep liburma +``` + +The two hosts must have mutually reachable EIDs and TCP reachability for the +TENT P2P control plane. Use the vendor-supported tool to verify each active UB +port, for example: + +```bash +urma_admin -p urma0 +``` + +If UMDK is installed outside the system search paths, export its library path +before configuring and running Mooncake: + +```bash +export LD_LIBRARY_PATH=/path/to/umdk/lib:${LD_LIBRARY_PATH} +``` + +Unset `MC_TENT_CONF` for the commands below unless the referenced file +explicitly enables `transports/ub`. The environment config replaces, rather +than merges with, the configuration assembled by `tebench`. + +```bash +unset MC_TENT_CONF +``` + +## 2. Build and Verify Real URMA Linkage + +Configure a Debug build first. CMake discovers standard UMDK installations +automatically; pass `URMA_SYSTEM_INCLUDE_DIR` and `URMA_LIBRARY` only for a +non-standard installation. + +```bash +cmake -S . -B build-ub-debug \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DUSE_CUDA=OFF \ + -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BENCHMARK=ON \ + -DCMAKE_BUILD_TYPE=Debug + +cmake --build build-ub-debug -j"$(nproc)" --target \ + tent_ub_core_test \ + tent_ub_teardown_test \ + tent_ub_native_data_path_test \ + tent_transport_selector_test \ + tebench +``` + +Confirm that CMake found a real library and that `tebench` has a runtime +dependency on it: + +```bash +grep '^URMA_LIBRARY:' build-ub-debug/CMakeCache.txt +ldd build-ub-debug/mooncake-transfer-engine/benchmark/tebench | grep urma +``` + +Both commands must show the intended UMDK installation. A header-only fallback +is sufficient for mock tests but **not** for hardware validation. + +Run the deterministic mock suite before touching the devices: + +```bash +ctest --test-dir build-ub-debug --output-on-failure \ + -R 'tent_(ub_core|ub_teardown|ub_native_data_path|transport_selector)_test' +``` + +Expected result: all selected tests pass. In particular, the teardown test +must demonstrate that a failed unregister/delete retains ownership and that a +second cleanup call releases the same handle. It also verifies that a failed +device unpublishes every old endpoint and cannot recover until every JFC has +reported healthy in the current failure epoch. + +## 3. Two-Node Correctness + +Use `HOST_A` as the target and `HOST_B` as the initiator. The examples use a +fixed control-plane port so peer-restart tests can reuse the same segment name. + +On `HOST_A`: + +```bash +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --rpc_server_port=18001 \ + --xport_type=ub \ + --seg_type=DRAM \ + --total_buffer_size=1073741824 2>&1 | tee target-ub.log +``` + +Copy the printed segment name. On `HOST_B`, run a consistency-checked mixed +transfer first: + +```bash +TARGET_SEG='' + +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --target_seg_name="${TARGET_SEG}" \ + --xport_type=ub \ + --tent_transport_hint=ub \ + --seg_type=DRAM \ + --op_type=mix \ + --check_consistency=true \ + --total_buffer_size=1073741824 \ + --start_block_size=4096 \ + --max_block_size=1048576 \ + --start_batch_size=1 \ + --max_batch_size=8 \ + --start_num_threads=1 \ + --max_num_threads=4 \ + --duration=10 2>&1 | tee initiator-mix.log +``` + +Then repeat with `--op_type=read` and `--op_type=write`. All rows must finish, +consistency checking must report no mismatch, and neither host may log a +fallback to a non-UB transport. This covers both upward-facing TENT APIs and +native URMA READ/WRITE without converting requests to the Classic TE path. + +## 4. Multi-Device and Pressure Scheduling + +Leave `transports/ub/device_filter` empty so every discovered active EID is +eligible. Confirm that the startup topology lists at least two `ub::eid` +entries on each host, then run a fixed high-concurrency case: + +```bash +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --target_seg_name="${TARGET_SEG}" \ + --xport_type=ub \ + --tent_transport_hint=ub \ + --seg_type=DRAM \ + --op_type=write \ + --total_buffer_size=1073741824 \ + --start_block_size=65536 \ + --max_block_size=65536 \ + --start_batch_size=16 \ + --max_batch_size=16 \ + --start_num_threads=16 \ + --max_num_threads=16 \ + --duration=60 2>&1 | tee initiator-multirail.log +``` + +Monitor per-device traffic with the platform's supported counters. Every +healthy eligible rail should make progress; saturating one path must not stop +posting on another path with free quota. Record aggregate bandwidth and the +per-device byte deltas. + +## 5. Fault and Recovery Cases + +Run these cases with a long-lived embedding harness that keeps the initiator's +TENT engine alive across each fault. Preserve both host logs, note exact fault +timestamps, and record endpoint generations before and after recovery. A new +`tebench` process may be used for basic reconnect smoke testing, but it is not +sufficient evidence for generation isolation. + +### Peer restart and endpoint rebuild + +1. Stop `HOST_A` with `SIGINT` while transfers are active. +2. Start it again with the same `--rpc_server_port=18001` command. +3. In the still-running `HOST_B` harness, discover the republished segment and + submit READ and WRITE again without reinstalling the local TENT transport. + +The old endpoint generation must not be reused. The live engine must establish +a new generation and complete new transfers; stale completions must not change +a new task. Save both generation values as part of the PASS evidence. + +### Remote segment republish + +Using the embedding application's TENT API, unregister and re-register the +target buffer at the same virtual address while the engines remain alive. +Publish the returned `BufferDesc`, then submit another READ and WRITE. The new +metadata generation must be imported and the stale imported segment must be +released. Do not use `tebench` for this case because it registers its buffers +only once at process startup. + +### UB port down/up + +1. Disable one UB port using the deployment's vendor-approved administration + procedure; do not power-cycle a host as the first test. +2. Confirm that traffic continues on another healthy rail. +3. Re-enable the port and wait at least the configured endpoint cooldown. +4. Confirm that the recovered rail resumes traffic and that a new endpoint is + built before posting. + +The current recovery loop first unpublishes every endpoint backed by the failed +device. It reactivates the context only after every provider JFC has polled +successfully in the latest failure epoch, the cooldown has elapsed, and all old +WRs have drained. This supports a transient port/link failure when the provider +keeps Context/JFC handles valid. A destructive device reset that invalidates +those handles still requires transport reinstall and must be reported as +unsupported for in-place recovery. + +### Cancellation and timeout + +Exercise cancellation through the embedding application's TENT API while a +large multi-slice transfer is active. For timeout testing, drop the UB path +without stopping the control-plane TCP connection. The transfer must either +complete on another rail or reach a terminal timeout/failure; it must not hang +and shutdown must still drain safely. + +### Transport fallback + +The `tebench --xport_type=ub` commands above deliberately disable every other +transport, so they validate UB path failover but **cannot** validate transport +fallback. Use an embedding test harness with both UB and TCP (or RDMA) enabled, +and a selector policy that prefers UB but permits the second transport. Submit +the request with `transport_hint=UNSPEC`; do not pin it to UB. + +1. Verify the first consistency-checked READ and WRITE select UB. +2. Disable every usable UB path while requests remain active. +3. Verify the runtime reaches a terminal result for the failed UB attempt and + resubmits through TCP/RDMA rather than leaving the task pending. +4. Restore UB, wait for its cooldown, and verify later policy-selected traffic + can use UB again. + +Capture selector/failover logs or transport metrics that identify both the +failed UB attempt and the replacement transport. A successful copy without +this evidence does not prove transport fallback. + +## 6. Repeated Shutdown and Sanitizers + +Configure a separate sanitizer build: + +```bash +cmake -S . -B build-ub-asan \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DUSE_CUDA=OFF \ + -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BENCHMARK=ON \ + -DCMAKE_BUILD_TYPE=Debug \ + -DCMAKE_CXX_FLAGS='-fsanitize=address,undefined -fno-omit-frame-pointer' \ + -DCMAKE_EXE_LINKER_FLAGS='-fsanitize=address,undefined' \ + -DCMAKE_SHARED_LINKER_FLAGS='-fsanitize=address,undefined' + +cmake --build build-ub-asan -j"$(nproc)" --target tebench +export ASAN_OPTIONS=detect_leaks=1:halt_on_error=1 +export UBSAN_OPTIONS=halt_on_error=1:print_stacktrace=1 +``` + +Use an embedding test harness to perform at least 100 +install-register-transfer-unregister-uninstall cycles **in one process**. +Include ten cycles where a UB port or peer is unavailable during uninstall. If +unregister or uninstall returns a provider `busy`/error, keep the engine alive +and retry the same cleanup operation; the retry must act on the retained native +handle and eventually succeed after the fault is removed. Also run 100 +target/initiator start-transfer-`SIGINT` process cycles as a supplementary +check. + +A cycle is successful only if shutdown returns, the next cycle can reuse the +same buffers and control-plane port, and ASan/UBSan reports no UAF, double free, +or leaked native ownership graph. Record the first and retry return codes plus +the provider log around every injected cleanup failure. + +## 7. Performance Baseline + +Use a `RelWithDebInfo` build and keep CPU affinity, NUMA placement, block size, +batch size, concurrency, UMDK version, and device set fixed. Record: + +| Case | Required measurements | +| --- | --- | +| 4 KiB, batch 1, one thread | average, P99, and P999 latency | +| 64 KiB and 1 MiB, one thread | single-stream bandwidth and CPU usage | +| 64 KiB and 1 MiB, 16 threads | aggregate bandwidth and CPU usage | +| Multi-device, 16 threads | aggregate and per-device bandwidth | +| First transfer to a peer | endpoint connection time | + +Run the same matrix against Classic TE UB using the existing Classic UB +benchmark on the same hosts. Any material regression needs a profile and an +explanation; do not compare results from different UMDK versions or topology. + +## 8. Result Checklist + +Attach the following to the review: + +```text +Mooncake commit: +Host/OS/kernel: +UMDK package and liburma path: +Devices/EIDs: +Debug USE_TENT+USE_UB build: PASS/FAIL +Mock UB tests: PASS/FAIL +Two-node READ: PASS/FAIL +Two-node WRITE: PASS/FAIL +Consistency-checked mix: PASS/FAIL +Multi-device pressure scheduling: PASS/FAIL +Peer restart / endpoint rebuild: PASS/FAIL +Segment republish: PASS/FAIL +Port down/up: PASS/FAIL +Cancellation and timeout: PASS/FAIL +UB to TCP/RDMA transport fallback: PASS/FAIL +100-cycle ASan/UBSan shutdown: PASS/FAIL +Classic UB baseline artifact: +TENT UB baseline artifact: +Known provider errors or unsupported cases: +``` + +Do not mark the hardware gate complete from bandwidth alone. Correctness, +generation isolation, retryable resource release, and sanitizer-clean shutdown +are independent acceptance requirements. diff --git a/docs/source/performance/mooncake/tebench.md b/docs/source/performance/mooncake/tebench.md index 47180d8abf..e39cec0fc9 100644 --- a/docs/source/performance/mooncake/tebench.md +++ b/docs/source/performance/mooncake/tebench.md @@ -397,7 +397,8 @@ gpu_id + thread_id **Transport (TENT only)** -* `--xport_type` : `rdma | shm | mnnvl | gds | iouring`. Selects a single +* `--xport_type` : `rdma | tcp | shm | mnnvl | nvlink | gds | iouring | ub | + sunrise_link`. Selects a single transport to enable (all others are disabled). Empty means no transport is explicitly enabled or disabled by tebench — the engine reads the transport enable list from the `MC_TENT_CONF` config file (see Section diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json index 01880b843e..2ba9fcd185 100644 --- a/mooncake-transfer-engine/tent/config/transfer-engine.json +++ b/mooncake-transfer-engine/tent/config/transfer-engine.json @@ -12,6 +12,16 @@ "max_failover_attempts": 3, "enable_auto_failover_on_poll": true, "enable_progress_worker": false, + "receiver_credit": { + "enable": false, + "data_bytes": 1073741824, + "request_slots": 1024, + "freshness_ttl_ms": 1000, + "max_freshness_ttl_ms": 60000, + "max_senders": 1024, + "max_authority_entries": 4096, + "max_sender_entries": 4096 + }, "metrics": { "enabled": true, "http_port": 9100, diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index b9d758180d..64562bf9d1 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -177,6 +177,12 @@ struct Request { // and docs/source/design/tent/deadline-scheduling.md. uint64_t deadline_ns = 0; IntentType intent_type = IntentType::INTENT_UNSPEC; + // Internal runtime fence for staging sub-transfers that are covered by + // their parent request's receiver-credit reservation. Public callers + // leave these fields zero. + uint64_t receiver_credit_session_high = 0; + uint64_t receiver_credit_session_low = 0; + uint64_t receiver_credit_epoch = 0; }; enum TransferStatusEnum { diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index daeac4916d..a19f0224ff 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -49,6 +49,7 @@ enum RpcFuncID { NotifySegmentUpdated, // Appended to preserve the numeric values of the existing RPCs. BootstrapUb, + ExchangeReceiverCredit, }; class ClientPool; 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 5623e9194e..786686988e 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -147,6 +147,12 @@ class LocalTransferAdmissionQueue { size_t max_owners, size_t max_bytes, std::vector* dropped_owner_ids = nullptr); + // Return an owner that could not be handed off to its selected transport + // from the dispatch window to the queue. Only a Dispatching owner can be + // deferred. Outstanding owner/byte accounting remains charged while the + // owner waits for a later dispatch attempt. + Status defer(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, @@ -192,10 +198,16 @@ class LocalTransferAdmissionQueue { QueueOwnerKind kind{QueueOwnerKind::User}; TransportType transport{UNSPEC}; bool degradation_eligible{false}; + // EDF owners that yield because a downstream resource is unavailable + // move to the next scheduling round. This preserves progress without + // relying on a rotated deque remaining globally deadline-sorted. + uint64_t scheduling_round{0}; QueueState state{QueueState::Queued}; TransferStatusEnum terminal_status{TransferStatusEnum::PENDING}; }; + void enqueueOwner(QueueOwnerId owner_id); + QueueLimits limits_; Status limits_status_; QueueOwnerId next_owner_id_{1}; 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 ab1ce5f3e1..93b439a996 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -77,25 +78,32 @@ struct UbBootstrapDesc { std::vector jetty_uasids; uint64_t endpoint_generation = 0; uint64_t segment_generation = 0; + uint64_t receiver_credit_session_high = 0; + uint64_t receiver_credit_session_low = 0; + uint64_t receiver_credit_epoch = 0; std::vector capabilities; std::string reply_msg; }; inline void to_json(nlohmann::json& j, const UbBootstrapDesc& desc) { - j = nlohmann::json{{"protocol_version", desc.protocol_version}, - {"segment_name", desc.segment_name}, - {"local_nic_path", desc.local_nic_path}, - {"peer_nic_path", desc.peer_nic_path}, - {"local_device_name", desc.local_device_name}, - {"local_device_id", desc.local_device_id}, - {"local_eid_index", desc.local_eid_index}, - {"local_eid", desc.local_eid}, - {"jetty_ids", desc.jetty_ids}, - {"jetty_uasids", desc.jetty_uasids}, - {"endpoint_generation", desc.endpoint_generation}, - {"segment_generation", desc.segment_generation}, - {"capabilities", desc.capabilities}, - {"reply_msg", desc.reply_msg}}; + j = nlohmann::json{ + {"protocol_version", desc.protocol_version}, + {"segment_name", desc.segment_name}, + {"local_nic_path", desc.local_nic_path}, + {"peer_nic_path", desc.peer_nic_path}, + {"local_device_name", desc.local_device_name}, + {"local_device_id", desc.local_device_id}, + {"local_eid_index", desc.local_eid_index}, + {"local_eid", desc.local_eid}, + {"jetty_ids", desc.jetty_ids}, + {"jetty_uasids", desc.jetty_uasids}, + {"endpoint_generation", desc.endpoint_generation}, + {"segment_generation", desc.segment_generation}, + {"receiver_credit_session_high", desc.receiver_credit_session_high}, + {"receiver_credit_session_low", desc.receiver_credit_session_low}, + {"receiver_credit_epoch", desc.receiver_credit_epoch}, + {"capabilities", desc.capabilities}, + {"reply_msg", desc.reply_msg}}; } inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { @@ -114,14 +122,44 @@ inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { desc.jetty_uasids = j.value("jetty_uasids", std::vector{}); desc.endpoint_generation = j.value("endpoint_generation", uint64_t{0}); desc.segment_generation = j.value("segment_generation", uint64_t{0}); + desc.receiver_credit_session_high = + j.value("receiver_credit_session_high", uint64_t{0}); + desc.receiver_credit_session_low = + j.value("receiver_credit_session_low", uint64_t{0}); + desc.receiver_credit_epoch = j.value("receiver_credit_epoch", uint64_t{0}); desc.capabilities = j.value("capabilities", std::vector{}); desc.reply_msg = j.value("reply_msg", ""); } +// SendData/RecvData v2 keeps the legacy address/length prefix, then adds a +// magic discriminator and the receiver-credit session fence. A new receiver +// can still parse legacy non-credit requests, while credit-required receivers +// reject any request without an exact current fence. +constexpr uint64_t kXferDataProtocolMagic = 0x54454e5458465232ULL; +constexpr uint64_t kXferDataV2LegacyRejectAddress = + std::numeric_limits::max(); +struct LegacyXferDataDesc { + uint64_t peer_mem_addr; + uint64_t length; +}; struct XferDataDesc { + // An old receiver interprets the first two fields as the legacy header. + // The impossible address makes it reject v2 instead of copying the + // appended fence bytes into user memory. + uint64_t legacy_reject_addr; + uint64_t length; + uint64_t protocol_magic; uint64_t peer_mem_addr; - size_t length; + uint64_t receiver_credit_session_high; + uint64_t receiver_credit_session_low; + uint64_t receiver_credit_epoch; }; +struct XferDataReplyHeader { + uint64_t protocol_magic; +}; +static_assert(sizeof(LegacyXferDataDesc) == 16); +static_assert(sizeof(XferDataDesc) == 56); +static_assert(sizeof(XferDataReplyHeader) == 8); // Standard TCP RPC payload boundaries. READ retains the server's 1 GiB cap; // WRITE uses a 32-bit RPC attachment that includes XferDataDesc. @@ -139,6 +177,10 @@ using OnReceiveBootstrap = using OnReceiveUbBootstrap = std::function; +using OnExchangeReceiverCredit = + std::function; + using OnNotify = std::function; class ControlClient { @@ -164,13 +206,24 @@ class ControlClient { const UbBootstrapDesc& request, UbBootstrapDesc& response); + static Status exchangeReceiverCredit( + const std::string& server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& response); + static Status sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length); + size_t length, + uint64_t receiver_credit_session_high = 0, + uint64_t receiver_credit_session_low = 0, + uint64_t receiver_credit_epoch = 0); static Status recvData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length); + size_t length, + uint64_t receiver_credit_session_high = 0, + uint64_t receiver_credit_session_low = 0, + uint64_t receiver_credit_epoch = 0); static Status notify(const std::string& server_addr, const Notification& message); @@ -224,6 +277,11 @@ class ControlService { ub_bootstrap_callback_ = callback; } + void setReceiverCreditCallback(const OnExchangeReceiverCredit& callback) { + std::lock_guard lock(receiver_credit_callback_mutex_); + receiver_credit_callback_ = callback; + } + void setNotifyCallback(const OnNotify& callback); Status start(uint16_t& port, bool ipv6_ = false, size_t threads = 1); @@ -237,6 +295,9 @@ class ControlService { void onBootstrapUb(const std::string_view& request, std::string& response); + void onExchangeReceiverCredit(const std::string_view& request, + std::string& response); + void onSendData(const std::string_view& request, std::string& response); void onRecvData(const std::string_view& request, std::string& response); @@ -280,6 +341,8 @@ class ControlService { std::mutex notify_cb_mutex_; std::condition_variable notify_cb_cv_; size_t notify_callbacks_in_flight_ = 0; + std::mutex receiver_credit_callback_mutex_; + OnExchangeReceiverCredit receiver_credit_callback_; OnNotify notify_callback_; static thread_local const ControlService* active_notify_service_; 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..76ad717327 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h @@ -6,93 +6,352 @@ #include #include +#include #include +#include #include #include #include #include "tent/common/status.h" +#include "tent/thirdparty/nlohmann/json.h" namespace mooncake::tent { +// A receiver session changes whenever a receiver process restarts. A sender +// instance has the same representation, but remains a distinct type so the +// two identities cannot be accidentally exchanged at an API boundary. struct ReceiverSessionId { uint64_t high{0}, low{0}; - bool operator==(const ReceiverSessionId& o) const { - return high == o.high && low == o.low; - } + bool operator==(const ReceiverSessionId&) const = default; + [[nodiscard]] bool empty() const noexcept { return high == 0 && low == 0; } }; + +struct SenderInstanceId { + uint64_t high{0}, low{0}; + bool operator==(const SenderInstanceId&) const = default; + [[nodiscard]] bool empty() const noexcept { return high == 0 && low == 0; } +}; + struct CreditKey { ReceiverSessionId receiver_session; - uint64_t sender_peer{0}; + SenderInstanceId sender_instance; uint32_t qos_class{0}; - bool operator==(const CreditKey& o) const { - return receiver_session == o.receiver_session && - sender_peer == o.sender_peer && qos_class == o.qos_class; - } + bool operator==(const CreditKey&) const = default; }; + struct CreditKeyHash { size_t operator()(const CreditKey&) const noexcept; }; + enum class CreditResource : uint16_t { DataBytes = 1, RequestSlots, StagingSlots, ConsumerSlots }; + constexpr size_t kCreditResourceCount = 4; +constexpr uint16_t kReceiverCreditProtocolVersion = 1; +constexpr uint16_t kReceiverCreditRequired = 1U << 0; + +// A cumulative grant. grant_total never decreases within a session/epoch. struct CreditAmount { - CreditResource resource; + CreditResource resource{CreditResource::DataBytes}; uint64_t grant_total{0}; }; + +// A cumulative counter used for release reports and desired grant totals. +struct CreditCounter { + CreditResource resource{CreditResource::DataBytes}; + uint64_t total{0}; +}; + struct CreditCharge { std::vector> resources; }; + +struct ReceiverCreditAdvertV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}; + uint16_t flags{kReceiverCreditRequired}; + ReceiverSessionId receiver_session_id; + uint64_t epoch{0}; + uint32_t freshness_ttl_ms{0}; + std::vector resources; + // Explicit capacities let a sender reject a request that can never fit + // instead of leaving it indefinitely deferred behind the credit gate. + std::vector capacities; +}; + struct ReceiverCreditUpdateV1 { - uint16_t schema_version{1}, flags{0}; + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; uint32_t qos_class{0}; ReceiverSessionId receiver_session_id; uint64_t epoch{0}, sequence{0}; uint32_t freshness_ttl_ms{0}; std::vector grants; }; + +// One cumulative exchange is sufficient for initial credit, replenishment, +// replay recovery, and terminal release reporting. desired_grant_totals is +// absolute (not a delta), making retries idempotent. +struct ReceiverCreditExchangeRequestV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; + ReceiverSessionId expected_receiver_session_id; + uint64_t epoch{0}; + SenderInstanceId sender_instance_id; + uint32_t qos_class{0}; + uint64_t known_update_sequence{0}; + uint64_t report_sequence{0}; + std::vector released_totals; + std::vector desired_grant_totals; +}; + +struct ReceiverCreditExchangeReplyV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; + SenderInstanceId sender_instance_id; + ReceiverCreditUpdateV1 update; + std::string reply_msg; +}; + enum class CreditUpdateDisposition : uint8_t { Applied, DuplicateOrOld, SequenceGap }; -// Private, sender-side state model. It has no network or Admission integration. +using CreditNowProvider = std::function; + +// Sender-side cumulative ledger. A successful transport handoff consumes a +// grant permanently; terminal completion reports a cumulative release to the +// receiver, which is the only authority allowed to mint a replacement grant. class SenderCreditLedger { public: - explicit SenderCreditLedger(size_t max_entries = 1024) - : max_entries_(max_entries) {} + explicit SenderCreditLedger(size_t max_entries = 1024, + uint32_t max_freshness_ttl_ms = 60000, + CreditNowProvider now_provider = {}); + Status activate(const CreditKey&, uint64_t epoch); - // Removes an exactly matched epoch after the caller has fenced and drained - // (or failed) its transport-owned work. An old cleanup cannot erase a - // reactivated, newer epoch. Status deactivate(const CreditKey&, uint64_t epoch); + Status deactivateSession(const ReceiverSessionId&); Status applyUpdate(const CreditKey&, const ReceiverCreditUpdateV1&, CreditUpdateDisposition&); + Status tryReserve(const CreditKey&, const CreditCharge&); - // Only for work not yet handed to a transport; completions need a new - // grant. + // Only for a failed submit that is proven to have handed no work to a + // transport. Consumption remains cumulative and the same amount is added + // to released, so the receiver (rather than this sender) reclaims it. Status rollbackReservation(const CreditKey&, const CreditCharge&); + // Records terminal completion of committed work. It deliberately does + // not reduce consumed; the receiver must authorize replacement credit. + Status releaseCommitted(const CreditKey&, const CreditCharge&); + + // Builds the next idempotent exchange. The desired totals request enough + // credit for `next_charge` in addition to all locally consumed credit. + Status prepareExchange(const CreditKey&, const CreditCharge& next_charge, + ReceiverCreditExchangeRequestV1&); + // Reports terminal releases without asking the receiver to replenish an + // idle sender. A later prepareExchange() requests credit on demand. + Status prepareReleaseExchange(const CreditKey&, + ReceiverCreditExchangeRequestV1&); + Status available(const CreditKey&, CreditResource, uint64_t&) const; Status consumed(const CreditKey&, CreditResource, uint64_t&) const; + Status released(const CreditKey&, CreditResource, uint64_t&) const; private: struct Entry { - uint64_t epoch{0}, last_sequence{0}; + uint64_t epoch{0}, last_sequence{0}, report_sequence{0}; + uint64_t expires_at_ns{0}; bool has_update{false}; - std::array grants{}, consumed{}; + std::array grants{}, consumed{}, + released{}; }; + static Status resourceIndex(CreditResource, size_t&); static Status normalize(const CreditCharge&, std::array&); + uint64_t nowNanos() const; + Status requireFresh(const Entry&) const; + mutable std::mutex mutex_; const size_t max_entries_; + const uint32_t max_freshness_ttl_ms_; + CreditNowProvider now_provider_; + std::unordered_map entries_; +}; + +struct ReceiverCreditLimits { + std::array capacity{}; + uint32_t freshness_ttl_ms{1000}; + size_t max_senders{1024}; + // Replay history is retained for the receiver session so an old initial + // demand cannot be accepted again after all of its grants are released. + // The bound is fail-closed; entries are reset only by a new receiver + // session, which also fences every old request on the wire. + size_t max_entries{4096}; +}; + +// Receiver-side grant authority. All resources in one desired request are +// granted atomically, and aggregate outstanding authorization can never exceed +// the configured capacity. Expiry is sender-side freshness only: this class +// conservatively does not recycle an unreachable sender without a transport +// fence proving its old DMA is gone. +class ReceiverCreditAuthority { + public: + ReceiverCreditAuthority(ReceiverSessionId session, uint64_t epoch, + ReceiverCreditLimits limits); + + Status status() const { return status_; } + ReceiverCreditAdvertV1 advert() const; + Status exchange(const ReceiverCreditExchangeRequestV1&, + ReceiverCreditExchangeReplyV1&); + Status outstanding(CreditResource, uint64_t&) const; + + private: + struct Entry { + uint64_t last_report_sequence{0}; + uint64_t update_sequence{0}; + std::array grants{}, released{}; + }; + + static Status validateLimits(const ReceiverCreditLimits&); + static Status normalizeCounters(const std::vector&, + std::array&, + std::array&); + void fillReply(const CreditKey&, const Entry&, + ReceiverCreditExchangeReplyV1&) const; + + ReceiverSessionId session_; + uint64_t epoch_{0}; + ReceiverCreditLimits limits_; + Status status_; + mutable std::mutex mutex_; + std::array outstanding_{}; std::unordered_map entries_; }; +inline void to_json(nlohmann::json& j, const ReceiverSessionId& id) { + j = nlohmann::json{{"high", id.high}, {"low", id.low}}; +} +inline void from_json(const nlohmann::json& j, ReceiverSessionId& id) { + j.at("high").get_to(id.high); + j.at("low").get_to(id.low); +} +inline void to_json(nlohmann::json& j, const SenderInstanceId& id) { + j = nlohmann::json{{"high", id.high}, {"low", id.low}}; +} +inline void from_json(const nlohmann::json& j, SenderInstanceId& id) { + j.at("high").get_to(id.high); + j.at("low").get_to(id.low); +} +inline void to_json(nlohmann::json& j, const CreditResource& resource) { + j = static_cast(resource); +} +inline void from_json(const nlohmann::json& j, CreditResource& resource) { + resource = static_cast(j.get()); +} +inline void to_json(nlohmann::json& j, const CreditAmount& amount) { + j = nlohmann::json{{"resource", amount.resource}, + {"grant_total", amount.grant_total}}; +} +inline void from_json(const nlohmann::json& j, CreditAmount& amount) { + j.at("resource").get_to(amount.resource); + j.at("grant_total").get_to(amount.grant_total); +} +inline void to_json(nlohmann::json& j, const CreditCounter& counter) { + j = nlohmann::json{{"resource", counter.resource}, + {"total", counter.total}}; +} +inline void from_json(const nlohmann::json& j, CreditCounter& counter) { + j.at("resource").get_to(counter.resource); + j.at("total").get_to(counter.total); +} +inline void to_json(nlohmann::json& j, const ReceiverCreditAdvertV1& advert) { + j = nlohmann::json{{"schema_version", advert.schema_version}, + {"flags", advert.flags}, + {"receiver_session_id", advert.receiver_session_id}, + {"epoch", advert.epoch}, + {"freshness_ttl_ms", advert.freshness_ttl_ms}, + {"resources", advert.resources}, + {"capacities", advert.capacities}}; +} +inline void from_json(const nlohmann::json& j, ReceiverCreditAdvertV1& advert) { + advert.schema_version = j.value("schema_version", uint16_t{0}); + advert.flags = j.value("flags", uint16_t{0}); + j.at("receiver_session_id").get_to(advert.receiver_session_id); + advert.epoch = j.value("epoch", uint64_t{0}); + advert.freshness_ttl_ms = j.value("freshness_ttl_ms", uint32_t{0}); + advert.resources = j.value("resources", std::vector{}); + advert.capacities = j.value("capacities", std::vector{}); +} +inline void to_json(nlohmann::json& j, const ReceiverCreditUpdateV1& update) { + j = nlohmann::json{{"schema_version", update.schema_version}, + {"flags", update.flags}, + {"qos_class", update.qos_class}, + {"receiver_session_id", update.receiver_session_id}, + {"epoch", update.epoch}, + {"sequence", update.sequence}, + {"freshness_ttl_ms", update.freshness_ttl_ms}, + {"grants", update.grants}}; +} +inline void from_json(const nlohmann::json& j, ReceiverCreditUpdateV1& update) { + update.schema_version = j.value("schema_version", uint16_t{0}); + update.flags = j.value("flags", uint16_t{0}); + update.qos_class = j.value("qos_class", uint32_t{0}); + j.at("receiver_session_id").get_to(update.receiver_session_id); + update.epoch = j.value("epoch", uint64_t{0}); + update.sequence = j.value("sequence", uint64_t{0}); + update.freshness_ttl_ms = j.value("freshness_ttl_ms", uint32_t{0}); + update.grants = j.value("grants", std::vector{}); +} +inline void to_json(nlohmann::json& j, + const ReceiverCreditExchangeRequestV1& request) { + j = nlohmann::json{ + {"schema_version", request.schema_version}, + {"flags", request.flags}, + {"expected_receiver_session_id", request.expected_receiver_session_id}, + {"epoch", request.epoch}, + {"sender_instance_id", request.sender_instance_id}, + {"qos_class", request.qos_class}, + {"known_update_sequence", request.known_update_sequence}, + {"report_sequence", request.report_sequence}, + {"released_totals", request.released_totals}, + {"desired_grant_totals", request.desired_grant_totals}}; +} +inline void from_json(const nlohmann::json& j, + ReceiverCreditExchangeRequestV1& request) { + request.schema_version = j.value("schema_version", uint16_t{0}); + request.flags = j.value("flags", uint16_t{0}); + j.at("expected_receiver_session_id") + .get_to(request.expected_receiver_session_id); + request.epoch = j.value("epoch", uint64_t{0}); + j.at("sender_instance_id").get_to(request.sender_instance_id); + request.qos_class = j.value("qos_class", uint32_t{0}); + request.known_update_sequence = + j.value("known_update_sequence", uint64_t{0}); + request.report_sequence = j.value("report_sequence", uint64_t{0}); + request.released_totals = + j.value("released_totals", std::vector{}); + request.desired_grant_totals = + j.value("desired_grant_totals", std::vector{}); +} +inline void to_json(nlohmann::json& j, + const ReceiverCreditExchangeReplyV1& reply) { + j = nlohmann::json{{"schema_version", reply.schema_version}, + {"flags", reply.flags}, + {"sender_instance_id", reply.sender_instance_id}, + {"update", reply.update}, + {"reply_msg", reply.reply_msg}}; +} +inline void from_json(const nlohmann::json& j, + ReceiverCreditExchangeReplyV1& reply) { + reply.schema_version = j.value("schema_version", uint16_t{0}); + reply.flags = j.value("flags", uint16_t{0}); + j.at("sender_instance_id").get_to(reply.sender_instance_id); + j.at("update").get_to(reply.update); + reply.reply_msg = j.value("reply_msg", ""); +} + } // namespace mooncake::tent #endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h index 63d73df312..1f6d8f846a 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,7 @@ #include "tent/common/concurrent/rw_spinlock.h" #include "tent/common/concurrent/thread_local_storage.h" #include "tent/runtime/metastore.h" +#include "tent/runtime/receiver_credit.h" #include "tent/rpc/rpc.h" #include "tent/runtime/topology.h" #include "tent/thirdparty/nlohmann/json.h" @@ -107,6 +109,11 @@ struct MemorySegmentDesc { // int maps to TransportType enum std::unordered_map transport_attrs; + // Optional runtime-level receiver-credit advertisement. Per-sender + // grants are never published in SegmentDesc; they are obtained through + // the credit exchange RPC. + std::optional receiver_credit; + public: // Get transport attributes for a specific transport type std::string& getTransportAttrs(int transport_type) { @@ -129,6 +136,7 @@ inline void to_json(json& j, const MemorySegmentDesc& m) { {"devices", m.devices}, {"topology", m.topology.toString()}, {"transport_attrs", m.transport_attrs}}; + if (m.receiver_credit) j["receiver_credit"] = *m.receiver_credit; } inline void from_json(const json& j, MemorySegmentDesc& m) { @@ -142,6 +150,11 @@ inline void from_json(const json& j, MemorySegmentDesc& m) { if (j.contains("transport_attrs")) { j.at("transport_attrs").get_to(m.transport_attrs); } + m.receiver_credit.reset(); + if (j.contains("receiver_credit")) { + m.receiver_credit = + j.at("receiver_credit").get(); + } } struct FileSegmentDesc { 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 6251927bda..3e4ac91f2c 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 @@ -33,6 +33,7 @@ #include "tent/common/status.h" #include "tent/common/types.h" #include "tent/runtime/admission_queue.h" +#include "tent/runtime/receiver_credit.h" #include "tent/runtime/transport.h" #include "tent/runtime/transport_selector.h" #include "tent/runtime/hp_tcp_transport_config.h" @@ -463,7 +464,7 @@ class TransferEngineImpl { QueueOwnerKind owner_kind) const; Status prepareSubmit(Batch* batch, const std::vector& request_list, - PreparedSubmit& prepared, + QueueOwnerKind owner_kind, PreparedSubmit& prepared, bool require_post_submit_cancellation = false); Status commitPreparedSubmit(Batch* batch, const PreparedSubmit& prepared); @@ -482,6 +483,43 @@ class TransferEngineImpl { Status dispatchQueuedOwner(QueueOwnerId owner_id); + struct ReceiverCreditTarget; + struct ReceiverCreditReservation; + + Status resolveReceiverCreditTarget(const Request& request, + ReceiverCreditTarget& target) const; + + Status acquireReceiverCredit(const Request& request, + ReceiverCreditReservation& reservation, + bool& acquired); + + Status validateReceiverCreditCapacity( + const Request& request, const ReceiverCreditTarget& target) const; + + void bindReceiverCreditFence(Transport::SubBatchRef, + const ReceiverCreditReservation&) const; + + void bindReceiverCreditFence(Transport::SubBatchRef, const Request&) const; + + void rollbackReceiverCredit(ReceiverCreditReservation& reservation); + + void releaseReceiverCredit(ReceiverCreditReservation& reservation); + + void scheduleReceiverCreditRelease(const CreditKey&, SegmentID, + const std::string&, uint64_t epoch); + + bool tryReportReceiverCreditRelease(const CreditKey&); + + void progressReceiverCreditReleases(); + + void observeReceiverCreditSession(SegmentID, + const ReceiverSessionId& session); + + Status exchangeReceiverCredit( + SegmentID target_id, const std::string& rpc_server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply); + Status markQueuedOwnerSubmitted(QueueOwnerId owner_id); Status finishQueuedOwner(QueueOwnerId owner_id, @@ -578,12 +616,46 @@ class TransferEngineImpl { std::chrono::microseconds progress_fallback_interval{50000}; }; + struct ReceiverCreditConfig { + bool sender_enabled{false}; + bool receiver_enabled{false}; + ReceiverCreditLimits limits{}; + size_t max_sender_entries{4096}; + uint32_t max_freshness_ttl_ms{60000}; + }; + + struct ReceiverCreditTarget { + bool required{false}; + ReceiverCreditAdvertV1 advert{}; + std::string rpc_server_addr; + }; + + struct ReceiverCreditReservation { + bool active{false}; + bool committed{false}; + CreditKey key{}; + CreditCharge charge{}; + SegmentID target_id{LOCAL_SEGMENT_ID}; + std::string rpc_server_addr; + uint64_t epoch{0}; + }; + + struct PendingReceiverCreditRelease { + SegmentID target_id{LOCAL_SEGMENT_ID}; + std::string rpc_server_addr; + uint64_t epoch{0}; + uint32_t attempts{0}; + std::chrono::steady_clock::time_point retry_at{}; + }; + struct QueuedOwnerState { Batch* batch{nullptr}; size_t owner_task_id{0}; std::vector public_task_ids; size_t byte_charge{0}; bool in_dispatch_window{false}; + bool receiver_credit_exempt{false}; + ReceiverCreditReservation credit; }; private: @@ -617,7 +689,14 @@ class TransferEngineImpl { bool notify_rpc_fallback_{true}; bool enable_progress_worker_{false}; RuntimeQueueConfig runtime_queue_config_; + ReceiverCreditConfig receiver_credit_config_; std::unique_ptr runtime_queue_; + SenderInstanceId sender_instance_id_; + std::unique_ptr sender_credit_ledger_; + std::unique_ptr receiver_credit_authority_; + std::unordered_map + pending_credit_releases_; + std::unordered_map receiver_credit_sessions_; std::unordered_map queued_owners_; size_t dispatch_inflight_owners_{0}; size_t dispatch_inflight_bytes_{0}; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index 0299475a7a..00bef306af 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -59,6 +59,13 @@ class Transport { // no pool (default spray). Carried like device_mask, from the matched // SelectionPolicy down to each RdmaTask. std::string qp_pool; + // Runtime receiver-credit authorization bound to the next submitted + // task. Native transports copy this fence into their asynchronous + // work so a refreshed segment cannot post under a different receiver + // session than the one that granted capacity. + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; BatchID progress_batch_id{0}; std::function notify_progress; }; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index de1f7a558b..1c97e6167d 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -107,6 +107,9 @@ struct RdmaTask { // no pool selected: slices spray across all data QPs as before. Resolved // from SelectionResult.qp_pool at task creation. std::string qp_pool; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; std::atomic success_slices{0}; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h index 1f0fc77d9a..d6ed97e340 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h @@ -55,6 +55,9 @@ struct TcpTask { std::atomic transferred_bytes{0}; bool non_replayable_failure{false}; // Published by status_word. uint64_t target_addr = 0; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; TcpTask() = default; TcpTask(TcpTask &&other) noexcept @@ -65,7 +68,10 @@ struct TcpTask { transferred_bytes( other.transferred_bytes.load(std::memory_order_relaxed)), non_replayable_failure(other.non_replayable_failure), - target_addr(other.target_addr) {} + target_addr(other.target_addr), + receiver_credit_session_high(other.receiver_credit_session_high), + receiver_credit_session_low(other.receiver_credit_session_low), + receiver_credit_epoch(other.receiver_credit_epoch) {} TcpTask(const TcpTask &) = delete; TcpTask &operator=(const TcpTask &) = delete; }; @@ -122,8 +128,7 @@ class TcpTransport : public Transport { Status doTransferWithRetry(TcpTask *task); - Status findRemoteSegment(uint64_t dest_addr, uint64_t length, - uint64_t target_id, std::string &rpc_server_addr); + Status findRemoteSegment(const TcpTask &task, std::string &rpc_server_addr); private: bool installed_; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h index 3dfe4533ba..3679bffaf3 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h @@ -10,6 +10,7 @@ #include #include #include +#include #include #include "tent/common/status.h" @@ -116,6 +117,11 @@ class UbEndpoint final : public std::enable_shared_from_this { [[nodiscard]] uint64_t peerGeneration() const noexcept { return peer_generation_.load(std::memory_order_acquire); } + [[nodiscard]] bool peerSupportsCapability( + std::string_view capability) const; + [[nodiscard]] bool peerMatchesReceiverCredit(uint64_t session_high, + uint64_t session_low, + uint64_t epoch) const; [[nodiscard]] uint64_t outstandingWrs() const noexcept { return outstanding_wrs_.load(std::memory_order_relaxed); } diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h index 3b70c3dc21..7428c5b0d8 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h @@ -147,6 +147,9 @@ struct UbSliceSpec { size_t length{0}; size_t request_offset{0}; uint32_t max_retries{0}; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; }; struct UbAttemptToken { diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h index a2e9031fe8..7f3fc4e1dd 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -147,6 +147,8 @@ class UbWorkers final { Status start(); Status stop(); + Status submitBatch(const std::vector& tasks, + uint64_t device_mask = ~0ULL); Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL); Status cancel(const UbTask::Ptr& task); diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 74c19416ad..3a273bf953 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -78,6 +78,30 @@ Status validateLimits(const QueueLimits& limits) { LocalTransferAdmissionQueue::LocalTransferAdmissionQueue(QueueLimits limits) : limits_(limits), limits_status_(validateLimits(limits)) {} +void LocalTransferAdmissionQueue::enqueueOwner(QueueOwnerId owner_id) { + fifo_.push_back(owner_id); + if (!limits_.deadline_aware) return; + + // Promotion and defer can both disturb the deque's prior global order. + // Re-establish the explicit (round, EDF) ordering on admission instead of + // passing a potentially rotated range to an ordered insertion algorithm. + // stable_sort preserves FIFO order for equal deadlines. + std::stable_sort( + fifo_.begin(), fifo_.end(), [this](QueueOwnerId lhs, QueueOwnerId rhs) { + const auto lhs_it = owners_.find(lhs); + const auto rhs_it = owners_.find(rhs); + if (lhs_it == owners_.end()) return false; + if (rhs_it == owners_.end()) return true; + const auto& lhs_owner = lhs_it->second; + const auto& rhs_owner = rhs_it->second; + if (lhs_owner.scheduling_round != rhs_owner.scheduling_round) { + return lhs_owner.scheduling_round < rhs_owner.scheduling_round; + } + return deadlineKey(lhs_owner.request.deadline_ns) < + deadlineKey(rhs_owner.request.deadline_ns); + }); +} + Status LocalTransferAdmissionQueue::tryAdmit( const QueueSubmit& submit, std::vector& admitted_owner_ids) { admitted_owner_ids.clear(); @@ -181,6 +205,17 @@ Status LocalTransferAdmissionQueue::tryAdmit( owner.kind = owner_input.kind; owner.transport = owner_input.transport; owner.degradation_eligible = owner_input.degradation_eligible; + if (limits_.deadline_aware) { + for (const auto queued_id : fifo_) { + auto queued_it = owners_.find(queued_id); + if (queued_it != owners_.end() && + queued_it->second.state == QueueState::Queued) { + owner.scheduling_round = + std::max(owner.scheduling_round, + queued_it->second.scheduling_round); + } + } + } owners_.emplace(owner_id, owner); public_to_owner_[{submit.batch_token, owner_input.owner_task_id}] = @@ -189,26 +224,8 @@ Status LocalTransferAdmissionQueue::tryAdmit( public_to_owner_[{submit.batch_token, derived_task_id}] = owner_id; } // RFC #2519 step 2: keep fifo_ ordered on admission so pickForDispatch - // never has to re-sort. Default (deadline_aware == false) appends in - // strict FIFO. When deadline-aware, insert at the earliest-deadline- - // first position; upper_bound places a new owner *after* existing - // owners with the same deadline, preserving FIFO order among ties. - if (limits_.deadline_aware) { - const uint64_t key = deadlineKey(owner.request.deadline_ns); - auto pos = std::upper_bound( - fifo_.begin(), fifo_.end(), key, - [this](uint64_t k, QueueOwnerId id) { - auto it = owners_.find(id); - uint64_t d = - (it == owners_.end()) - ? std::numeric_limits::max() - : deadlineKey(it->second.request.deadline_ns); - return k < d; - }); - fifo_.insert(pos, owner_id); - } else { - fifo_.push_back(owner_id); - } + // never has to re-sort. + enqueueOwner(owner_id); admitted_owner_ids.push_back(owner_id); } @@ -271,16 +288,26 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( .count())) : 0; - // Deadline proximity promotion: partition fifo_ so owners with critical - // slack (deadline approaching within promotion_slack_ns) appear before - // owners with comfortable slack or no deadline. stable_partition preserves - // relative EDF order within each group. + // Deadline proximity promotion applies only inside the earliest + // scheduling round. A resource-starved owner that defer() moved to the + // next round must not jump back ahead of current-round waiters merely + // because its deadline is urgent. if (promotion_enabled) { + uint64_t current_round = std::numeric_limits::max(); + for (const auto owner_id : fifo_) { + auto owner_it = owners_.find(owner_id); + if (owner_it != owners_.end() && + owner_it->second.state == QueueState::Queued) { + current_round = + std::min(current_round, owner_it->second.scheduling_round); + } + } std::stable_partition(fifo_.begin(), fifo_.end(), [&](QueueOwnerId id) { auto it = owners_.find(id); if (it == owners_.end() || it->second.state != QueueState::Queued) { return false; } + if (it->second.scheduling_round != current_round) return false; const uint64_t dl = it->second.request.deadline_ns; if (dl == 0 || dl <= now_ns) return false; return (dl - now_ns) < limits_.promotion_slack_ns; @@ -383,6 +410,35 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( return picked; } +Status LocalTransferAdmissionQueue::defer(QueueOwnerId owner_id) { + if (owner_id == 0) { + return Status::InvalidArgument("invalid queue owner id" LOC_MARK); + } + + 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("queue owner is not dispatching" LOC_MARK); + } + + // A resource-deferred EDF owner advances one scheduling round. It yields + // to every owner already waiting in this round while retaining EDF order + // among owners that yielded into the next one. + if (limits_.deadline_aware) { + if (owner_it->second.scheduling_round == + std::numeric_limits::max()) { + return Status::InternalError( + "admission scheduling round exhausted" LOC_MARK); + } + ++owner_it->second.scheduling_round; + } + owner_it->second.state = QueueState::Queued; + enqueueOwner(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 0d1624f19b..71cf343bf6 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -52,6 +53,113 @@ thread_local const ControlService* ControlService::active_notify_service_ = nullptr; thread_local CoroRpcAgent tl_rpc_agent; +namespace { + +constexpr uint64_t kDelegatedCreditProtocolMagic = 0x54454e54444c4732ULL; + +struct ParsedXferData { + uint64_t peer_mem_addr{0}; + size_t length{0}; + size_t header_size{0}; + bool version_two{false}; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; +}; + +Status parseXferData(const std::string_view& request, bool carries_payload, + ParsedXferData& parsed) { + if (request.size() < sizeof(LegacyXferDataDesc)) { + return Status::InvalidArgument("request too short" LOC_MARK); + } + LegacyXferDataDesc legacy; + memcpy(&legacy, request.data(), sizeof(legacy)); + const uint64_t length = le64toh(legacy.length); + if (length > std::numeric_limits::max()) { + return Status::InvalidArgument("transfer length overflow" LOC_MARK); + } + + parsed = ParsedXferData{}; + parsed.peer_mem_addr = le64toh(legacy.peer_mem_addr); + parsed.length = static_cast(length); + parsed.header_size = sizeof(LegacyXferDataDesc); + + if (request.size() >= sizeof(XferDataDesc)) { + XferDataDesc current; + memcpy(¤t, request.data(), sizeof(current)); + const bool magic_matches = + le64toh(current.protocol_magic) == kXferDataProtocolMagic; + const bool size_matches = + carries_payload + ? parsed.length <= std::numeric_limits::max() - + sizeof(XferDataDesc) && + request.size() == sizeof(XferDataDesc) + parsed.length + : request.size() == sizeof(XferDataDesc); + if (magic_matches && size_matches) { + parsed.header_size = sizeof(XferDataDesc); + parsed.version_two = true; + parsed.peer_mem_addr = le64toh(current.peer_mem_addr); + parsed.receiver_credit_session_high = + le64toh(current.receiver_credit_session_high); + parsed.receiver_credit_session_low = + le64toh(current.receiver_credit_session_low); + parsed.receiver_credit_epoch = + le64toh(current.receiver_credit_epoch); + } + } + + if (carries_payload) { + if (parsed.length > + std::numeric_limits::max() - parsed.header_size || + request.size() != parsed.header_size + parsed.length) { + return Status::InvalidArgument( + "invalid transfer request size" LOC_MARK); + } + } else if (request.size() != parsed.header_size) { + return Status::InvalidArgument( + "invalid transfer request header size" LOC_MARK); + } + return Status::OK(); +} + +Status validateXferReceiverCreditFence(const SegmentDescRef& local_desc, + const ParsedXferData& request) { + const bool has_fence = request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + if (!local_desc || local_desc->type != SegmentType::Memory) { + return has_fence ? Status::InvalidEntry( + "receiver-credit target is not memory" LOC_MARK) + : Status::OK(); + } + + const auto& advert = local_desc->getMemory().receiver_credit; + if (!advert) { + return has_fence + ? Status::InvalidEntry( + "receiver-credit session is unavailable" LOC_MARK) + : Status::OK(); + } + if (advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.empty() || advert->epoch == 0) { + return Status::InvalidEntry( + "invalid local receiver-credit advertisement" LOC_MARK); + } + if (!request.version_two || + request.receiver_credit_session_high != + advert->receiver_session_id.high || + request.receiver_credit_session_low != + advert->receiver_session_id.low || + request.receiver_credit_epoch != advert->epoch) { + return Status::InvalidEntry( + "receiver-credit session fence mismatch" LOC_MARK); + } + return Status::OK(); +} + +} // namespace + Status ControlClient::getSegmentDesc(const std::string& server_addr, std::string& response) { std::string request; @@ -113,29 +221,83 @@ Status ControlClient::bootstrapUb(const std::string& server_addr, return Status::OK(); } +Status ControlClient::exchangeReceiverCredit( + const std::string& server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& response) { + std::string request_raw, response_raw; + json encoded = request; + request_raw = encoded.dump(); + CHECK_STATUS(tl_rpc_agent.call(server_addr, ExchangeReceiverCredit, + request_raw, response_raw)); + try { + response = + json::parse(response_raw).get(); + } catch (const std::exception& error) { + return Status::MalformedJson( + std::string("Malformed receiver-credit response: ") + error.what() + + LOC_MARK); + } + if (response.schema_version != kReceiverCreditProtocolVersion || + response.flags != 0) { + return Status::RpcServiceError( + "Unsupported receiver-credit response version" LOC_MARK); + } + if (!response.reply_msg.empty()) { + return Status::RpcServiceError(response.reply_msg); + } + return Status::OK(); +} + Status ControlClient::sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length) { - // Check before addition/allocation or reading the caller's source buffer. + size_t length, + uint64_t receiver_credit_session_high, + uint64_t receiver_credit_session_low, + uint64_t receiver_credit_epoch) { if (length > kTcpMaxWriteBytes) return Status::InvalidArgument("TCP WRITE exceeds RPC payload limit"); - std::string response; - XferDataDesc desc{htole64(peer_mem_addr), htole64(length)}; - std::string request; + std::string request, response; + const bool has_any_fence = receiver_credit_session_high != 0 || + receiver_credit_session_low != 0 || + receiver_credit_epoch != 0; + const bool has_complete_fence = (receiver_credit_session_high != 0 || + receiver_credit_session_low != 0) && + receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial SendData receiver-credit fence" LOC_MARK); + } + if (length > std::numeric_limits::max() - sizeof(XferDataDesc)) { + return Status::InvalidArgument( + "SendData request size overflow" LOC_MARK); + } + size_t header_size = sizeof(LegacyXferDataDesc); request.reserve(sizeof(XferDataDesc) + length); - request.resize(sizeof(XferDataDesc)); - memcpy(request.data(), &desc, sizeof(desc)); + if (has_complete_fence) { + const XferDataDesc desc{htole64(kXferDataV2LegacyRejectAddress), + htole64(static_cast(length)), + htole64(kXferDataProtocolMagic), + htole64(peer_mem_addr), + htole64(receiver_credit_session_high), + htole64(receiver_credit_session_low), + htole64(receiver_credit_epoch)}; + header_size = sizeof(desc); + request.resize(header_size); + memcpy(request.data(), &desc, sizeof(desc)); + } else { + const LegacyXferDataDesc desc{htole64(peer_mem_addr), + htole64(static_cast(length))}; + request.resize(header_size); + memcpy(request.data(), &desc, sizeof(desc)); + } auto& loader = Platform::getLoader(); if (loader.getMemoryType(local_mem_addr) == MTYPE_CPU) { - // Single copy into the RPC attachment; avoids the resize() zero-fill - // and the extra copy in call(). request.append(reinterpret_cast(local_mem_addr), length); } else { - // resize() zero-fills the payload, so an unchecked copy failure would - // ship zeros that the peer stores successfully and reports COMPLETED. - request.resize(sizeof(XferDataDesc) + length); + request.resize(header_size + length); CHECK_STATUS( - loader.copy(request.data() + sizeof(desc), local_mem_addr, length)); + loader.copy(request.data() + header_size, local_mem_addr, length)); } auto status = tl_rpc_agent.callOwned(server_addr, SendData, std::move(request), response); @@ -146,19 +308,63 @@ Status ControlClient::sendData(const std::string& server_addr, Status ControlClient::recvData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length) { + size_t length, + uint64_t receiver_credit_session_high, + uint64_t receiver_credit_session_low, + uint64_t receiver_credit_epoch) { if (length > kTcpMaxReadBytes) return Status::InvalidArgument("TCP READ exceeds RPC payload limit"); std::string request, response; - XferDataDesc desc{htole64(peer_mem_addr), htole64(length)}; - request.resize(sizeof(XferDataDesc)); - memcpy(&request[0], &desc, sizeof(desc)); + const bool has_any_fence = receiver_credit_session_high != 0 || + receiver_credit_session_low != 0 || + receiver_credit_epoch != 0; + const bool has_complete_fence = (receiver_credit_session_high != 0 || + receiver_credit_session_low != 0) && + receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial RecvData receiver-credit fence" LOC_MARK); + } + if (has_complete_fence) { + const XferDataDesc desc{htole64(kXferDataV2LegacyRejectAddress), + htole64(static_cast(length)), + htole64(kXferDataProtocolMagic), + htole64(peer_mem_addr), + htole64(receiver_credit_session_high), + htole64(receiver_credit_session_low), + htole64(receiver_credit_epoch)}; + request.resize(sizeof(desc)); + memcpy(request.data(), &desc, sizeof(desc)); + } else { + const LegacyXferDataDesc desc{htole64(peer_mem_addr), + htole64(static_cast(length))}; + request.resize(sizeof(desc)); + memcpy(request.data(), &desc, sizeof(desc)); + } auto status = tl_rpc_agent.call(server_addr, RecvData, request, response); if (!status.ok()) return status; - if (response.size() != length) + if (!has_complete_fence) { + if (response.size() != length) { + return Status::RpcServiceError( + "RecvData failed: invalid or rejected legacy response"); + } + return Platform::getLoader().copy(local_mem_addr, response.data(), + length); + } + if (length > + std::numeric_limits::max() - sizeof(XferDataReplyHeader) || + response.size() != sizeof(XferDataReplyHeader) + length) { + return Status::RpcServiceError( + "RecvData failed: invalid or rejected response"); + } + XferDataReplyHeader header; + memcpy(&header, response.data(), sizeof(header)); + if (le64toh(header.protocol_magic) != kXferDataProtocolMagic) { return Status::RpcServiceError( - response.empty() ? "RecvData failed: empty response" : response); - return Platform::getLoader().copy(local_mem_addr, response.data(), length); + "RecvData failed: unsupported response version"); + } + return Platform::getLoader().copy( + local_mem_addr, response.data() + sizeof(XferDataReplyHeader), length); } inline void to_json(nlohmann::json& j, const Notification& n) { @@ -184,11 +390,26 @@ Status ControlClient::probe(const std::string& server_addr) { } inline void to_json(json& j, const Request& r) { - j = json{{"opcode", r.opcode == Request::READ ? "READ" : "WRITE"}, - {"source", reinterpret_cast(r.source)}, - {"target_id", r.target_id}, - {"target_offset", r.target_offset}, - {"length", r.length}}; + const bool has_complete_fence = (r.receiver_credit_session_high != 0 || + r.receiver_credit_session_low != 0) && + r.receiver_credit_epoch != 0; + j = json{ + {"opcode", r.opcode == Request::READ ? "READ" : "WRITE"}, + {"source", reinterpret_cast(r.source)}, + {"target_id", has_complete_fence ? std::numeric_limits::max() + : r.target_id}, + {"target_offset", has_complete_fence + ? std::numeric_limits::max() + : r.target_offset}, + {"length", r.length}, + {"receiver_credit_session_high", r.receiver_credit_session_high}, + {"receiver_credit_session_low", r.receiver_credit_session_low}, + {"receiver_credit_epoch", r.receiver_credit_epoch}}; + if (has_complete_fence) { + j["delegate_protocol_magic"] = kDelegatedCreditProtocolMagic; + j["actual_target_id"] = r.target_id; + j["actual_target_offset"] = r.target_offset; + } } inline void from_json(const json& j, Request& r) { @@ -201,13 +422,35 @@ inline void from_json(const json& j, Request& r) { throw std::runtime_error("Invalid opcode"); r.source = reinterpret_cast(j.at("source").get()); - r.target_id = j.at("target_id").get(); - r.target_offset = j.at("target_offset").get(); + const bool fenced_delegate = + j.value("delegate_protocol_magic", uint64_t{0}) == + kDelegatedCreditProtocolMagic; + r.target_id = fenced_delegate ? j.at("actual_target_id").get() + : j.at("target_id").get(); + r.target_offset = fenced_delegate + ? j.at("actual_target_offset").get() + : j.at("target_offset").get(); r.length = j.at("length").get(); + r.receiver_credit_session_high = + j.value("receiver_credit_session_high", uint64_t{0}); + r.receiver_credit_session_low = + j.value("receiver_credit_session_low", uint64_t{0}); + r.receiver_credit_epoch = j.value("receiver_credit_epoch", uint64_t{0}); } Status ControlClient::delegate(const std::string& server_addr, const Request& request) { + const bool has_any_fence = request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + const bool has_complete_fence = + (request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0) && + request.receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial delegated receiver-credit fence" LOC_MARK); + } std::string request_raw, response_raw; json j = request; request_raw = j.dump(); @@ -309,6 +552,11 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onBootstrapUb(request, response); }); + rpc_server_->registerFunction( + ExchangeReceiverCredit, + [this](const std::string_view& request, std::string& response) { + onExchangeReceiverCredit(request, response); + }); rpc_server_->registerFunction( SendData, [this](const std::string_view& request, std::string& response) { @@ -364,6 +612,10 @@ ControlService::~ControlService() { std::lock_guard lock(ub_bootstrap_callback_mutex_); ub_bootstrap_callback_ = {}; } + { + std::lock_guard lock(receiver_credit_callback_mutex_); + receiver_credit_callback_ = {}; + } rpc_server_.reset(); } @@ -507,37 +759,62 @@ void ControlService::onBootstrapUb(const std::string_view& request, response = j.dump(); } +void ControlService::onExchangeReceiverCredit(const std::string_view& request, + std::string& response) { + ReceiverCreditExchangeReplyV1 reply; + try { + auto parsed = json::parse(std::string(request)) + .get(); + Status status = Status::InvalidEntry( + "receiver credit callback is not registered" LOC_MARK); + { + // Serialize callback replacement with invocation so teardown + // cannot destroy the authority while an exchange is in flight. + std::lock_guard lock(receiver_credit_callback_mutex_); + if (receiver_credit_callback_) { + status = receiver_credit_callback_(parsed, reply); + } + } + if (!status.ok()) reply.reply_msg = status.ToString(); + } catch (const std::exception& error) { + reply.reply_msg = + std::string("Malformed receiver-credit request: ") + error.what(); + } + json encoded = reply; + response = encoded.dump(); +} + void ControlService::onSendData(const std::string_view& request, std::string& response) { - if (request.size() < sizeof(XferDataDesc)) { - response = "SendData failed: request too short"; + ParsedXferData parsed; + auto status = parseXferData(request, true, parsed); + if (!status.ok()) { + response = "SendData failed: " + status.ToString(); return; } - XferDataDesc* desc = (XferDataDesc*)request.data(); auto local_desc = manager_->getLocal(); - auto peer_mem_addr = le64toh(desc->peer_mem_addr); - auto length = le64toh(desc->length); + status = validateXferReceiverCreditFence(local_desc, parsed); + if (!status.ok()) { + response = "SendData failed: " + status.ToString(); + return; + } - // Validate request size to prevent buffer over-read - if (length > kTcpMaxWriteBytes || - length > request.size() - sizeof(XferDataDesc)) { + if (parsed.length > kTcpMaxWriteBytes) { response = "SendData failed: invalid request size"; return; } - - if (auto* buffer = local_desc->findBuffer(peer_mem_addr, length)) { + if (auto* buffer = + local_desc->findBuffer(parsed.peer_mem_addr, parsed.length)) { if (buffer->permission != kGlobalReadWrite) { response = "SendData failed: remote write permission denied"; return; } - auto status = - Platform::getLoader().copy((void*)peer_mem_addr, &desc[1], length); - if (!status.ok()) { - // A non-empty response is interpreted as an RPC error by the - // client (see ControlClient::sendData). Without this the sender's - // transfer would be reported COMPLETED even though the destination - // buffer was never written. - response = "SendData failed: copy: " + status.ToString(); + auto copy_status = Platform::getLoader().copy( + reinterpret_cast(parsed.peer_mem_addr), + const_cast(request.data() + parsed.header_size), + parsed.length); + if (!copy_status.ok()) { + response = "SendData failed: copy: " + copy_status.ToString(); } } else { response = "SendData failed: target address not in registered buffer"; @@ -546,46 +823,59 @@ void ControlService::onSendData(const std::string_view& request, void ControlService::onRecvData(const std::string_view& request, std::string& response) { - if (request.size() < sizeof(XferDataDesc)) { - response = "RecvData failed: request too short"; + ParsedXferData parsed; + auto status = parseXferData(request, false, parsed); + // Legacy readers accept any response of exactly the requested length. + // Fenced readers also require the reply magic; pad errors to keep both + // response formats unambiguous, including parse and session failures. + const size_t reply_header_size = + parsed.version_two ? sizeof(XferDataReplyHeader) : 0; + auto fail = [&response, &parsed, reply_header_size](std::string message) { + response = std::move(message); + if (parsed.length <= + std::numeric_limits::max() - reply_header_size && + response.size() == reply_header_size + parsed.length) { + response.push_back(' '); + } + }; + if (!status.ok()) { + fail("RecvData failed: " + status.ToString()); return; } - XferDataDesc* desc = (XferDataDesc*)request.data(); auto local_desc = manager_->getLocal(); - auto peer_mem_addr = le64toh(desc->peer_mem_addr); - auto length = le64toh(desc->length); - - // The client accepts any response of exactly `length` bytes as payload (see - // ControlClient::recvData), so an error of that size must be padded or it - // would be copied into the caller's buffer and reported as success. - auto fail = [&response, length](std::string message) { - response = std::move(message); - if (response.size() == length) response.push_back(' '); - }; - - // Validate length to prevent DoS via excessive memory allocation - if (length > kTcpMaxReadBytes) { + status = validateXferReceiverCreditFence(local_desc, parsed); + if (!status.ok()) { + fail("RecvData failed: " + status.ToString()); + return; + } + if (parsed.length > kTcpMaxReadBytes) { fail("RecvData failed: length exceeds maximum allowed"); return; } - - if (auto* buffer = local_desc->findBuffer(peer_mem_addr, length)) { + if (auto* buffer = + local_desc->findBuffer(parsed.peer_mem_addr, parsed.length)) { if (buffer->permission != kGlobalReadOnly && buffer->permission != kGlobalReadWrite) { fail("RecvData failed: remote read permission denied"); return; } auto& loader = Platform::getLoader(); - if (loader.getMemoryType((void*)peer_mem_addr) == MTYPE_CPU) { - // assign() skips the resize() zero-fill pass. - response.assign(reinterpret_cast(peer_mem_addr), - length); + auto* source = reinterpret_cast(parsed.peer_mem_addr); + response.clear(); + if (parsed.version_two) { + const XferDataReplyHeader reply{htole64(kXferDataProtocolMagic)}; + response.assign(reinterpret_cast(&reply), + sizeof(reply)); + } + if (loader.getMemoryType(source) == MTYPE_CPU) { + response.append(reinterpret_cast(source), + parsed.length); } else { - response.resize(length); - auto status = - loader.copy(response.data(), (void*)peer_mem_addr, length); - if (!status.ok()) { - fail("RecvData failed: copy: " + status.ToString()); + response.resize(reply_header_size + parsed.length); + auto copy_status = loader.copy(response.data() + reply_header_size, + source, parsed.length); + if (!copy_status.ok()) { + fail("RecvData failed: copy: " + copy_status.ToString()); } } } else { diff --git a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp index 545eca753d..3936f37361 100644 --- a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp @@ -233,6 +233,11 @@ BatchID ProxyManager::submitCrossStage(const Request& request, inter_stage.length = chunk_length; inter_stage.target_id = request.target_id; inter_stage.target_offset = remote_stage_buffer; + inter_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + inter_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + inter_stage.receiver_credit_epoch = request.receiver_credit_epoch; auto batch = impl_->allocateBatch(1); auto status = impl_->submitStagingTransfer(batch, {inter_stage}); if (!status.ok()) { @@ -283,6 +288,11 @@ Status ProxyManager::waitRemoteStage(const std::string& server_addr, remote_stage.length = chunk_length; remote_stage.target_id = LOCAL_SEGMENT_ID; remote_stage.target_offset = request.target_offset + offset; + remote_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + remote_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + remote_stage.receiver_credit_epoch = request.receiver_credit_epoch; return ControlClient::delegate(server_addr, remote_stage); } @@ -296,6 +306,11 @@ void ProxyManager::submitRemoteStage( remote_stage.length = chunk_length; remote_stage.target_id = LOCAL_SEGMENT_ID; remote_stage.target_offset = request.target_offset + offset; + remote_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + remote_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + remote_stage.receiver_credit_epoch = request.receiver_credit_epoch; operation = std::make_shared( server_addr, remote_stage_buffer, [server_addr, remote_stage_buffer] { ControlClient::unpinStageBufferAsync( diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp index 5ae6c14a2e..00bf199c6e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp @@ -3,176 +3,622 @@ #include "tent/runtime/receiver_credit.h" +#include +#include +#include +#include + namespace mooncake::tent { +namespace { + +uint64_t steadyNowNs() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} -size_t CreditKeyHash::operator()(const CreditKey& k) const noexcept { - size_t h = std::hash{}(k.receiver_session.high); - auto mix = [&h](uint64_t v) { - h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + - (h >> 2); - }; - mix(k.receiver_session.low); - mix(k.sender_peer); - mix(k.qos_class); - return h; -} - -Status SenderCreditLedger::resourceIndex(CreditResource r, size_t& i) { - auto raw = static_cast(r); - if (raw < 1 || raw > kCreditResourceCount) +Status indexFor(CreditResource resource, size_t& index) { + const auto raw = static_cast(resource); + if (raw < 1 || raw > kCreditResourceCount) { return Status::InvalidArgument("unknown credit resource" LOC_MARK); - i = raw - 1; + } + index = raw - 1; return Status::OK(); } +void mixHash(size_t& hash, uint64_t value) { + hash ^= std::hash{}(value) + 0x9e3779b97f4a7c15ULL + (hash << 6) + + (hash >> 2); +} + +struct SenderInstanceIdHash { + size_t operator()(const SenderInstanceId& id) const noexcept { + size_t hash = std::hash{}(id.high); + mixHash(hash, id.low); + return hash; + } +}; + +} // namespace + +size_t CreditKeyHash::operator()(const CreditKey& key) const noexcept { + size_t hash = std::hash{}(key.receiver_session.high); + mixHash(hash, key.receiver_session.low); + mixHash(hash, key.sender_instance.high); + mixHash(hash, key.sender_instance.low); + mixHash(hash, key.qos_class); + return hash; +} + +SenderCreditLedger::SenderCreditLedger(size_t max_entries, + uint32_t max_freshness_ttl_ms, + CreditNowProvider now_provider) + : max_entries_(max_entries), + max_freshness_ttl_ms_(max_freshness_ttl_ms), + now_provider_(std::move(now_provider)) {} + +Status SenderCreditLedger::resourceIndex(CreditResource resource, + size_t& index) { + return indexFor(resource, index); +} + Status SenderCreditLedger::normalize( - const CreditCharge& c, std::array& out) { - out.fill(0); - if (c.resources.empty()) + const CreditCharge& charge, + std::array& normalized) { + normalized.fill(0); + if (charge.resources.empty()) { return Status::InvalidArgument("empty credit charge" LOC_MARK); - for (auto [r, amount] : c.resources) { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); - if (!amount || out[i]) + } + for (const auto& [resource, amount] : charge.resources) { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); + if (amount == 0 || normalized[index] != 0) { return Status::InvalidArgument( "zero or duplicate credit charge" LOC_MARK); - out[i] = amount; + } + normalized[index] = amount; } return Status::OK(); } -Status SenderCreditLedger::activate(const CreditKey& k, uint64_t epoch) { - if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); +uint64_t SenderCreditLedger::nowNanos() const { + return now_provider_ ? now_provider_() : steadyNowNs(); +} + +Status SenderCreditLedger::requireFresh(const Entry& entry) const { + if (!entry.has_update || entry.expires_at_ns == 0 || + nowNanos() >= entry.expires_at_ns) { + return Status::InvalidEntry("receiver credit is stale" LOC_MARK); + } + return Status::OK(); +} + +Status SenderCreditLedger::activate(const CreditKey& key, uint64_t epoch) { + if (key.receiver_session.empty() || key.sender_instance.empty() || + epoch == 0) { + return Status::InvalidArgument( + "invalid credit activation identity" LOC_MARK); + } std::lock_guard lock(mutex_); - auto existing = entries_.find(k); + auto existing = entries_.find(key); if (existing != entries_.end()) { - if (epoch < existing->second.epoch) + if (epoch < existing->second.epoch) { return Status::InvalidEntry("stale credit activation" LOC_MARK); + } if (epoch == existing->second.epoch) return Status::OK(); Entry replacement; replacement.epoch = epoch; existing->second = replacement; return Status::OK(); - } else if (entries_.size() >= max_entries_) { + } + if (entries_.size() >= max_entries_) { return Status::TooManyRequests("credit ledger entry limit" LOC_MARK); } - Entry e; - e.epoch = epoch; - entries_.emplace(k, e); + Entry entry; + entry.epoch = epoch; + entries_.emplace(key, entry); return Status::OK(); } -Status SenderCreditLedger::deactivate(const CreditKey& k, uint64_t epoch) { - if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); +Status SenderCreditLedger::deactivate(const CreditKey& key, uint64_t epoch) { + if (epoch == 0) { + return Status::InvalidArgument("zero credit epoch" LOC_MARK); + } std::lock_guard lock(mutex_); - auto existing = entries_.find(k); - if (existing == entries_.end()) return Status::OK(); // idempotent cleanup - if (existing->second.epoch != epoch) + auto existing = entries_.find(key); + if (existing == entries_.end()) return Status::OK(); + if (existing->second.epoch != epoch) { return Status::InvalidEntry("credit cleanup epoch mismatch" LOC_MARK); + } entries_.erase(existing); return Status::OK(); } -Status SenderCreditLedger::applyUpdate(const CreditKey& k, - const ReceiverCreditUpdateV1& u, +Status SenderCreditLedger::deactivateSession(const ReceiverSessionId& session) { + if (session.empty()) { + return Status::InvalidArgument( + "empty receiver credit session cleanup" LOC_MARK); + } + std::lock_guard lock(mutex_); + for (auto it = entries_.begin(); it != entries_.end();) { + if (it->first.receiver_session == session) { + it = entries_.erase(it); + } else { + ++it; + } + } + return Status::OK(); +} + +Status SenderCreditLedger::applyUpdate(const CreditKey& key, + const ReceiverCreditUpdateV1& update, CreditUpdateDisposition& disposition) { - if (u.schema_version != 1 || !u.epoch || !u.sequence) + if (update.schema_version != kReceiverCreditProtocolVersion || + update.flags != 0 || update.epoch == 0 || update.sequence == 0 || + update.freshness_ttl_ms == 0 || + update.freshness_ttl_ms > max_freshness_ttl_ms_) { return Status::InvalidArgument("invalid credit update header" LOC_MARK); - if (!(u.receiver_session_id == k.receiver_session) || - u.qos_class != k.qos_class || u.grants.size() > kCreditResourceCount) + } + if (!(update.receiver_session_id == key.receiver_session) || + update.qos_class != key.qos_class || + update.grants.size() > kCreditResourceCount) { return Status::InvalidArgument("credit update identity/size" LOC_MARK); + } + std::array proposed{}; std::array present{}; - for (auto a : u.grants) { - size_t i = 0; - CHECK_STATUS(resourceIndex(a.resource, i)); - if (present[i]) + for (const auto& amount : update.grants) { + size_t index = 0; + CHECK_STATUS(resourceIndex(amount.resource, index)); + if (present[index]) { return Status::InvalidArgument("duplicate grant resource" LOC_MARK); - present[i] = true; - proposed[i] = a.grant_total; + } + present[index] = true; + proposed[index] = amount.grant_total; } + std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || it->second.epoch != u.epoch) + auto it = entries_.find(key); + if (it == entries_.end() || it->second.epoch != update.epoch) { return Status::InvalidEntry("inactive or stale credit epoch" LOC_MARK); - auto& e = it->second; - if (e.has_update && u.sequence <= e.last_sequence) { + } + auto& entry = it->second; + if (entry.has_update && update.sequence <= entry.last_sequence) { disposition = CreditUpdateDisposition::DuplicateOrOld; return Status::OK(); } - // `grants` is a partial cumulative update: each resource present in this - // message replaces that resource's cumulative grant total, while omitted - // resources retain their previous totals. This lets the receiver refresh - // only the resources whose available capacity changed. - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (present[i] && - (proposed[i] < e.grants[i] || proposed[i] < e.consumed[i])) + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (present[i] && (proposed[i] < entry.grants[i] || + proposed[i] < entry.consumed[i])) { return Status::InvalidArgument( "decreasing or under-consumed grant" LOC_MARK); - bool gap = e.has_update && u.sequence > e.last_sequence + 1; - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (present[i]) e.grants[i] = proposed[i]; - e.last_sequence = u.sequence; - e.has_update = true; + } + } + + const bool gap = + entry.has_update && update.sequence > entry.last_sequence + 1; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (present[i]) entry.grants[i] = proposed[i]; + } + entry.last_sequence = update.sequence; + entry.has_update = true; + const uint64_t ttl_ns = + static_cast(update.freshness_ttl_ms) * 1000000ULL; + const uint64_t now = nowNanos(); + entry.expires_at_ns = ttl_ns > std::numeric_limits::max() - now + ? std::numeric_limits::max() + : now + ttl_ns; disposition = gap ? CreditUpdateDisposition::SequenceGap : CreditUpdateDisposition::Applied; return Status::OK(); } -Status SenderCreditLedger::tryReserve(const CreditKey& k, - const CreditCharge& c) { - std::array n; - CHECK_STATUS(normalize(c, n)); +Status SenderCreditLedger::tryReserve(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || !it->second.has_update) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit unavailable" LOC_MARK); - auto& e = it->second; - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (e.consumed[i] > e.grants[i] || n[i] > e.grants[i] - e.consumed[i]) + } + CHECK_STATUS(requireFresh(it->second)); + auto& entry = it->second; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (entry.consumed[i] > entry.grants[i] || + normalized[i] > entry.grants[i] - entry.consumed[i]) { return Status::TooManyRequests("insufficient credit" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) e.consumed[i] += n[i]; + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + entry.consumed[i] += normalized[i]; + } return Status::OK(); } -Status SenderCreditLedger::rollbackReservation(const CreditKey& k, - const CreditCharge& c) { - std::array n; - CHECK_STATUS(normalize(c, n)); +Status SenderCreditLedger::rollbackReservation(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end()) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit session inactive" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (n[i] > it->second.consumed[i]) + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (it->second.released[i] > it->second.consumed[i] || + normalized[i] > it->second.consumed[i] - it->second.released[i]) { return Status::InvalidArgument( "credit rollback underflow" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) - it->second.consumed[i] -= n[i]; + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (normalized[i] > + std::numeric_limits::max() - it->second.released[i]) { + return Status::InvalidArgument("credit rollback overflow" LOC_MARK); + } + it->second.released[i] += normalized[i]; + } + return Status::OK(); +} + +Status SenderCreditLedger::releaseCommitted(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (it->second.released[i] > it->second.consumed[i] || + normalized[i] > it->second.consumed[i] - it->second.released[i]) { + return Status::InvalidArgument( + "credit release exceeds committed consumption" LOC_MARK); + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (normalized[i] > + std::numeric_limits::max() - it->second.released[i]) { + return Status::InvalidArgument("credit release overflow" LOC_MARK); + } + it->second.released[i] += normalized[i]; + } + return Status::OK(); +} + +Status SenderCreditLedger::prepareExchange( + const CreditKey& key, const CreditCharge& next_charge, + ReceiverCreditExchangeRequestV1& request) { + std::array normalized; + CHECK_STATUS(normalize(next_charge, normalized)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + auto& entry = it->second; + if (entry.report_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "credit report sequence exhausted" LOC_MARK); + } + + request = ReceiverCreditExchangeRequestV1{}; + request.expected_receiver_session_id = key.receiver_session; + request.epoch = entry.epoch; + request.sender_instance_id = key.sender_instance; + request.qos_class = key.qos_class; + request.known_update_sequence = entry.last_sequence; + request.report_sequence = ++entry.report_sequence; + request.released_totals.reserve(kCreditResourceCount); + request.desired_grant_totals.reserve(next_charge.resources.size()); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const auto resource = static_cast(i + 1); + request.released_totals.push_back({resource, entry.released[i]}); + if (normalized[i] == 0) continue; + if (normalized[i] > + std::numeric_limits::max() - entry.consumed[i]) { + return Status::InvalidArgument( + "desired credit total overflow" LOC_MARK); + } + request.desired_grant_totals.push_back( + {resource, entry.consumed[i] + normalized[i]}); + } + return Status::OK(); +} + +Status SenderCreditLedger::prepareReleaseExchange( + const CreditKey& key, ReceiverCreditExchangeRequestV1& request) { + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + auto& entry = it->second; + if (entry.report_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "credit report sequence exhausted" LOC_MARK); + } + + request = ReceiverCreditExchangeRequestV1{}; + request.expected_receiver_session_id = key.receiver_session; + request.epoch = entry.epoch; + request.sender_instance_id = key.sender_instance; + request.qos_class = key.qos_class; + request.known_update_sequence = entry.last_sequence; + request.report_sequence = ++entry.report_sequence; + request.released_totals.reserve(kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const auto resource = static_cast(i + 1); + request.released_totals.push_back({resource, entry.released[i]}); + if (entry.consumed[i] != 0) { + request.desired_grant_totals.push_back( + {resource, entry.consumed[i]}); + } + } + if (request.desired_grant_totals.empty()) { + return Status::InvalidEntry( + "credit release exchange has no consumption" LOC_MARK); + } return Status::OK(); } -Status SenderCreditLedger::available(const CreditKey& k, CreditResource r, - uint64_t& v) const { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); +Status SenderCreditLedger::available(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || !it->second.has_update) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit unavailable" LOC_MARK); - v = it->second.grants[i] - it->second.consumed[i]; + } + CHECK_STATUS(requireFresh(it->second)); + value = it->second.grants[index] - it->second.consumed[index]; + return Status::OK(); +} + +Status SenderCreditLedger::consumed(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + value = it->second.consumed[index]; return Status::OK(); } -Status SenderCreditLedger::consumed(const CreditKey& k, CreditResource r, - uint64_t& v) const { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); +Status SenderCreditLedger::released(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end()) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit session inactive" LOC_MARK); - v = it->second.consumed[i]; + } + value = it->second.released[index]; + return Status::OK(); +} + +ReceiverCreditAuthority::ReceiverCreditAuthority(ReceiverSessionId session, + uint64_t epoch, + ReceiverCreditLimits limits) + : session_(session), epoch_(epoch), limits_(std::move(limits)) { + if (session_.empty() || epoch_ == 0) { + status_ = Status::InvalidArgument( + "invalid receiver credit authority identity" LOC_MARK); + } else { + status_ = validateLimits(limits_); + } +} + +Status ReceiverCreditAuthority::validateLimits( + const ReceiverCreditLimits& limits) { + if (limits.freshness_ttl_ms == 0 || limits.max_senders == 0 || + limits.max_entries < limits.max_senders) { + return Status::InvalidArgument( + "invalid receiver credit limits" LOC_MARK); + } + bool any_capacity = false; + for (const auto capacity : limits.capacity) any_capacity |= capacity != 0; + if (!any_capacity) { + return Status::InvalidArgument( + "receiver credit capacity is empty" LOC_MARK); + } + return Status::OK(); +} + +ReceiverCreditAdvertV1 ReceiverCreditAuthority::advert() const { + ReceiverCreditAdvertV1 result; + result.receiver_session_id = session_; + result.epoch = epoch_; + result.freshness_ttl_ms = limits_.freshness_ttl_ms; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (limits_.capacity[i] != 0) { + const auto resource = static_cast(i + 1); + result.resources.push_back(resource); + result.capacities.push_back({resource, limits_.capacity[i]}); + } + } + return result; +} + +Status ReceiverCreditAuthority::normalizeCounters( + const std::vector& counters, + std::array& normalized, + std::array& present) { + normalized.fill(0); + present.fill(false); + if (counters.size() > kCreditResourceCount) { + return Status::InvalidArgument("too many credit counters" LOC_MARK); + } + for (const auto& counter : counters) { + size_t index = 0; + CHECK_STATUS(indexFor(counter.resource, index)); + if (present[index]) { + return Status::InvalidArgument( + "duplicate credit counter resource" LOC_MARK); + } + present[index] = true; + normalized[index] = counter.total; + } + return Status::OK(); +} + +void ReceiverCreditAuthority::fillReply( + const CreditKey& key, const Entry& entry, + ReceiverCreditExchangeReplyV1& reply) const { + reply = ReceiverCreditExchangeReplyV1{}; + reply.sender_instance_id = key.sender_instance; + reply.update.qos_class = key.qos_class; + reply.update.receiver_session_id = session_; + reply.update.epoch = epoch_; + reply.update.sequence = entry.update_sequence; + reply.update.freshness_ttl_ms = limits_.freshness_ttl_ms; + reply.update.grants.reserve(kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + reply.update.grants.push_back( + {static_cast(i + 1), entry.grants[i]}); + } +} + +Status ReceiverCreditAuthority::exchange( + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + CHECK_STATUS(status_); + if (request.schema_version != kReceiverCreditProtocolVersion || + request.flags != 0 || request.epoch != epoch_ || + !(request.expected_receiver_session_id == session_) || + request.sender_instance_id.empty() || request.qos_class > 2 || + request.report_sequence == 0 || request.desired_grant_totals.empty()) { + return Status::InvalidArgument( + "invalid receiver credit exchange identity/header" LOC_MARK); + } + + std::array released{}, desired{}; + std::array released_present{}, + desired_present{}; + CHECK_STATUS( + normalizeCounters(request.released_totals, released, released_present)); + CHECK_STATUS(normalizeCounters(request.desired_grant_totals, desired, + desired_present)); + + CreditKey key{request.expected_receiver_session_id, + request.sender_instance_id, request.qos_class}; + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + if (request.known_update_sequence != 0) { + return Status::InvalidArgument( + "new sender knows an impossible credit update" LOC_MARK); + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (released_present[i] && released[i] != 0) { + return Status::InvalidArgument( + "new sender reports impossible credit release" LOC_MARK); + } + } + if (entries_.size() >= limits_.max_entries) { + return Status::TooManyRequests( + "receiver credit replay-history limit" LOC_MARK); + } + it = entries_.emplace(key, Entry{}).first; + } + auto& entry = it->second; + if (request.known_update_sequence > entry.update_sequence) { + return Status::InvalidArgument( + "sender knows an impossible credit update" LOC_MARK); + } + if (request.report_sequence <= entry.last_report_sequence) { + fillReply(key, entry, reply); + return Status::OK(); + } + if (entry.update_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "receiver credit update sequence exhausted" LOC_MARK); + } + + // Validate the entire cumulative release before mutating any resource. + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + if (released[i] < entry.released[i] || released[i] > entry.grants[i]) { + return Status::InvalidArgument( + "invalid cumulative credit release" LOC_MARK); + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + const uint64_t delta = released[i] - entry.released[i]; + if (delta > outstanding_[i]) { + return Status::InternalError( + "receiver credit accounting underflow" LOC_MARK); + } + } + + // Apply releases first, then atomically decide whether every requested + // resource can reach its absolute desired total. + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + outstanding_[i] -= released[i] - entry.released[i]; + entry.released[i] = released[i]; + } + + std::array grant_delta{}; + bool can_grant_all = true; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!desired_present[i] || desired[i] <= entry.grants[i]) continue; + grant_delta[i] = desired[i] - entry.grants[i]; + const uint64_t free = limits_.capacity[i] - outstanding_[i]; + if (grant_delta[i] > free) can_grant_all = false; + } + const bool requests_new_grant = + std::any_of(grant_delta.begin(), grant_delta.end(), + [](uint64_t delta) { return delta != 0; }); + if (can_grant_all && requests_new_grant) { + std::unordered_set + active_senders; + active_senders.reserve(std::min(entries_.size(), limits_.max_senders)); + for (const auto& [entry_key, tracked] : entries_) { + bool has_outstanding = false; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + has_outstanding |= tracked.grants[i] > tracked.released[i]; + } + if (has_outstanding) { + active_senders.insert(entry_key.sender_instance); + } + } + if (!active_senders.count(request.sender_instance_id) && + active_senders.size() >= limits_.max_senders) { + can_grant_all = false; + } + } + if (can_grant_all) { + for (size_t i = 0; i < kCreditResourceCount; ++i) { + entry.grants[i] += grant_delta[i]; + outstanding_[i] += grant_delta[i]; + } + } + + entry.last_report_sequence = request.report_sequence; + ++entry.update_sequence; + fillReply(key, entry, reply); + return Status::OK(); +} + +Status ReceiverCreditAuthority::outstanding(CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(indexFor(resource, index)); + std::lock_guard lock(mutex_); + value = outstanding_[index]; return Status::OK(); } 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 305844d94f..4528e527d7 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -57,6 +57,33 @@ constexpr uint8_t kRedisDefaultDbIndex = 0; // the counter), so only a permanently failing poll or queue retire gets here. constexpr size_t kMaxReclaimAttempts = 4096; constexpr int kMaxHpTcpMetadataRefreshRetries = 1; +uint64_t randomIdentityWord() { + try { + static thread_local std::random_device random; + return (static_cast(random()) << 32) ^ random() ^ + getCurrentTimeInNano(); + } catch (const std::exception&) { + // random_device is allowed to be unavailable on unusual libstdc++ + // targets. Preserve startup while still mixing a process-local nonce + // into the clock-seeded fallback generator. + static std::atomic nonce{1}; + auto& random = SimpleRandom::Get(); + return (static_cast(random.next()) << 32) ^ random.next() ^ + nonce.fetch_add(1, std::memory_order_relaxed); + } +} + +SenderInstanceId makeSenderInstanceId() { + SenderInstanceId id{randomIdentityWord(), randomIdentityWord()}; + if (id.empty()) id.low = 1; + return id; +} + +ReceiverSessionId makeReceiverSessionId() { + ReceiverSessionId id{randomIdentityWord(), randomIdentityWord()}; + if (id.empty()) id.low = 1; + return id; +} } // namespace struct Batch { @@ -360,6 +387,11 @@ Status TransferEngineImpl::setupLocalSegment() { segment.rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); auto& detail = std::get(segment.detail); detail.topology = *(topology_.get()); + if (receiver_credit_authority_) { + detail.receiver_credit = receiver_credit_authority_->advert(); + } else { + detail.receiver_credit.reset(); + } return Status::OK(); })); local_segment_tracker_ = std::make_unique(manager); @@ -395,6 +427,15 @@ Status TransferEngineImpl::construct() { notify_rpc_fallback_ = conf_->get("notification/rpc_fallback", true); enable_progress_worker_ = conf_->get("enable_progress_worker", false); runtime_queue_config_.enabled = conf_->get("enable_runtime_queue", false); + const bool receiver_credit_legacy_enable = + conf_->get("receiver_credit/enable", false); + receiver_credit_config_.sender_enabled = conf_->get( + "receiver_credit/sender_enable", receiver_credit_legacy_enable); + receiver_credit_config_.receiver_enabled = conf_->get( + "receiver_credit/receiver_enable", receiver_credit_legacy_enable); + if (receiver_credit_config_.sender_enabled) { + runtime_queue_config_.enabled = true; + } if (runtime_queue_config_.enabled) enable_progress_worker_ = true; runtime_queue_config_.limits.max_outstanding_owners = conf_->get("runtime_queue/max_outstanding_owners", 1024UL); @@ -417,6 +458,59 @@ Status TransferEngineImpl::construct() { runtime_queue_config_.progress_fallback_interval = std::chrono::microseconds( conf_->get("runtime_queue/progress_fallback_interval_us", 50000UL)); + + receiver_credit_config_.limits.capacity[0] = + conf_->get("receiver_credit/data_bytes", + static_cast( + runtime_queue_config_.limits.max_outstanding_bytes)); + receiver_credit_config_.limits.capacity[1] = + conf_->get("receiver_credit/request_slots", + static_cast( + runtime_queue_config_.limits.max_outstanding_owners)); + receiver_credit_config_.limits.capacity[2] = 0; + receiver_credit_config_.limits.capacity[3] = 0; + receiver_credit_config_.limits.freshness_ttl_ms = + conf_->get("receiver_credit/freshness_ttl_ms", uint32_t{1000}); + receiver_credit_config_.limits.max_senders = + conf_->get("receiver_credit/max_senders", size_t{1024}); + receiver_credit_config_.limits.max_entries = + conf_->get("receiver_credit/max_authority_entries", size_t{4096}); + receiver_credit_config_.max_sender_entries = + conf_->get("receiver_credit/max_sender_entries", size_t{4096}); + receiver_credit_config_.max_freshness_ttl_ms = + conf_->get("receiver_credit/max_freshness_ttl_ms", uint32_t{60000}); + if (conf_->get("receiver_credit/staging_slots", uint64_t{0}) != 0 || + conf_->get("receiver_credit/consumer_slots", uint64_t{0}) != 0) { + return Status::InvalidArgument( + "staging/consumer receiver-credit resources are not supported " + "by the PR6 runtime gate" LOC_MARK); + } + if (receiver_credit_config_.sender_enabled) { + if (receiver_credit_config_.max_sender_entries == 0 || + receiver_credit_config_.max_freshness_ttl_ms == 0 || + runtime_queue_config_.progress_fallback_interval.count() <= 0) { + return Status::InvalidArgument( + "receiver-credit sender requires non-zero entries, TTL " + "limit, and progress fallback interval" LOC_MARK); + } + } + if (receiver_credit_config_.receiver_enabled) { + if (receiver_credit_config_.limits.capacity[0] == 0 || + receiver_credit_config_.limits.capacity[1] == 0 || + receiver_credit_config_.limits.max_senders == 0 || + receiver_credit_config_.limits.max_entries < + receiver_credit_config_.limits.max_senders) { + return Status::InvalidArgument( + "receiver credit requires non-zero bytes, request slots, " + "and a replay-history limit no smaller than the sender " + "limit" LOC_MARK); + } + if (receiver_credit_config_.limits.freshness_ttl_ms > + receiver_credit_config_.max_freshness_ttl_ms) { + return Status::InvalidArgument( + "receiver credit TTL exceeds sender safety limit" LOC_MARK); + } + } if (runtime_queue_config_.enabled && (runtime_queue_config_.max_dispatch_owners == 0 || runtime_queue_config_.max_dispatch_bytes == 0)) { @@ -425,6 +519,15 @@ Status TransferEngineImpl::construct() { } runtime_queue_ = std::make_unique( runtime_queue_config_.limits); + sender_instance_id_ = makeSenderInstanceId(); + sender_credit_ledger_ = std::make_unique( + std::max(receiver_credit_config_.max_sender_entries, 1), + std::max(receiver_credit_config_.max_freshness_ttl_ms, 1)); + if (receiver_credit_config_.receiver_enabled) { + receiver_credit_authority_ = std::make_unique( + makeReceiverSessionId(), 1, receiver_credit_config_.limits); + CHECK_STATUS(receiver_credit_authority_->status()); + } if (!hostname_.empty()) CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_)); else @@ -437,6 +540,14 @@ Status TransferEngineImpl::construct() { metadata_ = std::make_shared(metadata_type, metadata_servers, this); + if (receiver_credit_authority_) { + metadata_->setReceiverCreditCallback( + [this](const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + return receiver_credit_authority_->exchange(request, reply); + }); + } + CHECK_STATUS(metadata_->start(port_, ipv6_, rpc_server_threads)); if (metadata_type == "p2p") @@ -610,6 +721,28 @@ Status TransferEngineImpl::deconstruct() { << status.ToString(); } } + // Only a reservation that never reached a transport can be returned here. + // A remaining committed owner is not terminal, and the UB DMA fence is + // established later while its transport is destroyed. Conservatively + // abandon that authorization in the old receiver session instead of + // making it available while DMA may still be active. + size_t abandoned_receiver_credits = 0; + for (auto& [_, queued] : queued_owners_) { + if (!queued.credit.active) continue; + if (!queued.credit.committed) { + rollbackReceiverCredit(queued.credit); + } else { + queued.credit = ReceiverCreditReservation{}; + ++abandoned_receiver_credits; + } + } + if (abandoned_receiver_credits != 0) { + LOG(WARNING) << "Conservatively abandoned " + << abandoned_receiver_credits + << " in-flight receiver-credit reservations during " + "shutdown"; + } + progressReceiverCreditReleases(); // Destroy staging_proxy_ first: its destructor calls back into // unregisterLocalMemory/freeLocalMemory, which require @@ -1698,6 +1831,10 @@ struct RequestBoundaryInfo { std::optional source_key; std::optional target_key; uint64_t max_merge_bytes{std::numeric_limits::max()}; + // A required receiver-credit advert limits the size of one physical + // owner. Request merging must not turn individually admissible requests + // into an owner that can never acquire credit. + std::optional receiver_credit_data_capacity; }; struct MergeResult { @@ -1720,6 +1857,8 @@ struct TransferEngineImpl::PreparedSubmit { std::vector derived_task_ids; Request request{}; SelectionResult route{}; + bool receiver_credit_required{false}; + bool receiver_credit_exempt{false}; bool staging{false}; std::vector staging_params; }; @@ -1820,6 +1959,15 @@ MergeResult mergeRequests(const std::vector& requests, if (last.req.length + curr.req.length > last.boundary.max_merge_bytes) { return false; } + const size_t merged_length = last.req.length + curr.req.length; + if (last.boundary.receiver_credit_data_capacity != + curr.boundary.receiver_credit_data_capacity) { + return false; + } + if (last.boundary.receiver_credit_data_capacity && + merged_length > *last.boundary.receiver_credit_data_capacity) { + return false; + } uint64_t last_source_end = 0; uint64_t last_target_end = 0; @@ -1898,6 +2046,21 @@ std::vector resolveRequestBoundaries( std::min(hp_tcp_limit, endpoint.max_transfer_bytes); } } + std::optional receiver_credit_data_capacity; + if (target_desc->type == SegmentType::Memory && + target_desc->getMemory().receiver_credit) { + // Zero deliberately disables merging for malformed + // adverts; prepareSubmit validates and rejects them with + // the protocol-specific error below. + receiver_credit_data_capacity = 0; + for (const auto& capacity : + target_desc->getMemory().receiver_credit->capacities) { + if (capacity.resource == CreditResource::DataBytes) { + receiver_credit_data_capacity = capacity.total; + break; + } + } + } bool any_missing = false; for (size_t i : idxs) { const auto& r = requests[i]; @@ -1907,6 +2070,8 @@ std::vector resolveRequestBoundaries( r.transport_hint == UNSPEC)) { boundaries[i].max_merge_bytes = hp_tcp_limit; } + boundaries[i].receiver_credit_data_capacity = + receiver_credit_data_capacity; auto* buffer = target_desc->findBuffer(r.target_offset, r.length); if (!buffer) { @@ -2090,7 +2255,8 @@ SelectionResult TransferEngineImpl::resolveTransport(const Request& req, Status TransferEngineImpl::prepareSubmit( Batch* batch, const std::vector& request_list, - PreparedSubmit& prepared, bool require_post_submit_cancellation) { + QueueOwnerKind owner_kind, PreparedSubmit& prepared, + bool require_post_submit_cancellation) { if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); for (size_t i = 0; i < request_list.size(); ++i) { auto st = validateTransportHint(request_list[i], i); @@ -2171,6 +2337,49 @@ Status TransferEngineImpl::prepareSubmit( "transport" LOC_MARK); } } + owner.receiver_credit_exempt = + owner_kind == QueueOwnerKind::StagingInternal; + if (owner_kind != QueueOwnerKind::StagingInternal) { + ReceiverCreditTarget credit_target; + CHECK_STATUS( + resolveReceiverCreditTarget(owner.request, credit_target)); + owner.receiver_credit_required = credit_target.required; + // Scatter bypasses the runtime queue to preserve synchronous + // cancellation. It cannot bypass a peer's required credit gate. + if (credit_target.required && require_post_submit_cancellation) { + return Status::NotImplemented( + "scatter transfer does not support required receiver " + "credit" LOC_MARK); + } + if (credit_target.required && + !receiver_credit_config_.sender_enabled) { + return Status::InvalidArgument( + "peer requires receiver credit but the sender gate is " + "disabled" LOC_MARK); + } + CHECK_STATUS( + validateReceiverCreditCapacity(owner.request, credit_target)); + } else { + const bool has_credit_fence = + owner.request.receiver_credit_session_high != 0 || + owner.request.receiver_credit_session_low != 0 || + owner.request.receiver_credit_epoch != 0; + if (has_credit_fence) { + ReceiverCreditTarget credit_target; + CHECK_STATUS( + resolveReceiverCreditTarget(owner.request, credit_target)); + if (!credit_target.required || + credit_target.advert.receiver_session_id.high != + owner.request.receiver_credit_session_high || + credit_target.advert.receiver_session_id.low != + owner.request.receiver_credit_session_low || + credit_target.advert.epoch != + owner.request.receiver_credit_epoch) { + return Status::InvalidEntry( + "staging receiver-credit session changed" LOC_MARK); + } + } + } owner.staging_params = owner.route.staging_params; if (owner.staging_params.empty() && (owner.route.transport == TCP || owner.route.transport == HP_TCP)) { @@ -2214,6 +2423,354 @@ void TransferEngineImpl::attachProgressNotifier( }; } +Status TransferEngineImpl::resolveReceiverCreditTarget( + const Request& request, ReceiverCreditTarget& target) const { + target = ReceiverCreditTarget{}; + SegmentDescRef segment; + if (request.target_id == LOCAL_SEGMENT_ID) { + segment = metadata_->segmentManager().getLocal(); + } else { + CHECK_STATUS(metadata_->segmentManager().getRemoteCached( + segment, request.target_id)); + } + if (!segment || segment->type != SegmentType::Memory) return Status::OK(); + const auto& memory = segment->getMemory(); + if (!memory.receiver_credit) return Status::OK(); + + const auto& advert = *memory.receiver_credit; + if (advert.schema_version != kReceiverCreditProtocolVersion || + advert.flags != kReceiverCreditRequired || + advert.receiver_session_id.empty() || advert.epoch == 0 || + advert.freshness_ttl_ms == 0 || + advert.freshness_ttl_ms > + receiver_credit_config_.max_freshness_ttl_ms || + advert.resources.empty() || + advert.resources.size() > kCreditResourceCount || + advert.capacities.size() != advert.resources.size()) { + return Status::InvalidArgument( + "invalid receiver-credit advertisement" LOC_MARK); + } + + std::array resources{}; + for (const auto resource : advert.resources) { + const auto raw = static_cast(resource); + if (raw < 1 || raw > kCreditResourceCount || resources[raw - 1]) { + return Status::InvalidArgument( + "invalid receiver-credit resource advertisement" LOC_MARK); + } + resources[raw - 1] = true; + } + std::array capacities{}; + for (const auto& capacity : advert.capacities) { + const auto raw = static_cast(capacity.resource); + if (raw < 1 || raw > kCreditResourceCount || capacity.total == 0 || + capacities[raw - 1] || !resources[raw - 1]) { + return Status::InvalidArgument( + "invalid receiver-credit capacity advertisement" LOC_MARK); + } + capacities[raw - 1] = true; + } + if (capacities != resources) { + return Status::InvalidArgument( + "receiver-credit resources/capacities mismatch" LOC_MARK); + } + if (!resources[static_cast(CreditResource::DataBytes) - 1] || + !resources[static_cast(CreditResource::RequestSlots) - 1]) { + return Status::InvalidArgument( + "receiver credit lacks bytes/request resources" LOC_MARK); + } + if (request.target_id != LOCAL_SEGMENT_ID && + segment->rpc_server_addr.empty()) { + return Status::InvalidArgument( + "receiver credit peer has no RPC address" LOC_MARK); + } + target.required = true; + target.advert = advert; + target.rpc_server_addr = segment->rpc_server_addr; + return Status::OK(); +} + +Status TransferEngineImpl::validateReceiverCreditCapacity( + const Request& request, const ReceiverCreditTarget& target) const { + if (!target.required) return Status::OK(); + uint64_t data_bytes = 0; + uint64_t request_slots = 0; + for (const auto& capacity : target.advert.capacities) { + if (capacity.resource == CreditResource::DataBytes) { + data_bytes = capacity.total; + } else if (capacity.resource == CreditResource::RequestSlots) { + request_slots = capacity.total; + } + } + if (request_slots == 0 || request.length > data_bytes) { + return Status::TooManyRequests( + "request exceeds receiver-credit capacity" LOC_MARK); + } + return Status::OK(); +} + +Status TransferEngineImpl::exchangeReceiverCredit( + SegmentID target_id, const std::string& rpc_server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + if (target_id == LOCAL_SEGMENT_ID) { + if (!receiver_credit_authority_) { + return Status::InvalidEntry( + "local receiver-credit authority is unavailable" LOC_MARK); + } + return receiver_credit_authority_->exchange(request, reply); + } + return ControlClient::exchangeReceiverCredit(rpc_server_addr, request, + reply); +} + +Status TransferEngineImpl::acquireReceiverCredit( + const Request& request, ReceiverCreditReservation& reservation, + bool& acquired) { + acquired = false; + if (reservation.active) { + return Status::InternalError( + "receiver credit attempt already has a reservation" LOC_MARK); + } + + ReceiverCreditTarget target; + CHECK_STATUS(resolveReceiverCreditTarget(request, target)); + if (!target.required) { + acquired = true; + return Status::OK(); + } + CHECK_STATUS(validateReceiverCreditCapacity(request, target)); + if (!receiver_credit_config_.sender_enabled || !sender_credit_ledger_) { + return Status::InvalidEntry( + "receiver-credit sender gate is disabled" LOC_MARK); + } + + const uint32_t qos_class = static_cast( + std::clamp(request.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))); + CreditKey key{target.advert.receiver_session_id, sender_instance_id_, + qos_class}; + CreditCharge charge{{{CreditResource::DataBytes, request.length}, + {CreditResource::RequestSlots, 1}}}; + observeReceiverCreditSession(request.target_id, + target.advert.receiver_session_id); + CHECK_STATUS(sender_credit_ledger_->activate(key, target.advert.epoch)); + + auto reserve_status = sender_credit_ledger_->tryReserve(key, charge); + if (!reserve_status.ok()) { + if (!reserve_status.IsTooManyRequests() && + !reserve_status.IsInvalidEntry()) { + return reserve_status; + } + + ReceiverCreditExchangeRequestV1 exchange_request; + CHECK_STATUS(sender_credit_ledger_->prepareExchange(key, charge, + exchange_request)); + ReceiverCreditExchangeReplyV1 exchange_reply; + auto exchange_status = + exchangeReceiverCredit(request.target_id, target.rpc_server_addr, + exchange_request, exchange_reply); + if (!exchange_status.ok()) { + if (request.target_id != LOCAL_SEGMENT_ID) { + (void)metadata_->segmentManager().invalidateRemote( + request.target_id); + } + LOG(WARNING) << "Receiver-credit exchange deferred: " + << exchange_status.ToString(); + return Status::OK(); + } + if (!(exchange_reply.sender_instance_id == sender_instance_id_)) { + return Status::InvalidArgument( + "receiver-credit reply sender mismatch" LOC_MARK); + } + if (exchange_reply.flags != 0) { + return Status::InvalidArgument( + "receiver-credit demand reply flags" LOC_MARK); + } + CreditUpdateDisposition disposition; + CHECK_STATUS(sender_credit_ledger_->applyUpdate( + key, exchange_reply.update, disposition)); + pending_credit_releases_.erase(key); + reserve_status = sender_credit_ledger_->tryReserve(key, charge); + if (reserve_status.IsTooManyRequests() || + reserve_status.IsInvalidEntry()) { + return Status::OK(); + } + CHECK_STATUS(reserve_status); + } + + reservation.active = true; + reservation.committed = false; + reservation.key = key; + reservation.charge = std::move(charge); + reservation.target_id = request.target_id; + reservation.rpc_server_addr = std::move(target.rpc_server_addr); + reservation.epoch = target.advert.epoch; + acquired = true; + return Status::OK(); +} + +void TransferEngineImpl::rollbackReceiverCredit( + ReceiverCreditReservation& reservation) { + if (!reservation.active) return; + if (!reservation.committed) { + const auto key = reservation.key; + const auto target_id = reservation.target_id; + const auto rpc_server_addr = reservation.rpc_server_addr; + const auto epoch = reservation.epoch; + auto status = sender_credit_ledger_->rollbackReservation( + reservation.key, reservation.charge); + reservation = ReceiverCreditReservation{}; + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit rollback failed: " + << status.ToString(); + return; + } + scheduleReceiverCreditRelease(key, target_id, rpc_server_addr, epoch); + (void)tryReportReceiverCreditRelease(key); + return; + } + reservation = ReceiverCreditReservation{}; +} + +void TransferEngineImpl::releaseReceiverCredit( + ReceiverCreditReservation& reservation) { + if (!reservation.active) return; + if (!reservation.committed) { + rollbackReceiverCredit(reservation); + return; + } + + const auto key = reservation.key; + const auto target_id = reservation.target_id; + const auto rpc_server_addr = reservation.rpc_server_addr; + const auto epoch = reservation.epoch; + auto status = sender_credit_ledger_->releaseCommitted(reservation.key, + reservation.charge); + reservation = ReceiverCreditReservation{}; + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit terminal release failed: " + << status.ToString(); + return; + } + + scheduleReceiverCreditRelease(key, target_id, rpc_server_addr, epoch); + (void)tryReportReceiverCreditRelease(key); +} + +void TransferEngineImpl::scheduleReceiverCreditRelease( + const CreditKey& key, SegmentID target_id, + const std::string& rpc_server_addr, uint64_t epoch) { + auto& pending = pending_credit_releases_[key]; + pending.target_id = target_id; + pending.rpc_server_addr = rpc_server_addr; + pending.epoch = epoch; + pending.retry_at = std::chrono::steady_clock::now(); +} + +bool TransferEngineImpl::tryReportReceiverCreditRelease(const CreditKey& key) { + auto pending_it = pending_credit_releases_.find(key); + if (pending_it == pending_credit_releases_.end()) return true; + auto& pending = pending_it->second; + if (std::chrono::steady_clock::now() < pending.retry_at) return false; + + auto defer_retry = [&](const Status& status) { + ++pending.attempts; + const uint32_t shift = std::min(pending.attempts - 1, 4); + const auto delay = std::chrono::milliseconds(50U << shift); + pending.retry_at = std::chrono::steady_clock::now() + delay; + LOG(WARNING) << "Receiver-credit release exchange deferred (attempt " + << pending.attempts << "): " << status.ToString(); + return false; + }; + + ReceiverCreditExchangeRequestV1 request; + auto status = sender_credit_ledger_->prepareReleaseExchange(key, request); + if (!status.ok()) { + return defer_retry(status); + } + ReceiverCreditExchangeReplyV1 reply; + status = exchangeReceiverCredit(pending.target_id, pending.rpc_server_addr, + request, reply); + if (!status.ok()) { + return defer_retry(status); + } + if (!(reply.sender_instance_id == sender_instance_id_)) { + return defer_retry(Status::InvalidArgument( + "receiver-credit release reply sender mismatch" LOC_MARK)); + } + if (reply.flags != 0) { + return defer_retry(Status::InvalidArgument( + "receiver-credit release reply flags" LOC_MARK)); + } + CreditUpdateDisposition disposition; + status = sender_credit_ledger_->applyUpdate(key, reply.update, disposition); + if (!status.ok()) { + return defer_retry(status); + } + pending_credit_releases_.erase(pending_it); + notifyRuntimeQueueReady(); + return true; +} + +void TransferEngineImpl::progressReceiverCreditReleases() { + const auto now = std::chrono::steady_clock::now(); + std::vector due; + due.reserve(pending_credit_releases_.size()); + for (const auto& [key, pending] : pending_credit_releases_) { + if (pending.retry_at <= now) due.push_back(key); + } + for (const auto& key : due) (void)tryReportReceiverCreditRelease(key); +} + +void TransferEngineImpl::observeReceiverCreditSession( + SegmentID target_id, const ReceiverSessionId& session) { + auto [it, inserted] = + receiver_credit_sessions_.try_emplace(target_id, session); + if (inserted || it->second == session) return; + + const auto stale_session = it->second; + auto status = sender_credit_ledger_->deactivateSession(stale_session); + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit stale session cleanup failed: " + << status.ToString(); + } + for (auto pending = pending_credit_releases_.begin(); + pending != pending_credit_releases_.end();) { + if (pending->first.receiver_session == stale_session) { + pending = pending_credit_releases_.erase(pending); + } else { + ++pending; + } + } + it->second = session; +} + +void TransferEngineImpl::bindReceiverCreditFence( + Transport::SubBatchRef sub_batch, + const ReceiverCreditReservation& reservation) const { + if (!sub_batch) return; + sub_batch->receiver_credit_session_high = 0; + sub_batch->receiver_credit_session_low = 0; + sub_batch->receiver_credit_epoch = 0; + if (!reservation.active) return; + sub_batch->receiver_credit_session_high = + reservation.key.receiver_session.high; + sub_batch->receiver_credit_session_low = + reservation.key.receiver_session.low; + sub_batch->receiver_credit_epoch = reservation.epoch; +} + +void TransferEngineImpl::bindReceiverCreditFence( + Transport::SubBatchRef sub_batch, const Request& request) const { + if (!sub_batch) return; + sub_batch->receiver_credit_session_high = + request.receiver_credit_session_high; + sub_batch->receiver_credit_session_low = + request.receiver_credit_session_low; + sub_batch->receiver_credit_epoch = request.receiver_credit_epoch; +} + Status TransferEngineImpl::commitPreparedSubmit( Batch* batch, const PreparedSubmit& prepared) { if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); @@ -2487,6 +3044,8 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, queued.batch = batch; queued.owner_task_id = prepared.owners[i].owner_task_id; queued.byte_charge = prepared.owners[i].request.length; + queued.receiver_credit_exempt = + prepared.owners[i].receiver_credit_exempt; queued.public_task_ids.push_back(prepared.owners[i].owner_task_id); queued.public_task_ids.insert( queued.public_task_ids.end(), @@ -2511,6 +3070,7 @@ Status TransferEngineImpl::finishQueuedOwner( "runtime dispatch window accounting underflow" LOC_MARK); } } + releaseReceiverCredit(queued.credit); CHECK_STATUS(runtime_queue_->complete(owner_id, terminal_status)); if (queued.in_dispatch_window) { --dispatch_inflight_owners_; @@ -2611,7 +3171,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(); @@ -2629,13 +3189,43 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (staging_params.empty()) findStagingPolicy(task.request, staging_params); if (!staging_params.empty() && staging_proxy_) { + bool acquired = true; + auto credit_status = Status::OK(); + if (!queued.receiver_credit_exempt) { + acquired = false; + credit_status = acquireReceiverCredit(task.request, + queued.credit, acquired); + } + if (!credit_status.ok()) { + LOG(WARNING) << "Receiver-credit gate rejected staging " + "dispatch: " + << credit_status.ToString(); + return finishQueuedOwner(owner_id, FAILED); + } + if (!acquired) { + task.type = UNSPEC; + task.sub_task_id = -1; + CHECK_STATUS(runtime_queue_->defer(owner_id)); + return Status::OK(); + } + if (queued.credit.active) { + task.request.receiver_credit_session_high = + queued.credit.key.receiver_session.high; + task.request.receiver_credit_session_low = + queued.credit.key.receiver_session.low; + task.request.receiver_credit_epoch = queued.credit.epoch; + } task.staging = true; // Orchestration only; the real transport submissions issued by // ProxyManager are counted where they recurse through the // non-staging path below. auto status = staging_proxy_->submit(&task, (BatchID)batch, staging_params); - if (!status.ok()) return finishQueuedOwner(owner_id, FAILED); + if (!status.ok()) { + rollbackReceiverCredit(queued.credit); + return finishQueuedOwner(owner_id, FAILED); + } + if (queued.credit.active) queued.credit.committed = true; task.post_time = std::chrono::steady_clock::now(); return markQueuedOwnerSubmitted(owner_id); } @@ -2657,6 +3247,30 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { sub_batch->device_mask = task.device_mask; if (task.type == RDMA) sub_batch->qp_pool = task.qp_pool; } + + bool acquired = true; + auto credit_status = Status::OK(); + if (!queued.receiver_credit_exempt) { + acquired = false; + credit_status = + acquireReceiverCredit(task.request, queued.credit, acquired); + } + if (!credit_status.ok()) { + LOG(WARNING) << "Receiver-credit gate rejected dispatch: " + << credit_status.ToString(); + return finishQueuedOwner(owner_id, FAILED); + } + if (!acquired) { + task.type = UNSPEC; + task.sub_task_id = -1; + CHECK_STATUS(runtime_queue_->defer(owner_id)); + return Status::OK(); + } + if (queued.receiver_credit_exempt) { + bindReceiverCreditFence(sub_batch, task.request); + } else { + bindReceiverCreditFence(sub_batch, queued.credit); + } task.sub_task_id = sub_batch->size(); startTransportAttempt(task, task.type, std::chrono::steady_clock::now()); auto status = transport->submitTransferTasks(sub_batch, {task.request}); @@ -2669,8 +3283,12 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (attemptSubmitStageFailover(batch, queued.owner_task_id)) { return markQueuedOwnerSubmitted(owner_id); } + rollbackReceiverCredit(queued.credit); + task.type = UNSPEC; return finishQueuedOwner(owner_id, FAILED); } + if (queued.credit.active) queued.credit.committed = true; + task.post_time = std::chrono::steady_clock::now(); return markQueuedOwnerSubmitted(owner_id); } @@ -2712,6 +3330,7 @@ Status TransferEngineImpl::progressRuntimeQueue() { std::lock_guard lk(progress_mutex_); if (!runtime_queue_config_.enabled) return Status::OK(); + progressReceiverCreditReleases(); CHECK_STATUS(refillDispatchWindow()); std::vector owner_ids; @@ -2756,7 +3375,8 @@ Status TransferEngineImpl::progressRuntimeQueue() { bool TransferEngineImpl::hasActiveRuntimeQueue() { std::lock_guard lk(progress_mutex_); - return runtime_queue_config_.enabled && !queued_owners_.empty(); + return runtime_queue_config_.enabled && + (!queued_owners_.empty() || !pending_credit_releases_.empty()); } bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared, @@ -2764,6 +3384,12 @@ bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared, if (!runtime_queue_config_.enabled) return false; if (prepared.require_post_submit_cancellation) return false; if (owner_kind == QueueOwnerKind::StagingInternal) return true; + if (std::any_of(prepared.owners.begin(), prepared.owners.end(), + [](const PreparedSubmit::Owner& owner) { + return owner.receiver_credit_required; + })) { + return true; + } return std::none_of( prepared.owners.begin(), prepared.owners.end(), [](const PreparedSubmit::Owner& owner) { return owner.staging; }); @@ -2778,8 +3404,8 @@ Status TransferEngineImpl::submitTransfer( BatchRef batch_ref(*this, batch); const size_t start_task_id = batch_ref.get()->task_list.size(); PreparedSubmit prepared; - CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, prepared, - require_post_submit_cancellation)); + CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, owner_kind, + prepared, require_post_submit_cancellation)); if (shouldQueueSubmit(prepared, owner_kind)) { CHECK_STATUS( @@ -2967,6 +3593,18 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { auto& task = batch->task_list[task_id]; auto prev_type = task.type; + QueueOwnerId queue_owner_id = 0; + QueuedOwnerState* queued_owner = nullptr; + if (runtime_queue_config_.enabled && batch->queue_token != 0 && + runtime_queue_ + ->resolveOwner(batch->queue_token, task_id, queue_owner_id) + .ok()) { + auto queued_it = queued_owners_.find(queue_owner_id); + if (queued_it != queued_owners_.end()) { + queued_owner = &queued_it->second; + } + } + if (++task.failover_count > task.runtime_policy.max_failover_attempts) { LOG(WARNING) << "Task failover limit reached (" << task.runtime_policy.max_failover_attempts @@ -3027,13 +3665,72 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { } task.sub_task_id = sub_batch->size(); task.type = type; + ReceiverCreditTarget credit_target; + const bool credit_exempt = + queued_owner && queued_owner->receiver_credit_exempt; + if (!credit_exempt) { + CHECK_STATUS(resolveReceiverCreditTarget(task.request, credit_target)); + if (credit_target.required && !queued_owner) { + return Status::InvalidEntry( + "credit-required failover needs runtime queue " + "ownership" LOC_MARK); + } + } + if (queued_owner && queued_owner->credit.active && + (!credit_target.required || + queued_owner->credit.key.receiver_session != + credit_target.advert.receiver_session_id || + queued_owner->credit.epoch != credit_target.advert.epoch)) { + // A reservation from another receiver incarnation cannot authorize + // this attempt. Do not report it to a possibly replaced authority; + // abandon it and acquire against the current session below. + queued_owner->credit = ReceiverCreditReservation{}; + } + bool acquired_this_attempt = false; + if (!credit_exempt && queued_owner && !queued_owner->credit.active) { + bool acquired = false; + CHECK_STATUS(acquireReceiverCredit(task.request, queued_owner->credit, + acquired)); + if (!acquired) { + if (!queued_owner->in_dispatch_window || + dispatch_inflight_owners_ == 0 || + dispatch_inflight_bytes_ < queued_owner->byte_charge) { + return Status::InternalError( + "runtime dispatch window accounting underflow" LOC_MARK); + } + --dispatch_inflight_owners_; + dispatch_inflight_bytes_ -= queued_owner->byte_charge; + queued_owner->in_dispatch_window = false; + CHECK_STATUS(runtime_queue_->defer(queue_owner_id)); + task.type = UNSPEC; + task.sub_task_id = -1; + return Status::OK(); + } + acquired_this_attempt = queued_owner->credit.active; + } + if (queued_owner) { + if (queued_owner->receiver_credit_exempt) { + bindReceiverCreditFence(sub_batch, task.request); + } else { + bindReceiverCreditFence(sub_batch, queued_owner->credit); + } + } startTransportAttempt(task, type, std::chrono::steady_clock::now()); - auto status = transport->submitTransferTasks(sub_batch, {task.request}); - if (!status.ok()) { + auto submit_status = + transport->submitTransferTasks(sub_batch, {task.request}); + if (!submit_status.ok()) { if (task.failure_stage < 0) task.failure_stage = 0; finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + if (queued_owner && acquired_this_attempt) { + rollbackReceiverCredit(queued_owner->credit); + } + return submit_status; } - return status; + if (queued_owner && queued_owner->credit.active) { + queued_owner->credit.committed = true; + } + task.post_time = std::chrono::steady_clock::now(); + return Status::OK(); } bool TransferEngineImpl::attemptSubmitStageFailover(Batch* batch, @@ -3641,9 +4338,32 @@ Status TransferEngineImpl::waitTransferCompletion(BatchID batch_id) { Status TransferEngineImpl::transferSync( const std::vector& request_list) { + auto has_any_fence = [](const Request& request) { + return request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + }; + auto has_complete_fence = [](const Request& request) { + return (request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0) && + request.receiver_credit_epoch != 0; + }; + const bool delegated_credit_fence = + !request_list.empty() && + std::all_of(request_list.begin(), request_list.end(), + has_complete_fence); + if (!delegated_credit_fence && + std::any_of(request_list.begin(), request_list.end(), has_any_fence)) { + return Status::InvalidArgument( + "mixed or partial delegated receiver-credit fence" LOC_MARK); + } + auto batch_id = allocateBatch(request_list.size()); BatchGuard batch_guard(*this, batch_id); - CHECK_STATUS(submitTransfer(batch_id, request_list)); + CHECK_STATUS(submitTransfer(batch_id, request_list, nullptr, + delegated_credit_fence + ? QueueOwnerKind::StagingInternal + : QueueOwnerKind::User)); for (uint64_t poll_count = 0;; ++poll_count) { TransferStatus xfer_status; CHECK_STATUS(progressBatch(batch_id, xfer_status)); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index 4b64333b6c..0331a05cba 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -540,6 +540,11 @@ Status RdmaTransport::submitTransferTasks( task->request = request; task->device_mask = rdma_batch->device_mask; task->qp_pool = rdma_batch->qp_pool; // RFC #2568 step 3 + task->receiver_credit_session_high = + rdma_batch->receiver_credit_session_high; + task->receiver_credit_session_low = + rdma_batch->receiver_credit_session_low; + task->receiver_credit_epoch = rdma_batch->receiver_credit_epoch; task->num_slices = 0; task->status_word = PENDING; task->transferred_bytes = 0; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 468cddc2b2..0914c05dd8 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -1835,6 +1835,24 @@ Status Workers::generatePostPath(RdmaSlice* slice) { auto target_id = slice->task->request.target_id; CHECK_STATUS(getRouteHint(target, target_id, (uint64_t)slice->target_addr, slice->length)); + const bool has_credit_fence = + slice->task->receiver_credit_session_high != 0 || + slice->task->receiver_credit_session_low != 0 || + slice->task->receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = target.segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + slice->task->receiver_credit_session_high || + advert->receiver_session_id.low != + slice->task->receiver_credit_session_low || + advert->epoch != slice->task->receiver_credit_epoch) { + return Status::InvalidEntry( + "RDMA receiver-credit session changed before post" LOC_MARK); + } + } if (slice->retry_count == 0) CHECK_STATUS(selectOptimalDevice(source, target, slice)); diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp index ef6162dd07..d7288c5376 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp @@ -176,6 +176,9 @@ Status TcpTransport::submitTransferTasks( task.request = request; task.progress_batch_id = batch->progress_batch_id; task.notify_progress = batch->notify_progress; + task.receiver_credit_session_high = batch->receiver_credit_session_high; + task.receiver_credit_session_low = batch->receiver_credit_session_low; + task.receiver_credit_epoch = batch->receiver_credit_epoch; task.status_word.store(TransferStatusEnum::PENDING, std::memory_order_release); } @@ -262,9 +265,7 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { return Status::InternalError("Transport shutting down"); std::string rpc_server_addr; - auto status = - findRemoteSegment(task->request.target_offset, task->request.length, - task->request.target_id, rpc_server_addr); + auto status = findRemoteSegment(*task, rpc_server_addr); if (!status.ok()) return status; Status last_error; @@ -294,11 +295,15 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { if (task->request.opcode == Request::WRITE) { status = ControlClient::sendData( rpc_server_addr, task->request.target_offset, - task->request.source, task->request.length); + task->request.source, task->request.length, + task->receiver_credit_session_high, + task->receiver_credit_session_low, task->receiver_credit_epoch); } else { status = ControlClient::recvData( rpc_server_addr, task->request.target_offset, - task->request.source, task->request.length); + task->request.source, task->request.length, + task->receiver_credit_session_high, + task->receiver_credit_session_low, task->receiver_credit_epoch); } if (status.ok()) return Status::OK(); @@ -325,9 +330,7 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { CHECK_STATUS(metadata_->segmentManager().invalidateRemote( task->request.target_id)); if (attempt == params_.max_retry_count) break; - auto resolve = findRemoteSegment( - task->request.target_offset, task->request.length, - task->request.target_id, rpc_server_addr); + auto resolve = findRemoteSegment(*task, rpc_server_addr); if (!resolve.ok()) return resolve; } } @@ -335,12 +338,35 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { return last_error; } -Status TcpTransport::findRemoteSegment(uint64_t dest_addr, uint64_t length, - uint64_t target_id, +Status TcpTransport::findRemoteSegment(const TcpTask &task, std::string &rpc_server_addr) { return metadata_->segmentManager().withCachedSegment( - target_id, [&](SegmentDesc *segment) { - auto buffer = segment->findBuffer(dest_addr, length); + task.request.target_id, [&](SegmentDesc *segment) { + const bool has_credit_fence = + task.receiver_credit_session_high != 0 || + task.receiver_credit_session_low != 0 || + task.receiver_credit_epoch != 0; + if (has_credit_fence) { + if (segment->type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "TCP receiver-credit target is not memory" LOC_MARK); + } + const auto &advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + task.receiver_credit_session_high || + advert->receiver_session_id.low != + task.receiver_credit_session_low || + advert->epoch != task.receiver_credit_epoch) { + return Status::InvalidEntry( + "TCP receiver-credit session changed before " + "send" LOC_MARK); + } + } + auto buffer = segment->findBuffer(task.request.target_offset, + task.request.length); rpc_server_addr = segment->rpc_server_addr; if (!buffer) { return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp index 665e9204d1..e887bbf158 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp @@ -34,7 +34,11 @@ bool samePeer(const UbBootstrapDesc& lhs, const UbBootstrapDesc& rhs) { return lhs.protocol_version == rhs.protocol_version && lhs.local_eid == rhs.local_eid && lhs.jetty_ids == rhs.jetty_ids && lhs.jetty_uasids == rhs.jetty_uasids && - lhs.endpoint_generation == rhs.endpoint_generation; + lhs.endpoint_generation == rhs.endpoint_generation && + lhs.receiver_credit_session_high == + rhs.receiver_credit_session_high && + lhs.receiver_credit_session_low == rhs.receiver_credit_session_low && + lhs.receiver_credit_epoch == rhs.receiver_credit_epoch; } } // namespace @@ -340,10 +344,26 @@ Status UbEndpoint::makeBootstrapDesc(const std::string& segment_name, } output.endpoint_generation = generation_; output.segment_generation = segment_generation; - output.capabilities = {"read", "write", "endpoint_generation"}; + output.capabilities = {"read", "write", "endpoint_generation", + "receiver_credit_v1"}; return Status::OK(); } +bool UbEndpoint::peerSupportsCapability(std::string_view capability) const { + std::lock_guard lock(lifecycle_mutex_); + return std::find(peer_.capabilities.begin(), peer_.capabilities.end(), + capability) != peer_.capabilities.end(); +} + +bool UbEndpoint::peerMatchesReceiverCredit(uint64_t session_high, + uint64_t session_low, + uint64_t epoch) const { + std::lock_guard lock(lifecycle_mutex_); + return peer_.receiver_credit_session_high == session_high && + peer_.receiver_credit_session_low == session_low && + peer_.receiver_credit_epoch == epoch; +} + bool UbEndpoint::tryAcquireOutstanding(uint64_t bytes) noexcept { std::lock_guard lock(lifecycle_mutex_); if (state_.load(std::memory_order_relaxed) != State::kReady || !context_ || diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index 652b18c6d3..59c9c14422 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -36,6 +37,8 @@ namespace mooncake::tent { namespace { +constexpr std::string_view kReceiverCreditCapability = "receiver_credit_v1"; + bool filterAllows(const std::vector& filter, const ub::DeviceInfo& device) { if (filter.empty()) return true; @@ -45,6 +48,43 @@ bool filterAllows(const std::vector& filter, filter.end(); } +bool receiverCreditRequired(const SegmentDesc* segment) { + if (!segment || segment->type != SegmentType::Memory) return false; + const auto& advert = segment->getMemory().receiver_credit; + return advert && advert->schema_version == kReceiverCreditProtocolVersion && + advert->flags == kReceiverCreditRequired && + !advert->receiver_session_id.empty() && advert->epoch != 0; +} + +const ReceiverCreditAdvertV1* receiverCreditAdvert(const SegmentDesc* segment) { + if (!receiverCreditRequired(segment)) return nullptr; + return &*segment->getMemory().receiver_credit; +} + +void bindReceiverCreditIdentity(const SegmentDesc* segment, + UbBootstrapDesc& bootstrap) { + const auto* advert = receiverCreditAdvert(segment); + if (!advert) return; + bootstrap.receiver_credit_session_high = advert->receiver_session_id.high; + bootstrap.receiver_credit_session_low = advert->receiver_session_id.low; + bootstrap.receiver_credit_epoch = advert->epoch; +} + +bool matchesReceiverCredit(const UbBootstrapDesc& bootstrap, + const ReceiverCreditAdvertV1& advert) { + return bootstrap.receiver_credit_session_high == + advert.receiver_session_id.high && + bootstrap.receiver_credit_session_low == + advert.receiver_session_id.low && + bootstrap.receiver_credit_epoch == advert.epoch; +} + +bool supportsReceiverCredit(const UbBootstrapDesc& bootstrap) { + return std::find(bootstrap.capabilities.begin(), + bootstrap.capabilities.end(), + kReceiverCreditCapability) != bootstrap.capabilities.end(); +} + uint64_t generationSeed() { const auto wall = static_cast( std::chrono::duration_cast( @@ -510,7 +550,22 @@ struct UbTransport::Impl { request.remote_topology_id, peer_path}; CHECK_STATUS( endpoints->getOrCreate(key, request.local_context, endpoint)); + const auto* credit_advert = + receiverCreditAdvert(request.remote_segment); + const bool credit_required = credit_advert != nullptr; if (endpoint->ready()) { + if (credit_required && + (!endpoint->peerSupportsCapability(kReceiverCreditCapability) || + !endpoint->peerMatchesReceiverCredit( + credit_advert->receiver_session_id.high, + credit_advert->receiver_session_id.low, + credit_advert->epoch))) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return Status::InvalidArgument( + "Remote UB endpoint lacks required receiver-credit " + "support" LOC_MARK); + } recordReadyEndpoint(endpoint); return Status::OK(); } @@ -527,9 +582,20 @@ struct UbTransport::Impl { endpoint.reset(); return status; } + auto local_segment = + metadata ? metadata->segmentManager().getLocal() : SegmentDescRef{}; + bindReceiverCreditIdentity(local_segment.get(), bootstrap); UbBootstrapDesc response; status = ControlClient::bootstrapUb( request.remote_segment->rpc_server_addr, bootstrap, response); + if (status.ok() && credit_required) { + if (!supportsReceiverCredit(response) || + !matchesReceiverCredit(response, *credit_advert)) { + status = Status::InvalidArgument( + "Remote UB bootstrap receiver-credit identity " + "mismatch" LOC_MARK); + } + } if (status.ok()) status = endpoint->bind(response); if (!status.ok()) { (void)endpoints->retire(endpoint); @@ -545,6 +611,14 @@ struct UbTransport::Impl { response.reply_msg = "UB transport is shutting down"; return -1; } + auto local_segment = + metadata ? metadata->segmentManager().getLocal() : SegmentDescRef{}; + if (receiverCreditRequired(local_segment.get()) && + !supportsReceiverCredit(request)) { + response.reply_msg = + "UB bootstrap peer lacks required receiver-credit support"; + return -1; + } const std::string local_name = getNicNameFromNicPath(request.peer_nic_path); auto context_it = context_by_topology_name.find(local_name); @@ -571,6 +645,7 @@ struct UbTransport::Impl { request.local_nic_path, currentSegmentGeneration(), response); } if (status.ok()) { + bindReceiverCreditIdentity(local_segment.get(), response); recordReadyEndpoint(endpoint); } if (!status.ok()) { @@ -705,6 +780,7 @@ Status UbTransport::submitTransferTasks( const auto notify_progress = ub_batch->notify_progress; const auto progress_batch_id = ub_batch->progress_batch_id; + auto progress_armed = std::make_shared>(false); std::vector new_tasks; new_tasks.reserve(request_list.size()); for (const auto& request : request_list) { @@ -717,9 +793,12 @@ Status UbTransport::submitTransferTasks( "transports/ub/max_slices_per_task" LOC_MARK); } auto task = ub::UbTask::create( - request, - [notify_progress, progress_batch_id](const TransferStatus&) { - if (notify_progress) notify_progress(progress_batch_id); + request, [notify_progress, progress_batch_id, + progress_armed](const TransferStatus&) { + if (progress_armed->load(std::memory_order_acquire) && + notify_progress) { + notify_progress(progress_batch_id); + } }); size_t offset = 0; while (offset < request.length) { @@ -731,6 +810,11 @@ Status UbTransport::submitTransferTasks( spec.length = length; spec.request_offset = offset; spec.max_retries = impl_->params.max_retries; + spec.receiver_credit_session_high = + ub_batch->receiver_credit_session_high; + spec.receiver_credit_session_low = + ub_batch->receiver_credit_session_low; + spec.receiver_credit_epoch = ub_batch->receiver_credit_epoch; if (!task->addSlice(spec)) { return Status::InternalError( "Unable to construct UB slice" LOC_MARK); @@ -741,15 +825,15 @@ Status UbTransport::submitTransferTasks( new_tasks.push_back(std::move(task)); } - for (auto& task : new_tasks) { - ub_batch->task_list.push_back(task); - auto status = impl_->workers->submit(task, ub_batch->device_mask); - if (!status.ok()) { - for (auto& queued : new_tasks) { - if (queued) queued->requestCancellation(); - } - return status; - } + CHECK_STATUS(impl_->workers->submitBatch(new_tasks, ub_batch->device_mask)); + ub_batch->task_list.insert(ub_batch->task_list.end(), new_tasks.begin(), + new_tasks.end()); + progress_armed->store(true, std::memory_order_release); + // A very small task can finish between worker publication and the batch + // append above. Its callback intentionally stays disarmed until the task + // is queryable; this notification closes that race. + if (!new_tasks.empty() && notify_progress) { + notify_progress(progress_batch_id); } return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp index effe77fff2..28e3dc0e79 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -275,44 +276,101 @@ Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { if (!task) { return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK); } + return submitBatch({task}, device_mask); +} + +Status UbWorkers::submitBatch(const std::vector& tasks, + uint64_t device_mask) { if (!accepting_.load(std::memory_order_acquire)) { return Status::InvalidArgument( "UB workers are not accepting work" LOC_MARK); } - const auto snapshot = task->snapshot(); - if (!snapshot.sealed) { - return Status::InvalidArgument( - "UB task must be sealed before submission" LOC_MARK); - } - const int priority = - std::clamp(task->request().priority, static_cast(PRIO_HIGH), - static_cast(PRIO_LOW)); std::vector pending; - for (const auto& slice : task->slices()) { - if (slice && slice->markQueued()) { + std::unordered_set unique_tasks; + for (const auto& task : tasks) { + if (!task) { + return Status::InvalidArgument( + "Cannot submit a null UB task" LOC_MARK); + } + if (!unique_tasks.insert(task.get()).second) { + return Status::InvalidArgument( + "Cannot submit the same UB task twice in one batch" LOC_MARK); + } + const auto snapshot = task->snapshot(); + if (!snapshot.sealed) { + return Status::InvalidArgument( + "UB task must be sealed before submission" LOC_MARK); + } + const int priority = + std::clamp(task->request().priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW)); + for (const auto& slice : task->slices()) { + if (!slice || slice->snapshot().state != UbSliceState::kInitial) { + return Status::InvalidArgument( + "UB task contains a slice that cannot be " + "submitted" LOC_MARK); + } pending.push_back(PendingSlice{task, slice, device_mask, priority, task->request().target_id, task->request().opcode}); } } - bool stopped_during_submit = false; + + Status failure = Status::OK(); + bool notify = false; + size_t marked_count = 0; { std::lock_guard lock(queue_mutex_); if (!accepting_.load(std::memory_order_relaxed)) { - stopped_during_submit = true; - } else { - for (auto& item : pending) { - queues_[item.priority].push_back(std::move(item)); + failure = Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + + if (failure.ok()) { + for (const auto& item : pending) { + if (!item.slice->markQueued()) { + failure = Status::InvalidArgument( + "UB task changed while it was being " + "submitted" LOC_MARK); + break; + } + ++marked_count; + } + } + + if (failure.ok()) { + std::array inserted_counts{}; + try { + for (const auto& item : pending) { + queues_[item.priority].push_back(item); + ++inserted_counts[item.priority]; + } + notify = !pending.empty(); + } catch (const std::bad_alloc&) { + for (size_t priority = 0; priority < queues_.size(); + ++priority) { + while (inserted_counts[priority] != 0) { + queues_[priority].pop_back(); + --inserted_counts[priority]; + } + } + failure = Status::InternalError( + "Unable to enqueue UB task batch" LOC_MARK); } } } - if (stopped_during_submit) { - for (auto& item : pending) item.slice->requestCancellation(); - return Status::InvalidArgument( - "UB workers stopped during submission" LOC_MARK); + + if (!failure.ok()) { + // A markQueued() race or allocation failure may occur after earlier + // slices changed state. None of them are visible to posting workers; + // cancel the private tasks so no slice remains orphaned in kQueued. + for (size_t index = 0; index < marked_count; ++index) { + (void)pending[index].task->requestCancellation(); + } + return failure; } - if (!pending.empty()) queue_cv_.notify_all(); + if (notify) queue_cv_.notify_all(); return Status::OK(); } @@ -472,9 +530,28 @@ Status UbWorkers::buildRoute(const PendingSlice& pending, Route& route) { return Status::InvalidMetadataType( "UB target is not a memory segment" LOC_MARK); } + const auto& spec = pending.slice->spec(); + const bool has_credit_fence = + spec.receiver_credit_session_high != 0 || + spec.receiver_credit_session_low != 0 || + spec.receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + spec.receiver_credit_session_high || + advert->receiver_session_id.low != + spec.receiver_credit_session_low || + advert->epoch != spec.receiver_credit_epoch) { + return Status::InvalidEntry( + "UB receiver-credit session changed before " + "post" LOC_MARK); + } + } auto* buffer = - segment->findBuffer(pending.slice->spec().remote_address, - pending.slice->spec().length); + segment->findBuffer(spec.remote_address, spec.length); if (!buffer) { return Status::NeedsRefreshCache( "UB target range is not registered" LOC_MARK); diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index ab404fff3c..438bb3057e 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -252,6 +252,193 @@ TEST(AdmissionQueueTest, KeepsAdmissionOrderForDispatch) { EXPECT_EQ(picked, expected_ids); } +TEST(AdmissionQueueTest, DefersDispatchingOwnerToFifoTailWithoutReaccounting) { + LocalTransferAdmissionQueue queue({4, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {makeOwner(0, 16), makeOwner(1, 24), + makeOwner(2, 32)}), + admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 3u); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + const size_t owners_before = queue.outstandingOwners(); + const size_t bytes_before = queue.outstandingBytes(); + + EXPECT_TRUE(queue.defer(picked[0]).ok()); + EXPECT_EQ(queue.outstandingOwners(), owners_before); + EXPECT_EQ(queue.outstandingBytes(), bytes_before); + + picked = queue.pickForDispatch(3, 128); + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfOwnerYieldsToCurrentWaiters) { + QueueLimits limits{4, 128, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 100; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 200; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + EXPECT_TRUE(queue.defer(picked[0]).ok()); + + picked = queue.pickForDispatch(3, 128); + // A temporarily blocked earliest-deadline owner yields one round so it + // cannot starve later-deadline work that may have available resources. + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfOwnerKeepsRoundAheadOfNewAdmission) { + QueueLimits limits{4, 128, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 200; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 300; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto newcomer = makeOwner(0, 16); + newcomer.request.deadline_ns = 150; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(2, 1, {std::move(newcomer)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids, std::vector({4})); + + picked = queue.pickForDispatch(4, 128); + // The deferred owner starts the next EDF round. A newly admitted owner is + // ordered within that round instead of jumping ahead of owners 2 and 3, + // which were already waiting when owner 1 yielded. + const std::vector expected{2, 3, 1, 4}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfRoundsStayOrderedAcrossAdmissions) { + QueueLimits limits{6, 256, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 200; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 300; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 256); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto deadline_250 = makeOwner(0, 16); + deadline_250.request.deadline_ns = 250; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(2, 1, {std::move(deadline_250)}), admitted_ids) + .ok()); + + picked = queue.pickForDispatch(1, 256); + ASSERT_EQ(picked, std::vector({2})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto deadline_225 = makeOwner(0, 16); + deadline_225.request.deadline_ns = 225; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(3, 1, {std::move(deadline_225)}), admitted_ids) + .ok()); + + picked = queue.pickForDispatch(5, 256); + const std::vector expected{3, 1, 2, 5, 4}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionCannotUndoCreditDeferralRound) { + QueueLimits limits{2, 128, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = 10; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{100}; }); + + auto starved = makeOwner(0, 16); + starved.request.deadline_ns = 105; + auto waiter = makeOwner(1, 16); + waiter.request.deadline_ns = 1000; + std::vector admitted_ids; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(1, 2, {std::move(starved), std::move(waiter)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + // Owner 1 remains urgent, but it yielded into round 1. Promotion cannot + // move it ahead of owner 2, which is still waiting in round 0. + picked = queue.pickForDispatch(1, 128); + EXPECT_EQ(picked, std::vector({2})); +} + +TEST(AdmissionQueueTest, RejectsDeferForNonDispatchingOwner) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 1u); + + EXPECT_TRUE(queue.defer(0).IsInvalidArgument()); + EXPECT_TRUE(queue.defer(99).IsInvalidEntry()); + EXPECT_TRUE(queue.defer(admitted_ids[0]).IsInvalidEntry()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked.size(), 1u); + ASSERT_TRUE(queue.complete(picked[0], COMPLETED).ok()); + EXPECT_TRUE(queue.defer(picked[0]).IsInvalidEntry()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); +} + TEST(AdmissionQueueTest, RequiresDispatchBeforeTerminalCompletion) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; @@ -828,6 +1015,29 @@ TEST(AdmissionQueueTest, EXPECT_EQ(queue.dispatchingBytes(), 0u); } +TEST(AdmissionQueueTest, RejectsDeferForOwnerDroppedDuringDispatchPick) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(1, 1, + {makeDegradationEligibleOwnerWithDeadline( + 0, 16, 1'000'000'010)}), + admitted_ids) + .ok()); + + std::vector dropped; + EXPECT_TRUE(queue.pickForDispatch(1, 128, &dropped).empty()); + ASSERT_EQ(dropped, std::vector({1})); + EXPECT_TRUE(queue.defer(dropped[0]).IsInvalidEntry()); + EXPECT_TRUE(queue.pickForDispatch(1, 128).empty()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); +} + TEST(AdmissionQueueTest, Step3DropsAlreadyExpiredDeadline) { LocalTransferAdmissionQueue queue(step3Limits(1.5)); queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp index 33faa04a2c..86b74670a1 100644 --- a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp @@ -10,7 +10,7 @@ namespace mooncake::tent { namespace { -CreditKey key() { return {{1, 2}, 3, 4}; } +CreditKey key() { return {{1, 2}, {3, 4}, 1}; } ReceiverCreditUpdateV1 update(uint64_t epoch, uint64_t seq, std::vector grants) { ReceiverCreditUpdateV1 u; @@ -18,6 +18,7 @@ ReceiverCreditUpdateV1 update(uint64_t epoch, uint64_t seq, u.qos_class = key().qos_class; u.epoch = epoch; u.sequence = seq; + u.freshness_ttl_ms = 1000; u.grants = std::move(grants); return u; } @@ -126,7 +127,7 @@ TEST(ReceiverCredit, LedgerEntryCountIsBounded) { SenderCreditLedger l(1); ASSERT_TRUE(l.activate(key(), 7).ok()); auto other = key(); - ++other.sender_peer; + ++other.sender_instance.low; EXPECT_TRUE(l.activate(other, 7).IsTooManyRequests()); // A new epoch for an existing key does not consume another entry. EXPECT_TRUE(l.activate(key(), 8).ok()); @@ -138,7 +139,7 @@ TEST(ReceiverCredit, DeactivationReleasesCapacityAfterExactEpochFence) { grant(l, 1); ASSERT_TRUE(l.tryReserve(key(), charge(80, 1)).ok()); auto other = key(); - ++other.sender_peer; + ++other.sender_instance.low; EXPECT_TRUE(l.activate(other, 1).IsTooManyRequests()); EXPECT_TRUE(l.deactivate(key(), 6).IsInvalidEntry()); @@ -201,6 +202,27 @@ TEST(ReceiverCredit, RollbackChecksUnderflowAtomically) { 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.consumed(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 60); // cumulative consumption never moves backwards + ASSERT_TRUE(l.released(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 60); + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 40); // returned authorization is not locally re-minted +} + +TEST(ReceiverCredit, DeactivateSessionRemovesAllQosEntries) { + SenderCreditLedger ledger(2); + auto high = key(); + auto low = key(); + low.qos_class = 2; + ASSERT_TRUE(ledger.activate(high, 7).ok()); + ASSERT_TRUE(ledger.activate(low, 7).ok()); + ASSERT_TRUE(ledger.deactivateSession(key().receiver_session).ok()); + + auto replacement = key(); + replacement.receiver_session = {9, 10}; + replacement.sender_instance = {11, 12}; + ASSERT_TRUE(ledger.activate(replacement, 1).ok()); } TEST(ReceiverCredit, GrantCannotDecreaseOrFallBelowConsumption) { @@ -232,5 +254,335 @@ TEST(ReceiverCredit, ConcurrentReservationsNeverExceedGrant) { ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, consumed).ok()); EXPECT_EQ(consumed, 100); } + +TEST(ReceiverCredit, FreshnessExpiryFailsClosedAndReplayDoesNotRenew) { + uint64_t now_ns = 100; + SenderCreditLedger ledger(8, 1000, [&] { return now_ns; }); + ASSERT_TRUE(ledger.activate(key(), 7).ok()); + auto first = update( + 7, 1, + {{CreditResource::DataBytes, 100}, {CreditResource::RequestSlots, 2}}); + first.freshness_ttl_ms = 1; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger.applyUpdate(key(), first, disposition).ok()); + now_ns += 999999; + EXPECT_TRUE(ledger.tryReserve(key(), charge(1, 1)).ok()); + + // A replay is accepted idempotently, but must not extend the lease. + ASSERT_TRUE(ledger.applyUpdate(key(), first, disposition).ok()); + EXPECT_EQ(disposition, CreditUpdateDisposition::DuplicateOrOld); + ++now_ns; + EXPECT_TRUE(ledger.tryReserve(key(), charge(1, 1)).IsInvalidEntry()); +} + +TEST(ReceiverCredit, TerminalReleaseRequiresReceiverReplenishment) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 2; + limits.freshness_ttl_ms = 1000; + ReceiverCreditAuthority authority({9, 10}, 3, limits); + ASSERT_TRUE(authority.status().ok()); + + CreditKey sender{{9, 10}, {11, 12}, 1}; + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(sender, 3).ok()); + + ReceiverCreditExchangeRequestV1 request; + ASSERT_TRUE(ledger.prepareExchange(sender, charge(60, 1), request).ok()); + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger.applyUpdate(sender, reply.update, disposition).ok()); + ASSERT_TRUE(ledger.tryReserve(sender, charge(60, 1)).ok()); + ASSERT_TRUE(ledger.releaseCommitted(sender, charge(60, 1)).ok()); + + uint64_t available = 0; + ASSERT_TRUE( + ledger.available(sender, CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, 0); // local release never mints credit + + ASSERT_TRUE(ledger.prepareExchange(sender, charge(60, 1), request).ok()); + ASSERT_TRUE(authority.exchange(request, reply).ok()); + ASSERT_TRUE(ledger.applyUpdate(sender, reply.update, disposition).ok()); + EXPECT_TRUE(ledger.tryReserve(sender, charge(60, 1)).ok()); + + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 60); +} + +TEST(ReceiverCredit, AuthorityGrantsMultiResourceDemandAtomically) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender, uint64_t bytes) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = { + {CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}, + }; + request.desired_grant_totals = { + {CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, 1}, + }; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 60), reply).ok()); + ASSERT_EQ(reply.update.grants.size(), kCreditResourceCount); + EXPECT_EQ(reply.update.grants[0].grant_total, 60); + EXPECT_EQ(reply.update.grants[1].grant_total, 1); + + ASSERT_TRUE(authority.exchange(make_request({5, 6}, 40), reply).ok()); + // The second sender cannot get a request slot, so bytes are not stranded + // in a partial grant either. + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + EXPECT_EQ(reply.update.grants[1].grant_total, 0); + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 60); +} + +TEST(ReceiverCredit, RollbackReleaseLetsAnotherSenderAcquireCapacity) { + ReceiverCreditLimits limits; + limits.capacity[0] = 60; + limits.capacity[1] = 1; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + SenderCreditLedger first; + CreditKey first_key{{1, 2}, {3, 4}, 0}; + ASSERT_TRUE(first.activate(first_key, 1).ok()); + + ReceiverCreditExchangeRequestV1 request; + ASSERT_TRUE(first.prepareExchange(first_key, charge(60, 1), request).ok()); + const auto initial_demand = request; + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE(first.applyUpdate(first_key, reply.update, disposition).ok()); + ASSERT_TRUE(first.tryReserve(first_key, charge(60, 1)).ok()); + ASSERT_TRUE(first.rollbackReservation(first_key, charge(60, 1)).ok()); + ASSERT_TRUE(first.prepareReleaseExchange(first_key, request).ok()); + const auto release_report = request; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + EXPECT_EQ(reply.flags, 0); + + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 0); + + // Both a delayed initial demand and a replayed release receive retained + // cumulative state; neither can re-mint released capacity. + ASSERT_TRUE(authority.exchange(initial_demand, reply).ok()); + EXPECT_EQ(reply.flags, 0); + EXPECT_EQ(reply.update.grants[0].grant_total, 60); + ASSERT_TRUE(authority.exchange(release_report, reply).ok()); + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 0); + + SenderCreditLedger second; + CreditKey second_key{{1, 2}, {5, 6}, 0}; + ASSERT_TRUE(second.activate(second_key, 1).ok()); + ASSERT_TRUE( + second.prepareExchange(second_key, charge(60, 1), request).ok()); + ASSERT_TRUE(authority.exchange(request, reply).ok()); + ASSERT_TRUE(second.applyUpdate(second_key, reply.update, disposition).ok()); + EXPECT_TRUE(second.tryReserve(second_key, charge(60, 1)).ok()); +} + +TEST(ReceiverCredit, SenderLimitCountsSenderNotQosEntries) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 2; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender, uint32_t qos) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = qos; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 0), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 2), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({5, 6}, 0), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + EXPECT_EQ(reply.update.grants[1].grant_total, 0); +} + +TEST(ReceiverCredit, ZeroGrantHistoryDoesNotConsumeActiveSenderLimit) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 1; + limits.max_entries = 4; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + auto first = make_request({3, 4}); + ReceiverCreditExchangeReplyV1 first_reply; + ASSERT_TRUE(authority.exchange(first, first_reply).ok()); + EXPECT_EQ(first_reply.update.grants[0].grant_total, 10); + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({5, 6}), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + + first.known_update_sequence = first_reply.update.sequence; + first.report_sequence = 2; + first.released_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + ASSERT_TRUE(authority.exchange(first, reply).ok()); + + ASSERT_TRUE(authority.exchange(make_request({7, 8}), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 10); + EXPECT_EQ(reply.update.grants[1].grant_total, 1); +} + +TEST(ReceiverCredit, ReplayHistoryLimitFailsClosed) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 1; + limits.max_entries = 2; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({5, 6}), reply).ok()); + EXPECT_TRUE( + authority.exchange(make_request({7, 8}), reply).IsTooManyRequests()); + + // Existing identities remain replayable after the cap is reached. + EXPECT_TRUE(authority.exchange(make_request({3, 4}), reply).ok()); +} + +TEST(ReceiverCredit, ReplayHistoryLimitMustCoverActiveSenders) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 2; + limits.max_entries = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + EXPECT_TRUE(authority.status().IsInvalidArgument()); +} + +TEST(ReceiverCredit, AdvertPublishesHardCapacityBounds) { + ReceiverCreditLimits limits; + limits.capacity[0] = 4096; + limits.capacity[1] = 8; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + const auto advert = authority.advert(); + ASSERT_EQ(advert.capacities.size(), 2u); + EXPECT_EQ(advert.capacities[0].resource, CreditResource::DataBytes); + EXPECT_EQ(advert.capacities[0].total, 4096u); + EXPECT_EQ(advert.capacities[1].resource, CreditResource::RequestSlots); + EXPECT_EQ(advert.capacities[1].total, 8u); + + const auto round_trip = + nlohmann::json(advert).get(); + EXPECT_EQ(round_trip.capacities.size(), advert.capacities.size()); + EXPECT_EQ(round_trip.capacities[0].total, 4096u); +} + +TEST(ReceiverCredit, ExchangeReplayIsIdempotent) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = {3, 4}; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 50}, + {CreditResource::RequestSlots, 1}}; + + ReceiverCreditExchangeReplyV1 first, replay; + ASSERT_TRUE(authority.exchange(request, first).ok()); + ASSERT_TRUE(authority.exchange(request, replay).ok()); + EXPECT_EQ(replay.update.sequence, first.update.sequence); + EXPECT_EQ(replay.update.grants[0].grant_total, 50); + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 50); +} + +TEST(ReceiverCredit, InvalidNewSenderCannotConsumeAuthorityEntry) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = {3, 4}; + request.qos_class = 0; + request.known_update_sequence = 99; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 50}, + {CreditResource::RequestSlots, 1}}; + ReceiverCreditExchangeReplyV1 reply; + EXPECT_TRUE(authority.exchange(request, reply).IsInvalidArgument()); + + request.sender_instance_id = {5, 6}; + request.known_update_sequence = 0; + EXPECT_TRUE(authority.exchange(request, reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 50); +} } // namespace } // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp index 116616e30e..638dd2ff7d 100644 --- a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp +++ b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp @@ -24,6 +24,7 @@ struct RequestBoundaryInfo { std::optional source_key; std::optional target_key; uint64_t max_merge_bytes{std::numeric_limits::max()}; + std::optional receiver_credit_data_capacity; }; struct MergeResult { 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 41149ed4e1..dbaddebfc5 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -37,6 +37,12 @@ namespace mooncake { // left the normal lifecycle still holds. class TransferEngineImplTestPeer { public: + static tent::Status submitRequiringPostSubmitCancellation( + tent::TransferEngineImpl& engine, tent::BatchID batch_id, + const std::vector& requests) { + return engine.submitTransferRequiringPostSubmitCancellation(batch_id, + requests); + } static size_t queuedOwners(tent::TransferEngineImpl& engine) { return engine.queued_owners_.size(); } @@ -101,6 +107,10 @@ class FakeTransport : public Transport { std::atomic submit_calls{0}; std::atomic status_calls{0}; std::atomic cancel_calls{0}; + std::atomic fail_next_submits{0}; + std::atomic last_credit_session_high{0}; + std::atomic last_credit_session_low{0}; + std::atomic last_credit_epoch{0}; bool cancellation_supported{true}; Status install(std::string&, std::shared_ptr, @@ -123,6 +133,21 @@ class FakeTransport : public Transport { Status submitTransferTasks(SubBatchRef batch, const std::vector& requests) override { ++submit_calls; + int failures = fail_next_submits.load(std::memory_order_relaxed); + while (failures > 0 && + !fail_next_submits.compare_exchange_weak( + failures, failures - 1, std::memory_order_relaxed)) { + } + if (failures > 0) { + return Status::InternalError( + "injected zero-handoff submit failure" LOC_MARK); + } + last_credit_session_high.store(batch->receiver_credit_session_high, + std::memory_order_relaxed); + last_credit_session_low.store(batch->receiver_credit_session_low, + std::memory_order_relaxed); + last_credit_epoch.store(batch->receiver_credit_epoch, + std::memory_order_relaxed); auto* fake = static_cast(batch); for (const auto& request : requests) { fake->requests.push_back(request); @@ -252,6 +277,13 @@ void installFakeRdma(TransferEngineImpl& engine, engine.swapTransportForTest(RDMA, fake_rdma); } +void installFakeUb(TransferEngineImpl& engine, + const std::shared_ptr& fake_ub) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake_ub->install(seg_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(UB, fake_ub); +} + Request makeLocalWrite(uint8_t* ptr, size_t length) { Request request; request.opcode = Request::WRITE; @@ -263,6 +295,26 @@ Request makeLocalWrite(uint8_t* ptr, size_t length) { return request; } +Request makeLocalUbWrite(uint8_t* ptr, size_t length) { + auto request = makeLocalWrite(ptr, length); + request.transport_hint = UB; + return request; +} + +Request makeRemoteUbWrite(uint8_t* source, SegmentID target_id, uint8_t* target, + size_t length) { + auto request = makeLocalUbWrite(source, length); + request.target_id = target_id; + request.target_offset = reinterpret_cast(target); + return request; +} + +Request makeLocalNetworkWrite(uint8_t* ptr, size_t length) { + auto request = makeLocalWrite(ptr, length); + request.transport_hint = UNSPEC; + return request; +} + TEST(RuntimeQueueDispatch, RejectsOverfullBatchBeforePublishingTasks) { auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); TransferEngineImpl engine(cfg); @@ -339,14 +391,13 @@ TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) { TransferEngineImpl engine(cfg); ASSERT_TRUE(engine.available()); - std::atomic complete_first{false}; + std::atomic complete{false}; auto fake_rdma = std::make_shared( - RDMA, [&complete_first](const Request& request, int) { - if (!complete_first.load()) { - return TransferStatus{TransferStatusEnum::PENDING, 0}; - } - return TransferStatus{TransferStatusEnum::COMPLETED, - request.length}; + RDMA, [&complete](const Request& request, int) { + return complete.load(std::memory_order_acquire) + ? TransferStatus{TransferStatusEnum::COMPLETED, + request.length} + : TransferStatus{TransferStatusEnum::PENDING, 0}; }); installFakeRdma(engine, fake_rdma); @@ -365,8 +416,7 @@ TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) { .ok()); EXPECT_EQ(fake_rdma->submit_calls.load(), 1); - complete_first.store(true); - + complete.store(true, std::memory_order_release); TransferStatus first{}; ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); @@ -1104,6 +1154,390 @@ TEST(RuntimeQueueDispatch, TcpCapsPreservePublicMappingAndOwnerBytes) { } } +TEST(RuntimeQueueDispatch, ScatterCannotBypassRequiredReceiverCredit) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_ub = std::make_shared(UB); + installFakeUb(engine, fake_ub); + std::vector buffer(kReqLen, 0x71); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + auto status = + TransferEngineImplTestPeer::submitRequiringPostSubmitCancellation( + engine, batch, {makeLocalUbWrite(buffer.data(), kReqLen)}); + EXPECT_TRUE(status.IsNotImplemented()); + EXPECT_EQ(fake_ub->submit_calls.load(), 0); + EXPECT_EQ(TransferEngineImplTestPeer::queuedOwners(engine), 0u); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ReceiverCreditDefersUbUntilTerminalRelease) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20, true); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + cfg->set("receiver_credit/freshness_ttl_ms", uint32_t{1000}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_ub = std::make_shared( + UB, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeUb(engine, fake_ub); + + std::vector buffer(kReqLen * 2, 0x71); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), kReqLen), + makeLocalUbWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + // Both owners fit the local dispatch window, but receiver credit permits + // only one physical UB request. Merging must stop at the advertised hard + // byte capacity, so the second owner remains queued instead of making an + // inadmissible 8 KiB owner. + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_ub->submit_calls.load(), 2); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ZeroHandoffUbFailureRollsCreditBack) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_ub = std::make_shared(UB); + fake_ub->fail_next_submits.store(1); + installFakeUb(engine, fake_ub); + std::vector buffer(kReqLen * 2, 0x72); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), kReqLen), + makeLocalUbWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + // The first submit failed before handoff. Its cumulative rollback is + // reported to the receiver before the next owner gets a fresh grant. + EXPECT_EQ(fake_ub->submit_calls.load(), 2); + TransferStatus first{}, second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::FAILED); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RemoteReceiverCreditExchangeGatesUbDispatch) { + constexpr size_t kReqLen = 4096; + auto receiver_cfg = makeRuntimeQueueConfig(2, 1UL << 20); + receiver_cfg->set("receiver_credit/receiver_enable", true); + receiver_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + receiver_cfg->set("receiver_credit/request_slots", uint64_t{1}); + auto sender_cfg = makeRuntimeQueueConfig(2, 1UL << 20); + sender_cfg->set("receiver_credit/sender_enable", true); + sender_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + sender_cfg->set("receiver_credit/request_slots", uint64_t{1}); + + TransferEngineImpl receiver(receiver_cfg); + TransferEngineImpl sender(sender_cfg); + ASSERT_TRUE(receiver.available()); + ASSERT_TRUE(sender.available()); + + auto receiver_ub = std::make_shared(UB); + installFakeUb(receiver, receiver_ub); + std::atomic complete_first{false}; + auto sender_ub = std::make_shared( + UB, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeUb(sender, sender_ub); + + std::vector source(kReqLen * 2, 0x81); + std::vector target(kReqLen * 2, 0); + ASSERT_TRUE(sender.registerLocalMemory(source.data(), source.size()).ok()); + ASSERT_TRUE( + receiver.registerLocalMemory(target.data(), target.size()).ok()); + SegmentID remote = LOCAL_SEGMENT_ID; + ASSERT_TRUE(sender.openSegment(remote, receiver.getSegmentName()).ok()); + ASSERT_NE(remote, LOCAL_SEGMENT_ID); + + BatchID batch = sender.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + sender + .submitTransfer( + batch, {makeRemoteUbWrite(source.data(), remote, target.data(), + kReqLen), + makeRemoteUbWrite(source.data() + kReqLen, remote, + target.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(sender_ub->submit_calls.load(), 1); + + // TCP fallback carries the same receiver-session fence on the wire. The + // receiver rejects an unfenced or stale request even if the sender's + // metadata cache still describes the old process incarnation. + const auto receiver_addr = receiver.getSegmentName(); + const uint64_t session_high = sender_ub->last_credit_session_high.load(); + const uint64_t session_low = sender_ub->last_credit_session_low.load(); + const uint64_t epoch = sender_ub->last_credit_epoch.load(); + ASSERT_NE(session_high | session_low, 0u); + ASSERT_NE(epoch, 0u); + + auto delegated = makeLocalUbWrite(target.data(), kReqLen); + delegated.receiver_credit_session_high = session_high; + delegated.receiver_credit_session_low = session_low; + delegated.receiver_credit_epoch = epoch + 1; + EXPECT_TRUE( + ControlClient::delegate(receiver_addr, delegated).IsRpcServiceError()); + delegated.receiver_credit_epoch = epoch; + ASSERT_TRUE(ControlClient::delegate(receiver_addr, delegated).ok()); + EXPECT_EQ(receiver_ub->submit_calls.load(), 1); + + delegated.receiver_credit_epoch = 0; + EXPECT_TRUE( + ControlClient::delegate(receiver_addr, delegated).IsInvalidArgument()); + + EXPECT_TRUE(ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), + source.data(), kReqLen) + .IsRpcServiceError()); + EXPECT_EQ(target[0], 0); + EXPECT_TRUE(ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), + source.data(), kReqLen, session_high, session_low, + epoch + 1) + .IsRpcServiceError()); + EXPECT_EQ(target[0], 0); + const auto fenced_send = ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), source.data(), + kReqLen, session_high, session_low, epoch); + ASSERT_TRUE(fenced_send.ok()) << fenced_send.ToString(); + EXPECT_EQ(target[0], source[0]); + + std::vector readback(kReqLen, 0); + EXPECT_TRUE(ControlClient::recvData( + receiver_addr, reinterpret_cast(target.data()), + readback.data(), kReqLen, session_high, session_low, + epoch + 1) + .IsRpcServiceError()); + const auto fenced_recv = ControlClient::recvData( + receiver_addr, reinterpret_cast(target.data()), + readback.data(), kReqLen, session_high, session_low, epoch); + ASSERT_TRUE(fenced_recv.ok()) << fenced_recv.ToString(); + EXPECT_EQ(readback, std::vector(kReqLen, 0x81)); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}, second{}; + ASSERT_TRUE(sender.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(sender_ub->submit_calls.load(), 2); + ASSERT_TRUE(sender.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(sender.freeBatch(batch).ok()); + EXPECT_TRUE(sender.closeSegment(remote).ok()); + EXPECT_TRUE( + sender.unregisterLocalMemory(source.data(), source.size()).ok()); + EXPECT_TRUE( + receiver.unregisterLocalMemory(target.data(), target.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RdmaFailureGatesUbFailoverWithReceiverCredit) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = + std::make_shared(RDMA, [](const Request&, int) { + return TransferStatus{TransferStatusEnum::FAILED, 0}; + }); + auto fake_ub = std::make_shared(UB); + installFakeRdma(engine, fake_rdma); + installFakeUb(engine, fake_ub); + + std::vector buffer(kReqLen, 0x91); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE(engine + .submitTransfer( + batch, {makeLocalNetworkWrite(buffer.data(), kReqLen)}) + .ok()); + ASSERT_EQ(fake_rdma->submit_calls.load(), 1); + ASSERT_EQ(fake_ub->submit_calls.load(), 0); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ReceiverCreditAlsoGatesInitialRdmaPath) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_rdma = std::make_shared( + RDMA, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + std::vector buffer(kReqLen * 2, 0x92); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_NE(fake_rdma->last_credit_session_high.load() | + fake_rdma->last_credit_session_low.load(), + 0u); + EXPECT_EQ(fake_rdma->last_credit_epoch.load(), 1u); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}, second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsRequestLargerThanReceiverCreditCapacity) { + constexpr size_t kCapacity = 4096; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kCapacity}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_ub = std::make_shared(UB); + installFakeUb(engine, fake_ub); + std::vector buffer(kCapacity * 2, 0x93); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + + const auto status = engine.submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), buffer.size())}); + EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString(); + EXPECT_EQ(fake_ub->submit_calls.load(), 0); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, TcpControlDataUsesLegacyFramingWithoutCredit) { + constexpr size_t kLength = 256; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + std::vector source(kLength, 0xa5); + std::vector target(kLength, 0); + std::vector readback(kLength, 0); + ASSERT_TRUE(engine.registerLocalMemory(target.data(), target.size()).ok()); + const auto server_addr = engine.getSegmentName(); + + ASSERT_TRUE(ControlClient::sendData( + server_addr, reinterpret_cast(target.data()), + source.data(), source.size()) + .ok()); + EXPECT_EQ(target, source); + ASSERT_TRUE(ControlClient::recvData( + server_addr, reinterpret_cast(target.data()), + readback.data(), readback.size()) + .ok()); + EXPECT_EQ(readback, source); + + EXPECT_TRUE(ControlClient::sendData( + server_addr, reinterpret_cast(target.data()), + source.data(), source.size(), 1, 0, 0) + .IsInvalidArgument()); + EXPECT_TRUE( + engine.unregisterLocalMemory(target.data(), target.size()).ok()); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index dffe707bc0..625133e097 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -411,6 +411,9 @@ TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { source.jetty_uasids = {21, 22}; source.endpoint_generation = 41; source.segment_generation = 42; + source.receiver_credit_session_high = 43; + source.receiver_credit_session_low = 44; + source.receiver_credit_epoch = 45; source.capabilities = {"read", "write"}; source.reply_msg = "ok"; @@ -427,6 +430,11 @@ TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { EXPECT_EQ(parsed.jetty_uasids, source.jetty_uasids); EXPECT_EQ(parsed.endpoint_generation, source.endpoint_generation); EXPECT_EQ(parsed.segment_generation, source.segment_generation); + EXPECT_EQ(parsed.receiver_credit_session_high, + source.receiver_credit_session_high); + EXPECT_EQ(parsed.receiver_credit_session_low, + source.receiver_credit_session_low); + EXPECT_EQ(parsed.receiver_credit_epoch, source.receiver_credit_epoch); EXPECT_EQ(parsed.capabilities, source.capabilities); EXPECT_EQ(parsed.reply_msg, source.reply_msg); } @@ -445,6 +453,9 @@ TEST(ControlPlaneTest, UbBootstrapJsonDefaultsOptionalFields) { EXPECT_TRUE(parsed.jetty_uasids.empty()); EXPECT_EQ(parsed.endpoint_generation, 0u); EXPECT_EQ(parsed.segment_generation, 0u); + EXPECT_EQ(parsed.receiver_credit_session_high, 0u); + EXPECT_EQ(parsed.receiver_credit_session_low, 0u); + EXPECT_EQ(parsed.receiver_credit_epoch, 0u); EXPECT_TRUE(parsed.capabilities.empty()); EXPECT_TRUE(parsed.reply_msg.empty()); } @@ -470,6 +481,7 @@ TEST(ControlPlaneTest, UbBootstrapRpcIdIsAppendedWithoutRenumbering) { EXPECT_EQ(static_cast(SubscribeSegmentUpdate), 10); EXPECT_EQ(static_cast(NotifySegmentUpdated), 11); EXPECT_EQ(static_cast(BootstrapUb), 12); + EXPECT_EQ(static_cast(ExchangeReceiverCredit), 13); } // --------------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp index 4dc2694de6..2ad2b862fe 100644 --- a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -3,6 +3,7 @@ #include +#include #include #include #include @@ -547,6 +548,30 @@ TEST(UbNativeDataPathTest, (void)endpoints.retire(endpoint); }); ASSERT_TRUE(workers.start().ok()); + + Request invalid_batch_request{}; + invalid_batch_request.opcode = Request::WRITE; + invalid_batch_request.source = source.data(); + invalid_batch_request.target_id = LOCAL_SEGMENT_ID; + invalid_batch_request.target_offset = + reinterpret_cast(target.data()); + invalid_batch_request.length = 16; + auto valid_task = UbTask::create(invalid_batch_request); + auto valid_slice = valid_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, 0, 1}); + ASSERT_NE(valid_slice, nullptr); + ASSERT_TRUE(valid_task->seal()); + auto unsealed_task = UbTask::create(invalid_batch_request); + ASSERT_NE(unsealed_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, + 0, 1}), + nullptr); + + EXPECT_FALSE(workers.submitBatch({valid_task, unsealed_task}).ok()); + EXPECT_EQ(workers.queuedCount(), 0U); + EXPECT_EQ(workers.inflightCount(), 0U); + EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kInitial); + adapter->holdNextCompletion(); adapter->dropNextQuiescedCompletion(); @@ -762,6 +787,15 @@ TEST(UbNativeDataPathTest, EndpointStoreNeverReusesRetiredGeneration) { EXPECT_EQ(concurrent[i], concurrent[0]); } auto first = concurrent[0]; + UbBootstrapDesc local_bootstrap; + ASSERT_TRUE(first + ->makeBootstrapDesc("local", "local@ub:fake0:eid0", + key.peer_nic_path, 1, local_bootstrap) + .ok()); + EXPECT_NE( + std::find(local_bootstrap.capabilities.begin(), + local_bootstrap.capabilities.end(), "receiver_credit_v1"), + local_bootstrap.capabilities.end()); const uint64_t old_generation = first->generation(); EXPECT_TRUE(store.retire(key, old_generation)); std::shared_ptr replacement; @@ -885,8 +919,13 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { ASSERT_TRUE(transport.addMemoryBuffer(descriptors, options).ok()); ASSERT_TRUE(control->segmentManager() .updateLocal([&](SegmentDesc& segment) { - std::get(segment.detail).buffers = - descriptors; + auto& memory = + std::get(segment.detail); + memory.buffers = descriptors; + ReceiverCreditAdvertV1 advert; + advert.receiver_session_id = {1, 2}; + advert.epoch = 1; + memory.receiver_credit = std::move(advert); return Status::OK(); }) .ok()); @@ -904,6 +943,13 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { transport.submitTransferTasks(rejected_batch, {oversized_request}) .IsInvalidArgument()); EXPECT_TRUE(transport.freeSubBatch(rejected_batch).ok()); + UbBootstrapDesc unsupported_bootstrap; + UbBootstrapDesc rejected_bootstrap; + ASSERT_TRUE(ControlClient::bootstrapUb(segment_name, unsupported_bootstrap, + rejected_bootstrap) + .ok()); + EXPECT_NE(rejected_bootstrap.reply_msg.find("receiver-credit"), + std::string::npos); Transport::SubBatchRef batch = nullptr; ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok()); @@ -913,6 +959,11 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { request.target_id = LOCAL_SEGMENT_ID; request.target_offset = reinterpret_cast(target.data()); request.length = source.size(); + auto invalid_request = request; + invalid_request.length = 0; + EXPECT_FALSE( + transport.submitTransferTasks(batch, {request, invalid_request}).ok()); + EXPECT_EQ(batch->size(), 0U); ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); TransferStatus transfer{PENDING, 0}; From c0ba30f3a85f4b664f441889799e663fbd98fb96 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:29:31 +0800 Subject: [PATCH 2/8] [TENT] Preserve receiver credit fencing on HP TCP after migration --- .../src/transport/hp_tcp/hp_tcp_transport.cpp | 19 +++++ .../tent/tests/hp_tcp_transport_test.cpp | 79 +++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp index 978bd7d934..335360c70f 100644 --- a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp @@ -500,6 +500,25 @@ Status HighPerformanceTcpTransport::planTask( if (segment == nullptr || segment->type != SegmentType::Memory) { return NeedsRefresh("HP TCP target is not a memory segment"); } + const bool has_credit_fence = + batch.receiver_credit_session_high != 0 || + batch.receiver_credit_session_low != 0 || + batch.receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + batch.receiver_credit_session_high || + advert->receiver_session_id.low != + batch.receiver_credit_session_low || + advert->epoch != batch.receiver_credit_epoch) { + return Status::InvalidEntry( + "HP TCP receiver-credit session changed before " + "post" LOC_MARK); + } + } BufferDesc* buffer = segment->findBuffer(request.target_offset, request.length); if (buffer == nullptr || diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp index 452b2e3f04..abca485bc6 100644 --- a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp @@ -251,6 +251,85 @@ Status WaitForTransportResult(HighPerformanceTcpTransport& transport, return Status::InternalError("HP TCP test transfer did not finish"); } +TEST(HighPerformanceTcpTransportTest, ReceiverCreditFenceTracksPinnedMetadata) { + auto metadata = MakeLocalMetadata(); + HighPerformanceTcpTransport transport(MakeParams()); + std::string name = "hp_transport_test"; + ASSERT_TRUE(transport.install(name, metadata, nullptr, nullptr).ok()); + std::array source{}; + std::array target{}; + source.fill(0x5a); + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = "cpu:0"; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = "cpu:0"; + MemoryOptions options; + options.type = HP_TCP; + options.perm = kGlobalReadWrite; + ASSERT_TRUE(transport.addMemoryBuffer(source_desc, options).ok()); + ASSERT_TRUE(transport.addMemoryBuffer(target_desc, options).ok()); + ASSERT_TRUE(PublishBuffers(metadata, {source_desc, target_desc}).ok()); + + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = target_desc.addr; + request.length = source.size(); + ReceiverCreditAdvertV1 advert; + advert.receiver_session_id = {1, 2}; + advert.epoch = 7; + // No advert, wrong session, wrong epoch, and unsupported protocol must + // reject before worker handoff; a matching pinned advert must still work. + for (int scenario = 0; scenario < 6; ++scenario) { + SCOPED_TRACE(scenario); + auto current = advert; + if (scenario == 1) ++current.receiver_session_id.low; + if (scenario == 2) ++current.epoch; + if (scenario == 3) ++current.schema_version; + if (scenario == 4) current.flags = 0; + ASSERT_TRUE(metadata->segmentManager() + .updateLocal([&](SegmentDesc& segment) -> Status { + auto& published = + segment.getMemory().receiver_credit; + if (scenario == 0) + published.reset(); + else + published = current; + return Status::OK(); + }) + .ok()); + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 1).ok()); + batch->receiver_credit_session_high = advert.receiver_session_id.high; + batch->receiver_credit_session_low = advert.receiver_session_id.low; + batch->receiver_credit_epoch = advert.epoch; + const Status submitted = + transport.submitTransferTasks(batch, {request}); + if (scenario < 5) { + EXPECT_TRUE(submitted.IsInvalidEntry()) << submitted.ToString(); + EXPECT_EQ(batch->size(), 0u); + EXPECT_TRUE(std::all_of(target.begin(), target.end(), + [](uint8_t byte) { return byte == 0; })); + } else { + ASSERT_TRUE(submitted.ok()) << submitted.ToString(); + TransferStatus status{}; + ASSERT_TRUE(WaitForTransportResult(transport, batch, status).ok()); + EXPECT_EQ(status.s, COMPLETED); + EXPECT_EQ(source, target); + } + EXPECT_TRUE(transport.freeSubBatch(batch).ok()); + } + EXPECT_TRUE(transport.quiesce().ok()); + EXPECT_TRUE(transport.removeMemoryBuffer(source_desc).ok()); + EXPECT_TRUE(transport.removeMemoryBuffer(target_desc).ok()); + EXPECT_TRUE(transport.uninstall().ok()); +} + // Test-only relay: observe the wire request and forward it to the real server. // Pumped by the test thread; no transport worker waits for another callback. class RailRelay { From 999edc7f440cd9d8c3d0c322ea1ae048814239d3 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:08:24 +0800 Subject: [PATCH 3/8] [TENT] Fix TCP regression test after receiver-credit RPC changes ControlClient::sendData/recvData gained receiver-credit session/epoch parameters. tcp_regression_test calls them through function pointers, where default arguments do not apply, so the calls failed to compile. Pass explicit zero fences for the non-credit paths, and build the malformed-length probe from the legacy header layout so the server rejects it as an invalid transfer request size. --- .../tent/tests/tcp_regression_test.cpp | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/mooncake-transfer-engine/tent/tests/tcp_regression_test.cpp b/mooncake-transfer-engine/tent/tests/tcp_regression_test.cpp index 716ca8eaee..94b7f50c34 100644 --- a/mooncake-transfer-engine/tent/tests/tcp_regression_test.cpp +++ b/mooncake-transfer-engine/tent/tests/tcp_regression_test.cpp @@ -254,7 +254,7 @@ TEST_F(TcpRegression, ServerChecksEveryRegistrationPermission) { // Direct RPC bypasses client-side registration/routing policy. auto status = call(engine->getSegmentName(), reinterpret_cast(data.data()), - local.data(), local.size()); + local.data(), local.size(), 0, 0, 0); EXPECT_EQ(status.ok(), allowed) << status.ToString(); if (!allowed) { EXPECT_NE(status.ToString().find("permission denied"), @@ -298,7 +298,7 @@ TEST_F(TcpRegression, PayloadCapsRejectBeforeAllocationOrDataRpc) { EXPECT_NO_THROW({ EXPECT_TRUE(call(address, reinterpret_cast(data.data()), - data.data(), length) + data.data(), length, 0, 0, 0) .IsInvalidArgument()); }); } @@ -316,14 +316,20 @@ TEST_F(TcpRegression, MalformedWriteLengthCannotOverflowOrModifyTarget) { CoroRpcAgent rpc; for (size_t length : {size_t{1}, kTcpMaxWriteBytes + 1, std::numeric_limits::max()}) { - XferDataDesc header{htole64(reinterpret_cast(data.data())), - htole64(length)}; + LegacyXferDataDesc header{ + htole64(reinterpret_cast(data.data())), htole64(length)}; std::string request(reinterpret_cast(&header), sizeof(header)); std::string response; ASSERT_TRUE( rpc.call(engine->getSegmentName(), SendData, request, response) .ok()); - EXPECT_EQ(response, "SendData failed: invalid request size"); + // The v2 parser rejects the legacy header before any copy: the + // declared length can never match a request that carries no payload. + EXPECT_EQ(response.rfind("SendData failed: InvalidArgument: invalid " + "transfer request size", + 0), + 0u) + << response; EXPECT_EQ(data, original); } } From f1c63bddfacb5171d16f276ae39e9253fdba3c05 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:27:22 +0800 Subject: [PATCH 4/8] [TENT] Fix mutable receiver-credit metadata setup in HP TCP test --- mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp index abca485bc6..2271ccab2a 100644 --- a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp @@ -295,7 +295,8 @@ TEST(HighPerformanceTcpTransportTest, ReceiverCreditFenceTracksPinnedMetadata) { ASSERT_TRUE(metadata->segmentManager() .updateLocal([&](SegmentDesc& segment) -> Status { auto& published = - segment.getMemory().receiver_credit; + std::get(segment.detail) + .receiver_credit; if (scenario == 0) published.reset(); else From cd561d9a5fd67128f4f50c0c8b2b7a005cf6e597 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:26:10 +0800 Subject: [PATCH 5/8] [TENT] Surface RecvData server errors on TCP legacy replies ControlClient::recvData replaced the server's rejection text with a generic 'invalid or rejected legacy response' string for every short reply, so a denied READ lost 'permission denied' and the upstream tcp_regression permission test started failing. Propagate the server message when the reply cannot be payload, keeping a descriptive fallback for empty replies; apply the same treatment to the v2 reply path so fence rejections stay diagnosable. --- .../tent/src/runtime/control_plane.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 71cf343bf6..dc861e89fc 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -343,10 +343,16 @@ Status ControlClient::recvData(const std::string& server_addr, } auto status = tl_rpc_agent.call(server_addr, RecvData, request, response); if (!status.ok()) return status; + // A short reply is an error, not payload: surface the server's message so + // callers keep the rejection reason (e.g. "permission denied") instead of + // a generic string, matching the pre-v2 client behavior. An empty or + // unpadded reply still gets a descriptive fallback. if (!has_complete_fence) { if (response.size() != length) { return Status::RpcServiceError( - "RecvData failed: invalid or rejected legacy response"); + response.empty() + ? "RecvData failed: invalid or rejected legacy response" + : response); } return Platform::getLoader().copy(local_mem_addr, response.data(), length); @@ -355,7 +361,8 @@ Status ControlClient::recvData(const std::string& server_addr, std::numeric_limits::max() - sizeof(XferDataReplyHeader) || response.size() != sizeof(XferDataReplyHeader) + length) { return Status::RpcServiceError( - "RecvData failed: invalid or rejected response"); + response.empty() ? "RecvData failed: invalid or rejected response" + : response); } XferDataReplyHeader header; memcpy(&header, response.data(), sizeof(header)); From 68f1b9c6d858867e208f775aee5d138ed9539d95 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:08:24 +0800 Subject: [PATCH 6/8] [Store][Test] Restore DFS writes before retrying revoked key DfsOnlyBatchFailureOnlyRevokesFailedKey leaves the persistent write fault injection armed, so the retry Put it asserts on can never succeed. Disarm the injection after the rollback assertions and before the retry. Co-Authored-By: Claude Sonnet 5.5 --- mooncake-store/tests/dfs_sync_client_test.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mooncake-store/tests/dfs_sync_client_test.cpp b/mooncake-store/tests/dfs_sync_client_test.cpp index 119b1b9b6d..00046dc8e4 100644 --- a/mooncake-store/tests/dfs_sync_client_test.cpp +++ b/mooncake-store/tests/dfs_sync_client_test.cpp @@ -630,6 +630,9 @@ TEST_F(DfsSyncClientTest, DfsOnlyBatchFailureOnlyRevokesFailedKey) { ExpectDfsOnlyValue(keys[0], values[0]); ExpectNotFound(keys[1]); + // Restore writes after the persistent failure before retrying the key. + adapter_->FailWriteCall(-1); + // The revoked key and its extent can be reused. std::vector retry{{values[1].data(), values[1].size()}}; ASSERT_TRUE(writer_->Put(keys[1], retry, DfsOnlyConfig()).has_value()); From 9eb5b2fe9d80bffd4add95450a1f75936fc57d4f Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:28:37 +0800 Subject: [PATCH 7/8] [TENT] Test receiver credit is held until failed staged chunks drain A staged owner holds the receiver credit for all of its chunks, and the chunks themselves are credit-exempt. If one chunk fails while a sibling's cross stage is still writing into the receiver, returning the credit at that point would let the receiver grant the same capacity to another sender while the old write is still landing. Upstream #4720 made ProxyManager keep the parent PENDING until in-flight siblings drain. Add an end-to-end regression through the runtime queue and receiver-credit gate: the staged owner stays PENDING and a second credit-gated write never reaches the transport until the held cross-stage write drains; only then does the owner fail, return the credit, and let the second write complete. Reverting the upstream drain makes the test fail on all three checks. CI hosts have no device memory, so add a test-only setStagingPolicyForTest hook that forces host staging for selected requests. --- .../tent/runtime/transfer_engine_impl.h | 12 ++ .../tent/src/runtime/transfer_engine_impl.cpp | 2 + .../tests/runtime_queue_dispatch_test.cpp | 142 ++++++++++++++++++ 3 files changed, 156 insertions(+) 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 4e8e7addbd..66face97fe 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,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -361,6 +362,16 @@ class TransferEngineImpl { } } + // Test-only hook: CI hosts have no device memory, so no request is ever + // staged. A test can force a staging policy for the requests it selects + // (return true and fill `policy`) to drive a staged owner end to end. + // Set before the first submit; not thread-safe with in-flight transfers. + using StagingPolicyForTest = + std::function&)>; + void setStagingPolicyForTest(StagingPolicyForTest hook) { + staging_policy_for_test_ = std::move(hook); + } + Status publishRuntimeConfigForTest( std::shared_ptr snapshot) { if (!snapshot) { @@ -720,6 +731,7 @@ class TransferEngineImpl { std::string local_segment_name_; std::unique_ptr staging_proxy_; + StagingPolicyForTest staging_policy_for_test_; bool merge_requests_; std::shared_ptr runtime_config_snapshot_; bool notify_rpc_fallback_{true}; 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 56e551f866..14e55d8eea 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -2261,6 +2261,8 @@ bool TransferEngineImpl::xpuIpcReady(const Request& request, void TransferEngineImpl::findStagingPolicy(const Request& request, std::vector& policy) { + if (staging_policy_for_test_ && staging_policy_for_test_(request, policy)) + return; if (request.target_id == LOCAL_SEGMENT_ID) return; SegmentDesc* desc = nullptr; 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 dbaddebfc5..d9001023d4 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,7 @@ #include "tent/common/types.h" #include "tent/runtime/segment.h" #include "tent/runtime/control_plane.h" +#include "tent/runtime/proxy_manager.h" #include "tent/runtime/transfer_engine_impl.h" #include "tent/runtime/transport.h" @@ -69,6 +71,12 @@ class TransferEngineImplTestPeer { [bytes_per_second]() { return bytes_per_second; }, tent::DegradationHooks{}, nullptr); } + // Small chunks so a short request splits into several staged chunks. + static void useSmallStagePool(tent::TransferEngineImpl& engine, + size_t chunk_size, size_t chunk_count) { + engine.staging_proxy_ = std::make_unique( + &engine, chunk_size, chunk_count); + } static void adoptDeferredBatches(tent::TransferEngineImpl& engine, std::vector batches) { tent::TransferEngineImpl::DeferredStageTeardown deferred; @@ -1538,6 +1546,140 @@ TEST(RuntimeQueueDispatch, TcpControlDataUsesLegacyFramingWithoutCredit) { engine.unregisterLocalMemory(target.data(), target.size()).ok()); } +// A staged owner holds the receiver credit for all of its chunks; the chunks +// themselves are credit-exempt. If one chunk fails while a sibling's cross +// stage is still writing into the receiver, the owner must not turn terminal +// (and so must not return the credit) until that write has drained. +// Otherwise the receiver could grant the same capacity to another sender +// while the old write is still landing. +TEST(RuntimeQueueDispatch, StagedFailureKeepsReceiverCreditUntilChunksDrain) { + constexpr size_t kChunk = 64; + constexpr size_t kStagedLen = 2 * kChunk; + constexpr size_t kSmallLen = 32; + auto cfg = makeRuntimeQueueConfig(8, 1UL << 20); + cfg->set("max_failover_attempts", 0); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kStagedLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + cfg->set("receiver_credit/freshness_ttl_ms", uint32_t{1000}); + std::vector source(kStagedLen, 0x5a); + std::vector target(kStagedLen, 0); + std::vector small_source(kSmallLen, 0x6b); + std::vector small_target(kSmallLen, 0); + // Declared after the buffers the fake transport and hook capture. + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto in = [](const void* addr, const std::vector& buf) { + auto p = reinterpret_cast(addr); + auto base = reinterpret_cast(buf.data()); + return p >= base && p < base + buf.size(); + }; + + std::atomic cross_inflight{false}; + std::atomic release_cross{false}; + std::atomic failed_polls{0}; + std::atomic small_polls{0}; + auto fake_tcp = + std::make_shared(TCP, [&](const Request& request, int) { + const auto* src = static_cast(request.source); + const auto* dst = + reinterpret_cast(request.target_offset); + // Second chunk's local stage: fail only once the first chunk's + // cross stage is writing into the receiver. + if (src == source.data() + kChunk) { + if (!cross_inflight.load()) return TransferStatus{PENDING, 0}; + ++failed_polls; + return TransferStatus{FAILED, 0}; + } + // Cross stage into the receiver: held until the test lets go. + if (in(dst, target)) { + cross_inflight = true; + if (!release_cross.load()) return TransferStatus{PENDING, 0}; + } + if (in(src, small_source)) ++small_polls; + std::memcpy(const_cast(dst), request.source, request.length); + return TransferStatus{COMPLETED, request.length}; + }); + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(TCP, fake_tcp); + TransferEngineImplTestPeer::useSmallStagePool(engine, kChunk, 2); + // CI has no device memory, so force host staging for the large write. + engine.setStagingPolicyForTest( + [&](const Request& request, std::vector& policy) { + if (request.source != source.data() || request.length != kStagedLen) + return false; + policy = {"", kWildcardLocation, ""}; + return true; + }); + + for (auto* buf : {&source, &target, &small_source, &small_target}) + ASSERT_TRUE(engine.registerLocalMemory(buf->data(), buf->size()).ok()); + + auto make_write = [](std::vector& from, std::vector& to) { + Request request; + request.opcode = Request::WRITE; + request.source = from.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(to.data()); + request.length = from.size(); + request.transport_hint = TCP; + return request; + }; + auto wait_until = [](const std::function& done) { + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!done() && std::chrono::steady_clock::now() < deadline) + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + return done(); + }; + + BatchID staged = engine.allocateBatch(1); + ASSERT_NE(staged, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(staged, {make_write(source, target)}).ok()); + ASSERT_TRUE(wait_until([&] { return failed_polls.load() > 0; })) + << "second chunk never failed while the first was in its cross stage"; + // Let the failure path run while the first chunk's write is held. + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + TransferStatus staged_status{}; + ASSERT_TRUE(engine.getTransferStatus(staged, 0, staged_status).ok()); + EXPECT_EQ(staged_status.s, PENDING); + + // The staged owner still holds the only request slot, so a second + // credit-gated write must wait instead of reaching the transport. + BatchID small = engine.allocateBatch(1); + ASSERT_NE(small, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(small, {make_write(small_source, small_target)}) + .ok()); + TransferStatus small_status{}; + ASSERT_TRUE(engine.getTransferStatus(small, 0, small_status).ok()); + EXPECT_EQ(small_status.s, PENDING); + EXPECT_EQ(small_polls.load(), 0); + + // Once the held write drains, the owner fails and returns the credit. + release_cross = true; + ASSERT_TRUE(wait_until([&] { + return engine.getTransferStatus(staged, 0, staged_status).ok() && + staged_status.s != PENDING; + })); + EXPECT_EQ(staged_status.s, FAILED); + ASSERT_TRUE(wait_until([&] { + return engine.getTransferStatus(small, 0, small_status).ok() && + small_status.s != PENDING; + })); + EXPECT_EQ(small_status.s, COMPLETED); + EXPECT_EQ(small_target, small_source); + + EXPECT_TRUE(engine.freeBatch(staged).ok()); + EXPECT_TRUE(engine.freeBatch(small).ok()); + for (auto* buf : {&source, &target, &small_source, &small_target}) + EXPECT_TRUE( + engine.unregisterLocalMemory(buf->data(), buf->size()).ok()); +} + } // namespace } // namespace tent } // namespace mooncake From 461b462e80896b0b6a232a359bfa14de92814d28 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:03:42 +0800 Subject: [PATCH 8/8] [TENT] Back off refused receiver-credit exchanges instead of retrying per poll On two UB nodes, tebench throughput collapsed ~1000x (6 GB/s -> 0.004 GB/s, multi-second latency) whenever concurrent demand exceeded the receiver's data_bytes credit, and recovered as soon as capacity covered the demand. A deferred owner is retried on every refill, and refill runs on every status poll. Each retry re-sent the demand exchange, a synchronous RPC to the receiver made under progress_mutex_. With credit exhausted, the waiting threads issued thousands of refused RPCs per second and kept the lock busy, so the in-flight owners could not be polled to completion and never returned the credit they held. A local receiver did not show it because a local exchange is a plain function call. Remember a refused (or failed) exchange per credit key and requeue the owner without another RPC until this sender releases credit for the key, asks for less, or a capped exponential backoff (1 ms .. 32 ms) expires. A local release retries immediately; the backoff bounds the wait for capacity that another sender returned, which this sender is not told about. A cross-engine reproduction of the tebench workload (4 threads x 8 x 1 MiB, 1 ms simulated completion) goes from 4 to ~1500 batches/s at 16 MiB of credit. New tests check that a refused owner exchanges at most a handful of times while busy-polled (686 times in 50 ms without the fix) and is dispatched immediately on a local release, and that a refused sender recovers when another sender releases. --- .../tent/runtime/transfer_engine_impl.h | 16 ++ .../tent/src/runtime/transfer_engine_impl.cpp | 43 +++++ .../tests/runtime_queue_dispatch_test.cpp | 179 ++++++++++++++++++ 3 files changed, 238 insertions(+) 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 66face97fe..c4f5fcf682 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 @@ -695,6 +695,19 @@ class TransferEngineImpl { std::chrono::steady_clock::time_point retry_at{}; }; + // A demand exchange the receiver did not satisfy (or that failed). A + // deferred owner is retried on every refill, and the exchange is an RPC + // made under progress_mutex_; repeating it on every poll can starve the + // in-flight owners whose completion is what returns the credit. Until + // this sender releases credit for the key, asks for less, or the backoff + // expires, the owner is requeued without another RPC. + struct ReceiverCreditDenial { + std::array released{}; // DataBytes, RequestSlots + std::array charge{}; + uint32_t attempts{0}; + std::chrono::steady_clock::time_point retry_at{}; + }; + struct QueuedOwnerState { Batch* batch{nullptr}; size_t owner_task_id{0}; @@ -745,6 +758,9 @@ class TransferEngineImpl { std::unordered_map pending_credit_releases_; std::unordered_map receiver_credit_sessions_; + std::unordered_map + receiver_credit_denials_; + uint64_t receiver_credit_demand_exchanges_{0}; std::unordered_map queued_owners_; size_t dispatch_inflight_owners_{0}; size_t dispatch_inflight_bytes_{0}; 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 14e55d8eea..35a0c6acb7 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -2725,10 +2725,42 @@ Status TransferEngineImpl::acquireReceiverCredit( return reserve_status; } + std::array released_now{}; + (void)sender_credit_ledger_->released(key, CreditResource::DataBytes, + released_now[0]); + (void)sender_credit_ledger_->released(key, CreditResource::RequestSlots, + released_now[1]); + const std::array charge_now{request.length, 1}; + const auto now = std::chrono::steady_clock::now(); + auto denial_it = receiver_credit_denials_.find(key); + if (denial_it != receiver_credit_denials_.end()) { + const auto& denial = denial_it->second; + const bool asks_less = charge_now[0] < denial.charge[0] || + charge_now[1] < denial.charge[1]; + if (denial.released == released_now && !asks_less && + now < denial.retry_at) { + return Status::OK(); // Nothing changed: requeue, no RPC. + } + } + // Back off exponentially (1 ms .. 32 ms) while the receiver keeps + // refusing. A local release retries at once; the cap bounds the + // wait for capacity that another sender returned. + auto note_denial = [&] { + auto& denial = receiver_credit_denials_[key]; + if (denial.released != released_now) denial.attempts = 0; + denial.released = released_now; + denial.charge = charge_now; + const uint32_t shift = std::min(denial.attempts, 5); + ++denial.attempts; + denial.retry_at = std::chrono::steady_clock::now() + + std::chrono::milliseconds(1U << shift); + }; + ReceiverCreditExchangeRequestV1 exchange_request; CHECK_STATUS(sender_credit_ledger_->prepareExchange(key, charge, exchange_request)); ReceiverCreditExchangeReplyV1 exchange_reply; + ++receiver_credit_demand_exchanges_; auto exchange_status = exchangeReceiverCredit(request.target_id, target.rpc_server_addr, exchange_request, exchange_reply); @@ -2739,6 +2771,7 @@ Status TransferEngineImpl::acquireReceiverCredit( } LOG(WARNING) << "Receiver-credit exchange deferred: " << exchange_status.ToString(); + note_denial(); return Status::OK(); } if (!(exchange_reply.sender_instance_id == sender_instance_id_)) { @@ -2756,10 +2789,12 @@ Status TransferEngineImpl::acquireReceiverCredit( reserve_status = sender_credit_ledger_->tryReserve(key, charge); if (reserve_status.IsTooManyRequests() || reserve_status.IsInvalidEntry()) { + note_denial(); return Status::OK(); } CHECK_STATUS(reserve_status); } + receiver_credit_denials_.erase(key); reservation.active = true; reservation.committed = false; @@ -2905,6 +2940,14 @@ void TransferEngineImpl::observeReceiverCreditSession( ++pending; } } + for (auto denial = receiver_credit_denials_.begin(); + denial != receiver_credit_denials_.end();) { + if (denial->first.receiver_session == stale_session) { + denial = receiver_credit_denials_.erase(denial); + } else { + ++denial; + } + } it->second = session; } 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 d9001023d4..07b3092204 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -77,6 +77,10 @@ class TransferEngineImplTestPeer { engine.staging_proxy_ = std::make_unique( &engine, chunk_size, chunk_count); } + static uint64_t creditDemandExchanges(tent::TransferEngineImpl& engine) { + std::lock_guard lk(engine.progress_mutex_); + return engine.receiver_credit_demand_exchanges_; + } static void adoptDeferredBatches(tent::TransferEngineImpl& engine, std::vector batches) { tent::TransferEngineImpl::DeferredStageTeardown deferred; @@ -1680,6 +1684,181 @@ TEST(RuntimeQueueDispatch, StagedFailureKeepsReceiverCreditUntilChunksDrain) { engine.unregisterLocalMemory(buf->data(), buf->size()).ok()); } +// Receiver credit lives in another engine, so each demand exchange is an RPC +// made under progress_mutex_. A deferred owner must not repeat a refused +// exchange on every poll: back-to-back RPCs hold the lock and starve the +// in-flight owner whose completion is what returns the credit (on two UB +// nodes this collapsed throughput by ~1000x once demand exceeded capacity). +TEST(RuntimeQueueDispatch, RefusedCreditIsNotReexchangedOnEveryPoll) { + constexpr size_t kReqLen = 4096; + auto receiver_cfg = makeRuntimeQueueConfig(4, 1UL << 20); + receiver_cfg->set("receiver_credit/receiver_enable", true); + receiver_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + receiver_cfg->set("receiver_credit/request_slots", uint64_t{1}); + auto sender_cfg = makeRuntimeQueueConfig(4, 1UL << 20); + sender_cfg->set("receiver_credit/sender_enable", true); + std::vector source(kReqLen * 2, 0x91); + std::vector target(kReqLen * 2, 0); + TransferEngineImpl receiver(receiver_cfg); + TransferEngineImpl sender(sender_cfg); + ASSERT_TRUE(receiver.available()); + ASSERT_TRUE(sender.available()); + installFakeUb(receiver, std::make_shared(UB)); + std::atomic complete_first{false}; + auto sender_ub = + std::make_shared(UB, [&](const Request& request, int) { + if (request.source == source.data() && !complete_first.load()) + return TransferStatus{PENDING, 0}; + return TransferStatus{COMPLETED, request.length}; + }); + installFakeUb(sender, sender_ub); + ASSERT_TRUE(sender.registerLocalMemory(source.data(), source.size()).ok()); + ASSERT_TRUE( + receiver.registerLocalMemory(target.data(), target.size()).ok()); + SegmentID remote = LOCAL_SEGMENT_ID; + ASSERT_TRUE(sender.openSegment(remote, receiver.getSegmentName()).ok()); + + BatchID holder = sender.allocateBatch(1); + BatchID waiter = sender.allocateBatch(1); + ASSERT_NE(holder, (BatchID)0); + ASSERT_NE(waiter, (BatchID)0); + ASSERT_TRUE( + sender + .submitTransfer(holder, {makeRemoteUbWrite(source.data(), remote, + target.data(), kReqLen)}) + .ok()); + ASSERT_TRUE( + sender + .submitTransfer( + waiter, {makeRemoteUbWrite(source.data() + kReqLen, remote, + target.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(sender_ub->submit_calls.load(), 1); + + // Busy-poll the refused owner. Without the backoff every poll is another + // exchange (thousands here); with it, a handful (1, 2, 4, ... ms). + const auto before = + TransferEngineImplTestPeer::creditDemandExchanges(sender); + TransferStatus waiter_status{}; + const auto until = + std::chrono::steady_clock::now() + std::chrono::milliseconds(50); + while (std::chrono::steady_clock::now() < until) { + ASSERT_TRUE(sender.getTransferStatus(waiter, 0, waiter_status).ok()); + ASSERT_EQ(waiter_status.s, PENDING); + } + const auto refused = + TransferEngineImplTestPeer::creditDemandExchanges(sender) - before; + EXPECT_GE(refused, 1u); + EXPECT_LE(refused, 12u); + EXPECT_EQ(sender_ub->submit_calls.load(), 1); + + // Returning the credit retries at once, without waiting out the backoff. + complete_first = true; + TransferStatus holder_status{}; + ASSERT_TRUE(sender.getTransferStatus(holder, 0, holder_status).ok()); + EXPECT_EQ(holder_status.s, COMPLETED); + EXPECT_EQ(sender_ub->submit_calls.load(), 2); + ASSERT_TRUE(sender.getTransferStatus(waiter, 0, waiter_status).ok()); + EXPECT_EQ(waiter_status.s, COMPLETED); + + EXPECT_TRUE(sender.freeBatch(holder).ok()); + EXPECT_TRUE(sender.freeBatch(waiter).ok()); + EXPECT_TRUE(sender.closeSegment(remote).ok()); + EXPECT_TRUE( + sender.unregisterLocalMemory(source.data(), source.size()).ok()); + EXPECT_TRUE( + receiver.unregisterLocalMemory(target.data(), target.size()).ok()); +} + +// The refused sender is not told when another sender returns credit. The +// capped backoff must still let it through without any local release. +TEST(RuntimeQueueDispatch, RefusedCreditRecoversWhenAnotherSenderReleases) { + constexpr size_t kReqLen = 4096; + auto receiver_cfg = makeRuntimeQueueConfig(4, 1UL << 20); + receiver_cfg->set("receiver_credit/receiver_enable", true); + receiver_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + receiver_cfg->set("receiver_credit/request_slots", uint64_t{1}); + auto sender_cfg = makeRuntimeQueueConfig(4, 1UL << 20); + sender_cfg->set("receiver_credit/sender_enable", true); + std::vector source_a(kReqLen, 0xa1), source_b(kReqLen, 0xb2); + std::vector target(kReqLen * 2, 0); + TransferEngineImpl receiver(receiver_cfg); + TransferEngineImpl sender_a(sender_cfg); + TransferEngineImpl sender_b(sender_cfg); + ASSERT_TRUE(receiver.available()); + ASSERT_TRUE(sender_a.available()); + ASSERT_TRUE(sender_b.available()); + installFakeUb(receiver, std::make_shared(UB)); + std::atomic complete_a{false}; + auto ub_a = + std::make_shared(UB, [&](const Request& request, int) { + if (!complete_a.load()) return TransferStatus{PENDING, 0}; + return TransferStatus{COMPLETED, request.length}; + }); + auto ub_b = std::make_shared(UB); + installFakeUb(sender_a, ub_a); + installFakeUb(sender_b, ub_b); + ASSERT_TRUE( + sender_a.registerLocalMemory(source_a.data(), source_a.size()).ok()); + ASSERT_TRUE( + sender_b.registerLocalMemory(source_b.data(), source_b.size()).ok()); + ASSERT_TRUE( + receiver.registerLocalMemory(target.data(), target.size()).ok()); + SegmentID remote_a = LOCAL_SEGMENT_ID, remote_b = LOCAL_SEGMENT_ID; + ASSERT_TRUE(sender_a.openSegment(remote_a, receiver.getSegmentName()).ok()); + ASSERT_TRUE(sender_b.openSegment(remote_b, receiver.getSegmentName()).ok()); + + BatchID batch_a = sender_a.allocateBatch(1); + BatchID batch_b = sender_b.allocateBatch(1); + ASSERT_TRUE(sender_a + .submitTransfer( + batch_a, {makeRemoteUbWrite(source_a.data(), remote_a, + target.data(), kReqLen)}) + .ok()); + ASSERT_TRUE( + sender_b + .submitTransfer( + batch_b, {makeRemoteUbWrite(source_b.data(), remote_b, + target.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(ub_a->submit_calls.load(), 1); + EXPECT_EQ(ub_b->submit_calls.load(), 0); + // Let sender B back off a few rounds while A holds the only slot. + TransferStatus status_b{}; + const auto hold = + std::chrono::steady_clock::now() + std::chrono::milliseconds(40); + while (std::chrono::steady_clock::now() < hold) { + ASSERT_TRUE(sender_b.getTransferStatus(batch_b, 0, status_b).ok()); + ASSERT_EQ(status_b.s, PENDING); + } + + complete_a = true; + TransferStatus status_a{}; + ASSERT_TRUE(sender_a.getTransferStatus(batch_a, 0, status_a).ok()); + EXPECT_EQ(status_a.s, COMPLETED); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (std::chrono::steady_clock::now() < deadline) { + ASSERT_TRUE(sender_b.getTransferStatus(batch_b, 0, status_b).ok()); + if (status_b.s != PENDING) break; + std::this_thread::sleep_for(std::chrono::microseconds(200)); + } + EXPECT_EQ(status_b.s, COMPLETED); + EXPECT_EQ(ub_b->submit_calls.load(), 1); + + EXPECT_TRUE(sender_a.freeBatch(batch_a).ok()); + EXPECT_TRUE(sender_b.freeBatch(batch_b).ok()); + EXPECT_TRUE(sender_a.closeSegment(remote_a).ok()); + EXPECT_TRUE(sender_b.closeSegment(remote_b).ok()); + EXPECT_TRUE( + sender_a.unregisterLocalMemory(source_a.data(), source_a.size()).ok()); + EXPECT_TRUE( + sender_b.unregisterLocalMemory(source_b.data(), source_b.size()).ok()); + EXPECT_TRUE( + receiver.unregisterLocalMemory(target.data(), target.size()).ok()); +} + } // namespace } // namespace tent } // namespace mooncake