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 1023992b84..a65c979ca8 100644 --- a/docs/source/performance/mooncake/tebench.md +++ b/docs/source/performance/mooncake/tebench.md @@ -433,9 +433,9 @@ gpu_id + thread_id **Transport (TENT only)** * `--xport_type` : selects the benchmark transport, such as `tcp`, `hp_tcp`, - `rdma`, `shm`, or `nvlink`, and disables the other selectable transports. - MLU VRAM runs also enable the local-copy transport needed for host staging. - Empty means no transport + `rdma`, `shm`, `nvlink`, `ub`, or `sunrise_link`, and disables the other + selectable transports. MLU VRAM runs also enable the local-copy transport + needed for host staging. 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 5.8 for multi-transport scenarios). diff --git a/mooncake-store/tests/dfs_sync_client_test.cpp b/mooncake-store/tests/dfs_sync_client_test.cpp index 0be16160c9..165fa827a5 100644 --- a/mooncake-store/tests/dfs_sync_client_test.cpp +++ b/mooncake-store/tests/dfs_sync_client_test.cpp @@ -627,6 +627,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()); diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json index feadc511f2..12ea6bf57b 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 1ae7654de1..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 @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,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" @@ -360,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) { @@ -477,7 +489,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); @@ -496,6 +508,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, @@ -614,12 +663,59 @@ 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{}; + }; + + // 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}; std::vector public_task_ids; size_t byte_charge{0}; bool in_dispatch_window{false}; + bool receiver_credit_exempt{false}; + ReceiverCreditReservation credit; }; private: @@ -648,12 +744,23 @@ 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}; 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 + 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/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index f81fe66ec2..daa54da4da 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -61,6 +61,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 fe378fea55..60bfb22d4f 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 #include @@ -58,6 +59,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; @@ -119,29 +227,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); @@ -152,19 +314,70 @@ 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) + // 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( + response.empty() + ? "RecvData failed: invalid or rejected legacy response" + : 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( - response.empty() ? "RecvData failed: empty response" : response); - return Platform::getLoader().copy(local_mem_addr, response.data(), length); + response.empty() ? "RecvData failed: invalid or rejected response" + : response); + } + XferDataReplyHeader header; + memcpy(&header, response.data(), sizeof(header)); + if (le64toh(header.protocol_magic) != kXferDataProtocolMagic) { + return Status::RpcServiceError( + "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) { @@ -196,11 +409,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) { @@ -213,13 +441,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(); @@ -321,6 +571,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) { @@ -376,6 +631,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(); } @@ -519,37 +778,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"; @@ -558,46 +842,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 64cad52563..55851c94ca 100644 --- a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp @@ -238,6 +238,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()) { @@ -288,6 +293,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); } @@ -301,6 +311,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 e00c70f55d..6d6471127e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -60,6 +60,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 { @@ -363,6 +390,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); @@ -398,6 +430,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); @@ -420,6 +461,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)) { @@ -428,6 +522,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 @@ -440,6 +543,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") @@ -611,6 +722,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 @@ -1790,6 +1923,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 { @@ -1812,6 +1949,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; }; @@ -1912,6 +2051,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; @@ -1990,6 +2138,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]; @@ -1999,6 +2162,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) { @@ -2096,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; @@ -2249,7 +2416,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); @@ -2330,6 +2498,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)) { @@ -2374,6 +2585,397 @@ 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; + } + + 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); + 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(); + note_denial(); + 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()) { + note_denial(); + return Status::OK(); + } + CHECK_STATUS(reserve_status); + } + receiver_credit_denials_.erase(key); + + 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; + } + } + 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; +} + +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); @@ -2648,6 +3250,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(), @@ -2672,6 +3276,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_; @@ -2772,7 +3377,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(); @@ -2790,13 +3395,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); } @@ -2818,6 +3453,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}); @@ -2830,8 +3489,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); } @@ -2873,6 +3536,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; @@ -2917,7 +3581,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, @@ -2925,6 +3590,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; }); @@ -2944,8 +3615,8 @@ Status TransferEngineImpl::submitTransfer( std::lock_guard lk(progressLockFor(batch_id)); 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( @@ -3138,6 +3809,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 @@ -3198,13 +3881,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, @@ -3814,9 +4556,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/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/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index b147400054..2a13b43279 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -548,6 +548,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 5b6a4db38d..2078762cfb 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -1842,6 +1842,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 9def334c40..6a84520236 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp @@ -178,6 +178,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); } @@ -264,9 +267,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; @@ -296,11 +297,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(); @@ -327,9 +332,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; } } @@ -337,12 +340,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/hp_tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp index 452b2e3f04..2271ccab2a 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,86 @@ 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 = + std::get(segment.detail) + .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 { 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..07b3092204 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" @@ -37,6 +39,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(); } @@ -63,6 +71,16 @@ 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 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; @@ -101,6 +119,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 +145,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 +289,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 +307,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 +403,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 +428,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 +1166,699 @@ 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()); +} + +// 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()); +} + +// 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 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); } } diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index 7f7ad40dd8..d2709d2d30 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};