diff --git a/.github/workflows/_build-wheel.yaml b/.github/workflows/_build-wheel.yaml index be35662247..21a0ed9e53 100644 --- a/.github/workflows/_build-wheel.yaml +++ b/.github/workflows/_build-wheel.yaml @@ -214,6 +214,7 @@ jobs: VERSION: ${{ env.VERSION }} - name: Smoke test repaired wheel + shell: bash run: | smoke_venv=$(mktemp -d) python -m venv "$smoke_venv" @@ -223,11 +224,20 @@ jobs: if [ "${VARIANT_FLAG:-}" = "NON_CUDA_BUILD" ]; then site_packages=$("$smoke_venv/bin/python" -c \ 'import sysconfig; print(sysconfig.get_paths()["purelib"])') - master="$site_packages/mooncake/mooncake_master" - if readelf -d "$master" | grep -Eq \ - 'Shared library: \[(libcuda|libcudart)\.so'; then - echo "Non-CUDA mooncake_master depends on CUDA" - readelf -d "$master" | grep 'Shared library:' + cuda_dependency_found=false + for package_path in "$site_packages"/mooncake*; do + [ -e "$package_path" ] || continue + while IFS= read -r -d '' file; do + cuda_dependencies=$(readelf -d "$file" 2>/dev/null | grep -E \ + 'Shared library: \[(libcuda|libcudart|libcublas|libcufft|libcurand|libcusolver|libcusparse|libcufile|libcupti|libnvrtc|libnvJitLink|libnvToolsExt|libnvfatbin|libnvidia|libnccl)\.so' || true) + if [ -n "$cuda_dependencies" ]; then + echo "::error file=$file::Non-CUDA wheel artifact depends on CUDA" + echo "$cuda_dependencies" + cuda_dependency_found=true + fi + done < <(find "$package_path" -type f -print0) + done + if [ "$cuda_dependency_found" = true ]; then exit 1 fi fi diff --git a/.github/workflows/ci-on-label.yml b/.github/workflows/ci-on-label.yml new file mode 100644 index 0000000000..f4d5929ca6 --- /dev/null +++ b/.github/workflows/ci-on-label.yml @@ -0,0 +1,53 @@ +name: Retrigger CI on run-ci label + +# Same-SHA retrigger for Build & Test. This is a separate workflow so labels +# other than `run-ci` (especially `run-e2e-ci`) cannot start or cancel +# `.github/workflows/ci.yml`. +# +# pull_request_target is required so fork PRs can rerun Actions. This +# workflow only calls the GitHub API; it does not check out PR code. +on: + pull_request_target: + branches: + - "main" + - "release/**" + types: [labeled] + +permissions: + actions: write + contents: read + +jobs: + retrigger: + if: > + github.event.label.name == 'run-ci' && + github.actor != 'github-actions[bot]' + runs-on: ubuntu-latest + steps: + - name: Re-run Build & Test for this SHA + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + SHA: ${{ github.event.pull_request.head.sha }} + shell: bash + run: | + set -euo pipefail + + run_json=$(gh api \ + "repos/${REPO}/actions/workflows/ci.yml/runs?head_sha=${SHA}&per_page=20") + run_id=$(echo "$run_json" | jq -r '.workflow_runs[0].id // empty') + status=$(echo "$run_json" | jq -r '.workflow_runs[0].status // empty') + + if [ -z "$run_id" ]; then + echo "No Build & Test run found for SHA ${SHA}." + echo "Open or push to the PR first so ci.yml has a run to rerun." + exit 1 + fi + + echo "Matched workflow run ${run_id} (status=${status})" + if [ "$status" != "completed" ]; then + echo "Build & Test is still ${status}; not starting a duplicate." + exit 0 + fi + + gh run rerun "$run_id" --repo "$REPO" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 095abf2f7c..739c519aee 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,7 +9,11 @@ on: branches: - "main" - "release/**" - types: [opened, synchronize, reopened, labeled] + # `labeled` is intentionally omitted. Auto-labeler already applies + # `run-ci`, so any new label (including `run-e2e-ci`) would retrigger + # this whole workflow and cancel the in-progress run. Same-SHA + # retrigger via the `run-ci` label lives in ci-on-label.yml. + types: [opened, synchronize, reopened] workflow_dispatch: {} permissions: @@ -945,6 +949,17 @@ jobs: run: ./scripts/ci/run_transfer_engine_rust_smoke.sh shell: bash + - name: Smoke test TENT UB benchmark CLI + if: matrix.name == 'ub-mock' + run: | + cd build-tent + help_output="$(./mooncake-transfer-engine/benchmark/tebench \ + --backend=tent --xport_type=ub --tent_transport_hint=ub \ + --help 2>&1 || true)" + grep -q 'iouring|ub|sunrise_link' <<< "${help_output}" + grep -q 'ascend|ub|sunrise_link' <<< "${help_output}" + shell: bash + - name: Run sccache stat for check if: ${{ env.SCCACHE_PATH != '' }} shell: bash diff --git a/.github/workflows/ci_ascend.yml b/.github/workflows/ci_ascend.yml index 13f79a244f..1f2b6d1ea2 100644 --- a/.github/workflows/ci_ascend.yml +++ b/.github/workflows/ci_ascend.yml @@ -45,12 +45,13 @@ jobs: fetch-depth: 1 persist-credentials: false - - name: Retry checkout via GitHub mirror + # Mirror is only a git transport fallback. Keep retry on actions/checkout so + # pull_request_target fork checks still run; do not fetch with raw git. + - name: Configure GitHub mirror rewrite if: steps.checkout_code.outcome == 'failure' shell: bash env: ASCEND_GITHUB_MIRROR_URLS: 'https://ghfast.top/' - CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }} run: | set -euo pipefail @@ -83,27 +84,21 @@ jobs: workdir="${GITHUB_WORKSPACE}" git config --global --add safe.directory "$workdir" + find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} + - for base in "${candidates[@]}"; do - mirror_url="${base}https://github.com/${GITHUB_REPOSITORY}.git" - echo "Retrying checkout with ${mirror_url}" - - find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} + - git init "$workdir" - git -C "$workdir" remote add origin "$mirror_url" - - if git -C "$workdir" fetch --depth=1 origin "$CHECKOUT_REF" && \ - git -C "$workdir" checkout --force --detach FETCH_HEAD; then - echo "Mirror checkout succeeded via ${base}" - exit 0 - fi - - echo "Mirror checkout failed via ${base}" - rm -rf "$workdir/.git" - done + # insteadOf only rewrites github.com fetches; checkout still runs + # assertSafePrCheckout before any git network I/O. + mirror_base="${candidates[0]}" + echo "Rewriting https://github.com/ to ${mirror_base}https://github.com/" + git config --global url."${mirror_base}https://github.com/".insteadOf "https://github.com/" - echo "Direct GitHub checkout failed and all mirror retries failed" - exit 1 + - name: Retry checkout via GitHub mirror + if: steps.checkout_code.outcome == 'failure' + uses: actions/checkout@v4 + with: + ref: ${{ inputs.checkout_ref || github.sha }} + fetch-depth: 1 + persist-credentials: false - name: Configure CMake shell: bash diff --git a/docs/source/api-reference/cpp/index.md b/docs/source/api-reference/cpp/index.md index 089c657c8b..022c3a08b8 100644 --- a/docs/source/api-reference/cpp/index.md +++ b/docs/source/api-reference/cpp/index.md @@ -4,7 +4,7 @@ |--------|-------------| | [Transfer Engine C++ API](transfer-engine) | `TransferEngine` class — memory registration, batch transfer, segment management, RDMA transport | | [TENT C++ API](tent) | `mooncake::tent::TransferEngine` — next-gen transfer engine with automatic transport selection and fault tolerance | -| [Mooncake Store Client C++ API](mooncake-store) | `Client` class — `Put`/`Get`/`Remove`/`Replicate` operations, `BufferAllocatorBase` interface | +| [Mooncake Store Client C++ API](mooncake-store) | `Client` class — `Put`/`Get`/`Remove`/`Replicate` operations | :::{toctree} :maxdepth: 1 diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 5d8419a9d7..3556e42e8a 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -929,7 +929,23 @@ allocation_strategy: "local_first" When enabled, the master applies local-first allocation only for memory replicas with `replica_num == 1`. Explicit `preferred_segment` or `preferred_segments` are tried first; if they are unavailable or full, Mooncake falls back through active hosts in cyclic lexicographic host-id order, starting from the writer host when it has active segments, or otherwise from the next greater active host id. Within the same host, segment names are sorted and rotated by key hash so multiple segments on one host do not always receive the first allocation attempt. -The client derives the host id from `local_hostname` by removing the port. For example, `host-a:50051` and `host-a:50052` map to the same host id, `host-a`. For local-first allocation to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique host part in `local_hostname`. In deployments with multiple NIC IPs, hostname aliases, or container/pod networking, choose one canonical host name or IP and use it consistently across processes on that host. Empty, loopback, and wildcard values such as `localhost`, `127.0.0.1`, `0.0.0.0`, `::1`, and `::` are treated as unknown and do not trigger automatic local-first placement for that client. +By default, the client derives the host id from `local_hostname` by removing the port. For example, `host-a:50051` and `host-a:50052` map to the same host id, `host-a`. Set `MOONCAKE_HOST_ID` to override this derived value with a stable, globally unique node identifier. The override is read directly by the C++ client, so it applies to every client initialization method. It must be set before creating the client, and all writer and store processes on the same physical or logical host must use the same value. An empty or whitespace-only override falls back to `local_hostname`. Loopback and wildcard values such as `localhost`, `127.0.0.1`, `0.0.0.0`, `::1`, and `::` are treated as unknown and do not trigger automatic local-first placement. + +In Kubernetes, keep `MOONCAKE_LOCAL_HOSTNAME` as the routable pod IP for the transfer endpoint and use `spec.nodeName` as the shared placement identity: + +```yaml +env: + - name: MOONCAKE_LOCAL_HOSTNAME + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: MOONCAKE_HOST_ID + valueFrom: + fieldRef: + fieldPath: spec.nodeName +``` + +Apply the same `MOONCAKE_HOST_ID` mapping to every writer and store pod. This separates the per-pod network address from the node-level placement identity, allowing colocated pods with different IPs to match for local-first allocation. --- @@ -1186,7 +1202,7 @@ Do not run binaries from before and after checksum support was introduced in the | Variable | Default | Description | |----------|---------|-------------| | `MC_STORE_USE_HUGEPAGE` | unset | Set `1` to request HugeTLB-backed `mmap()` | -| `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `1GB` | +| `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `512MB`, `1GB` | | `MC_MMAP_ARENA_POOL_SIZE` | unset | Pre-allocated arena pool size (e.g., `8gb`). Explicitly set to enable the arena | | `MC_DISABLE_MMAP_ARENA` | unset | Disable arena, fall back to per-call `mmap()`. Accepts `1`/`true`/`yes`/`on` (or `0`/`false`/`no`/`off`) | diff --git a/docs/source/design/store/mooncake-store.md b/docs/source/design/store/mooncake-store.md index 56c006ac52..f8521557d6 100644 --- a/docs/source/design/store/mooncake-store.md +++ b/docs/source/design/store/mooncake-store.md @@ -554,7 +554,7 @@ Valid values are: `random` (default), `free_ratio_first`, `ssd_free_ratio_first` | `free_ratio_first` | Balanced utilization, dynamic scaling | Slightly lower throughput due to sampling and sorting overhead | | `ssd_free_ratio_first` | SSD-aware memory allocation when SSD offloading is enabled | Depends on SSD usage metrics; falls back to random allocation when needed | | `cxl` | CXL memory hardware | CXL-specific; single-replica only | -| `local_first` | Colocated inference workers and memory store segments | Requires stable host identity in `local_hostname`; single memory replica only | +| `local_first` | Colocated inference workers and memory store segments | Requires stable host identity from `MOONCAKE_HOST_ID` or `local_hostname`; single memory replica only | **Use `random`** (default) when your cluster is relatively stable (segments rarely join or leave) and you want the highest possible allocation throughput. @@ -566,7 +566,7 @@ Valid values are: `random` (default), `free_ratio_first`, `ssd_free_ratio_first` **Use `cxl`** only when your hardware includes CXL (Compute Express Link) memory devices and you want to allocate data exclusively on CXL segments. -**Use `local_first`** when inference workers and Mooncake Store memory segments are colocated and you want writes to prefer the writer's host before falling back to other hosts. For this strategy to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique host part in `local_hostname`. +**Use `local_first`** when inference workers and Mooncake Store memory segments are colocated and you want writes to prefer the writer's host before falling back to other hosts. For this strategy to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique `MOONCAKE_HOST_ID`. When the variable is unset or empty, Mooncake derives the host identity from `local_hostname` by removing the port. For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../../performance/mooncake/allocation-strategy-benchmark-result.md). @@ -605,6 +605,8 @@ An SSD-aware variant of the free-ratio-first strategy. It first tries preferred Host-aware local-first allocation reuses the normal preferred-segment flow. The master derives the writer host id from the request's client host identity and builds an ordered preferred segment list: active hosts are visited in cyclic lexicographic host-id order, starting from the writer host when it has active segments, or otherwise from the next greater active host id. Within the same host, segment names are sorted and rotated by key hash so multiple local segments do not always receive the first allocation attempt. +The C++ client reads `MOONCAKE_HOST_ID` as an explicit deployment override for the client identity carried in allocation requests and the identity recorded for mounted segments. This lets containerized deployments keep `local_hostname` as a routable per-pod transfer endpoint while using a shared node-level placement identity. Loopback and wildcard overrides are rejected; an empty override preserves the derived-hostname behavior. + This strategy currently applies to memory allocation with `replica_num == 1`. Explicit `preferred_segment` or `preferred_segments` in `ReplicateConfig` are still tried first; if they are unavailable or full, allocation continues with the local-first ordered fallback list. **`cxl` — CxlAllocationStrategy** diff --git a/docs/source/design/tent/hp-tcp.md b/docs/source/design/tent/hp-tcp.md new file mode 100644 index 0000000000..1b4e6c1b65 --- /dev/null +++ b/docs/source/design/tent/hp-tcp.md @@ -0,0 +1,95 @@ +# TENT High-Performance TCP + +`hp_tcp` is a standalone TENT transport for CPU DRAM transfers over +data-center TCP. Standard `tcp` remains the RPC-based compatibility path. +The first version intentionally excludes GPU memory, TLS, multi-endpoint +routing, multi-NIC striping, transparent replay after an ambiguous WRITE and +dynamic lane scheduling. + +## Architecture + +Each worker owns one `asio::io_context` and one thread. Each peer has a +configured number of persistent lanes, and request IDs distribute operations +across them. A stable hash of peer and lane selects the owner; socket state +never moves between workers, and operations on a lane are FIFO. ASIO provides +the event queue; process-wide task and byte admission limits bound all accepted +work, including callbacks waiting in that queue. + +The server uses the same worker pool. Accepted sockets are assigned to workers +and stored in worker-owned session sets. A global connection limit bounds live +sessions; closing a session removes it immediately rather than retaining one +thread per historical connection. + +```text +TENT request -> bounded admission -> owner worker -> persistent lane + -> versioned TCP protocol -> registered remote buffer +``` + +## Protocol and memory safety + +Requests contain a version, opcode, request ID, registration ID, remote +address and length. Responses contain the request ID, status and committed +byte count. A WRITE completes only after the target has copied the full payload +and returned an acknowledgement. A READ completes after the full response +payload arrives. + +Every registered buffer has an ID formed from a per-registry random namespace +and a monotonic sequence, plus a remote permission. This prevents a stale ID +from a previous server incarnation from becoming valid after restart. The +target validates the ID, range and permission before access. An operation holds +a lease until its final I/O callback retires; unregister hides the range from +new work and waits for existing leases. Stale registration metadata causes one +bounded metadata refresh and retry on the same transport. Permission and range +failures are terminal. + +If a WRITE request may have reached the peer but no valid acknowledgement is +received, the remote outcome is unknown. That failure is terminal and is not +replayed through another transport; otherwise a committed WRITE whose ACK was +lost could execute twice. + +## Timeouts and shutdown + +Resolve/connect use `connect_timeout_ms`. Header, payload and response progress +use `progress_timeout_ms` on both client and server. A newly accepted connection +must send its first header byte before the deadline, and every partial header +or payload must continue to make progress. After a valid request completes, +pure idle time on its persistent connection is not treated as stalled I/O; the +deadline resumes as soon as the next header begins. A timeout cancels the +resolver or socket; terminal completion is published only after the +corresponding callback retires. + +Shutdown closes admission and the listener, drains queued dispatch callbacks, +cancels every client lane and server session on its owner, waits for operations +and leases, then stops and joins worker threads. This makes shutdown bounded +even when a peer sends only part of a request. + +This ordering is a lifecycle invariant, not an incidental destructor detail: +the client and server are destroyed before the worker contexts they use. In a +debug build, normal teardown asserts that admission, client operations and +server sessions have all drained before their owners are destroyed. + +An exception escaping an ASIO handler marks the runtime failed and blocks +further admission. The owner event loop continues only to retire previously +committed work and process teardown cancellation with the same affinity. Once +those resources drain, shutdown joins the workers and reports the failure. +Likewise, admission-release underflow is fail-closed: counters are preserved, +new work is rejected, and drain returns an error instead of treating live work +as complete. + +## Configuration + +The transport is configured under `transports.hp_tcp`: + +| Field | Meaning | +| --- | --- | +| `enable` | Enable `hp_tcp`; set `transports.tcp.enable` to `false`. The two transports cannot be enabled together because control-plane notification ownership is singular. | +| `bind_address`, `advertise_address`, `port` | Listener and published endpoint. | +| `worker_count` | ASIO event-loop threads. | +| `connections_per_peer` | Persistent lanes per peer. | +| `max_outstanding_tasks`, `max_outstanding_bytes` | Global admission bounds. | +| `max_transfer_bytes` | Maximum request size. I/O progress is tracked in fixed internal steps. | +| `connect_timeout_ms`, `progress_timeout_ms` | Connection and I/O deadlines. | + +Tests cover wire validation, admission, buffer leases, connection reuse, +session reaping, client/server timeout, stale-registration recovery, ambiguous +WRITE completion and a two-process READ/WRITE smoke test. diff --git a/docs/source/design/tent/overview.md b/docs/source/design/tent/overview.md index 8fb8ce093d..4f7104d230 100644 --- a/docs/source/design/tent/overview.md +++ b/docs/source/design/tent/overview.md @@ -93,6 +93,14 @@ metrics transport-selector ::: +## TENT High-Performance TCP + +:::{toctree} +:maxdepth: 1 + +hp-tcp +::: + ## TENT Quality of Service :::{toctree} diff --git a/docs/source/design/tent/slice-spraying.md b/docs/source/design/tent/slice-spraying.md index 5c8db04587..ff942fdbc6 100644 --- a/docs/source/design/tent/slice-spraying.md +++ b/docs/source/design/tent/slice-spraying.md @@ -188,7 +188,8 @@ All slice spraying parameters are configurable via the configuration file: { "transports": { "rdma": { - "numa_penalties": [1.0, 5.0, 10.0] + "numa_penalties": [1.0, 5.0, 10.0], + "strict_local_numa": false } } } @@ -197,12 +198,44 @@ All slice spraying parameters are configurable via the configuration file: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `numa_penalties` | array[float] | `[1.0, 5.0, 10.0]` | Penalty multipliers for each NUMA tier | +| `strict_local_numa` | bool | `false` | Never select a cross-NUMA NIC instead of penalizing it | **Guidelines**: - Higher values = stronger preference for local devices - Set all to `1.0` to disable NUMA awareness - Increase remote penalties if cross-NUMA latency is high +### Strict Local NUMA + +`numa_penalties` makes a remote NIC expensive but still selectable, so a busy +local NIC eventually loses to a cross-NUMA one. Set `strict_local_numa` (or the +`MC_STRICT_LOCAL_NUMA` environment variable, which accepts `1`/`0` and +`true`/`false`) to remove those NICs from selection entirely. + +A NIC is only excluded when the memory location and the NIC **both** report a +NUMA node and the nodes differ. If either side is unknown the NIC keeps its +`numa_penalties` weight, because discovery reports `-1` in cases where excluding +everything would break otherwise working hosts: + +- virtual machines and some GPUs, where sysfs exposes no `numa_node` +- bonded NICs such as `mlx5_bond_0` +- classic priority-matrix topologies (`MC_CUSTOM_TOPO_JSON`, + `topology/priority_matrix`), which carry no NUMA information at all + +On those hosts the flag has no effect; a warning is logged at startup so this is +visible rather than silent. A second warning names any location left without a +same-NUMA NIC, since transfers from it will fail with `DeviceNotFound`. + +**Trade-off**: strict mode converts a performance problem into an availability +one. Without a local NIC an allocation fails instead of degrading, so enable it +only where every memory location provably has a same-NUMA rail. + +**Scope**: the exclusion is enforced on the local NIC for both the first +selection and the retry path. For the remote NIC it is only a preference — the +peer publishes its own topology and may not run this policy, so failing a slice +because another host has no local rail would turn a local setting into a +cross-node outage. + ### Bandwidth Estimation ```json @@ -270,7 +303,15 @@ All slice spraying parameters are configurable via the configuration file: **Notes**: - Each device's bandwidth is read from the speed and width its port negotiated (`ibv_query_port`), so a 100G and a 400G NIC in the same host - start from different theoretical rates + start from different theoretical rates. Where libibverbs provides + `ibv_query_port_speed()` (rdma-core >= 62) the *effective* speed it + reports is preferred: for a VF over LAG that is the bandwidth left after + a PF drops out of the bond, which the encoded link rate cannot express. + The verb is resolved as an optional symbol, so older libraries keep + working on the encoded rate. A query *error* keeps the last known + effective speed (falling back would briefly restore the higher encoded + rate on a degraded LAG); failures are counted per device and logged once + per episode - The theoretical rate seeds the EWMA and bounds it to `[ewma_min_multiplier, ewma_max_multiplier]` times that rate - If a device's port speed cannot be read or is outside [min, max], diff --git a/docs/source/getting_started/build.md b/docs/source/getting_started/build.md index f095beb60a..726770115f 100644 --- a/docs/source/getting_started/build.md +++ b/docs/source/getting_started/build.md @@ -75,6 +75,38 @@ sudo make install `-DUSE_NOF=ON` builds the NoF registration APIs and deployment tools. Use `-DUSE_NOF=OFF` or omit the option when the NVMe-oF SSD pool is not needed. +### RISC-V Build + +Mooncake supports native 64-bit RISC-V Linux builds with `USE_RISCV` enabled. +The option keeps regular Release optimizations but disables interprocedural +optimization for the Python extensions, avoiding the excessive memory use of +full GNU LTO on RISC-V build hosts. The build also detects whether 16-byte +atomic operations require `libatomic` and links it automatically. + +The following configuration builds the C++, Python, and Rust components while +disabling every component that requires Go: + +```bash +mkdir build-riscv +cd build-riscv +cmake -G Ninja .. \ + -DCMAKE_BUILD_TYPE=Release \ + -DUSE_RISCV=ON \ + -DWITH_STORE_GO=OFF \ + -DWITH_P2P_STORE=OFF \ + -DUSE_ETCD=OFF \ + -DSTORE_USE_ETCD=OFF \ + -DSTORE_USE_K8S_LEASE=OFF \ + -DBUILD_UNIT_TESTS=OFF \ + -DBUILD_EXAMPLES=OFF \ + -DBUILD_BENCHMARK=OFF +cmake --build . --parallel 4 +``` + +Adjust the parallel job count for the available memory. The example uses four +jobs because optimized C++ and Python binding translation units can each need +several gigabytes of memory on RISC-V. + ### Hardware Backend Setup Run `sudo bash dependencies.sh` before using any of these backend-specific build @@ -206,6 +238,7 @@ The following options can be passed to `cmake ..`. | `-DUSE_HYGON=ON/OFF` | `OFF` | Enable Hygon DCU support via DTK SDK. Uses a CUDA-compatible runtime. | | `-DUSE_COREX=ON/OFF` | `OFF` | Enable Iluvatar CoreX GPU support. Uses a CUDA-compatible runtime. | | `-DUSE_MLU=ON/OFF` | `OFF` | Enable Cambricon MLU memory support via Neuware, including memory detection, topology discovery, and RDMA registration. | +| `-DUSE_RISCV=ON/OFF` | `OFF` | Enable RISC-V build compatibility settings, including disabling full IPO/LTO for Python extensions. | | `-DUSE_ASCEND_DIRECT=ON/OFF` | `OFF` | Enable Ascend Direct transport and HCCS support via the ADXL engine. Recommended for Ascend builds. | | `-DUSE_UBSHMEM=ON/OFF` | `OFF` | Enable Huawei Ascend NPU shared memory transport via CANN VMM APIs. | | `-DUSE_INTRA_NVLINK=ON/OFF` | `OFF` | Enable intranode NVLink transport. | diff --git a/docs/source/performance/mooncake/tebench.md b/docs/source/performance/mooncake/tebench.md index 514113077b..b299331465 100644 --- a/docs/source/performance/mooncake/tebench.md +++ b/docs/source/performance/mooncake/tebench.md @@ -197,6 +197,20 @@ computed from each class's actual transfer size. `--qos_classes_json`, and the global `--tent_intent_type`. Non-default per-class intents and deadlines require the TENT backend. +### 4.3 Per-Target Metrics + +Multi-target runs print one `[target-summary]` line per target. Use +`--result_output_jsonl=` to also append a schema-versioned JSON record for +each benchmark configuration. The record keeps the aggregate operation, byte, +and throughput totals plus each target's segment name, assigned thread count, +completed operations, transferred bytes, throughput, and latency distribution. +The aggregate throughput uses the pooled average worker duration, matching the +existing `BW (GB/s)` table calculation. + +Targets with no assigned worker are retained with zero-valued metrics. This +makes an under-provisioned run (`threads < targets`) visible instead of silently +dropping targets from the result. + ## 5. Runtime Configuration This section summarizes the key runtime options that control workload behavior, @@ -371,6 +385,8 @@ gpu_id + thread_id * `--qos_link_capacity_gbps` : measured usable link capacity in decimal GB/s * `--qos_output_jsonl` : append one schema-versioned JSON object per benchmark configuration +* `--result_output_jsonl` : append aggregate and per-target metrics for each + benchmark configuration QoS mode intentionally requires a fixed thread count. Sweep offered load by running explicit cases with different class thread allocations so every output diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 3745b0a716..7b84ed7664 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -87,6 +87,20 @@ option(USE_HIP "option for enabling gpu features for AMD GPU" OFF) option(USE_HYGON "option for enabling gpu features for Hygon DCU with DTK" OFF) option(USE_COREX "option for enabling gpu features for Iluvatar CoreX" OFF) option(USE_SUPA "option for enabling gpu features for Biren GPU with SUPA" OFF) +option(USE_RISCV "Enable RISC-V build compatibility settings" OFF) +if(USE_RISCV) + if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^riscv") + message( + WARNING + "USE_RISCV is enabled, but CMAKE_SYSTEM_PROCESSOR is '${CMAKE_SYSTEM_PROCESSOR}'" + ) + endif() + # Define this before any pybind11 module is created. Otherwise pybind11 adds + # its default full-LTO target, which is prohibitively resource-intensive on + # RISC-V build hosts. + set(CMAKE_INTERPROCEDURAL_OPTIMIZATION OFF) + message(STATUS "RISC-V: IPO disabled for Mooncake Python extensions") +endif() option(USE_NVMEOF "option for using NVMe over Fabric" OFF) option(USE_TCP "option for using TCP transport" ON) option(USE_BAREX "option for using accl-barex transport" OFF) diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 552cd2b6bb..25439a0726 100644 --- a/mooncake-common/include/environment_variables.h +++ b/mooncake-common/include/environment_variables.h @@ -38,6 +38,52 @@ struct FileStorageEnvironmentVariables { MC_DEFINE_ENV_VAR(std::string, MOONCAKE_USE_URING); }; +struct ClientAutoPortEnvironmentVariables { + MC_DEFINE_ENV_VAR(int, MC_STORE_CLIENT_SETUP_RETRIES); + MC_DEFINE_ENV_VAR(int, MC_STORE_CLIENT_MIN_PORT); + MC_DEFINE_ENV_VAR(int, MC_STORE_CLIENT_MAX_PORT); +}; + +struct RegisteredPinnedMemoryEnvironmentVariables { + // Keep the raw string because the legacy parser rejects a leading '+', + // unlike the shared typed integer parser. + MC_DEFINE_ENV_VAR(std::string, MC_STORE_PIN_MEMORY_MAX_BYTES); +}; + +struct LocalHotCacheEnvironmentVariables { + // Keep these values as strings to preserve their existing per-setting + // parsing, fallback, and logging behavior. + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_CACHE_SIZE); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_BLOCK_SIZE); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_CACHE_USE_SHM); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD); +}; + +struct ClientMetricEnvironmentVariables { + // Keep these values as strings because ClientMetricConfig preserves the + // existing per-setting fallback and logging behavior. + MC_DEFINE_ENV_VAR(std::string, MC_STORE_CLIENT_METRIC); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_CLIENT_METRIC_INTERVAL); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_CLIENT_METRIC_BANDWIDTH); +}; + +struct DistributedStorageEnvironmentVariables { + MC_DEFINE_ENV_VAR(std::string, MOONCAKE_DFS_ROOT_DIR); + MC_DEFINE_ENV_VAR(std::string, MOONCAKE_DISTRIBUTED_ROOT_DIR); + MC_DEFINE_ENV_VAR(std::string, MOONCAKE_DFS_FS_ADAPTER); + MC_DEFINE_ENV_VAR(std::string, MOONCAKE_DISTRIBUTED_FS_TYPE); + MC_DEFINE_ENV_VAR(bool, MOONCAKE_DISTRIBUTED_HEALTH_CHECK); + MC_DEFINE_ENV_VAR(int, MOONCAKE_DFS_SHARD_COUNT); + MC_DEFINE_ENV_VAR(uint64_t, MOONCAKE_DFS_SHARD_CAPACITY); + MC_DEFINE_ENV_VAR(uint64_t, MOONCAKE_DFS_ALIGNMENT); + MC_DEFINE_ENV_VAR(bool, MOONCAKE_DFS_SINGLE_TENANT); + MC_DEFINE_ENV_VAR(bool, MOONCAKE_DFS_EVICTION_ENABLED); + MC_DEFINE_ENV_VAR(double, MOONCAKE_DFS_EVICTION_HIGH_WATERMARK); + MC_DEFINE_ENV_VAR(double, MOONCAKE_DFS_EVICTION_LOW_WATERMARK); + MC_DEFINE_ENV_VAR(int, MOONCAKE_DFS_DEFERRED_FREE_SECONDS); + MC_DEFINE_ENV_VAR(int, MOONCAKE_DFS_EVICTION_CHECK_INTERVAL); +}; + #undef MC_DEFINE_ENV_VAR } // namespace mooncake diff --git a/mooncake-common/include/ib_link_speed.h b/mooncake-common/include/ib_link_speed.h index 150039d3a2..7e9fba9261 100644 --- a/mooncake-common/include/ib_link_speed.h +++ b/mooncake-common/include/ib_link_speed.h @@ -70,6 +70,18 @@ inline double ibLinkSpeedGbps(int active_speed, int active_width) { return ibLaneSpeedGbps(active_speed) * ibLinkWidthLanes(active_width); } +// Port speed in Gbps, preferring the effective speed ibv_query_port_speed() +// reports (rdma-core >= 62, here in Mb/s) over the encoded link rate. +// The two differ for a VF over LAG: a PF dropping out of the bond halves +// the VF's bandwidth while its port stays ACTIVE at the same encoding, and +// only the effective speed reflects that. 0 for effective_mbps means the +// verb is unavailable or reported nothing, and the encodings decide. +inline double ibPortSpeedGbps(unsigned long long effective_mbps, + int active_speed, int active_width) { + if (effective_mbps > 0) return effective_mbps / 1000.0; + return ibLinkSpeedGbps(active_speed, active_width); +} + } // namespace mooncake #endif // MOONCAKE_IB_LINK_SPEED_H_ diff --git a/mooncake-common/tests/ib_link_speed_test.cpp b/mooncake-common/tests/ib_link_speed_test.cpp index 9798cce7c1..774dbcf687 100644 --- a/mooncake-common/tests/ib_link_speed_test.cpp +++ b/mooncake-common/tests/ib_link_speed_test.cpp @@ -52,6 +52,23 @@ TEST(IbLinkSpeedTest, ConvertsPortAttrEncodingsToGbps) { // An encoding the table does not know must not be guessed at: 0 tells the // caller the speed is unknown so it can fall back explicitly. +// ibv_query_port_speed() (rdma-core >= 62) reports the port's *effective* +// speed in Mb/s: for a VF over LAG that is the bandwidth left +// after a PF drops out, which the encoded link rate cannot express. When +// available it wins; otherwise the encodings decide as before. +TEST(IbLinkSpeedTest, EffectiveSpeedWinsOverEncodedRate) { + // 400G link, but the LAG under this VF is down to one 200G PF. + EXPECT_DOUBLE_EQ(ibPortSpeedGbps(200'000, 128, 2), 200.0); + // Effective speed known, encodings unknown: still usable. + EXPECT_DOUBLE_EQ(ibPortSpeedGbps(100'000, 0, 0), 100.0); +} + +TEST(IbLinkSpeedTest, EncodedRateWhenEffectiveSpeedIsUnavailable) { + // 0 = the library predates the verb or the driver reported nothing. + EXPECT_DOUBLE_EQ(ibPortSpeedGbps(0, 128, 2), 400.0); + EXPECT_DOUBLE_EQ(ibPortSpeedGbps(0, 0, 0), 0.0); +} + TEST(IbLinkSpeedTest, UnknownEncodingsReportZero) { EXPECT_DOUBLE_EQ(ibLinkSpeedGbps(0, 2), 0.0); // speed unset EXPECT_DOUBLE_EQ(ibLinkSpeedGbps(32, 0), 0.0); // width unset diff --git a/mooncake-store/AGENTS.md b/mooncake-store/AGENTS.md index 7e2480637e..2bf0ba0444 100644 --- a/mooncake-store/AGENTS.md +++ b/mooncake-store/AGENTS.md @@ -11,3 +11,12 @@ - Do not migrate `src/cachelib_memory_allocator/` unless explicitly requested. - The shared engine is not cryptographically secure; do not use it for secrets or authentication tokens. + +## Local-first Host Identity + +- Keep the placement host ID separate from the routable transfer endpoint. +- `MOONCAKE_HOST_ID` takes priority after ASCII whitespace trimming. Unset, + empty, or whitespace-only values fall back to the existing `local_hostname` + behavior. +- Normalize endpoint-shaped host IDs before rejecting loopback or wildcard + identities. diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h index 5b641db1d3..961008fb61 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -2,11 +2,14 @@ #define BUFFER_ALLOCATOR_H #include +#include #include #include #include #include +#include + #include "cachelib_memory_allocator/MemoryAllocator.h" #include "offset_allocator/offset_allocator.h" #include "storage_usage.h" @@ -29,12 +32,18 @@ enum class ReplicaType { DFS = 100, // Distributed filesystem page-offset replica }; +struct LiveAllocation { + uint64_t offset_from_base{0}; + uint64_t requested_size{0}; +}; + // Constant for unknown free space in allocators that don't track it precisely static constexpr size_t kAllocatorUnknownFreeSpace = std::numeric_limits::max(); // Forward declarations class BufferAllocatorBase; +class Replica; class AllocatedBuffer { public: @@ -70,6 +79,10 @@ class AllocatedBuffer { return !allocator_.expired(); } + [[nodiscard]] std::shared_ptr getAllocator() const { + return allocator_.lock(); + } + // Serialize the buffer into a descriptor for transfer [[nodiscard]] Descriptor get_descriptor() const; @@ -93,6 +106,8 @@ class AllocatedBuffer { void* get_vaddr_from_cxl(); private: + bool copyTransferProtocolFrom(const AllocatedBuffer& source); + std::weak_ptr allocator_; std::string segment_name_; void* buffer_ptr_{nullptr}; @@ -103,6 +118,7 @@ class AllocatedBuffer { std::nullopt}; friend class Serializer; + friend class Replica; }; /** @@ -117,6 +133,7 @@ class BufferAllocatorBase { virtual void deallocate(AllocatedBuffer* handle) = 0; virtual size_t capacity() const = 0; virtual size_t size() const = 0; + virtual uintptr_t base() const = 0; virtual std::string getSegmentName() const = 0; virtual std::string getTransportEndpoint() const = 0; @@ -175,6 +192,7 @@ class DummyBufferAllocator final : public BufferAllocatorBase { return kAllocatorUnknownFreeSpace; } size_t size() const override { return 0; } + uintptr_t base() const override { return 0; } std::string getSegmentName() const override { return segment_name_; } std::string getTransportEndpoint() const override { return transport_endpoint_; @@ -189,12 +207,7 @@ class DummyBufferAllocator final : public BufferAllocatorBase { * CachelibBufferAllocator manages memory allocation using CacheLib's slab * allocation strategy. * - * Important alignment requirements: - * 1. Base address must be at least 8-byte aligned (CacheLib requirement) - * 2. Base address should be 4MB aligned since the total size must be a multiple - * of 4MB - * 3. Use sufficiently high base addresses (e.g., 0x100000000 for 4GB) to avoid - * memory conflicts + * The base address and size must both be aligned to CacheLib's slab size. * * Example usage: * ```cpp @@ -202,7 +215,7 @@ class DummyBufferAllocator final : public BufferAllocatorBase { * const size_t base = 0x100000000; // 4GB aligned * const size_t base = 0x200000000; // 8GB aligned * - * // Bad - will likely crash + * // Bad - Create() returns ErrorCode::INVALID_PARAMS * const size_t base = 0x1234; // Too low, unaligned * const size_t base = 0x100000001; // Not 4MB aligned * ``` @@ -211,9 +224,10 @@ class CachelibBufferAllocator : public BufferAllocatorBase, public std::enable_shared_from_this { public: - CachelibBufferAllocator(std::string segment_name, size_t base, size_t size, - std::string transport_endpoint, - ReplicaType replica_type = ReplicaType::MEMORY); + static tl::expected, ErrorCode> + Create(std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + ReplicaType replica_type = ReplicaType::MEMORY); ~CachelibBufferAllocator() override; @@ -223,6 +237,7 @@ class CachelibBufferAllocator size_t capacity() const override { return total_size_; } size_t size() const override { return GetUsageBytes(); } + uintptr_t base() const override { return base_; } std::string getSegmentName() const override { return segment_name_; } std::string getTransportEndpoint() const override { return transport_endpoint_; @@ -238,8 +253,12 @@ class CachelibBufferAllocator } private: + CachelibBufferAllocator(std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + ReplicaType replica_type); + std::unique_ptr adoptImportedBuffer( - const AllocatedBuffer::Descriptor& descriptor); + const LiveAllocation& allocation); // metadata const std::string segment_name_; const size_t base_; @@ -257,10 +276,10 @@ class CachelibBufferAllocator friend struct RestoredCachelibBufferAllocator; friend std::optional - RestoreCachelibBufferAllocator( + ImportCachelibBufferAllocator( std::string segment_name, size_t base, size_t size, std::string transport_endpoint, - const std::vector& descriptors, + const std::vector& allocations, ReplicaType replica_type); }; @@ -269,10 +288,10 @@ struct RestoredCachelibBufferAllocator { std::vector> buffers; }; -std::optional RestoreCachelibBufferAllocator( +std::optional ImportCachelibBufferAllocator( std::string segment_name, size_t base, size_t size, std::string transport_endpoint, - const std::vector& descriptors, + const std::vector& allocations, ReplicaType replica_type = ReplicaType::MEMORY); /** @@ -296,6 +315,7 @@ class OffsetBufferAllocator size_t capacity() const override { return total_size_; } size_t size() const override { return GetUsageBytes(); } + uintptr_t base() const override { return base_; } std::string getSegmentName() const override { return segment_name_; } std::string getTransportEndpoint() const override { return transport_endpoint_; @@ -335,15 +355,21 @@ struct RestoredOffsetBufferAllocator { std::vector> buffers; }; -// Reconstructs an empty OffsetBufferAllocator from final live descriptors. -// The returned buffers follow descriptor input order. No state is exposed on +// Reconstructs an empty OffsetBufferAllocator from final live allocations. +// The returned buffers follow allocation input order. No state is exposed on // validation or allocation failure. -std::optional RestoreOffsetBufferAllocator( +std::optional ImportOffsetBufferAllocator( std::string segment_name, size_t base, size_t size, std::string transport_endpoint, - const std::vector& descriptors, + const std::vector& allocations, ReplicaType replica_type = ReplicaType::MEMORY); +tl::expected, ErrorCode> +CreateBufferAllocator(BufferAllocatorType allocator_type, + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + ReplicaType replica_type = ReplicaType::MEMORY); + // The main difference is that it allocates real memory and returns it, while // BufferAllocator allocates an address class SimpleAllocator { diff --git a/mooncake-store/include/client_auto_port_config.h b/mooncake-store/include/client_auto_port_config.h new file mode 100644 index 0000000000..1bbc44cdc5 --- /dev/null +++ b/mooncake-store/include/client_auto_port_config.h @@ -0,0 +1,13 @@ +#pragma once + +namespace mooncake { + +struct ClientAutoPortConfig { + int max_retries = 20; + int min_port = 12300; + int max_port = 14300; + + static ClientAutoPortConfig FromEnvironment(); +}; + +} // namespace mooncake diff --git a/mooncake-store/include/client_metric.h b/mooncake-store/include/client_metric.h index 9784600d3a..f1c75680e8 100644 --- a/mooncake-store/include/client_metric.h +++ b/mooncake-store/include/client_metric.h @@ -653,6 +653,14 @@ struct SsdMetric { } }; +struct ClientMetricConfig { + bool enabled = true; + std::chrono::milliseconds reporting_interval{0}; + bool bandwidth_reporting_enabled = true; + + static ClientMetricConfig FromEnvironment(); +}; + struct ClientMetric { TransferMetric transfer_metric; MasterClientMetric master_client_metric; diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index cdb156303c..4b753c0667 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -814,8 +814,8 @@ class Client { const DiskDescriptor& disk_descriptor); /** * @brief Initialize local hot cache - * @return ErrorCode::OK if use local hot cache, - * ErrorCode::INVALID_PARAMS if invalid MC_STORE_LOCAL_HOT_CACHE_SIZE config + * @return ErrorCode::OK if disabled or initialized successfully; + * ErrorCode::INVALID_PARAMS if cache allocation or registration fails */ ErrorCode InitLocalHotCache(); @@ -824,21 +824,6 @@ class Client { */ void UnregisterLocalHotCacheMemory(); - /** - * @brief Read MC_STORE_LOCAL_HOT_CACHE_SIZE from environment variable - * @return Cache size in bytes, or 0 if not set or invalid - */ - size_t GetLocalHotCacheSizeFromEnv(); - - /** - * @brief Read MC_STORE_LOCAL_HOT_BLOCK_SIZE from environment variable - * @param default_value Default block size to use if env var is not set or - * invalid - * @return Parsed block size from environment, or default_value if not - * set/invalid - */ - size_t GetLocalHotBlockSizeFromEnv(size_t default_value); - /** * @brief Redirect replica descriptor to local hot cache if cache hit * @param key Object key diff --git a/mooncake-store/include/config/distributed_storage_config.h b/mooncake-store/include/config/distributed_storage_config.h new file mode 100644 index 0000000000..bce5facebc --- /dev/null +++ b/mooncake-store/include/config/distributed_storage_config.h @@ -0,0 +1,29 @@ +#pragma once + +#include +#include +#include + +namespace mooncake { + +struct DistributedStorageConfig { + std::string fsdir = "/mnt/3fs/mooncake"; + std::string fs_adapter_type = "hf3fs"; + bool enable_health_check = false; + int shard_count = 64; + uint64_t shard_capacity = 4ULL * 1024 * 1024 * 1024; + uint64_t alignment = 4096; + bool single_tenant = true; + bool eviction_enabled = true; + double eviction_high_watermark = 0.9; + double eviction_low_watermark = 0.7; + std::chrono::seconds deferred_free_duration{30}; + std::chrono::seconds eviction_check_interval{5}; + + bool Validate() const; + bool ValidateForAllocator() const; + static DistributedStorageConfig FromEnvironment(); + std::string FormatStr() const; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h b/mooncake-store/include/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h new file mode 100644 index 0000000000..f8f7b4f95b --- /dev/null +++ b/mooncake-store/include/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h @@ -0,0 +1,126 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ha/oplog/oplog_batch_types.h" +#include "types.h" + +namespace mooncake { + +inline constexpr uint64_t kDefaultBatchOpLogSnapshotIntervalSeconds = 600; + +class HaKvBackend; +class HotStandbyService; +class SnapshotMaintenanceLease; +class SnapshotObjectStore; + +struct BatchOpLogSnapshotCoordinatorConfig { + uint64_t snapshot_interval_seconds{ + kDefaultBatchOpLogSnapshotIntervalSeconds}; + size_t chunk_object_count{1000000}; + std::string snapshot_root; + std::function clock; +}; + +struct BatchOpLogSnapshotCoordinatorStatus { + bool running{false}; + bool attempt_in_flight{false}; + bool promotion_requested{false}; + uint64_t attempts{0}; + ErrorCode last_error{ErrorCode::OK}; + std::optional catch_up_target; +}; + +// Coordinates the opt-in batch-OpLog snapshot path. Construction alone does +// not start a worker or change HotStandbyService behavior. +class BatchOpLogSnapshotCoordinator final { + public: + using LeaseFactory = + std::function()>; + + BatchOpLogSnapshotCoordinator(HotStandbyService& standby, + HaKvBackend& backend, + SnapshotObjectStore& object_store, + std::string cluster_id, + BatchOpLogSnapshotCoordinatorConfig config, + LeaseFactory lease_factory = {}); + BatchOpLogSnapshotCoordinator(HotStandbyService& standby, + HaKvBackend& backend, + SnapshotObjectStore& object_store, + std::string cluster_id, + std::string snapshot_root, + uint64_t snapshot_interval_seconds = + kDefaultBatchOpLogSnapshotIntervalSeconds, + size_t chunk_object_count = 1000000); + ~BatchOpLogSnapshotCoordinator(); + + BatchOpLogSnapshotCoordinator(const BatchOpLogSnapshotCoordinator&) = + delete; + BatchOpLogSnapshotCoordinator& operator=( + const BatchOpLogSnapshotCoordinator&) = delete; + + // Starts periodic scheduling. RunOnce() remains available for + // deterministic tests and callers that own the scheduling loop. + void Start(); + void Stop(); + + // Executes at most one attempt. An ineligible cycle returns OK and leaves + // the standby OpLog apply loop untouched. + ErrorCode RunOnce(); + ErrorCode PollOnce() { return RunOnce(); } + + // Called by HotStandbyService before promotion/stop. Promotion keeps a + // fully uploaded candidate eligible for the background publish step. + void NotifyPromotion(); + void OnPromotion() { NotifyPromotion(); } + + BatchOpLogSnapshotCoordinatorStatus GetStatus() const; + bool IsRunning() const; + bool IsAttemptInFlight() const; + ErrorCode last_error() const; + + private: + using Clock = std::chrono::steady_clock; + + void SchedulerLoop(); + ErrorCode RunAttempt(); + std::optional ReadLatestBatchId(ErrorCode& error) const; + bool CatchUpComplete(const DurablePrefix& target) const; + void FinishAttempt(ErrorCode error, bool count_attempt); + Clock::time_point Now() const; + void OnCaptureReleased(); + std::optional ReadDurablePrefix() const; + void RequestStop(); + + HotStandbyService& standby_; + HaKvBackend& backend_; + SnapshotObjectStore& object_store_; + std::string cluster_id_; + BatchOpLogSnapshotCoordinatorConfig config_; + LeaseFactory lease_factory_; + + mutable std::mutex mutex_; + std::condition_variable cv_; + std::thread worker_; + bool running_{false}; + bool stop_requested_{false}; + bool attempt_in_flight_{false}; + bool capture_active_{false}; + bool promotion_requested_{false}; + uint64_t attempts_{0}; + ErrorCode last_error_{ErrorCode::OK}; + std::optional last_attempt_complete_; + std::optional capture_cursor_; + std::optional catch_up_target_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/snapshot/batch_oplog_snapshot_coordinator.h b/mooncake-store/include/ha/snapshot/batch_oplog_snapshot_coordinator.h new file mode 100644 index 0000000000..428289d413 --- /dev/null +++ b/mooncake-store/include/ha/snapshot/batch_oplog_snapshot_coordinator.h @@ -0,0 +1,5 @@ +#pragma once + +// Compatibility include; the batch-OpLog implementation lives in the +// batch_oplog-specific directory. +#include "ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h" diff --git a/mooncake-store/include/hot_standby_service.h b/mooncake-store/include/hot_standby_service.h index 5e722c71e7..a1c671d5bf 100644 --- a/mooncake-store/include/hot_standby_service.h +++ b/mooncake-store/include/hot_standby_service.h @@ -178,6 +178,17 @@ class HotStandbyService { std::vector& out); void EndBatchOpLogSnapshotCapture(BatchOpLogSnapshotCapture& capture); + // N06 coordinator seams. These stay inert unless a coordinator is + // explicitly constructed by the caller. + std::optional GetLastAppliedBatchOpLogSnapshotPrefix() const; + void CancelBatchOpLogSnapshotCapture(); + using SnapshotLifecycleCallback = std::function; + void SetBatchOpLogSnapshotCaptureReleasedCallback( + SnapshotLifecycleCallback callback); + void SetBatchOpLogSnapshotPromotionCallback( + SnapshotLifecycleCallback callback); + void SetBatchOpLogSnapshotStopCallback(SnapshotLifecycleCallback callback); + // Inject a snapshot provider (from external snapshot implementation). void SetSnapshotProvider(std::unique_ptr provider); @@ -222,6 +233,8 @@ class HotStandbyService { void HandleSnapshotCaptureRequest( const OpLogBatchStandbyPollResult& result); void CancelSnapshotCapture(); + void NotifySnapshotPromotion(); + void NotifySnapshotStop(); // Shared body for Promote() and PromoteAndExportSnapshot(): runs the // promotion sequence machine transitions + gap resolution + final @@ -278,6 +291,8 @@ class HotStandbyService { std::atomic replication_loop_running_{false}; std::mutex replication_loop_mutex_; std::condition_variable replication_loop_cv_; + mutable std::mutex batch_snapshot_cursor_mutex_; + std::optional last_applied_batch_snapshot_prefix_; std::shared_ptr snapshot_capture_state_{ @@ -287,7 +302,11 @@ class HotStandbyService { // Synchronization mutable std::mutex mutex_; mutable std::mutex sync_status_callback_mutex_; + mutable std::mutex snapshot_lifecycle_callback_mutex_; SyncStatusCallback sync_status_callback_; + SnapshotLifecycleCallback snapshot_capture_released_callback_; + SnapshotLifecycleCallback snapshot_promotion_callback_; + SnapshotLifecycleCallback snapshot_stop_callback_; }; } // namespace mooncake diff --git a/mooncake-store/include/local_hot_cache.h b/mooncake-store/include/local_hot_cache.h index 3fa827f331..4ecd6bd953 100644 --- a/mooncake-store/include/local_hot_cache.h +++ b/mooncake-store/include/local_hot_cache.h @@ -19,6 +19,15 @@ namespace mooncake { +struct LocalHotCacheConfig { + size_t total_size_bytes = 0; + size_t block_size_bytes = 16 * 1024 * 1024; + bool use_shm = false; + uint8_t admission_threshold = 2; + + static LocalHotCacheConfig FromEnvironment(); +}; + /** * @brief Token captured at async hot cache fill submission time. * Invalidated when RemoveHotKey, BumpKeyGeneration, or Clear bumps diff --git a/mooncake-store/include/placement/target.h b/mooncake-store/include/placement/target.h new file mode 100644 index 0000000000..52c58b8b9a --- /dev/null +++ b/mooncake-store/include/placement/target.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include + +#include "allocator.h" + +namespace mooncake { + +// A stable allocation endpoint published to PlacementIndex. RegionResource +// owns the target and must outlive every placement reference to it. +class PlacementTarget { + public: + virtual ~PlacementTarget() = default; + + virtual std::unique_ptr Allocate(size_t size) const = 0; + + size_t Capacity() const { return allocator_->capacity(); } + size_t Used() const { return allocator_->size(); } + + protected: + explicit PlacementTarget(std::shared_ptr allocator) + : allocator_(std::move(allocator)) {} + + BufferAllocatorBase& allocator() const noexcept { return *allocator_; } + + private: + std::shared_ptr allocator_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index a5472cc83a..b790a70f31 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -454,7 +454,14 @@ class Replica { if (!buffer || !is_memory_replica()) { return false; } - std::get(data_).buffer = std::move(buffer); + auto& memory = std::get(data_); + if (!memory.buffer || + !buffer->copyTransferProtocolFrom(*memory.buffer)) { + return false; + } + // Allocator import rebuilds address ownership; the replica keeps the + // transfer protocol advertised before remount. + memory.buffer = std::move(buffer); return true; } diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h index 1b2e60894a..11e9a7e472 100644 --- a/mooncake-store/include/segment.h +++ b/mooncake-store/include/segment.h @@ -450,8 +450,8 @@ class SegmentManager { return usage_tracker_->GetUsage(); } - void initializeCxlAllocator(const std::string& cxl_path, - const size_t cxl_size); + ErrorCode initializeCxlAllocator(const std::string& cxl_path, + size_t cxl_size); // Endpoint-based segment queries (for standby restore) bool HasSegmentByEndpoint(const std::string& endpoint) const; diff --git a/mooncake-store/include/segment/region.h b/mooncake-store/include/segment/region.h new file mode 100644 index 0000000000..b1680f21f2 --- /dev/null +++ b/mooncake-store/include/segment/region.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include +#include + +#include "types.h" + +namespace mooncake { + +enum class RegionKind { + HOST_MEMORY = 0, + CXL, +}; + +struct RegionResourceSpec { + UUID id{0, 0}; + std::string name; + uintptr_t base{0}; + size_t size{0}; + std::string transport_endpoint; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/segment/region_driver.h b/mooncake-store/include/segment/region_driver.h new file mode 100644 index 0000000000..1766122e82 --- /dev/null +++ b/mooncake-store/include/segment/region_driver.h @@ -0,0 +1,120 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "allocator.h" +#include "placement/target.h" +#include "segment/region.h" + +namespace mooncake { + +struct RegionResource final { + explicit RegionResource(std::unique_ptr placement_target); + + std::unique_ptr target; + bool active{false}; +}; + +class RegionDriver; + +class PreparedRegionResource final { + public: + ~PreparedRegionResource(); + + PreparedRegionResource(PreparedRegionResource&& other) noexcept; + PreparedRegionResource& operator=(PreparedRegionResource&& other) noexcept; + PreparedRegionResource(const PreparedRegionResource&) = delete; + PreparedRegionResource& operator=(const PreparedRegionResource&) = delete; + + // The staged resource remains valid until Commit() or a move. + RegionResource& resource() const noexcept; + const std::vector>& imported_buffers() + const noexcept; + std::vector> TakeImportedBuffers(); + + void Commit() noexcept; + + private: + PreparedRegionResource( + RegionDriver& driver, const UUID& id, + std::unique_ptr resource, + std::vector> imported_buffers); + + struct State; + std::unique_ptr state_; + + friend class RegionDriver; +}; + +class RegionDriver { + public: + virtual ~RegionDriver() = default; + + std::optional allocator_type() const noexcept { + return allocator_type_; + } + + virtual tl::expected PrepareOpen( + const RegionResourceSpec& spec, + const std::vector& live_allocations) = 0; + virtual tl::expected PrepareAdopt( + const RegionResourceSpec& spec, + std::shared_ptr allocator) = 0; + + RegionResource* GetResource(const UUID& id); + const RegionResource* GetResource(const UUID& id) const; + bool Deactivate(const UUID& id); + bool Reactivate(const UUID& id); + bool Erase(const UUID& id); + + protected: + explicit RegionDriver( + std::optional allocator_type = std::nullopt) + : allocator_type_(allocator_type) {} + + PreparedRegionResource Stage( + const UUID& id, std::unique_ptr resource, + std::vector> imported_buffers = {}); + + private: + const std::optional allocator_type_; + void CommitPrepared(PreparedRegionResource& prepared) noexcept; + + std::map> resources_; + + friend class PreparedRegionResource; +}; + +using RegionDriverRegistry = + std::unordered_map>; + +struct CxlRegionDriverConfig { + std::string path; + size_t size{0}; +}; + +struct RegionDriverConfig { + BufferAllocatorType memory_allocator{BufferAllocatorType::CACHELIB}; + std::optional cxl; +}; + +tl::expected CreateRegionDrivers( + const RegionDriverConfig& config); + +// Converts descriptors that have already been canonicalized to +// spec.transport_endpoint. Segment-name aliases must be resolved by the +// recovery layer that owns the segment/catalog context before calling this +// helper. +tl::expected, ErrorCode> BuildRegionLiveAllocations( + const RegionResourceSpec& spec, + std::span descriptors); + +} // namespace mooncake diff --git a/mooncake-store/include/storage/distributed/distributed_storage_backend.h b/mooncake-store/include/storage/distributed/distributed_storage_backend.h index 7bc512625e..252cde5a2b 100644 --- a/mooncake-store/include/storage/distributed/distributed_storage_backend.h +++ b/mooncake-store/include/storage/distributed/distributed_storage_backend.h @@ -1,12 +1,12 @@ #pragma once -#include #include #include #include #include #include +#include "config/distributed_storage_config.h" #include "fs_adapter.h" #include "replica.h" #include "storage/distributed/object_storage_adapter.h" @@ -21,26 +21,6 @@ enum class DistributedStorageMode { kObjectStorage, }; -struct DistributedStorageConfig { - std::string fsdir = "/mnt/3fs/mooncake"; - std::string fs_adapter_type = "hf3fs"; - bool enable_health_check = false; - int shard_count = 64; - uint64_t shard_capacity = 4ULL * 1024 * 1024 * 1024; - uint64_t alignment = 4096; - bool single_tenant = true; - bool eviction_enabled = true; - double eviction_high_watermark = 0.9; - double eviction_low_watermark = 0.7; - std::chrono::seconds deferred_free_duration{30}; - std::chrono::seconds eviction_check_interval{5}; - - bool Validate() const; - bool ValidateForAllocator() const; - static DistributedStorageConfig FromEnvironment(); - std::string FormatStr() const; -}; - struct DfsWriteRequest { std::string key; DistributedFSDescriptor descriptor; diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 9943f71595..d357d0e339 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -262,9 +262,19 @@ std::string expected_to_str(const tl::expected& expected) { // Buffer allocator functions constexpr size_t SZ_2MB = 2 * 1024 * 1024; +constexpr size_t SZ_512MB = 512 * 1024 * 1024; constexpr size_t SZ_1GB = 1024 * 1024 * 1024; constexpr double BYTES_PER_GIB = static_cast(SZ_1GB); +// 512MiB hugepages (PMD size on arm64 kernels with 64K base pages) are not +// defined by older glibc/kernel headers; provide fallbacks. +#ifndef MAP_HUGE_512MB +#define MAP_HUGE_512MB (29 << 26) // MAP_HUGE_SHIFT = 26 +#endif +#ifndef MFD_HUGE_512MB +#define MFD_HUGE_512MB (29 << 26) // MFD_HUGE_SHIFT = 26 +#endif + /** * @brief Allocates memory for the `BufferAllocator` class. * @param total_size The total size of the memory to allocate. @@ -286,7 +296,7 @@ inline size_t align_up(size_t size, size_t alignment) { * @brief Get hugepage size from env and optionally set the corresponding memfd * flags. * * @param out_flags Optional pointer to an int. If provided, - * MAP_HUGETLB and MAP_HUGE_2MB/1GB will be OR-ed into it. + * MAP_HUGETLB and MAP_HUGE_2MB/512MB/1GB will be OR-ed into it. * @return size_t Hugepage size in bytes, or 0 if disabled. */ [[nodiscard]] inline size_t get_hugepage_size_from_env( @@ -302,11 +312,12 @@ inline size_t align_up(size_t size, size_t alignment) { if (size_env != nullptr) { size_t parsed_size = string_to_byte_size(size_env); - if (parsed_size == SZ_2MB || parsed_size == SZ_1GB) { + if (parsed_size == SZ_2MB || parsed_size == SZ_512MB || + parsed_size == SZ_1GB) { size = parsed_size; } else { LOG(WARNING) << "Invalid MC_STORE_HUGEPAGE_SIZE='" << size_env - << "'. Supported: 2MB, 1GB. Fallback to 2MB."; + << "'. Supported: 2MB, 512MB, 1GB. Fallback to 2MB."; size = SZ_2MB; } } @@ -325,6 +336,12 @@ inline size_t align_up(size_t size, size_t alignment) { } else { *out_flags |= MAP_HUGE_2MB; } + } else if (size == SZ_512MB) { + if (use_memfd) { + *out_flags |= MFD_HUGE_512MB; + } else { + *out_flags |= MAP_HUGE_512MB; + } } else if (size == SZ_1GB) { if (use_memfd) { *out_flags |= MFD_HUGE_1GB; @@ -333,7 +350,9 @@ inline size_t align_up(size_t size, size_t alignment) { } } LOG(INFO) << "Using hugepage size: " - << (size == SZ_2MB ? "2MB" : "1GB"); + << (size == SZ_2MB ? "2MB" + : size == SZ_512MB ? "512MB" + : "1GB"); } return size; diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 2c2d5f9089..1b933e1bcf 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -8,6 +8,7 @@ set(MOONCAKE_STORE_SHARED_SOURCES thread_pool.cpp etcd_helper.cpp uds_transport.cpp + config/distributed_storage_config.cpp storage/distributed/distributed_storage_backend.cpp storage/distributed/posix_fs_adapter.cpp utils/file_util.cpp) @@ -19,6 +20,7 @@ set(MOONCAKE_STORE_MASTER_SOURCES master_snapshot_manager.cpp master_snapshot_repository.cpp master_metric_manager.cpp + segment/region_driver.cpp segment.cpp tenant_quota.cpp tenant_quota_ledger.cpp @@ -43,6 +45,7 @@ set(MOONCAKE_STORE_MASTER_SOURCES ha/snapshot/batch_oplog/batch_oplog_snapshot_provider.cpp ha/snapshot/batch_oplog/writer.cpp ha/snapshot/batch_oplog/batch_oplog_snapshot_publisher.cpp + ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.cpp ha/snapshot/snapshot_maintenance_lease.cpp ha/snapshot/master_snapshot_codec.cpp ha/snapshot/local_ssd_codec.cpp @@ -84,11 +87,15 @@ set(MOONCAKE_STORE_CLIENT_SOURCES client_buffer.cpp aligned_client_buffer.cpp real_client.cpp + config/registered_pinned_memory_config.cpp registered_pinned_memory.cpp dummy_client.cpp shm_helper.cpp file_storage.cpp + config/client_metric_config.cpp config/file_storage_config.cpp + config/client_auto_port_config.cpp + config/local_hot_cache_config.cpp device/accelerator_device.cpp device/accelerator_registry.cpp device/runtime_accelerator.cpp @@ -106,10 +113,47 @@ set(MOONCAKE_STORE_CLIENT_SOURCES if(BUILD_UNIT_TESTS) list(APPEND MOONCAKE_STORE_CLIENT_SOURCES nvme_kv_executor_stub.cpp) endif() - set(EXTRA_LIBS "") set(SPDK_STATIC_LIBS "") +# RISC-V toolchains may implement 16-byte atomic operations in libatomic rather +# than emitting native instructions. Boost.Lockfree uses these operations in the +# Store master, so detect and propagate the runtime library in RISC-V mode. +if(USE_RISCV) + include(CheckCXXSourceCompiles) + check_cxx_source_compiles( + " +alignas(16) volatile unsigned __int128 value = 0; +int main(int argc, char**) { + __atomic_store_n(&value, static_cast(argc), + __ATOMIC_SEQ_CST); + return static_cast(__atomic_load_n(&value, __ATOMIC_SEQ_CST)); +} +" + MOONCAKE_HAS_NATIVE_16BYTE_ATOMICS) + if(NOT MOONCAKE_HAS_NATIVE_16BYTE_ATOMICS) + set(CMAKE_REQUIRED_LIBRARIES_SAVE ${CMAKE_REQUIRED_LIBRARIES}) + set(CMAKE_REQUIRED_LIBRARIES atomic) + check_cxx_source_compiles( + " +alignas(16) volatile unsigned __int128 value = 0; +int main(int argc, char**) { + __atomic_store_n(&value, static_cast(argc), + __ATOMIC_SEQ_CST); + return static_cast(__atomic_load_n(&value, __ATOMIC_SEQ_CST)); +} +" + MOONCAKE_HAS_LIBATOMIC_16BYTE_ATOMICS) + set(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES_SAVE}) + if(MOONCAKE_HAS_LIBATOMIC_16BYTE_ATOMICS) + list(APPEND EXTRA_LIBS atomic) + message(STATUS "16-byte atomics: using libatomic") + else() + message(FATAL_ERROR "16-byte atomic operations are not supported") + endif() + endif() +endif() + # Find AWS SDK find_package(AWSSDK QUIET COMPONENTS s3) if(AWSSDK_FOUND) @@ -299,12 +343,13 @@ check_pie_supported(LANGUAGES CXX) string(TOUPPER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_UPPER) -# Store Client sources call GPU runtime APIs directly. Detect the available -# runtime independently from the Transfer Engine feature flags so that each -# target below can declare its own compile and link requirements. -find_package(CUDAToolkit QUIET) -if(NOT CUDAToolkit_FOUND) - find_package(hip QUIET) +# Store Client accelerator staging follows the explicitly selected build +# backend. In particular, non-CUDA wheels are built in a CUDA toolchain image, +# so SDK presence must not enable CUDA support by itself. +if(USE_CUDA) + find_package(CUDAToolkit REQUIRED) +elseif(USE_HIP) + find_package(hip REQUIRED) endif() # Sources used by both the Master and Client sides of Store. @@ -413,13 +458,13 @@ endif() if(STORE_USE_ETCD) add_dependencies(mooncake_store_client_objects build_etcd_wrapper) endif() -if(CUDAToolkit_FOUND) - message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging") +if(USE_CUDA) + message(STATUS "mooncake_store: CUDA enabled, enabling D2H staging") target_compile_definitions(mooncake_store_client_objects PRIVATE USE_CUDA) target_include_directories(mooncake_store_client_objects PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) -elseif(hip_FOUND) - message(STATUS "mooncake_store: HIP detected, enabling D2H staging") +elseif(USE_HIP) + message(STATUS "mooncake_store: HIP enabled, enabling D2H staging") target_compile_definitions(mooncake_store_client_objects PRIVATE USE_HIP) endif() if(USE_ASCEND @@ -493,9 +538,9 @@ if(TARGET Mooncake::liburing) target_compile_definitions(mooncake_store PUBLIC USE_URING) target_link_libraries(mooncake_store PUBLIC Mooncake::liburing) endif() -if(CUDAToolkit_FOUND) +if(USE_CUDA) target_link_libraries(mooncake_store PRIVATE CUDA::cudart) -elseif(hip_FOUND) +elseif(USE_HIP) target_link_libraries(mooncake_store PRIVATE hip::host) endif() if(USE_ASCEND @@ -549,12 +594,12 @@ set_target_properties(mooncake_client PROPERTIES POSITION_INDEPENDENT_CODE ON) target_link_libraries( mooncake_client PRIVATE mooncake_store transfer_engine asio_shared gflags::gflags yalantinglibs::yalantinglibs) -if(CUDAToolkit_FOUND) +if(USE_CUDA) target_compile_definitions(mooncake_client PRIVATE USE_CUDA) target_include_directories(mooncake_client PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) target_link_libraries(mooncake_client PRIVATE CUDA::cudart) -elseif(hip_FOUND) +elseif(USE_HIP) target_compile_definitions(mooncake_client PRIVATE USE_HIP) target_link_libraries(mooncake_client PRIVATE hip::host) endif() diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp index cc7e0f2bb0..23405e8607 100644 --- a/mooncake-store/src/allocator.cpp +++ b/mooncake-store/src/allocator.cpp @@ -10,6 +10,25 @@ #include "master_metric_manager.h" namespace mooncake { +namespace { + +bool IsValidCachelibLayout(size_t base, size_t size) noexcept { + const size_t slab_count = size / sizeof(facebook::cachelib::Slab); + return base != 0 && base % facebook::cachelib::Slab::kSize == 0 && + size >= facebook::cachelib::Slab::kSize && + size % facebook::cachelib::Slab::kSize == 0 && + slab_count <= std::numeric_limits::max() && + base <= std::numeric_limits::max() - size; +} + +bool IsValidAllocation(const LiveAllocation& allocation, + size_t capacity) noexcept { + return allocation.requested_size != 0 && + allocation.offset_from_base < capacity && + allocation.requested_size <= capacity - allocation.offset_from_base; +} + +} // namespace void BufferAllocatorBase::AttachUsageTracker( const std::shared_ptr& usage_tracker) { @@ -64,6 +83,14 @@ AllocatedBuffer::AllocatedBuffer(std::shared_ptr allocator, } } +bool AllocatedBuffer::copyTransferProtocolFrom(const AllocatedBuffer& source) { + if (protocol == "cxl" || source.protocol == "cxl") { + return false; + } + protocol = source.protocol; + return true; +} + // Implementation of get_descriptor AllocatedBuffer::Descriptor AllocatedBuffer::get_descriptor() const { auto alloc = allocator_.lock(); @@ -106,7 +133,49 @@ std::ostream& operator<<(std::ostream& os, const AllocatedBuffer& buffer) { << "buffer_ptr: " << static_cast(buffer.data()) << " }"; } -// Removed allocated_bytes parameter and member initialization +tl::expected, ErrorCode> +CachelibBufferAllocator::Create(std::string segment_name, size_t base, + size_t size, std::string transport_endpoint, + ReplicaType replica_type) { + if (!IsValidCachelibLayout(base, size)) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + // CacheLib's parameter-dependent constructor failures are covered by the + // layout validation above. Do not catch allocation failures here: metadata + // exhaustion follows the process-level fail-fast policy. + return std::shared_ptr(new CachelibBufferAllocator( + std::move(segment_name), base, size, std::move(transport_endpoint), + replica_type)); +} + +tl::expected, ErrorCode> +CreateBufferAllocator(BufferAllocatorType allocator_type, + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + ReplicaType replica_type) { + switch (allocator_type) { + case BufferAllocatorType::CACHELIB: { + auto allocator = CachelibBufferAllocator::Create( + std::move(segment_name), base, size, + std::move(transport_endpoint), replica_type); + if (!allocator) { + return tl::make_unexpected(allocator.error()); + } + return std::shared_ptr(std::move(*allocator)); + } + case BufferAllocatorType::OFFSET: + // Offset construction has no parameter-dependent throwing path; + // metadata allocation failures intentionally follow fail-fast. + return std::shared_ptr( + std::make_shared( + std::move(segment_name), base, size, + std::move(transport_endpoint), replica_type)); + default: + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } +} + CachelibBufferAllocator::CachelibBufferAllocator(std::string segment_name, size_t base, size_t size, std::string transport_endpoint, @@ -216,58 +285,57 @@ void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) { } std::unique_ptr CachelibBufferAllocator::adoptImportedBuffer( - const AllocatedBuffer::Descriptor& descriptor) { - RecordAllocation(descriptor.size_); + const LiveAllocation& allocation) { + RecordAllocation(allocation.requested_size); if (replica_type_ == ReplicaType::MEMORY) { MasterMetricManager::instance().inc_allocated_mem_size( - segment_name_, descriptor.size_); + segment_name_, allocation.requested_size); } else if (replica_type_ == ReplicaType::NOF_SSD) { MasterMetricManager::instance().inc_allocated_nof_size( - segment_name_, descriptor.size_); + segment_name_, allocation.requested_size); } - return std::make_unique(shared_from_this(), descriptor); + return std::make_unique( + shared_from_this(), + reinterpret_cast(base_ + allocation.offset_from_base), + allocation.requested_size); } -std::optional RestoreCachelibBufferAllocator( +std::optional ImportCachelibBufferAllocator( std::string segment_name, size_t base, size_t size, std::string transport_endpoint, - const std::vector& descriptors, - ReplicaType replica_type) { + const std::vector& allocations, ReplicaType replica_type) { if (replica_type != ReplicaType::MEMORY || - base % facebook::cachelib::Slab::kSize != 0 || - size < facebook::cachelib::Slab::kSize || - size % facebook::cachelib::Slab::kSize != 0 || - base > std::numeric_limits::max() - size) { + !IsValidCachelibLayout(base, size)) { return std::nullopt; } - const size_t end = base + size; std::vector imports; - imports.reserve(descriptors.size()); - for (const auto& descriptor : descriptors) { - if (descriptor.protocol_ == "cxl" || - descriptor.transport_endpoint_ != transport_endpoint || - descriptor.size_ == 0 || descriptor.size_ > UINT32_MAX || - descriptor.buffer_address_ < base || - descriptor.buffer_address_ >= end || - descriptor.size_ > end - descriptor.buffer_address_) { + imports.reserve(allocations.size()); + for (const auto& allocation : allocations) { + if (!IsValidAllocation(allocation, size) || + allocation.requested_size > UINT32_MAX) { return std::nullopt; } - imports.push_back({reinterpret_cast(descriptor.buffer_address_), - static_cast(std::max( - descriptor.size_, kMinSliceSize))}); + imports.push_back( + {reinterpret_cast(base + allocation.offset_from_base), + static_cast(std::max(allocation.requested_size, + kMinSliceSize))}); } - auto allocator = std::make_shared( + auto created = CachelibBufferAllocator::Create( std::move(segment_name), base, size, transport_endpoint, replica_type); + if (!created) { + return std::nullopt; + } + auto allocator = std::move(*created); if (!allocator->memory_allocator_->importAllocations(allocator->pool_id_, imports)) { return std::nullopt; } std::vector> buffers; - buffers.reserve(descriptors.size()); - for (const auto& descriptor : descriptors) { - buffers.push_back(allocator->adoptImportedBuffer(descriptor)); + buffers.reserve(allocations.size()); + for (const auto& allocation : allocations) { + buffers.push_back(allocator->adoptImportedBuffer(allocation)); } return RestoredCachelibBufferAllocator{std::move(allocator), std::move(buffers)}; @@ -413,11 +481,10 @@ size_t OffsetBufferAllocator::getLargestFreeRegion() const { } } -std::optional RestoreOffsetBufferAllocator( +std::optional ImportOffsetBufferAllocator( std::string segment_name, size_t base, size_t size, std::string transport_endpoint, - const std::vector& descriptors, - ReplicaType replica_type) { + const std::vector& allocations, ReplicaType replica_type) { if (base > std::numeric_limits::max() - size) { return std::nullopt; } @@ -426,14 +493,14 @@ std::optional RestoreOffsetBufferAllocator( std::move(segment_name), base, size, transport_endpoint, replica_type); const auto offset_allocator = allocator->getOffsetAllocator(); - std::vector order(descriptors.size()); + std::vector order(allocations.size()); std::iota(order.begin(), order.end(), 0); std::sort(order.begin(), order.end(), [&](size_t lhs, size_t rhs) { - return descriptors[lhs].buffer_address_ < - descriptors[rhs].buffer_address_; + return allocations[lhs].offset_from_base < + allocations[rhs].offset_from_base; }); - std::vector> buffers(descriptors.size()); + std::vector> buffers(allocations.size()); std::vector> gaps; size_t cursor = base; @@ -473,26 +540,25 @@ std::optional RestoreOffsetBufferAllocator( }; for (const size_t index : order) { - const auto& descriptor = descriptors[index]; - if (descriptor.transport_endpoint_ != transport_endpoint || - descriptor.size_ == 0 || descriptor.buffer_address_ < cursor || - descriptor.buffer_address_ < base || - descriptor.buffer_address_ >= end || - descriptor.size_ > end - descriptor.buffer_address_) { + const auto& allocation = allocations[index]; + if (!IsValidAllocation(allocation, size)) { + return std::nullopt; + } + const size_t address = base + allocation.offset_from_base; + if (address < cursor) { return std::nullopt; } - const uint64_t occupied = - offset_allocator->normalizedAllocationSize(descriptor.size_); - if (occupied == 0 || occupied > end - descriptor.buffer_address_ || - !fill_gap(descriptor.buffer_address_ - cursor)) { + const uint64_t occupied = offset_allocator->normalizedAllocationSize( + allocation.requested_size); + if (occupied == 0 || occupied > end - address || + !fill_gap(address - cursor)) { return std::nullopt; } - auto buffer = allocator->allocate(descriptor.size_); - if (!buffer || reinterpret_cast(buffer->data()) != - descriptor.buffer_address_) { + auto buffer = allocator->allocate(allocation.requested_size); + if (!buffer || reinterpret_cast(buffer->data()) != address) { return std::nullopt; } - cursor = descriptor.buffer_address_ + occupied; + cursor = address + occupied; buffers[index] = std::move(buffer); } diff --git a/mooncake-store/src/client_metric.cpp b/mooncake-store/src/client_metric.cpp index 449417ab39..41f9fca5e4 100644 --- a/mooncake-store/src/client_metric.cpp +++ b/mooncake-store/src/client_metric.cpp @@ -2,11 +2,8 @@ #include #include -#include #include -#include "bool_parser.h" -#include "integer_parser.h" #include "version.h" namespace mooncake { @@ -22,55 +19,6 @@ std::map WithBuildInfoLabels( return labels; } -bool parseMetricsEnabled() { - const char* metric_env = std::getenv("MC_STORE_CLIENT_METRIC"); - if (!metric_env) { - return true; - } - return TryParseBool(metric_env).value_or(false); -} - -bool parseBoolEnv(const char* env_name, bool default_value) { - const char* env_value = std::getenv(env_name); - if (!env_value) { - return default_value; - } - - const auto parsed = TryParseBool(env_value); - if (parsed.has_value()) { - return *parsed; - } - - LOG(WARNING) << "Failed to parse " << env_name << ": " << env_value - << ", fallback to default=" << default_value; - return default_value; -} - -uint64_t parseMetricsInterval() { - const char* interval_env = std::getenv("MC_STORE_CLIENT_METRIC_INTERVAL"); - if (!interval_env) { - // Default to disabled - return 0; - } - - const auto interval = TryParseInteger( - interval_env, - {.trim_ascii_whitespace = true, .allow_leading_plus = true}); - if (!interval.has_value()) { - LOG(WARNING) << "Failed to parse MC_STORE_CLIENT_METRIC_INTERVAL: " - << interval_env << ", disabling metrics reporting"; - return 0; - } - if (*interval == 0) { - LOG(INFO) << "Client metrics reporting disabled (interval=0) via " - "MC_STORE_CLIENT_METRIC_INTERVAL"; - } else { - LOG(INFO) << "Client metrics interval set to " << *interval - << "s via MC_STORE_CLIENT_METRIC_INTERVAL"; - } - return *interval; -} - } // anonymous namespace ClientMetric::ClientMetric(uint64_t interval_seconds, @@ -105,24 +53,23 @@ ClientMetric::~ClientMetric() { StopMetricsReportingThread(); } std::unique_ptr ClientMetric::Create( const std::map& labels, bool master_rpc_metrics_enabled) { - if (!parseMetricsEnabled()) { + const auto config = ClientMetricConfig::FromEnvironment(); + if (!config.enabled) { LOG(INFO) << "Client metrics disabled (set MC_STORE_CLIENT_METRIC=0 to " "disable)"; return nullptr; } - uint64_t interval = parseMetricsInterval(); - bool bandwidth_reporting_enabled = - parseBoolEnv("MC_STORE_CLIENT_METRIC_BANDWIDTH", true); - LOG(INFO) << "Client metrics enabled (default enabled)"; LOG(INFO) << "Client bandwidth summary " - << (bandwidth_reporting_enabled ? "enabled" : "disabled") + << (config.bandwidth_reporting_enabled ? "enabled" : "disabled") << " via MC_STORE_CLIENT_METRIC_BANDWIDTH"; - return std::make_unique(interval, labels, - bandwidth_reporting_enabled, - master_rpc_metrics_enabled); + return std::make_unique( + std::chrono::duration_cast( + config.reporting_interval) + .count(), + labels, config.bandwidth_reporting_enabled, master_rpc_metrics_enabled); } void ClientMetric::serialize(std::string& str) { diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 77291f2913..4b076a4e12 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -4975,60 +4975,6 @@ tl::expected Client::GetPreferredReplica( return replica_list[0]; } -size_t Client::GetLocalHotCacheSizeFromEnv() { - if (const char* ev_size = std::getenv("MC_STORE_LOCAL_HOT_CACHE_SIZE")) { - std::string ev_size_str(ev_size); - std::string error_msg = "Invalid MC_STORE_LOCAL_HOT_CACHE_SIZE='" + - ev_size_str + "', disable local hot cache"; - // Check for negative values - if (!ev_size_str.empty() && ev_size_str[0] == '-') { - LOG(WARNING) << error_msg; - return 0; - } - try { - unsigned long long v = std::stoull(ev_size_str, nullptr, 10); - if (v > 0) { - return static_cast(v); - } else { - LOG(WARNING) << error_msg; - return 0; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - return 0; - } - } - return 0; -} - -size_t Client::GetLocalHotBlockSizeFromEnv(size_t default_value) { - if (const char* ev_block_size = - std::getenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE")) { - std::string ev_block_size_str(ev_block_size); - std::string error_msg = "Invalid MC_STORE_LOCAL_HOT_BLOCK_SIZE='" + - ev_block_size_str + - "', using default block size"; - // Check for negative values - if (!ev_block_size_str.empty() && ev_block_size_str[0] == '-') { - LOG(WARNING) << error_msg; - return default_value; - } - try { - unsigned long long v = std::stoull(ev_block_size_str, nullptr, 10); - if (v > 0) { - return static_cast(v); - } else { - LOG(WARNING) << error_msg; - return default_value; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - return default_value; - } - } - return default_value; -} - ErrorCode Client::InitLocalHotCache() { hot_cache_handler_.reset(); UnregisterLocalHotCacheMemory(); @@ -5039,38 +4985,22 @@ ErrorCode Client::InitLocalHotCache() { return ErrorCode::OK; } - // Defaults: hot cache is disabled unless MC_STORE_LOCAL_HOT_CACHE_SIZE is - // set to a positive value; when enabled, default block size is 16MB and - // thread_num is 2. - size_t block_size = 16 * 1024 * 1024; // 16MB default block size - size_t thread_num = 2; - - // Read MC_STORE_LOCAL_HOT_CACHE_SIZE from environment - size_t total_cache = GetLocalHotCacheSizeFromEnv(); - if (total_cache == 0) { - // Environment variable not set or invalid, disable cache + const auto config = LocalHotCacheConfig::FromEnvironment(); + if (config.total_size_bytes == 0) { return ErrorCode::OK; } - // Read MC_STORE_LOCAL_HOT_BLOCK_SIZE from environment - block_size = GetLocalHotBlockSizeFromEnv(block_size); - - // MC_STORE_LOCAL_HOT_CACHE_USE_SHM: "1" enables memfd-backed shm (default - // off). When enabled, hot cache is shareable with dummy clients via IPC. - bool use_shm = false; - if (const char* ev = std::getenv("MC_STORE_LOCAL_HOT_CACHE_USE_SHM")) { - use_shm = (std::string(ev) == "1"); - } + size_t thread_num = 2; // Enable hot cache { - hot_cache_ = - std::make_shared(total_cache, block_size, use_shm); + hot_cache_ = std::make_shared( + config.total_size_bytes, config.block_size_bytes, config.use_shm); // Check if cache initialization was successful if (hot_cache_->GetCacheSize() == 0) { LOG(ERROR) << "Local hot cache creation failed: no blocks allocated. " - << "total_cache=" << total_cache; + << "total_cache=" << config.total_size_bytes; hot_cache_.reset(); hot_cache_handler_.reset(); admission_sketch_.reset(); @@ -5094,35 +5024,17 @@ ErrorCode Client::InitLocalHotCache() { } hot_cache_memory_registered_ = true; - LOG(INFO) << "Local hot cache enabled with cache size=" << total_cache - << ", block size=" << block_size + LOG(INFO) << "Local hot cache enabled with cache size=" + << config.total_size_bytes + << ", block size=" << config.block_size_bytes << ", block amount=" << hot_cache_->GetCacheSize() - << ", shm=" << (use_shm ? "on" : "off") + << ", shm=" << (config.use_shm ? "on" : "off") << ", transfer engine registered=on"; // Create async handler with 2 worker threads hot_cache_handler_ = std::make_unique(hot_cache_, thread_num); admission_sketch_ = std::make_unique(); - - // MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD: minimum CMS count before a - // key is admitted to hot cache (default 2). - if (const char* ev = - std::getenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD")) { - std::string ev_str(ev); - std::string error_msg = - "Invalid MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD='" + ev_str + - "', using default"; - try { - unsigned long long v = std::stoull(ev_str, nullptr, 10); - if (v > 0 && v <= 255) { - admission_threshold_ = static_cast(v); - } else { - LOG(WARNING) << error_msg; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - } - } + admission_threshold_ = config.admission_threshold; } return ErrorCode::OK; } diff --git a/mooncake-store/src/config/client_auto_port_config.cpp b/mooncake-store/src/config/client_auto_port_config.cpp new file mode 100644 index 0000000000..00f2dc3bab --- /dev/null +++ b/mooncake-store/src/config/client_auto_port_config.cpp @@ -0,0 +1,27 @@ +#include "client_auto_port_config.h" + +#include "config.h" +#include "environ.h" +#include "environment_variables.h" + +namespace mooncake { + +ClientAutoPortConfig ClientAutoPortConfig::FromEnvironment() { + ClientAutoPortConfig config; + using Variables = ClientAutoPortEnvironmentVariables; + + config.max_retries = Environ::ReadOr( + Variables::MC_STORE_CLIENT_SETUP_RETRIES, config.max_retries); + const int raw_min_port = + Environ::ReadOr(Variables::MC_STORE_CLIENT_MIN_PORT, config.min_port); + const int raw_max_port = + Environ::ReadOr(Variables::MC_STORE_CLIENT_MAX_PORT, config.max_port); + + const auto [min_port, max_port] = ValidatePortRange( + raw_min_port, raw_max_port, config.min_port, config.max_port); + config.min_port = min_port; + config.max_port = max_port; + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/client_metric_config.cpp b/mooncake-store/src/config/client_metric_config.cpp new file mode 100644 index 0000000000..b63d970d67 --- /dev/null +++ b/mooncake-store/src/config/client_metric_config.cpp @@ -0,0 +1,67 @@ +#include "client_metric.h" + +#include +#include + +#include "bool_parser.h" +#include "environ.h" +#include "environment_variables.h" +#include "integer_parser.h" + +namespace mooncake { + +ClientMetricConfig ClientMetricConfig::FromEnvironment() { + ClientMetricConfig config; + using Variables = ClientMetricEnvironmentVariables; + + const auto enabled = Environ::Read(Variables::MC_STORE_CLIENT_METRIC); + if (enabled.has_value()) { + config.enabled = TryParseBool(*enabled).value_or(false); + } + if (!config.enabled) { + return config; + } + + const auto interval = + Environ::Read(Variables::MC_STORE_CLIENT_METRIC_INTERVAL); + if (interval.has_value()) { + const auto parsed = TryParseInteger( + *interval, + {.trim_ascii_whitespace = true, .allow_leading_plus = true}); + if (!parsed.has_value()) { + LOG(WARNING) << "Failed to parse " + << Variables::MC_STORE_CLIENT_METRIC_INTERVAL.name + << ": " << *interval + << ", disabling metrics reporting"; + } else { + config.reporting_interval = std::chrono::seconds(*parsed); + if (*parsed == 0) { + LOG(INFO) + << "Client metrics reporting disabled (interval=0) via " + << Variables::MC_STORE_CLIENT_METRIC_INTERVAL.name; + } else { + LOG(INFO) << "Client metrics interval set to " << *parsed + << "s via " + << Variables::MC_STORE_CLIENT_METRIC_INTERVAL.name; + } + } + } + + const auto bandwidth = + Environ::Read(Variables::MC_STORE_CLIENT_METRIC_BANDWIDTH); + if (bandwidth.has_value()) { + const auto parsed = TryParseBool(*bandwidth); + if (parsed.has_value()) { + config.bandwidth_reporting_enabled = *parsed; + } else { + LOG(WARNING) << "Failed to parse " + << Variables::MC_STORE_CLIENT_METRIC_BANDWIDTH.name + << ": " << *bandwidth << ", fallback to default=" + << config.bandwidth_reporting_enabled; + } + } + + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/distributed_storage_config.cpp b/mooncake-store/src/config/distributed_storage_config.cpp new file mode 100644 index 0000000000..e2e8ae14de --- /dev/null +++ b/mooncake-store/src/config/distributed_storage_config.cpp @@ -0,0 +1,144 @@ +#include "config/distributed_storage_config.h" + +#include +#include +#include +#include + +#include "environ.h" +#include "environment_variables.h" + +namespace mooncake { + +bool DistributedStorageConfig::Validate() const { + if (fsdir.empty()) { + LOG(ERROR) << "DistributedStorageConfig: fsdir is empty"; + return false; + } + if (!std::filesystem::path(fsdir).is_absolute()) { + LOG(ERROR) + << "DistributedStorageConfig: fsdir must be an absolute path: " + << fsdir; + return false; + } + if (fs_adapter_type != "hf3fs" && fs_adapter_type != "posix") { + LOG(ERROR) << "DistributedStorageConfig: unsupported fs_adapter_type: " + << fs_adapter_type; + return false; + } + if (shard_count <= 0) { + LOG(ERROR) << "DistributedStorageConfig: shard_count must > 0"; + return false; + } + if (shard_capacity == 0) { + LOG(ERROR) << "DistributedStorageConfig: shard_capacity must > 0"; + return false; + } + if (alignment == 0 || (alignment & (alignment - 1)) != 0) { + LOG(ERROR) << "DistributedStorageConfig: alignment must be power of 2"; + return false; + } + if (shard_capacity % alignment != 0) { + LOG(ERROR) << "DistributedStorageConfig: shard_capacity must align"; + return false; + } + if (!single_tenant) { + LOG(ERROR) << "DistributedStorageConfig: Currently, DFS requires " + "single_tenant=true"; + return false; + } + return true; +} + +bool DistributedStorageConfig::ValidateForAllocator() const { + if (!Validate()) return false; + + if (eviction_low_watermark < 0.0 || eviction_low_watermark > 1.0 || + eviction_high_watermark < 0.0 || eviction_high_watermark > 1.0 || + eviction_low_watermark >= eviction_high_watermark) { + LOG(ERROR) << "DistributedStorageConfig: eviction watermarks must " + "satisfy 0 <= low < high <= 1, low=" + << eviction_low_watermark + << ", high=" << eviction_high_watermark; + return false; + } + if (deferred_free_duration.count() < 0) { + LOG(ERROR) << "DistributedStorageConfig: deferred_free_duration must " + "be non-negative, seconds=" + << deferred_free_duration.count(); + return false; + } + if (eviction_enabled && eviction_check_interval.count() <= 0) { + LOG(ERROR) << "DistributedStorageConfig: eviction_check_interval must " + "be positive when eviction is enabled, seconds=" + << eviction_check_interval.count(); + return false; + } + return true; +} + +DistributedStorageConfig DistributedStorageConfig::FromEnvironment() { + DistributedStorageConfig config; + using Variables = DistributedStorageEnvironmentVariables; + + const auto legacy_root_dir = + Environ::ReadOr(Variables::MOONCAKE_DISTRIBUTED_ROOT_DIR, config.fsdir); + config.fsdir = + Environ::ReadOr(Variables::MOONCAKE_DFS_ROOT_DIR, legacy_root_dir); + if (!std::filesystem::path(config.fsdir).is_absolute()) { + config.fsdir = std::filesystem::absolute(config.fsdir).string(); + } + + const auto legacy_fs_adapter = Environ::ReadOr( + Variables::MOONCAKE_DISTRIBUTED_FS_TYPE, config.fs_adapter_type); + config.fs_adapter_type = + Environ::ReadOr(Variables::MOONCAKE_DFS_FS_ADAPTER, legacy_fs_adapter); + config.enable_health_check = + Environ::ReadOr(Variables::MOONCAKE_DISTRIBUTED_HEALTH_CHECK, + config.enable_health_check); + config.shard_count = Environ::ReadOr(Variables::MOONCAKE_DFS_SHARD_COUNT, + config.shard_count); + config.shard_capacity = Environ::ReadOr( + Variables::MOONCAKE_DFS_SHARD_CAPACITY, config.shard_capacity); + config.alignment = + Environ::ReadOr(Variables::MOONCAKE_DFS_ALIGNMENT, config.alignment); + config.single_tenant = Environ::ReadOr( + Variables::MOONCAKE_DFS_SINGLE_TENANT, config.single_tenant); + config.eviction_enabled = Environ::ReadOr( + Variables::MOONCAKE_DFS_EVICTION_ENABLED, config.eviction_enabled); + + // GetDouble silently falls back for an empty value; ReadOr emits a + // warning. Keep the existing diagnostics while this refactor is + // behavior-preserving. + config.eviction_high_watermark = + Environ::GetDouble(Variables::MOONCAKE_DFS_EVICTION_HIGH_WATERMARK.name, + config.eviction_high_watermark); + config.eviction_low_watermark = + Environ::GetDouble(Variables::MOONCAKE_DFS_EVICTION_LOW_WATERMARK.name, + config.eviction_low_watermark); + config.deferred_free_duration = std::chrono::seconds(Environ::ReadOr( + Variables::MOONCAKE_DFS_DEFERRED_FREE_SECONDS, + static_cast(config.deferred_free_duration.count()))); + config.eviction_check_interval = std::chrono::seconds(Environ::ReadOr( + Variables::MOONCAKE_DFS_EVICTION_CHECK_INTERVAL, + static_cast(config.eviction_check_interval.count()))); + return config; +} + +std::string DistributedStorageConfig::FormatStr() const { + std::ostringstream oss; + oss << "fsdir=" << fsdir << ", fs_adapter_type=" << fs_adapter_type + << ", enable_health_check=" << enable_health_check + << ", shard_count=" << shard_count + << ", shard_capacity=" << shard_capacity << ", alignment=" << alignment + << ", single_tenant=" << single_tenant + << ", eviction_enabled=" << eviction_enabled + << ", eviction_high_watermark=" << eviction_high_watermark + << ", eviction_low_watermark=" << eviction_low_watermark + << ", deferred_free_seconds=" << deferred_free_duration.count() + << ", eviction_check_interval_seconds=" + << eviction_check_interval.count(); + return oss.str(); +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/local_hot_cache_config.cpp b/mooncake-store/src/config/local_hot_cache_config.cpp new file mode 100644 index 0000000000..ab77f00332 --- /dev/null +++ b/mooncake-store/src/config/local_hot_cache_config.cpp @@ -0,0 +1,92 @@ +#include "local_hot_cache.h" + +#include + +#include +#include +#include + +#include "environ.h" +#include "environment_variables.h" + +namespace mooncake { +namespace { + +size_t ParseLegacyPositiveSizeOr(const std::optional& raw_value, + size_t default_value, + const char* variable_name, + const char* fallback_message) { + if (!raw_value.has_value()) { + return default_value; + } + const std::string error_message = "Invalid " + std::string(variable_name) + + "='" + *raw_value + "'" + + fallback_message; + if (!raw_value->empty() && raw_value->front() == '-') { + LOG(WARNING) << error_message; + return default_value; + } + + try { + // Keep the legacy numeric-prefix behavior of std::stoull. + const unsigned long long value = std::stoull(*raw_value, nullptr, 10); + if (value > 0) { + return static_cast(value); + } + } catch (const std::exception&) { + } + + LOG(WARNING) << error_message; + return default_value; +} + +} // namespace + +LocalHotCacheConfig LocalHotCacheConfig::FromEnvironment() { + LocalHotCacheConfig config; + using Variables = LocalHotCacheEnvironmentVariables; + + const auto total_size = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_CACHE_SIZE); + config.total_size_bytes = ParseLegacyPositiveSizeOr( + total_size, 0, Variables::MC_STORE_LOCAL_HOT_CACHE_SIZE.name, + ", disable local hot cache"); + if (config.total_size_bytes == 0) { + return config; + } + + const auto block_size = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_BLOCK_SIZE); + config.block_size_bytes = + ParseLegacyPositiveSizeOr(block_size, config.block_size_bytes, + Variables::MC_STORE_LOCAL_HOT_BLOCK_SIZE.name, + ", using default block size"); + + config.use_shm = Environ::Read(Variables::MC_STORE_LOCAL_HOT_CACHE_USE_SHM) + .value_or(std::string{}) == "1"; + + const auto admission_threshold = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD); + if (admission_threshold.has_value()) { + const std::string error_message = + "Invalid " + + std::string( + Variables::MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD.name) + + "='" + *admission_threshold + "', using default"; + try { + const unsigned long long value = + std::stoull(*admission_threshold, nullptr, 10); + if (value > 0 && value <= 255) { + config.admission_threshold = static_cast(value); + } else { + LOG(WARNING) << error_message; + } + } catch (const std::exception&) { + LOG(WARNING) << error_message; + } + } + + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/registered_pinned_memory_config.cpp b/mooncake-store/src/config/registered_pinned_memory_config.cpp new file mode 100644 index 0000000000..304ca2bbd6 --- /dev/null +++ b/mooncake-store/src/config/registered_pinned_memory_config.cpp @@ -0,0 +1,34 @@ +#include "registered_pinned_memory_config.h" + +#include + +#include "ascii_string.h" +#include "environ.h" +#include "environment_variables.h" +#include "integer_parser.h" + +namespace mooncake { + +RegisteredPinnedMemoryConfig RegisteredPinnedMemoryConfig::FromEnvironment() { + RegisteredPinnedMemoryConfig config; + using Variables = RegisteredPinnedMemoryEnvironmentVariables; + + const auto raw_value = + Environ::Read(Variables::MC_STORE_PIN_MEMORY_MAX_BYTES); + if (!raw_value.has_value() || raw_value->empty()) { + return config; + } + + const auto limit = + TryParseInteger(TrimAsciiWhitespace(*raw_value)); + if (!limit.has_value()) { + LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << *raw_value + << "', disabling Store segment pinning"; + return config; + } + + config.max_bytes = *limit; + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/registered_pinned_memory_config.h b/mooncake-store/src/config/registered_pinned_memory_config.h new file mode 100644 index 0000000000..f13fafb536 --- /dev/null +++ b/mooncake-store/src/config/registered_pinned_memory_config.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace mooncake { + +struct RegisteredPinnedMemoryConfig { + uint64_t max_bytes = 0; + + static RegisteredPinnedMemoryConfig FromEnvironment(); +}; + +} // namespace mooncake diff --git a/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.cpp b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.cpp new file mode 100644 index 0000000000..8999249e03 --- /dev/null +++ b/mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.cpp @@ -0,0 +1,462 @@ +#include "ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h" + +#include +#include +#include +#include + +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/snapshot/batch_oplog/batch_oplog_snapshot_publisher.h" +#include "ha/snapshot/batch_oplog/metadata.h" +#include "ha/snapshot/batch_oplog/writer.h" +#include "ha/snapshot/snapshot_maintenance_lease.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "hot_standby_service.h" + +namespace mooncake { +namespace { + +int64_t CurrentTimeMs() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} + +bool IsAtOrAfter(const DurablePrefix& current, const DurablePrefix& target) { + return !IsSequenceOlder(current.last_seq, target.last_seq) && + current.batch_id >= target.batch_id; +} + +} // namespace + +BatchOpLogSnapshotCoordinator::BatchOpLogSnapshotCoordinator( + HotStandbyService& standby, HaKvBackend& backend, + SnapshotObjectStore& object_store, std::string cluster_id, + BatchOpLogSnapshotCoordinatorConfig config, LeaseFactory lease_factory) + : standby_(standby), + backend_(backend), + object_store_(object_store), + cluster_id_(std::move(cluster_id)), + config_(std::move(config)), + lease_factory_(std::move(lease_factory)) { + if (!config_.clock) { + config_.clock = [] { return Clock::now(); }; + } + if (!lease_factory_) { + lease_factory_ = [this] { + return std::make_unique(cluster_id_); + }; + } + standby_.SetBatchOpLogSnapshotCaptureReleasedCallback( + [this] { OnCaptureReleased(); }); + standby_.SetBatchOpLogSnapshotPromotionCallback( + [this] { NotifyPromotion(); }); + standby_.SetBatchOpLogSnapshotStopCallback([this] { RequestStop(); }); +} + +BatchOpLogSnapshotCoordinator::BatchOpLogSnapshotCoordinator( + HotStandbyService& standby, HaKvBackend& backend, + SnapshotObjectStore& object_store, std::string cluster_id, + std::string snapshot_root, uint64_t snapshot_interval_seconds, + size_t chunk_object_count) + : BatchOpLogSnapshotCoordinator( + standby, backend, object_store, std::move(cluster_id), + BatchOpLogSnapshotCoordinatorConfig{ + .snapshot_interval_seconds = snapshot_interval_seconds, + .chunk_object_count = chunk_object_count, + .snapshot_root = std::move(snapshot_root), + .clock = {}}, + {}) {} + +BatchOpLogSnapshotCoordinator::~BatchOpLogSnapshotCoordinator() { + Stop(); + standby_.SetBatchOpLogSnapshotCaptureReleasedCallback(nullptr); + standby_.SetBatchOpLogSnapshotPromotionCallback(nullptr); + standby_.SetBatchOpLogSnapshotStopCallback(nullptr); +} + +void BatchOpLogSnapshotCoordinator::Start() { + std::thread stale_worker; + { + std::lock_guard lock(mutex_); + if (running_) { + return; + } + stop_requested_ = true; + stale_worker = std::move(worker_); + } + cv_.notify_all(); + if (stale_worker.joinable()) { + stale_worker.join(); + } + + std::lock_guard lock(mutex_); + stop_requested_ = false; + promotion_requested_ = false; + running_ = true; + worker_ = std::thread(&BatchOpLogSnapshotCoordinator::SchedulerLoop, this); +} + +void BatchOpLogSnapshotCoordinator::Stop() { + { + std::lock_guard lock(mutex_); + stop_requested_ = true; + running_ = false; + } + standby_.CancelBatchOpLogSnapshotCapture(); + cv_.notify_all(); + if (worker_.joinable()) { + worker_.join(); + } + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return !attempt_in_flight_; }); +} + +void BatchOpLogSnapshotCoordinator::NotifyPromotion() { + bool cancel_capture = false; + { + std::lock_guard lock(mutex_); + promotion_requested_ = true; + cancel_capture = capture_active_; + } + if (cancel_capture) { + standby_.CancelBatchOpLogSnapshotCapture(); + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return !capture_active_; }); + } + cv_.notify_all(); +} + +BatchOpLogSnapshotCoordinatorStatus BatchOpLogSnapshotCoordinator::GetStatus() + const { + std::lock_guard lock(mutex_); + return {.running = running_, + .attempt_in_flight = attempt_in_flight_, + .promotion_requested = promotion_requested_, + .attempts = attempts_, + .last_error = last_error_, + .catch_up_target = catch_up_target_}; +} + +bool BatchOpLogSnapshotCoordinator::IsRunning() const { + std::lock_guard lock(mutex_); + return running_; +} + +bool BatchOpLogSnapshotCoordinator::IsAttemptInFlight() const { + std::lock_guard lock(mutex_); + return attempt_in_flight_; +} + +ErrorCode BatchOpLogSnapshotCoordinator::last_error() const { + std::lock_guard lock(mutex_); + return last_error_; +} + +BatchOpLogSnapshotCoordinator::Clock::time_point +BatchOpLogSnapshotCoordinator::Now() const { + return config_.clock ? config_.clock() : Clock::now(); +} + +void BatchOpLogSnapshotCoordinator::OnCaptureReleased() { + const auto prefix = ReadDurablePrefix(); + std::lock_guard lock(mutex_); + capture_active_ = false; + if (prefix) { + catch_up_target_ = *prefix; + if (capture_cursor_ && IsSequenceOlder(catch_up_target_->last_seq, + capture_cursor_->last_seq)) { + catch_up_target_ = *capture_cursor_; + } + } else if (capture_cursor_) { + catch_up_target_ = *capture_cursor_; + } + cv_.notify_all(); +} + +std::optional BatchOpLogSnapshotCoordinator::ReadDurablePrefix() + const { + OpLogBatchStorage storage(cluster_id_, backend_); + DurablePrefix prefix; + if (storage.ReadDurablePrefix(prefix) != ErrorCode::OK) { + return std::nullopt; + } + return prefix; +} + +std::optional BatchOpLogSnapshotCoordinator::ReadLatestBatchId( + ErrorCode& error) const { + error = ErrorCode::OK; + uint64_t published_batch_id = 0; + for (const auto& key : + {ha::BuildBatchOpLogSnapshotLatestKey(cluster_id_), + ha::BuildBatchOpLogSnapshotFallbackKey(cluster_id_)}) { + std::string value; + const auto get_error = backend_.Get(key, value); + if (get_error == ErrorCode::ETCD_KEY_NOT_EXIST) { + continue; + } + if (get_error != ErrorCode::OK) { + error = get_error; + return std::nullopt; + } + auto descriptor = ha::DecodeBatchOpLogSnapshotDescriptor(value); + // Corrupt pointers are handled by the fenced publisher; they do not + // qualify a newer local cursor on their own. + if (descriptor) { + published_batch_id = std::max(published_batch_id, + descriptor->last_included_batch_id); + } + } + return published_batch_id; +} + +bool BatchOpLogSnapshotCoordinator::CatchUpComplete( + const DurablePrefix& target) const { + const auto current = standby_.GetLastAppliedBatchOpLogSnapshotPrefix(); + return current && IsAtOrAfter(*current, target); +} + +void BatchOpLogSnapshotCoordinator::FinishAttempt(ErrorCode error, + bool count_attempt) { + std::lock_guard lock(mutex_); + if (count_attempt) { + ++attempts_; + last_attempt_complete_ = Now(); + } + last_error_ = error; + attempt_in_flight_ = false; + capture_active_ = false; + cv_.notify_all(); +} + +ErrorCode BatchOpLogSnapshotCoordinator::RunOnce() { + try { + { + std::lock_guard lock(mutex_); + if (attempt_in_flight_) { + return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + if (stop_requested_) { + return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + if (last_attempt_complete_ && + Now() - *last_attempt_complete_ < + std::chrono::seconds(config_.snapshot_interval_seconds)) { + return ErrorCode::OK; + } + if (promotion_requested_) { + return ErrorCode::OK; + } + } + return RunAttempt(); + } catch (const std::exception& e) { + LOG(ERROR) << "Batch snapshot coordinator failed: " << e.what(); + FinishAttempt(ErrorCode::INTERNAL_ERROR, false); + return ErrorCode::INTERNAL_ERROR; + } catch (...) { + LOG(ERROR) << "Batch snapshot coordinator failed with unknown error"; + FinishAttempt(ErrorCode::INTERNAL_ERROR, false); + return ErrorCode::INTERNAL_ERROR; + } +} + +ErrorCode BatchOpLogSnapshotCoordinator::RunAttempt() { + if (config_.snapshot_root.empty() || config_.chunk_object_count == 0 || + !NormalizeAndValidateClusterId(cluster_id_) || cluster_id_.empty()) { + FinishAttempt(ErrorCode::INVALID_PARAMS, false); + return ErrorCode::INVALID_PARAMS; + } + + ErrorCode read_error = ErrorCode::OK; + const auto latest_batch_id = ReadLatestBatchId(read_error); + if (!latest_batch_id) { + FinishAttempt(read_error, false); + return read_error; + } + const auto local_prefix = standby_.GetLastAppliedBatchOpLogSnapshotPrefix(); + if (!local_prefix || local_prefix->batch_id <= *latest_batch_id) { + FinishAttempt(ErrorCode::OK, false); + return ErrorCode::OK; + } + bool catch_up_blocked = false; + { + std::lock_guard lock(mutex_); + if (attempt_in_flight_) { + return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + catch_up_blocked = catch_up_target_.has_value(); + attempt_in_flight_ = true; + } + if (catch_up_blocked) { + std::optional target; + { + std::lock_guard lock(mutex_); + target = catch_up_target_; + } + if (target && !CatchUpComplete(*target)) { + FinishAttempt(ErrorCode::OK, false); + return ErrorCode::OK; + } + } + + auto lease = lease_factory_(); + if (!lease) { + // A factory may return null to represent a busy maintenance lease. + FinishAttempt(ErrorCode::OK, false); + return ErrorCode::OK; + } + const ErrorCode lease_error = + lease->IsHeld() ? ErrorCode::OK : lease->Acquire(); + if (lease_error != ErrorCode::OK) { + // A busy maintenance lease is an ordinary skipped cycle. + const ErrorCode result = lease_error == ErrorCode::ETCD_TRANSACTION_FAIL + ? ErrorCode::OK + : lease_error; + FinishAttempt(result, false); + return result; + } + + auto release_lease = [&] { (void)lease->Release(); }; + bool cancel_before_capture = false; + { + std::lock_guard lock(mutex_); + cancel_before_capture = stop_requested_ || promotion_requested_; + } + if (cancel_before_capture) { + release_lease(); + FinishAttempt(ErrorCode::OK, true); + return ErrorCode::OK; + } + + // Re-read both sides after fencing the maintenance lease. + const auto reread_latest = ReadLatestBatchId(read_error); + const auto reread_local = standby_.GetLastAppliedBatchOpLogSnapshotPrefix(); + if (!reread_latest || !reread_local || + reread_local->batch_id <= *reread_latest) { + const ErrorCode result = + read_error == ErrorCode::OK ? ErrorCode::OK : read_error; + release_lease(); + FinishAttempt(result, true); + return result; + } + + auto capture = standby_.BeginBatchOpLogSnapshotCapture(); + { + std::lock_guard lock(mutex_); + capture_active_ = capture.has_value(); + if (capture && promotion_requested_) { + capture_active_ = false; + } + } + if (!capture || capture->last_included_batch_id <= *reread_latest) { + release_lease(); + FinishAttempt(ErrorCode::OK, true); + return ErrorCode::OK; + } + + bool cancel_after_promotion = false; + { + std::lock_guard lock(mutex_); + cancel_after_promotion = promotion_requested_; + } + if (cancel_after_promotion) { + standby_.CancelBatchOpLogSnapshotCapture(); + standby_.EndBatchOpLogSnapshotCapture(*capture); + release_lease(); + FinishAttempt(ErrorCode::OK, true); + return ErrorCode::OK; + } + + const std::string snapshot_id = + std::to_string(capture->last_included_batch_id) + "-" + + lease->owner_token(); + { + std::lock_guard lock(mutex_); + capture_cursor_ = + DurablePrefix{.batch_id = capture->last_included_batch_id, + .last_seq = capture->last_included_seq}; + } + const std::string artifact_prefix = + ha::BuildBatchOpLogSnapshotArtifactPrefix(config_.snapshot_root, + snapshot_id); + + BatchOpLogSnapshotWriter writer(object_store_); + auto descriptor = + writer.Write(standby_, *capture, config_.snapshot_root, snapshot_id, + config_.chunk_object_count, CurrentTimeMs()); + { + std::lock_guard lock(mutex_); + capture_active_ = false; + } + if (!descriptor) { + release_lease(); + FinishAttempt(ErrorCode::INTERNAL_ERROR, true); + return ErrorCode::INTERNAL_ERROR; + } + + bool stop_before_publish = false; + { + std::lock_guard lock(mutex_); + stop_before_publish = stop_requested_ && !promotion_requested_; + } + if (stop_before_publish) { + auto cleanup = object_store_.DeleteObjectsWithPrefix(artifact_prefix); + if (!cleanup) { + LOG(WARNING) << "Failed to clean snapshot candidate after stop: " + << cleanup.error(); + } + release_lease(); + FinishAttempt(ErrorCode::OK, true); + return ErrorCode::OK; + } + + BatchOpLogSnapshotPublisher publisher(backend_, cluster_id_); + ErrorCode publish_error = publisher.Publish(*lease, *descriptor); + if (publish_error != ErrorCode::OK) { + auto cleanup = object_store_.DeleteObjectsWithPrefix(artifact_prefix); + if (!cleanup) { + LOG(WARNING) << "Failed to clean unpublished snapshot candidate: " + << cleanup.error(); + } + } + release_lease(); + FinishAttempt(publish_error, true); + return publish_error; +} + +void BatchOpLogSnapshotCoordinator::SchedulerLoop() { + while (true) { + { + std::unique_lock lock(mutex_); + const auto delay = + std::chrono::seconds(config_.snapshot_interval_seconds == 0 + ? 1 + : config_.snapshot_interval_seconds); + if (cv_.wait_for(lock, delay, [this] { return stop_requested_; })) { + return; + } + } + if (RunOnce() == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS) { + std::lock_guard lock(mutex_); + if (stop_requested_) { + return; + } + } + } +} + +void BatchOpLogSnapshotCoordinator::RequestStop() { + { + std::lock_guard lock(mutex_); + stop_requested_ = true; + running_ = false; + } + standby_.CancelBatchOpLogSnapshotCapture(); + cv_.notify_all(); +} + +} // namespace mooncake diff --git a/mooncake-store/src/hot_standby_service.cpp b/mooncake-store/src/hot_standby_service.cpp index 2230f44237..9c9a9a2760 100644 --- a/mooncake-store/src/hot_standby_service.cpp +++ b/mooncake-store/src/hot_standby_service.cpp @@ -86,6 +86,10 @@ ErrorCode HotStandbyService::Start(const std::string& primary_address, batch_standby_reader_.reset(); batch_standby_kv_backend_.reset(); batch_snapshot_baseline_.reset(); + { + std::lock_guard cursor_lock(batch_snapshot_cursor_mutex_); + last_applied_batch_snapshot_prefix_.reset(); + } last_error_.store(ErrorCode::OK, std::memory_order_release); @@ -300,6 +304,8 @@ ErrorCode HotStandbyService::StartOplogFollowingLocked( if (cursor_error != ErrorCode::OK) { return cursor_error; } + std::lock_guard cursor_lock(batch_snapshot_cursor_mutex_); + last_applied_batch_snapshot_prefix_ = *batch_snapshot_baseline_; } state_machine_.ProcessEvent(StandbyEvent::SYNC_COMPLETE); @@ -367,6 +373,9 @@ void HotStandbyService::Stop() { return; } + if (current_state != StandbyState::PROMOTED) { + NotifySnapshotStop(); + } state_machine_.ProcessEvent(StandbyEvent::STOP); StopReplicationLoop(); @@ -379,6 +388,56 @@ void HotStandbyService::Stop() { << StandbyStateToString(GetState()); } +std::optional +HotStandbyService::GetLastAppliedBatchOpLogSnapshotPrefix() const { + std::lock_guard lock(batch_snapshot_cursor_mutex_); + return last_applied_batch_snapshot_prefix_; +} + +void HotStandbyService::CancelBatchOpLogSnapshotCapture() { + CancelSnapshotCapture(); +} + +void HotStandbyService::SetBatchOpLogSnapshotCaptureReleasedCallback( + SnapshotLifecycleCallback callback) { + std::lock_guard lock(snapshot_lifecycle_callback_mutex_); + snapshot_capture_released_callback_ = std::move(callback); +} + +void HotStandbyService::SetBatchOpLogSnapshotPromotionCallback( + SnapshotLifecycleCallback callback) { + std::lock_guard lock(snapshot_lifecycle_callback_mutex_); + snapshot_promotion_callback_ = std::move(callback); +} + +void HotStandbyService::SetBatchOpLogSnapshotStopCallback( + SnapshotLifecycleCallback callback) { + std::lock_guard lock(snapshot_lifecycle_callback_mutex_); + snapshot_stop_callback_ = std::move(callback); +} + +void HotStandbyService::NotifySnapshotPromotion() { + SnapshotLifecycleCallback callback; + { + std::lock_guard lock(snapshot_lifecycle_callback_mutex_); + callback = snapshot_promotion_callback_; + } + if (callback) { + callback(); + } +} + +void HotStandbyService::NotifySnapshotStop() { + SnapshotLifecycleCallback callback; + { + std::lock_guard lock(snapshot_lifecycle_callback_mutex_); + callback = snapshot_stop_callback_; + } + if (callback) { + callback(); + } +} + StandbySyncStatus HotStandbyService::GetSyncStatus() const { StandbySyncStatus status; @@ -560,6 +619,7 @@ ErrorCode HotStandbyService::Promote() { ErrorCode HotStandbyService::PromoteLockedInternal( uint64_t current_applied_seq_id) { + NotifySnapshotPromotion(); StopReplicationLoop(); ErrorCode catch_up_err = FinalCatchUpForPromotionLocked(current_applied_seq_id); @@ -735,6 +795,15 @@ void HotStandbyService::EndBatchOpLogSnapshotCapture( BatchOpLogSnapshotCapture& capture) { if (capture.lease_state_ == snapshot_capture_state_) { capture.Release(); + SnapshotLifecycleCallback callback; + { + std::lock_guard lock( + snapshot_lifecycle_callback_mutex_); + callback = snapshot_capture_released_callback_; + } + if (callback) { + callback(); + } } } @@ -838,6 +907,12 @@ void HotStandbyService::ReplicationLoop() { const uint64_t expected_before = oplog_applier_->GetExpectedSequenceId(); auto result = batch_standby_reader_->PollOnce(); + { + std::lock_guard cursor_lock( + batch_snapshot_cursor_mutex_); + last_applied_batch_snapshot_prefix_ = + batch_standby_reader_->GetLastAppliedDurablePrefix(); + } HandleSnapshotCaptureRequest(result); if (result.durable_prefix_present) { const uint64_t current_primary = primary_seq_id_.load(); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 7e299f6ebb..68f26b5ddd 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -33,6 +33,7 @@ #include "common.h" #include "environ.h" #include "segment.h" +#include "segment/region_driver.h" #ifdef USE_HTTP #include "transfer_metadata_plugin.h" #endif @@ -550,8 +551,11 @@ MasterService::MasterService(const MasterServiceConfig& config) if (config.enable_cxl) { allocation_strategy_ = std::make_shared(); - segment_manager_.initializeCxlAllocator(config.cxl_path, - config.cxl_size); + const auto result = segment_manager_.initializeCxlAllocator( + config.cxl_path, config.cxl_size); + LOG_IF(FATAL, result != ErrorCode::OK) + << "Failed to initialize CXL allocator: " + << static_cast(result); VLOG(1) << "action=start_cxl_global_allocator"; } } @@ -1176,28 +1180,36 @@ auto MasterService::ReMountSegment(const std::vector& segments, if (restore.descriptors.empty()) { continue; } + const RegionResourceSpec spec{ + restore.segment.id, restore.segment.name, restore.segment.base, + restore.segment.size, restore.segment.te_endpoint}; + auto allocations = + BuildRegionLiveAllocations(spec, restore.descriptors); + if (!allocations) { + return fail_remount(allocations.error()); + } if (std::dynamic_pointer_cast( restore.old_allocator)) { - auto restored = RestoreOffsetBufferAllocator( + auto imported = ImportOffsetBufferAllocator( restore.segment.name, restore.segment.base, restore.segment.size, restore.segment.te_endpoint, - restore.descriptors); - if (!restored) { + *allocations); + if (!imported) { return fail_remount(ErrorCode::INVALID_PARAMS); } - restore.restored_allocator = std::move(restored->allocator); - restore.buffers = std::move(restored->buffers); + restore.restored_allocator = std::move(imported->allocator); + restore.buffers = std::move(imported->buffers); } else if (std::dynamic_pointer_cast( restore.old_allocator)) { - auto restored = RestoreCachelibBufferAllocator( + auto imported = ImportCachelibBufferAllocator( restore.segment.name, restore.segment.base, restore.segment.size, restore.segment.te_endpoint, - restore.descriptors); - if (!restored) { + *allocations); + if (!imported) { return fail_remount(ErrorCode::INVALID_PARAMS); } - restore.restored_allocator = std::move(restored->allocator); - restore.buffers = std::move(restored->buffers); + restore.restored_allocator = std::move(imported->allocator); + restore.buffers = std::move(imported->buffers); } else { return fail_remount(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } diff --git a/mooncake-store/src/mmap_arena.cpp b/mooncake-store/src/mmap_arena.cpp index a0fc41e2cf..8ee216f9d5 100644 --- a/mooncake-store/src/mmap_arena.cpp +++ b/mooncake-store/src/mmap_arena.cpp @@ -76,9 +76,13 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment, const size_t actual_alignment = std::max(alignment, kMinAlignment); - // Align pool size to 2MB for huge pages with overflow protection + // Align pool size to the configured huge page size (2MB when hugepages + // are not explicitly requested) with overflow protection + unsigned int hugepage_flags = 0; + const size_t hugepage_size = get_hugepage_size_from_env(&hugepage_flags); + const size_t page_align = hugepage_size > 0 ? hugepage_size : SZ_2MB; size_t aligned_pool_size; - if (!safe_align_up(pool_size, SZ_2MB, &aligned_pool_size)) { + if (!safe_align_up(pool_size, page_align, &aligned_pool_size)) { LOG(ERROR) << "Arena pool size overflow: requested=" << pool_size; return false; } @@ -91,7 +95,12 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment, // Try huge pages for better TLB performance #ifdef MAP_HUGETLB - flags |= MAP_HUGETLB; + if (hugepage_size > 0) { + // MC_STORE_USE_HUGEPAGE set: honor the configured page size. + flags |= static_cast(hugepage_flags); + } else { + flags |= MAP_HUGETLB; + } #endif void* pool_base = @@ -111,7 +120,7 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment, } // Retry without huge pages - flags &= ~MAP_HUGETLB; + flags &= ~(MAP_HUGETLB | (0x3f << 26)); // strip MAP_HUGE_* size bits LOG(WARNING) << "Arena hugepage mmap failed for pool_size=" << aligned_pool_size << " bytes" << ", errno=" << mmap_errno << " (" << strerror(mmap_errno) diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index e764f49184..565e74e54b 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -24,7 +24,7 @@ #include "config.h" #include "store_rpc_client_io_context.h" #include "bool_parser.h" -#include "environ.h" +#include "client_auto_port_config.h" #include "integer_parser.h" #include "mutex.h" #include "types.h" @@ -815,25 +815,18 @@ tl::expected RealClient::setup_internal( client_ = *client_opt; } else { // Auto port binding with retry on metadata registration failure - const int kMaxRetries = - Environ::GetInt("MC_STORE_CLIENT_SETUP_RETRIES", 20); - const int rawMinPort = - Environ::GetInt("MC_STORE_CLIENT_MIN_PORT", 12300); - const int rawMaxPort = - Environ::GetInt("MC_STORE_CLIENT_MAX_PORT", 14300); - constexpr int kDefaultMinPort = 12300; - constexpr int kDefaultMaxPort = 14300; - auto [minPort, maxPort] = ValidatePortRange( - rawMinPort, rawMaxPort, kDefaultMinPort, kDefaultMaxPort); + const auto auto_port_config = ClientAutoPortConfig::FromEnvironment(); bool success = false; - for (int retry = 0; retry < kMaxRetries; ++retry) { + for (int retry = 0; retry < auto_port_config.max_retries; ++retry) { // Create port binder to hold a port - port_binder_ = std::make_unique(minPort, maxPort); + port_binder_ = std::make_unique( + auto_port_config.min_port, auto_port_config.max_port); int port = port_binder_->getPort(); if (port < 0) { - LOG(WARNING) << "Failed to bind available port, retry " - << (retry + 1) << "/" << kMaxRetries; + LOG(WARNING) + << "Failed to bind available port, retry " << (retry + 1) + << "/" << auto_port_config.max_retries; port_binder_.reset(); std::this_thread::sleep_for(std::chrono::milliseconds(100)); continue; @@ -857,14 +850,15 @@ tl::expected RealClient::setup_internal( // Failed to create client (possibly due to metadata registration // conflict), release port and retry with a different port LOG(WARNING) << "Failed to create client on port " << port - << ", retry " << (retry + 1) << "/" << kMaxRetries; + << ", retry " << (retry + 1) << "/" + << auto_port_config.max_retries; port_binder_.reset(); std::this_thread::sleep_for(std::chrono::milliseconds(100)); } if (!success) { - LOG(ERROR) << "Failed to create client after " << kMaxRetries - << " retries"; + LOG(ERROR) << "Failed to create client after " + << auto_port_config.max_retries << " retries"; return tl::unexpected(ErrorCode::INTERNAL_ERROR); } } diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 92aecb6917..b9b3351222 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -1,14 +1,9 @@ #include "registered_pinned_memory.h" -#include #include -#include #include -#include "ascii_string.h" -#include "integer_parser.h" - #if defined(USE_CUDA) #include #endif @@ -16,20 +11,6 @@ namespace mooncake { namespace { -std::pair ParsePinnedMemoryConfig() { - const char* raw_value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); - if (!raw_value || raw_value[0] == '\0') return {false, 0}; - - const auto limit = - TryParseInteger(TrimAsciiWhitespace(raw_value)); - if (!limit.has_value()) { - LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << raw_value - << "', disabling Store segment pinning"; - return {false, 0}; - } - return {*limit != 0, *limit}; -} - void LogPinSkip(const std::string& owner, const char* reason, size_t size) { LOG(WARNING) << "Skip cudaHostRegister for " << owner << ": " << reason << ", size=" << size; @@ -88,18 +69,18 @@ RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { } RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() - : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig(), - DefaultPinOps()) {} + : RegisteredPinnedMemoryManager( + RegisteredPinnedMemoryConfig::FromEnvironment(), DefaultPinOps()) {} RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( - std::pair config, PinOps pin_ops) - : enabled_(config.first), limit_bytes_(config.second), pin_ops_(pin_ops) { + RegisteredPinnedMemoryConfig config, PinOps pin_ops) + : limit_bytes_(config.max_bytes), pin_ops_(pin_ops) { #if defined(USE_CUDA) LOG(INFO) << "Store segment pinned memory is " - << (enabled_ ? "enabled" : "disabled") + << (limit_bytes_ != 0 ? "enabled" : "disabled") << ", max_bytes=" << limit_bytes_; #else - if (enabled_) { + if (limit_bytes_ != 0) { LOG(INFO) << "Store segment pinning requested but this build has no " "CUDA runtime support"; } @@ -108,7 +89,7 @@ RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( std::shared_ptr RegisteredPinnedMemoryManager::try_pin( void* addr, size_t size, const std::string& owner) { - if (!addr || size == 0 || !enabled_) return nullptr; + if (!addr || size == 0 || limit_bytes_ == 0) return nullptr; if (!pin_ops_.register_region || !pin_ops_.unregister_region) { return nullptr; } diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index cbc43c1388..f9da64f6b8 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -5,9 +5,10 @@ #include #include #include -#include #include +#include "config/registered_pinned_memory_config.h" + namespace mooncake { class RegisteredPinnedMemoryManager; @@ -62,7 +63,7 @@ class RegisteredPinnedMemoryManager { #if defined(MOONCAKE_STORE_TEST) public: #endif - RegisteredPinnedMemoryManager(std::pair config, + RegisteredPinnedMemoryManager(RegisteredPinnedMemoryConfig config, PinOps pin_ops); #if defined(MOONCAKE_STORE_TEST) private: @@ -71,7 +72,6 @@ class RegisteredPinnedMemoryManager { bool release(RegisteredPinnedRegion* region); void remove_inactive_region_locked(void* addr, size_t size); - const bool enabled_; const uint64_t limit_bytes_; const PinOps pin_ops_; diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp index 703509f258..0ecfefdd5a 100644 --- a/mooncake-store/src/segment.cpp +++ b/mooncake-store/src/segment.cpp @@ -236,38 +236,13 @@ ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment, } } - std::shared_ptr allocator; - // CachelibBufferAllocator may throw an exception if the size or base is - // invalid for the slab allocator. - try { - // Create allocator based on the configured type - switch (segment_manager_->memory_allocator_) { - case BufferAllocatorType::CACHELIB: - allocator = std::make_shared( - segment.name, buffer, size, segment.te_endpoint); - break; - case BufferAllocatorType::OFFSET: - allocator = std::make_shared( - segment.name, buffer, size, segment.te_endpoint); - break; - default: - LOG(ERROR) << "segment_name=" << segment.name - << ", error=unknown_memory_allocator=" - << static_cast( - segment_manager_->memory_allocator_); - return ErrorCode::INVALID_PARAMS; - } - - if (!allocator) { - LOG(ERROR) << "segment_name=" << segment.name - << ", error=failed_to_create_allocator"; - return ErrorCode::INVALID_PARAMS; - } - } catch (...) { - LOG(ERROR) << "segment_name=" << segment.name - << ", error=exception_during_allocator_creation"; - return ErrorCode::INVALID_PARAMS; + auto created = + CreateBufferAllocator(segment_manager_->memory_allocator_, segment.name, + buffer, size, segment.te_endpoint); + if (!created) { + return created.error(); } + auto allocator = std::move(*created); allocator->AttachUsageTracker(segment_manager_->usage_tracker_); segment_manager_->allocator_manager_.addAllocator(segment.name, allocator); @@ -1222,37 +1197,13 @@ ErrorCode ScopedNoFSegmentAccess::MountSegment(const NoFSegment& segment, } } - std::shared_ptr allocator; - try { - switch (nof_segment_manager_->memory_allocator_) { - case BufferAllocatorType::CACHELIB: - allocator = std::make_shared( - segment.name, buffer, size, segment.te_endpoint, - ReplicaType::NOF_SSD); - break; - case BufferAllocatorType::OFFSET: - allocator = std::make_shared( - segment.name, buffer, size, segment.te_endpoint, - ReplicaType::NOF_SSD); - break; - default: - LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name - << ", error=unknown_memory_allocator=" - << static_cast( - nof_segment_manager_->memory_allocator_); - return ErrorCode::INVALID_PARAMS; - } - - if (!allocator) { - LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name - << ", error=failed_to_create_allocator"; - return ErrorCode::INVALID_PARAMS; - } - } catch (...) { - LOG(ERROR) << "NoF segment mount: segment_name=" << segment.name - << ", error=exception_during_allocator_creation"; - return ErrorCode::INVALID_PARAMS; + auto created = CreateBufferAllocator( + nof_segment_manager_->memory_allocator_, segment.name, buffer, size, + segment.te_endpoint, ReplicaType::NOF_SSD); + if (!created) { + return created.error(); } + auto allocator = std::move(*created); allocator->AttachUsageTracker(nof_segment_manager_->usage_tracker_); nof_segment_manager_->allocator_manager_.addAllocator(segment.name, @@ -1469,8 +1420,8 @@ void SegmentManager::releaseCapacityMetrics() { } } -void SegmentManager::initializeCxlAllocator(const std::string& cxl_path, - const size_t cxl_size) { +ErrorCode SegmentManager::initializeCxlAllocator(const std::string& cxl_path, + size_t cxl_size) { LOG(INFO) << "Init CXL global allocator."; LOG(INFO) << "[CXL] create allocator with " << "path=" << cxl_path << " base=0x" << std::hex @@ -1478,14 +1429,20 @@ void SegmentManager::initializeCxlAllocator(const std::string& cxl_path, << std::fixed << std::setprecision(2) << cxl_size / (1024.0 * 1024 * 1024) << " GB)"; - auto allocator = std::make_shared( - cxl_path, DEFAULT_CXL_BASE, cxl_size, cxl_path); + auto created = + CreateBufferAllocator(BufferAllocatorType::CACHELIB, cxl_path, + DEFAULT_CXL_BASE, cxl_size, cxl_path); + if (!created) { + return created.error(); + } + auto allocator = std::move(*created); allocator->AttachUsageTracker(usage_tracker_); { std::unique_lock lock(segment_mutex_); cxl_global_allocator_ = std::move(allocator); } MasterMetricManager::instance().inc_total_mem_capacity(cxl_path, cxl_size); + return ErrorCode::OK; } bool SegmentManager::HasSegmentByEndpoint(const std::string& endpoint) const { diff --git a/mooncake-store/src/segment/region_driver.cpp b/mooncake-store/src/segment/region_driver.cpp new file mode 100644 index 0000000000..ae127f8c7e --- /dev/null +++ b/mooncake-store/src/segment/region_driver.cpp @@ -0,0 +1,336 @@ +#include "segment/region_driver.h" + +#include +#include +#include + +#include "master_metric_manager.h" + +namespace mooncake { +namespace { + +using RegionResourceMap = std::map>; + +bool IsValidSpec(const RegionResourceSpec& spec) { + return spec.id != UUID{0, 0} && !spec.name.empty() && spec.base != 0 && + spec.size != 0 && + spec.base <= std::numeric_limits::max() - spec.size; +} + +class NativePlacementTarget final : public PlacementTarget { + public: + explicit NativePlacementTarget( + std::shared_ptr allocator) + : PlacementTarget(std::move(allocator)) {} + + std::unique_ptr Allocate(size_t size) const override { + return allocator().allocate(size); + } +}; + +class CxlPlacementTarget final : public PlacementTarget { + public: + CxlPlacementTarget(std::shared_ptr allocator, + std::string binding_name) + : PlacementTarget(std::move(allocator)), + cxl_binding_name_(std::move(binding_name)) {} + + std::unique_ptr Allocate(size_t size) const override { + auto buffer = allocator().allocate(size); + if (buffer) { + buffer->change_to_cxl(cxl_binding_name_); + } + return buffer; + } + + private: + std::string cxl_binding_name_; +}; + +std::unique_ptr MakeNativeResource( + std::shared_ptr allocator) { + auto target = std::make_unique(std::move(allocator)); + return std::make_unique(std::move(target)); +} + +std::unique_ptr MakeCxlResource( + const RegionResourceSpec& spec, + std::shared_ptr allocator) { + auto target = + std::make_unique(std::move(allocator), spec.name); + return std::make_unique(std::move(target)); +} + +class MemoryRegionDriver final : public RegionDriver { + public: + explicit MemoryRegionDriver(BufferAllocatorType allocator_type) + : RegionDriver(allocator_type) {} + + tl::expected PrepareOpen( + const RegionResourceSpec& spec, + const std::vector& live_allocations) override; + tl::expected PrepareAdopt( + const RegionResourceSpec& spec, + std::shared_ptr allocator) override; +}; + +class CxlRegionDriver final : public RegionDriver { + public: + explicit CxlRegionDriver( + std::shared_ptr global_allocator); + ~CxlRegionDriver() override; + + tl::expected PrepareOpen( + const RegionResourceSpec& spec, + const std::vector& live_allocations) override; + tl::expected PrepareAdopt( + const RegionResourceSpec& spec, + std::shared_ptr allocator) override; + + private: + std::shared_ptr global_allocator_; +}; + +} // namespace + +RegionResource::RegionResource( + std::unique_ptr placement_target) + : target(std::move(placement_target)) {} + +struct PreparedRegionResource::State { + State(RegionDriver& resource_driver, const UUID& id, + std::unique_ptr staged_resource, + std::vector> buffers) + : driver(resource_driver), imported_buffers(std::move(buffers)) { + RegionResourceMap staged; + auto inserted = staged.emplace(id, std::move(staged_resource)); + resource = staged.extract(inserted.first); + } + + RegionDriver& driver; + RegionResourceMap::node_type resource; + std::vector> imported_buffers; + std::unique_ptr replaced_resource; +}; + +PreparedRegionResource::PreparedRegionResource( + RegionDriver& driver, const UUID& id, + std::unique_ptr resource, + std::vector> imported_buffers) + : state_(std::make_unique(driver, id, std::move(resource), + std::move(imported_buffers))) {} + +PreparedRegionResource::~PreparedRegionResource() = default; +PreparedRegionResource::PreparedRegionResource( + PreparedRegionResource&& other) noexcept = default; +PreparedRegionResource& PreparedRegionResource::operator=( + PreparedRegionResource&& other) noexcept = default; + +RegionResource& PreparedRegionResource::resource() const noexcept { + return *state_->resource.mapped(); +} + +const std::vector>& +PreparedRegionResource::imported_buffers() const noexcept { + static const std::vector> empty; + return state_ ? state_->imported_buffers : empty; +} + +std::vector> +PreparedRegionResource::TakeImportedBuffers() { + return state_ ? std::move(state_->imported_buffers) + : std::vector>{}; +} + +void PreparedRegionResource::Commit() noexcept { + if (!state_ || state_->resource.empty()) { + return; + } + state_->driver.CommitPrepared(*this); +} + +RegionResource* RegionDriver::GetResource(const UUID& id) { + auto it = resources_.find(id); + return it == resources_.end() ? nullptr : it->second.get(); +} + +const RegionResource* RegionDriver::GetResource(const UUID& id) const { + auto it = resources_.find(id); + return it == resources_.end() ? nullptr : it->second.get(); +} + +bool RegionDriver::Deactivate(const UUID& id) { + auto* resource = GetResource(id); + if (!resource || !resource->active) { + return false; + } + resource->active = false; + return true; +} + +bool RegionDriver::Reactivate(const UUID& id) { + auto* resource = GetResource(id); + if (!resource || resource->active) { + return false; + } + resource->active = true; + return true; +} + +bool RegionDriver::Erase(const UUID& id) { return resources_.erase(id) != 0; } + +PreparedRegionResource RegionDriver::Stage( + const UUID& id, std::unique_ptr resource, + std::vector> imported_buffers) { + return PreparedRegionResource(*this, id, std::move(resource), + std::move(imported_buffers)); +} + +void RegionDriver::CommitPrepared(PreparedRegionResource& prepared) noexcept { + auto existing = resources_.extract(prepared.state_->resource.key()); + if (!existing.empty()) { + prepared.state_->replaced_resource = std::move(existing.mapped()); + } + prepared.state_->resource.mapped()->active = true; + resources_.insert(std::move(prepared.state_->resource)); +} + +namespace { + +tl::expected MemoryRegionDriver::PrepareOpen( + const RegionResourceSpec& spec, + const std::vector& live_allocations) { + if (!IsValidSpec(spec)) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + if (live_allocations.empty()) { + auto allocator = + CreateBufferAllocator(*allocator_type(), spec.name, spec.base, + spec.size, spec.transport_endpoint); + if (!allocator) { + return tl::make_unexpected(allocator.error()); + } + return Stage(spec.id, MakeNativeResource(std::move(*allocator))); + } + + if (*allocator_type() == BufferAllocatorType::CACHELIB) { + auto restored = ImportCachelibBufferAllocator( + spec.name, spec.base, spec.size, spec.transport_endpoint, + live_allocations); + if (!restored) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + auto resource = MakeNativeResource(std::move(restored->allocator)); + return Stage(spec.id, std::move(resource), + std::move(restored->buffers)); + } + if (*allocator_type() == BufferAllocatorType::OFFSET) { + auto restored = ImportOffsetBufferAllocator( + spec.name, spec.base, spec.size, spec.transport_endpoint, + live_allocations); + if (!restored) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + auto resource = MakeNativeResource(std::move(restored->allocator)); + return Stage(spec.id, std::move(resource), + std::move(restored->buffers)); + } + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); +} + +tl::expected +MemoryRegionDriver::PrepareAdopt( + const RegionResourceSpec& spec, + std::shared_ptr allocator) { + if (!IsValidSpec(spec) || !allocator || + allocator->getSegmentName() != spec.name || + allocator->getTransportEndpoint() != spec.transport_endpoint || + allocator->base() != spec.base || allocator->capacity() != spec.size) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return Stage(spec.id, MakeNativeResource(std::move(allocator))); +} + +CxlRegionDriver::CxlRegionDriver( + std::shared_ptr global_allocator) + : global_allocator_(std::move(global_allocator)) { + MasterMetricManager::instance().inc_total_mem_capacity( + global_allocator_->getSegmentName(), global_allocator_->capacity()); +} + +CxlRegionDriver::~CxlRegionDriver() { + const std::string name = global_allocator_->getSegmentName(); + auto& metrics = MasterMetricManager::instance(); + metrics.dec_total_mem_capacity(name, global_allocator_->capacity()); + if (metrics.get_segment_total_mem_capacity(name) == 0) { + metrics.remove_segment_metrics(name); + } +} + +tl::expected CxlRegionDriver::PrepareOpen( + const RegionResourceSpec& spec, + const std::vector& live_allocations) { + if (!live_allocations.empty()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + if (spec.id == UUID{0, 0} || spec.name.empty() || spec.size == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return Stage(spec.id, MakeCxlResource(spec, global_allocator_)); +} + +tl::expected CxlRegionDriver::PrepareAdopt( + const RegionResourceSpec& spec, + std::shared_ptr allocator) { + (void)spec; + (void)allocator; + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); +} + +} // namespace + +tl::expected CreateRegionDrivers( + const RegionDriverConfig& config) { + RegionDriverRegistry drivers; + drivers.emplace( + RegionKind::HOST_MEMORY, + std::make_unique(config.memory_allocator)); + if (config.cxl) { + auto allocator = CreateBufferAllocator( + BufferAllocatorType::CACHELIB, config.cxl->path, DEFAULT_CXL_BASE, + config.cxl->size, config.cxl->path); + if (!allocator) { + return tl::make_unexpected(allocator.error()); + } + drivers.emplace(RegionKind::CXL, std::make_unique( + std::move(*allocator))); + } + return drivers; +} + +tl::expected, ErrorCode> BuildRegionLiveAllocations( + const RegionResourceSpec& spec, + std::span descriptors) { + if (spec.id == UUID{0, 0} || spec.name.empty() || spec.base == 0 || + spec.size == 0 || + spec.base > std::numeric_limits::max() - spec.size) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + const uintptr_t end = spec.base + spec.size; + std::vector allocations; + allocations.reserve(descriptors.size()); + for (const auto& descriptor : descriptors) { + if (descriptor.transport_endpoint_ != spec.transport_endpoint || + descriptor.size_ == 0 || descriptor.buffer_address_ < spec.base || + descriptor.buffer_address_ >= end || + descriptor.size_ > end - descriptor.buffer_address_) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + allocations.push_back( + {descriptor.buffer_address_ - spec.base, descriptor.size_}); + } + return allocations; +} + +} // namespace mooncake diff --git a/mooncake-store/src/storage/distributed/dfs_global_allocator.cpp b/mooncake-store/src/storage/distributed/dfs_global_allocator.cpp index 3554d704b3..99ae56427f 100644 --- a/mooncake-store/src/storage/distributed/dfs_global_allocator.cpp +++ b/mooncake-store/src/storage/distributed/dfs_global_allocator.cpp @@ -7,7 +7,7 @@ #include #include -#include "storage/distributed/distributed_storage_backend.h" +#include "config/distributed_storage_config.h" #include "storage/distributed/fs_adapter.h" #include "storage/distributed/posix_fs_adapter.h" #include "utils.h" diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp index d49f3adcd2..46b93d55ed 100644 --- a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp +++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp @@ -3,9 +3,7 @@ #include #include #include -#include -#include "environ.h" #include "storage/distributed/dfs_global_allocator.h" #include "types.h" #include "utils.h" @@ -35,126 +33,6 @@ bool IsDfsDescriptorRangeValid(const DistributedFSDescriptor& desc, } // namespace -bool DistributedStorageConfig::Validate() const { - if (fsdir.empty()) { - LOG(ERROR) << "DistributedStorageConfig: fsdir is empty"; - return false; - } - if (!std::filesystem::path(fsdir).is_absolute()) { - LOG(ERROR) - << "DistributedStorageConfig: fsdir must be an absolute path: " - << fsdir; - return false; - } - if (fs_adapter_type != "hf3fs" && fs_adapter_type != "posix") { - LOG(ERROR) << "DistributedStorageConfig: unsupported fs_adapter_type: " - << fs_adapter_type; - return false; - } - if (shard_count <= 0) { - LOG(ERROR) << "DistributedStorageConfig: shard_count must > 0"; - return false; - } - if (shard_capacity == 0) { - LOG(ERROR) << "DistributedStorageConfig: shard_capacity must > 0"; - return false; - } - if (alignment == 0 || (alignment & (alignment - 1)) != 0) { - LOG(ERROR) << "DistributedStorageConfig: alignment must be power of 2"; - return false; - } - if (shard_capacity % alignment != 0) { - LOG(ERROR) << "DistributedStorageConfig: shard_capacity must align"; - return false; - } - if (!single_tenant) { - LOG(ERROR) << "DistributedStorageConfig: Currently, DFS requires " - "single_tenant=true"; - return false; - } - return true; -} - -bool DistributedStorageConfig::ValidateForAllocator() const { - if (!Validate()) return false; - - if (eviction_low_watermark < 0.0 || eviction_low_watermark > 1.0 || - eviction_high_watermark < 0.0 || eviction_high_watermark > 1.0 || - eviction_low_watermark >= eviction_high_watermark) { - LOG(ERROR) << "DistributedStorageConfig: eviction watermarks must " - "satisfy 0 <= low < high <= 1, low=" - << eviction_low_watermark - << ", high=" << eviction_high_watermark; - return false; - } - if (deferred_free_duration.count() < 0) { - LOG(ERROR) << "DistributedStorageConfig: deferred_free_duration must " - "be non-negative, seconds=" - << deferred_free_duration.count(); - return false; - } - if (eviction_enabled && eviction_check_interval.count() <= 0) { - LOG(ERROR) << "DistributedStorageConfig: eviction_check_interval must " - "be positive when eviction is enabled, seconds=" - << eviction_check_interval.count(); - return false; - } - return true; -} - -DistributedStorageConfig DistributedStorageConfig::FromEnvironment() { - DistributedStorageConfig config; - config.fsdir = Environ::GetString( - "MOONCAKE_DFS_ROOT_DIR", - Environ::GetString("MOONCAKE_DISTRIBUTED_ROOT_DIR", config.fsdir)); - if (!std::filesystem::path(config.fsdir).is_absolute()) { - config.fsdir = std::filesystem::absolute(config.fsdir).string(); - } - config.fs_adapter_type = - Environ::GetString("MOONCAKE_DFS_FS_ADAPTER", - Environ::GetString("MOONCAKE_DISTRIBUTED_FS_TYPE", - config.fs_adapter_type)); - config.enable_health_check = - Environ::GetBool("MOONCAKE_DISTRIBUTED_HEALTH_CHECK", false); - config.shard_count = - Environ::GetInt("MOONCAKE_DFS_SHARD_COUNT", config.shard_count); - config.shard_capacity = Environ::GetUInt64("MOONCAKE_DFS_SHARD_CAPACITY", - config.shard_capacity); - config.alignment = - Environ::GetUInt64("MOONCAKE_DFS_ALIGNMENT", config.alignment); - config.single_tenant = - Environ::GetBool("MOONCAKE_DFS_SINGLE_TENANT", config.single_tenant); - config.eviction_enabled = Environ::GetBool("MOONCAKE_DFS_EVICTION_ENABLED", - config.eviction_enabled); - config.eviction_high_watermark = Environ::GetDouble( - "MOONCAKE_DFS_EVICTION_HIGH_WATERMARK", config.eviction_high_watermark); - config.eviction_low_watermark = Environ::GetDouble( - "MOONCAKE_DFS_EVICTION_LOW_WATERMARK", config.eviction_low_watermark); - config.deferred_free_duration = std::chrono::seconds(Environ::GetInt( - "MOONCAKE_DFS_DEFERRED_FREE_SECONDS", - static_cast(config.deferred_free_duration.count()))); - config.eviction_check_interval = std::chrono::seconds(Environ::GetInt( - "MOONCAKE_DFS_EVICTION_CHECK_INTERVAL", - static_cast(config.eviction_check_interval.count()))); - return config; -} - -std::string DistributedStorageConfig::FormatStr() const { - std::ostringstream oss; - oss << "fsdir=" << fsdir << ", fs_adapter_type=" << fs_adapter_type - << ", enable_health_check=" << enable_health_check - << ", shard_count=" << shard_count - << ", shard_capacity=" << shard_capacity << ", alignment=" << alignment - << ", single_tenant=" << single_tenant - << ", eviction_enabled=" << eviction_enabled - << ", eviction_high_watermark=" << eviction_high_watermark - << ", eviction_low_watermark=" << eviction_low_watermark - << ", deferred_free_seconds=" << deferred_free_duration.count() - << ", eviction_check_interval_seconds=" - << eviction_check_interval.count(); - return oss.str(); -} - DistributedStorageBackend::DistributedStorageBackend( const FileStorageConfig& file_storage_config, const DistributedStorageConfig& distributed_config, diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 61d69908ee..2670c62689 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -684,14 +684,9 @@ tl::expected StorageBackend::LoadObject( void StorageBackend::RemoveFile(const std::string& path) { namespace fs = std::filesystem; - // TODO: attention: this function is not thread-safe, need to add lock if - // used in multi-thread environment Check if the file exists before - // attempting to remove it - // TODO: add a sleep to ensure the write thread has time to create the - // corresponding file it will be fixed in the next version - std::this_thread::sleep_for( - std::chrono::microseconds(50)); // sleep for 50 us - + // StoreObject holds the same striped path lock across file creation, write, + // and queue insertion. Acquiring it here serializes deletion with those + // operations without relying on a timing delay. MutexLocker path_locker(&GetFilePathMutex(path)); // Eviction disabled, use simple delete (no queue tracking) diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index dfe174e8cc..12e00b56ea 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -425,6 +425,8 @@ void *allocate_buffer_numa_segments(size_t total_size, unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS; if (page_size == SZ_2MB) { flags |= MAP_HUGETLB | MAP_HUGE_2MB; + } else if (page_size == SZ_512MB) { + flags |= MAP_HUGETLB | MAP_HUGE_512MB; } else if (page_size == SZ_1GB) { flags |= MAP_HUGETLB | MAP_HUGE_1GB; } else if (page_size != static_cast(getpagesize())) { @@ -598,22 +600,30 @@ int64_t time_gen() { .count(); } -std::string ResolveMooncakeHostId(const std::string &local_hostname) { - const std::string hostname(TrimAsciiWhitespace(local_hostname)); +static bool IsUsableMooncakeHostId(std::string_view host_id) { + return !host_id.empty() && + !AsciiCaseInsensitiveEquals(host_id, "localhost") && + host_id != "127.0.0.1" && host_id != "0.0.0.0" && host_id != "::1" && + host_id != "[::1]" && host_id != "::" && host_id != "[::]"; +} + +static std::string NormalizeMooncakeHostId(std::string_view value) { + const std::string hostname(TrimAsciiWhitespace(value)); const std::string host_id = (hostname == "::1" || hostname == "::") ? hostname : std::string(TrimAsciiWhitespace( getHostNameWithoutPort(hostname))); - if (host_id.empty()) { - return ""; - } + return IsUsableMooncakeHostId(host_id) ? host_id : ""; +} - if (AsciiCaseInsensitiveEquals(host_id, "localhost") || - host_id == "127.0.0.1" || host_id == "0.0.0.0" || host_id == "::1" || - host_id == "[::1]" || host_id == "::" || host_id == "[::]") { - return ""; +std::string ResolveMooncakeHostId(const std::string &local_hostname) { + const std::string configured_host_id( + TrimAsciiWhitespace(Environ::GetString("MOONCAKE_HOST_ID", ""))); + if (!configured_host_id.empty()) { + return NormalizeMooncakeHostId(configured_host_id); } - return host_id; + + return NormalizeMooncakeHostId(local_hostname); } static std::string SanitizeKey(const std::string &key) { diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 66372aedf2..3a9e9fc92d 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -36,6 +36,7 @@ function(add_ha_test name) endfunction() add_store_test(buffer_allocator_test buffer_allocator_test.cpp) +add_store_test(region_driver_test region_driver_test.cpp) add_store_test(runtime_accelerator_test runtime_accelerator_test.cpp) add_store_test(registered_pinned_memory_test registered_pinned_memory_test.cpp) add_store_test(allocation_strategy_test allocation_strategy_test.cpp) @@ -143,6 +144,8 @@ add_store_test(local_ssd_codec_test ha/snapshot/local_ssd_codec_test.cpp) add_store_test(offset_allocator_test offset_allocator_test.cpp) add_store_test(utils_test utils_test.cpp) add_store_test(client_buffer_test client_buffer_test.cpp) +add_store_test(client_auto_port_config_test client_auto_port_config_test.cpp) +add_store_test(local_hot_cache_config_test local_hot_cache_config_test.cpp) add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp) add_store_test(client_tcp_local_memcpy_test client_tcp_local_memcpy_test.cpp) add_store_test(pybind_client_test pybind_client_test.cpp) @@ -153,9 +156,12 @@ endif() add_store_test(ipv6_client_test ipv6_client_test.cpp) add_store_test(host_port_fix_test host_port_fix_test.cpp) add_store_test(client_metrics_test client_metrics_test.cpp) +add_store_test(client_metric_config_test client_metric_config_test.cpp) add_store_test(ssd_metrics_test ssd_metrics_test.cpp) add_store_test(serializer_test serializer_test.cpp) add_store_test(dfs_posix_test dfs_posix_test.cpp) +add_store_test(distributed_storage_config_test + distributed_storage_config_test.cpp) add_store_test(dfs_sync_client_test dfs_sync_client_test.cpp) add_test(NAME dfs_batch_read_checksum_test COMMAND dfs_sync_client_test @@ -188,6 +194,8 @@ add_ha_test(batch_oplog_snapshot_provider_test ha/snapshot/batch_oplog/provider_test.cpp) add_ha_test(batch_oplog_snapshot_publisher_test ha/snapshot/batch_oplog/publisher_test.cpp) +add_ha_test(batch_oplog_snapshot_coordinator_test + ha/snapshot/batch_oplog/coordinator_test.cpp) add_store_test(master_service_test_for_snapshot ha/snapshot/master_service_test_for_snapshot.cpp) add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp) diff --git a/mooncake-store/tests/allocation_strategy_test.cpp b/mooncake-store/tests/allocation_strategy_test.cpp index 2fdef43dc6..98dc50cc84 100644 --- a/mooncake-store/tests/allocation_strategy_test.cpp +++ b/mooncake-store/tests/allocation_strategy_test.cpp @@ -59,16 +59,12 @@ class AllocationStrategyParameterizedTest const std::string& segment_name, size_t base_offset, size_t size = 64 * MiB) { const size_t base = 0x100000000ULL + base_offset; // 4GB + offset - switch (allocator_type_) { - case BufferAllocatorType::CACHELIB: - return std::make_shared( - segment_name, base, size, segment_name); - case BufferAllocatorType::OFFSET: - return std::make_shared( - segment_name, base, size, segment_name); - default: - throw std::invalid_argument("Invalid allocator type"); + auto allocator = CreateBufferAllocator(allocator_type_, segment_name, + base, size, segment_name); + if (!allocator) { + throw std::invalid_argument("Invalid allocator test parameters"); } + return std::move(*allocator); } BufferAllocatorType allocator_type_; diff --git a/mooncake-store/tests/buffer_allocator_test.cpp b/mooncake-store/tests/buffer_allocator_test.cpp index 3d8122bc01..3449f9d470 100644 --- a/mooncake-store/tests/buffer_allocator_test.cpp +++ b/mooncake-store/tests/buffer_allocator_test.cpp @@ -15,6 +15,15 @@ namespace mooncake { +namespace { + +LiveAllocation ToLiveAllocation(uintptr_t base, + const AllocatedBuffer::Descriptor& descriptor) { + return {descriptor.buffer_address_ - base, descriptor.size_}; +} + +} // namespace + // Test fixture for BufferAllocator tests class BufferAllocatorTest : public ::testing::Test { protected: @@ -34,16 +43,12 @@ class BufferAllocatorTest : public ::testing::Test { const std::string& segment_name, size_t base_offset, size_t size, BufferAllocatorType allocator_type) { const size_t base = 0x100000000ULL + base_offset; // 4GB + offset - switch (allocator_type) { - case BufferAllocatorType::CACHELIB: - return std::make_shared( - segment_name, base, size, segment_name); - case BufferAllocatorType::OFFSET: - return std::make_shared( - segment_name, base, size, segment_name); - default: - throw std::invalid_argument("Invalid allocator type"); + auto allocator = CreateBufferAllocator(allocator_type, segment_name, + base, size, segment_name); + if (!allocator) { + throw std::invalid_argument("Invalid allocator test parameters"); } + return std::move(*allocator); } void VerifyAllocatedBuffer(const AllocatedBuffer& bufHandle, @@ -134,7 +139,7 @@ TEST_F(BufferAllocatorTest, OffsetLargestFreeRegionRemainsExact) { EXPECT_EQ(allocator->getLargestFreeRegion(), CAPACITY); } -TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsAtOriginalAddresses) { +TEST_F(BufferAllocatorTest, ImportOffsetAllocationsAtOriginalAddresses) { constexpr uintptr_t kBase = 0x180000000ULL; constexpr size_t kCapacity = 16 * 1024 * 1024; const std::string segment = "restore-segment"; @@ -149,13 +154,16 @@ TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsAtOriginalAddresses) { ASSERT_NE(removed, nullptr); ASSERT_NE(last, nullptr); - std::vector descriptors = { + const std::vector descriptors = { first->get_descriptor(), last->get_descriptor()}; + std::vector allocations = { + ToLiveAllocation(kBase, descriptors[0]), + ToLiveAllocation(kBase, descriptors[1])}; const auto removed_descriptor = removed->get_descriptor(); removed.reset(); - auto restored = RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, descriptors); + auto restored = ImportOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, allocations); ASSERT_TRUE(restored.has_value()); ASSERT_EQ(restored->buffers.size(), descriptors.size()); EXPECT_EQ(restored->allocator->size(), @@ -169,111 +177,119 @@ TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsAtOriginalAddresses) { ASSERT_NE(new_buffer, nullptr); EXPECT_EQ(reinterpret_cast(new_buffer->data()), removed_descriptor.buffer_address_); - - auto wrong_endpoint = descriptors; - wrong_endpoint[0].transport_endpoint_ = "other-endpoint"; - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, wrong_endpoint) - .has_value()); - - auto duplicate = descriptors; - duplicate.push_back(descriptors.front()); - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, duplicate) - .has_value()); - - auto out_of_range = descriptors; - out_of_range[0].buffer_address_ = kBase + kCapacity; - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, out_of_range) - .has_value()); } -TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsValidatesRangesAndOrder) { +TEST_F(BufferAllocatorTest, ImportOffsetAllocationsValidatesRangesAndOrder) { constexpr uintptr_t kBase = 0x190000000ULL; constexpr size_t kCapacity = 4096; const std::string segment = "restore-validation"; const std::string endpoint = "restore-validation-endpoint"; - auto descriptor = [&](uintptr_t address, uint64_t size) { - return AllocatedBuffer::Descriptor{size, address, "tcp", endpoint}; + auto allocation = [&](uintptr_t address, uint64_t size) { + return LiveAllocation{address - kBase, size}; }; - std::vector unsorted = { - descriptor(kBase + 512, 64), descriptor(kBase + 128, 64)}; - auto restored = RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, unsorted); + std::vector unsorted = {allocation(kBase + 512, 64), + allocation(kBase + 128, 64)}; + auto restored = ImportOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, unsorted); ASSERT_TRUE(restored.has_value()); ASSERT_EQ(restored->buffers.size(), unsorted.size()); EXPECT_EQ(reinterpret_cast(restored->buffers[0]->data()), - unsorted[0].buffer_address_); + kBase + unsorted[0].offset_from_base); EXPECT_EQ(reinterpret_cast(restored->buffers[1]->data()), - unsorted[1].buffer_address_); + kBase + unsorted[1].offset_from_base); - std::vector overlapping = { - descriptor(kBase + 128, 100), descriptor(kBase + 200, 32)}; - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, overlapping) + std::vector overlapping = {allocation(kBase + 128, 100), + allocation(kBase + 200, 32)}; + EXPECT_FALSE(ImportOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, overlapping) .has_value()); - std::vector normalized_past_end = { - descriptor(kBase + kCapacity - 100, 100)}; - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, normalized_past_end) + std::vector normalized_past_end = { + allocation(kBase + kCapacity - 100, 100)}; + EXPECT_FALSE(ImportOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, normalized_past_end) .has_value()); - EXPECT_FALSE(RestoreOffsetBufferAllocator( + EXPECT_FALSE(ImportOffsetBufferAllocator( segment, std::numeric_limits::max() - 100, 200, endpoint, {}) .has_value()); - std::vector descriptor_overflow = { - descriptor(std::numeric_limits::max() - 10, 20)}; - EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, - endpoint, descriptor_overflow) + std::vector allocation_overflow = { + {std::numeric_limits::max() - kBase - 10, 20}}; + EXPECT_FALSE(ImportOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, allocation_overflow) .has_value()); } -TEST_F(BufferAllocatorTest, RestoredOffsetHandleReleasesItsExactAddress) { +TEST_F(BufferAllocatorTest, ImportedOffsetHandleReleasesItsExactAddress) { constexpr uintptr_t kBase = 0x1A0000000ULL; constexpr size_t kCapacity = 4096; const std::string endpoint = "restore-release"; - std::vector descriptors = { - {64, kBase, "tcp", endpoint}, {64, kBase + 512, "tcp", endpoint}}; - auto restored = RestoreOffsetBufferAllocator( - "restore-release", kBase, kCapacity, endpoint, descriptors); + std::vector allocations = {{0, 64}, {512, 64}}; + auto restored = ImportOffsetBufferAllocator( + "restore-release", kBase, kCapacity, endpoint, allocations); ASSERT_TRUE(restored.has_value()); restored->buffers[0].reset(); auto replacement = restored->allocator->allocate(64); ASSERT_NE(replacement, nullptr); EXPECT_EQ(reinterpret_cast(replacement->data()), - descriptors[0].buffer_address_); + kBase + allocations[0].offset_from_base); } -TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsHasNoArbitraryGapLimit) { +TEST_F(BufferAllocatorTest, ImportOffsetAllocationsHasNoArbitraryGapLimit) { constexpr uintptr_t kBase = 0x1B0000000ULL; constexpr size_t kGapCount = 65537; const std::string endpoint = "restore-many-gaps"; - std::vector descriptors; - descriptors.reserve(kGapCount); + std::vector allocations; + allocations.reserve(kGapCount); for (size_t i = 0; i < kGapCount; ++i) { - descriptors.push_back({1, kBase + 1 + i * 2, "tcp", endpoint}); + allocations.push_back({1 + i * 2, 1}); } - auto restored = RestoreOffsetBufferAllocator( - "restore-many-gaps", kBase, kGapCount * 2 + 1, endpoint, descriptors); + auto restored = ImportOffsetBufferAllocator( + "restore-many-gaps", kBase, kGapCount * 2 + 1, endpoint, allocations); ASSERT_TRUE(restored.has_value()); - EXPECT_EQ(restored->buffers.size(), descriptors.size()); + EXPECT_EQ(restored->buffers.size(), allocations.size()); EXPECT_EQ(reinterpret_cast(restored->buffers.back()->data()), - descriptors.back().buffer_address_); + kBase + allocations.back().offset_from_base); +} + +TEST_F(BufferAllocatorTest, CachelibCreateRejectsInvalidMemoryLayout) { + constexpr size_t kSlabSize = facebook::cachelib::Slab::kSize; + constexpr uintptr_t kBase = 0x1C0000000ULL; + + auto expect_invalid = [](size_t base, size_t size) { + auto result = CachelibBufferAllocator::Create("cachelib-invalid", base, + size, "endpoint"); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + }; + + expect_invalid(kBase + 1, kSlabSize); + expect_invalid(kBase, kSlabSize + 1); + expect_invalid(std::numeric_limits::max() - kSlabSize, + 2 * kSlabSize); + if constexpr (std::numeric_limits::max() / kSlabSize > + std::numeric_limits::max()) { + const size_t too_many_slabs = + (static_cast(std::numeric_limits::max()) + + 1) * + kSlabSize; + expect_invalid(kBase, too_many_slabs); + } } -TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsAtOriginalAddresses) { +TEST_F(BufferAllocatorTest, ImportCachelibAllocationsAtOriginalAddresses) { constexpr uintptr_t kBase = 0x1C0000000ULL; constexpr size_t kCapacity = 4 * facebook::cachelib::Slab::kSize; const std::string segment = "cachelib-restore"; const std::string endpoint = "cachelib-restore-endpoint"; - auto original = std::make_shared( - segment, kBase, kCapacity, endpoint); + auto created = + CachelibBufferAllocator::Create(segment, kBase, kCapacity, endpoint); + ASSERT_TRUE(created.has_value()); + auto original = std::move(*created); auto small_first = original->allocate(64); auto small_hole = original->allocate(64); @@ -291,11 +307,16 @@ TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsAtOriginalAddresses) { std::vector descriptors = { large_last->get_descriptor(), small_first->get_descriptor(), large_first->get_descriptor(), small_last->get_descriptor()}; + std::vector allocations; + allocations.reserve(descriptors.size()); + for (const auto& descriptor : descriptors) { + allocations.push_back(ToLiveAllocation(kBase, descriptor)); + } small_hole.reset(); large_hole.reset(); - auto restored = RestoreCachelibBufferAllocator(segment, kBase, kCapacity, - endpoint, descriptors); + auto restored = ImportCachelibBufferAllocator(segment, kBase, kCapacity, + endpoint, allocations); ASSERT_TRUE(restored.has_value()); ASSERT_EQ(restored->buffers.size(), descriptors.size()); for (size_t i = 0; i < descriptors.size(); ++i) { @@ -317,42 +338,24 @@ TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsAtOriginalAddresses) { EXPECT_EQ(reinterpret_cast(replacement->data()), released); } -TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsRejectsInvalidLayouts) { +TEST_F(BufferAllocatorTest, ImportCachelibAllocationsRejectsInvalidLayouts) { constexpr uintptr_t kBase = 0x1D0000000ULL; constexpr size_t kCapacity = 4 * facebook::cachelib::Slab::kSize; - constexpr size_t kSlabSize = facebook::cachelib::Slab::kSize; const std::string endpoint = "cachelib-invalid-endpoint"; - auto descriptor = [&](uintptr_t address, uint64_t size) { - return AllocatedBuffer::Descriptor{size, address, "tcp", endpoint}; + auto allocation = [&](uintptr_t address, uint64_t size) { + return LiveAllocation{address - kBase, size}; }; - auto restore = [&](const std::vector& descs) { - return RestoreCachelibBufferAllocator("cachelib-invalid", kBase, - kCapacity, endpoint, descs); + auto import = [&](const std::vector& allocations) { + return ImportCachelibBufferAllocator("cachelib-invalid", kBase, + kCapacity, endpoint, allocations); }; EXPECT_FALSE( - restore({descriptor(kBase, 64), descriptor(kBase, 4096)}).has_value()); - EXPECT_FALSE(restore({descriptor(kBase + 1, 64)}).has_value()); + import({allocation(kBase, 64), allocation(kBase, 4096)}).has_value()); + EXPECT_FALSE(import({allocation(kBase + 1, 64)}).has_value()); EXPECT_FALSE( - restore({descriptor(kBase, 64), descriptor(kBase, 64)}).has_value()); - - auto wrong_endpoint = descriptor(kBase, 64); - wrong_endpoint.transport_endpoint_ = "wrong"; - EXPECT_FALSE(restore({wrong_endpoint}).has_value()); - EXPECT_FALSE(restore({descriptor(kBase + kCapacity, 64)}).has_value()); - EXPECT_FALSE(RestoreCachelibBufferAllocator("cachelib-invalid", kBase + 1, - kCapacity, endpoint, {}) - .has_value()); - EXPECT_FALSE(RestoreCachelibBufferAllocator( - "cachelib-invalid", - std::numeric_limits::max() - kSlabSize, - 2 * kSlabSize, endpoint, {}) - .has_value()); - - auto valid_after_fail = restore({descriptor(kBase + kSlabSize, 4096)}); - ASSERT_TRUE(valid_after_fail.has_value()); - EXPECT_EQ(reinterpret_cast(valid_after_fail->buffers[0]->data()), - kBase + kSlabSize); + import({allocation(kBase, 64), allocation(kBase, 64)}).has_value()); + EXPECT_FALSE(import({allocation(kBase + kCapacity, 64)}).has_value()); } TEST_F(BufferAllocatorTest, CachelibImportRejectsChunkInSlabTail) { @@ -370,32 +373,16 @@ TEST_F(BufferAllocatorTest, CachelibImportRejectsChunkInSlabTail) { pool, {{reinterpret_cast(kBase + kAllocSize), kAllocSize}})); } -TEST_F(BufferAllocatorTest, RestoreCachelibRejectsNonMemoryDescriptors) { +TEST_F(BufferAllocatorTest, CachelibImportRejectsNonMemoryReplicaType) { constexpr uintptr_t kBase = 0x1F0000000ULL; constexpr size_t kCapacity = 2 * facebook::cachelib::Slab::kSize; const std::string endpoint = "cachelib-memory-only"; - std::vector descriptors = { - {64, kBase, "tcp", endpoint}}; - - EXPECT_FALSE(RestoreCachelibBufferAllocator( - "cachelib-memory-only", kBase, kCapacity, endpoint, - descriptors, ReplicaType::NOF_SSD) - .has_value()); + std::vector allocations = {{0, 64}}; - descriptors[0].protocol_ = "cxl"; - EXPECT_FALSE(RestoreCachelibBufferAllocator("cachelib-memory-only", kBase, - kCapacity, endpoint, - descriptors) + EXPECT_FALSE(ImportCachelibBufferAllocator("cachelib-memory-only", kBase, + kCapacity, endpoint, allocations, + ReplicaType::NOF_SSD) .has_value()); - - descriptors[0].protocol_ = "rdma"; - auto rdma = RestoreCachelibBufferAllocator( - "cachelib-memory-only", kBase, kCapacity, endpoint, descriptors); - ASSERT_TRUE(rdma.has_value()); - const auto restored = rdma->buffers[0]->get_descriptor(); - EXPECT_EQ(restored.protocol_, descriptors[0].protocol_); - EXPECT_EQ(restored.buffer_address_, descriptors[0].buffer_address_); - EXPECT_EQ(restored.transport_endpoint_, descriptors[0].transport_endpoint_); } // Test allocation request larger than available space diff --git a/mooncake-store/tests/client_auto_port_config_test.cpp b/mooncake-store/tests/client_auto_port_config_test.cpp new file mode 100644 index 0000000000..da37ece192 --- /dev/null +++ b/mooncake-store/tests/client_auto_port_config_test.cpp @@ -0,0 +1,143 @@ +#include + +#include +#include +#include + +#include "client_auto_port_config.h" + +namespace mooncake { +namespace { + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + if (const char* value = std::getenv(name)) { + original_ = value; + } + unsetenv(name); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +struct ClientAutoPortEnvironment { + ScopedEnvVar setup_retries{"MC_STORE_CLIENT_SETUP_RETRIES"}; + ScopedEnvVar min_port{"MC_STORE_CLIENT_MIN_PORT"}; + ScopedEnvVar max_port{"MC_STORE_CLIENT_MAX_PORT"}; +}; + +class ClientAutoPortConfigTest : public ::testing::Test { + protected: + ClientAutoPortEnvironment env; +}; + +TEST_F(ClientAutoPortConfigTest, UsesExistingDefaultsWhenEnvironmentIsUnset) { + const auto config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.max_retries, 20); + EXPECT_EQ(config.min_port, 12300); + EXPECT_EQ(config.max_port, 14300); +} + +TEST_F(ClientAutoPortConfigTest, ReadsValidValues) { + env.setup_retries.Set("7"); + env.min_port.Set("12000"); + env.max_port.Set("14000"); + + const auto config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.max_retries, 7); + EXPECT_EQ(config.min_port, 12000); + EXPECT_EQ(config.max_port, 14000); +} + +TEST_F(ClientAutoPortConfigTest, InvalidIntegersUseIndividualFieldDefaults) { + for (const char* value : {"", "invalid", "2147483648"}) { + env.setup_retries.Set(value); + env.min_port.Set(value); + env.max_port.Set("15000"); + + auto config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.max_retries, 20) << value; + EXPECT_EQ(config.min_port, 12300) << value; + EXPECT_EQ(config.max_port, 15000) << value; + + env.min_port.Set("13000"); + env.max_port.Set(value); + + config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.min_port, 13000) << value; + EXPECT_EQ(config.max_port, 14300) << value; + } +} + +TEST_F(ClientAutoPortConfigTest, SupportsIndependentEndpointOverrides) { + env.min_port.Set("13000"); + + auto config = ClientAutoPortConfig::FromEnvironment(); + EXPECT_EQ(config.min_port, 13000); + EXPECT_EQ(config.max_port, 14300); + + env.min_port.Set("12300"); + env.max_port.Set("15000"); + + config = ClientAutoPortConfig::FromEnvironment(); + EXPECT_EQ(config.min_port, 12300); + EXPECT_EQ(config.max_port, 15000); +} + +TEST_F(ClientAutoPortConfigTest, InvalidPortPairsRestoreBothDefaults) { + struct PortPair { + const char* min_port; + const char* max_port; + }; + for (const auto& value : + {PortPair{"14301", "14300"}, PortPair{"80", "443"}, + PortPair{"32768", "40000"}, PortPair{"61000", "65536"}}) { + env.min_port.Set(value.min_port); + env.max_port.Set(value.max_port); + + const auto config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.min_port, 12300); + EXPECT_EQ(config.max_port, 14300); + } +} + +TEST_F(ClientAutoPortConfigTest, PreservesNonPositiveRetryCounts) { + env.setup_retries.Set("0"); + EXPECT_EQ(ClientAutoPortConfig::FromEnvironment().max_retries, 0); + + env.setup_retries.Set("-1"); + EXPECT_EQ(ClientAutoPortConfig::FromEnvironment().max_retries, -1); +} + +TEST_F(ClientAutoPortConfigTest, PreservesAcceptedIntegerSyntax) { + env.setup_retries.Set(" +7 "); + env.min_port.Set(" +12000 "); + env.max_port.Set(" +14000 "); + + const auto config = ClientAutoPortConfig::FromEnvironment(); + + EXPECT_EQ(config.max_retries, 7); + EXPECT_EQ(config.min_port, 12000); + EXPECT_EQ(config.max_port, 14000); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/client_metric_config_test.cpp b/mooncake-store/tests/client_metric_config_test.cpp new file mode 100644 index 0000000000..c6d426a166 --- /dev/null +++ b/mooncake-store/tests/client_metric_config_test.cpp @@ -0,0 +1,162 @@ +#include +#include + +#include +#include +#include +#include + +#include "client_metric.h" +#include "environment_variables.h" + +namespace mooncake { +namespace { + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + const char* value = std::getenv(name); + if (value != nullptr) { + original_ = value; + } + unsetenv(name); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + ScopedEnvVar(const ScopedEnvVar&) = delete; + ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +struct ClientMetricEnvironment { + using Variables = ClientMetricEnvironmentVariables; + + ScopedEnvVar enabled{Variables::MC_STORE_CLIENT_METRIC.name}; + ScopedEnvVar interval{Variables::MC_STORE_CLIENT_METRIC_INTERVAL.name}; + ScopedEnvVar bandwidth{Variables::MC_STORE_CLIENT_METRIC_BANDWIDTH.name}; +}; + +class ClientMetricConfigTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("ClientMetricConfigTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } + + ClientMetricEnvironment env; +}; + +TEST_F(ClientMetricConfigTest, UsesDefaultsWhenEnvironmentIsUnset) { + const auto config = ClientMetricConfig::FromEnvironment(); + + EXPECT_TRUE(config.enabled); + EXPECT_EQ(config.reporting_interval, std::chrono::milliseconds::zero()); + EXPECT_TRUE(config.bandwidth_reporting_enabled); +} + +TEST_F(ClientMetricConfigTest, ReadsValidEnvironmentValues) { + env.enabled.Set("true"); + env.interval.Set(" +15 "); + env.bandwidth.Set("false"); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_TRUE(config.enabled); + EXPECT_EQ(config.reporting_interval, std::chrono::seconds(15)); + EXPECT_FALSE(config.bandwidth_reporting_enabled); + EXPECT_NE(logs.find("Client metrics interval set to 15s via " + "MC_STORE_CLIENT_METRIC_INTERVAL"), + std::string::npos); +} + +TEST_F(ClientMetricConfigTest, InvalidEnableValueSilentlyDisablesMetrics) { + env.enabled.Set("invalid"); + env.interval.Set("invalid"); + env.bandwidth.Set("invalid"); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_FALSE(config.enabled); + EXPECT_EQ(config.reporting_interval, std::chrono::milliseconds::zero()); + EXPECT_TRUE(config.bandwidth_reporting_enabled); + EXPECT_EQ(logs.find("MC_STORE_CLIENT_METRIC_INTERVAL"), std::string::npos); + EXPECT_EQ(logs.find("MC_STORE_CLIENT_METRIC_BANDWIDTH"), std::string::npos); +} + +TEST_F(ClientMetricConfigTest, EmptyEnableValueSilentlyDisablesMetrics) { + env.enabled.Set(""); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_FALSE(config.enabled); + EXPECT_TRUE(logs.empty()); +} + +TEST_F(ClientMetricConfigTest, InvalidIntervalValuesUseDefaultAndWarn) { + for (const char* value : {"invalid", "-1", "18446744073709551616", ""}) { + SCOPED_TRACE(value); + env.interval.Set(value); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_EQ(config.reporting_interval, std::chrono::milliseconds::zero()); + EXPECT_NE(logs.find("Failed to parse " + "MC_STORE_CLIENT_METRIC_INTERVAL"), + std::string::npos); + } +} + +TEST_F(ClientMetricConfigTest, InvalidBandwidthValuesUseDefaultAndWarn) { + for (const char* value : {"invalid", ""}) { + SCOPED_TRACE(value); + env.bandwidth.Set(value); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_TRUE(config.bandwidth_reporting_enabled); + EXPECT_NE(logs.find("Failed to parse " + "MC_STORE_CLIENT_METRIC_BANDWIDTH"), + std::string::npos); + } +} + +TEST_F(ClientMetricConfigTest, ZeroIntervalRemainsEnabled) { + env.interval.Set("0"); + + ::testing::internal::CaptureStderr(); + const auto config = ClientMetricConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_TRUE(config.enabled); + EXPECT_EQ(config.reporting_interval, std::chrono::milliseconds::zero()); + EXPECT_NE(logs.find("Client metrics reporting disabled (interval=0) via " + "MC_STORE_CLIENT_METRIC_INTERVAL"), + std::string::npos); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/client_metrics_test.cpp b/mooncake-store/tests/client_metrics_test.cpp index 9f777b17f5..5053067404 100644 --- a/mooncake-store/tests/client_metrics_test.cpp +++ b/mooncake-store/tests/client_metrics_test.cpp @@ -11,6 +11,7 @@ #include #include "client_metric.h" +#include "environment_variables.h" #include "real_client.h" #include "test_server_helpers.h" #include "utils.h" @@ -296,15 +297,18 @@ TEST_F(ClientMetricsTest, ZeroSumHybridHistogramPreservesExistingMetrics) { } TEST_F(ClientMetricsTest, BandwidthSummaryRespectsEnvFlag) { - setenv("MC_STORE_CLIENT_METRIC_BANDWIDTH", "0", 1); + ScopedEnv bandwidth_env( + ClientMetricEnvironmentVariables::MC_STORE_CLIENT_METRIC_BANDWIDTH + .name); + setenv( + ClientMetricEnvironmentVariables::MC_STORE_CLIENT_METRIC_BANDWIDTH.name, + "0", 1); auto metrics = ClientMetric::Create(); ASSERT_NE(metrics, nullptr); metrics->transfer_metric.total_read_bytes.inc(1024); std::string summary = metrics->summary_metrics(); EXPECT_TRUE(summary.find("Average Read Throughput:") == std::string::npos); - - unsetenv("MC_STORE_CLIENT_METRIC_BANDWIDTH"); } TEST_F(ClientMetricsTest, SummaryCanOmitMasterRpcMetrics) { diff --git a/mooncake-store/tests/dfs_posix_test.cpp b/mooncake-store/tests/dfs_posix_test.cpp index 1e5ec533d6..e792b81a08 100644 --- a/mooncake-store/tests/dfs_posix_test.cpp +++ b/mooncake-store/tests/dfs_posix_test.cpp @@ -331,47 +331,6 @@ TEST(DfsGlobalAllocatorTest, AllocateFreeAndFormatShardIdx) { EXPECT_EQ(DfsGlobalAllocator::FormatShardIdx(100, 1000), "100"); } -TEST(DistributedStorageConfigTest, ReadsValidatesAndFormatsEnvironment) { - EnvGuard env; - TempDir tmp("dfs_config"); - env.Set("MOONCAKE_DFS_ROOT_DIR", tmp.path().c_str()); - env.Set("MOONCAKE_DFS_FS_ADAPTER", "posix"); - env.Set("MOONCAKE_DFS_SHARD_COUNT", "8"); - env.Set("MOONCAKE_DFS_SHARD_CAPACITY", "1048576"); - env.Set("MOONCAKE_DFS_ALIGNMENT", "4096"); - env.Set("MOONCAKE_DFS_SINGLE_TENANT", "1"); - env.Set("MOONCAKE_DFS_EVICTION_ENABLED", "1"); - env.Set("MOONCAKE_DFS_EVICTION_HIGH_WATERMARK", "0.85"); - env.Set("MOONCAKE_DFS_EVICTION_LOW_WATERMARK", "0.65"); - env.Set("MOONCAKE_DFS_DEFERRED_FREE_SECONDS", "12"); - env.Set("MOONCAKE_DFS_EVICTION_CHECK_INTERVAL", "3"); - - const auto config = DistributedStorageConfig::FromEnvironment(); - EXPECT_EQ(config.fsdir, tmp.path()); - EXPECT_EQ(config.fs_adapter_type, "posix"); - EXPECT_EQ(config.shard_count, 8); - EXPECT_EQ(config.shard_capacity, 1048576); - EXPECT_EQ(config.alignment, 4096); - EXPECT_TRUE(config.eviction_enabled); - EXPECT_DOUBLE_EQ(config.eviction_high_watermark, 0.85); - EXPECT_DOUBLE_EQ(config.eviction_low_watermark, 0.65); - EXPECT_EQ(config.deferred_free_duration, std::chrono::seconds(12)); - EXPECT_EQ(config.eviction_check_interval, std::chrono::seconds(3)); - EXPECT_TRUE(config.Validate()); - EXPECT_TRUE(config.ValidateForAllocator()); - - const std::string formatted = config.FormatStr(); - EXPECT_NE(formatted.find("fs_adapter_type=posix"), std::string::npos); - EXPECT_NE(formatted.find("shard_count=8"), std::string::npos); - EXPECT_NE(formatted.find("eviction_high_watermark=0.85"), - std::string::npos); - - auto invalid_eviction = config; - invalid_eviction.eviction_low_watermark = 0.9; - EXPECT_TRUE(invalid_eviction.Validate()); - EXPECT_FALSE(invalid_eviction.ValidateForAllocator()); -} - TEST(DfsGlobalAllocatorTest, InitReturnsSpecificErrors) { EnvGuard env; ConfigurePosixDfs(env); diff --git a/mooncake-store/tests/distributed_storage_config_test.cpp b/mooncake-store/tests/distributed_storage_config_test.cpp new file mode 100644 index 0000000000..bdd1c53320 --- /dev/null +++ b/mooncake-store/tests/distributed_storage_config_test.cpp @@ -0,0 +1,297 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "config/distributed_storage_config.h" + +namespace mooncake { + +namespace { + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + const char* value = std::getenv(name); + if (value != nullptr) { + original_ = value; + } + unsetenv(name); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + ScopedEnvVar(const ScopedEnvVar&) = delete; + ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +struct DistributedStorageEnvironment { + ScopedEnvVar root_dir{"MOONCAKE_DFS_ROOT_DIR"}; + ScopedEnvVar legacy_root_dir{"MOONCAKE_DISTRIBUTED_ROOT_DIR"}; + ScopedEnvVar fs_adapter{"MOONCAKE_DFS_FS_ADAPTER"}; + ScopedEnvVar legacy_fs_adapter{"MOONCAKE_DISTRIBUTED_FS_TYPE"}; + ScopedEnvVar health_check{"MOONCAKE_DISTRIBUTED_HEALTH_CHECK"}; + ScopedEnvVar shard_count{"MOONCAKE_DFS_SHARD_COUNT"}; + ScopedEnvVar shard_capacity{"MOONCAKE_DFS_SHARD_CAPACITY"}; + ScopedEnvVar alignment{"MOONCAKE_DFS_ALIGNMENT"}; + ScopedEnvVar single_tenant{"MOONCAKE_DFS_SINGLE_TENANT"}; + ScopedEnvVar eviction_enabled{"MOONCAKE_DFS_EVICTION_ENABLED"}; + ScopedEnvVar eviction_high_watermark{ + "MOONCAKE_DFS_EVICTION_HIGH_WATERMARK"}; + ScopedEnvVar eviction_low_watermark{"MOONCAKE_DFS_EVICTION_LOW_WATERMARK"}; + ScopedEnvVar deferred_free_seconds{"MOONCAKE_DFS_DEFERRED_FREE_SECONDS"}; + ScopedEnvVar eviction_check_interval{ + "MOONCAKE_DFS_EVICTION_CHECK_INTERVAL"}; +}; + +void ExpectDefaultConfig(const DistributedStorageConfig& config) { + EXPECT_EQ(config.fsdir, "/mnt/3fs/mooncake"); + EXPECT_EQ(config.fs_adapter_type, "hf3fs"); + EXPECT_FALSE(config.enable_health_check); + EXPECT_EQ(config.shard_count, 64); + EXPECT_EQ(config.shard_capacity, 4ULL * 1024 * 1024 * 1024); + EXPECT_EQ(config.alignment, 4096); + EXPECT_TRUE(config.single_tenant); + EXPECT_TRUE(config.eviction_enabled); + EXPECT_DOUBLE_EQ(config.eviction_high_watermark, 0.9); + EXPECT_DOUBLE_EQ(config.eviction_low_watermark, 0.7); + EXPECT_EQ(config.deferred_free_duration, std::chrono::seconds(30)); + EXPECT_EQ(config.eviction_check_interval, std::chrono::seconds(5)); +} + +DistributedStorageConfig ValidConfig() { + DistributedStorageConfig config; + config.fsdir = "/tmp/mooncake-distributed-storage"; + config.fs_adapter_type = "posix"; + config.shard_count = 8; + config.shard_capacity = 1024 * 1024; + config.alignment = 4096; + config.single_tenant = true; + config.eviction_enabled = true; + config.eviction_high_watermark = 0.85; + config.eviction_low_watermark = 0.65; + config.deferred_free_duration = std::chrono::seconds(12); + config.eviction_check_interval = std::chrono::seconds(3); + return config; +} + +class DistributedStorageConfigTest : public ::testing::Test { + protected: + DistributedStorageEnvironment env; +}; + +TEST_F(DistributedStorageConfigTest, UsesDefaultsWhenEnvironmentIsUnset) { + const auto config = DistributedStorageConfig::FromEnvironment(); + + ExpectDefaultConfig(config); + EXPECT_TRUE(config.Validate()); + EXPECT_TRUE(config.ValidateForAllocator()); +} + +TEST_F(DistributedStorageConfigTest, ReadsValidEnvironmentValues) { + env.root_dir.Set("/tmp/mooncake-dfs"); + env.fs_adapter.Set("posix"); + env.health_check.Set("true"); + env.shard_count.Set("8"); + env.shard_capacity.Set("1048576"); + env.alignment.Set("4096"); + env.single_tenant.Set("1"); + env.eviction_enabled.Set("1"); + env.eviction_high_watermark.Set("0.85"); + env.eviction_low_watermark.Set("0.65"); + env.deferred_free_seconds.Set("12"); + env.eviction_check_interval.Set("3"); + + const auto config = DistributedStorageConfig::FromEnvironment(); + + EXPECT_EQ(config.fsdir, "/tmp/mooncake-dfs"); + EXPECT_EQ(config.fs_adapter_type, "posix"); + EXPECT_TRUE(config.enable_health_check); + EXPECT_EQ(config.shard_count, 8); + EXPECT_EQ(config.shard_capacity, 1048576); + EXPECT_EQ(config.alignment, 4096); + EXPECT_TRUE(config.single_tenant); + EXPECT_TRUE(config.eviction_enabled); + EXPECT_DOUBLE_EQ(config.eviction_high_watermark, 0.85); + EXPECT_DOUBLE_EQ(config.eviction_low_watermark, 0.65); + EXPECT_EQ(config.deferred_free_duration, std::chrono::seconds(12)); + EXPECT_EQ(config.eviction_check_interval, std::chrono::seconds(3)); + EXPECT_TRUE(config.Validate()); + EXPECT_TRUE(config.ValidateForAllocator()); + + const std::string formatted = config.FormatStr(); + EXPECT_NE(formatted.find("fs_adapter_type=posix"), std::string::npos); + EXPECT_NE(formatted.find("shard_count=8"), std::string::npos); + EXPECT_NE(formatted.find("eviction_high_watermark=0.85"), + std::string::npos); +} + +TEST_F(DistributedStorageConfigTest, PreservesAliasPrecedence) { + env.legacy_root_dir.Set("/tmp/legacy-dfs"); + env.legacy_fs_adapter.Set("posix"); + + const auto legacy = DistributedStorageConfig::FromEnvironment(); + EXPECT_EQ(legacy.fsdir, "/tmp/legacy-dfs"); + EXPECT_EQ(legacy.fs_adapter_type, "posix"); + + env.root_dir.Set("/tmp/preferred-dfs"); + env.fs_adapter.Set("hf3fs"); + + const auto preferred = DistributedStorageConfig::FromEnvironment(); + EXPECT_EQ(preferred.fsdir, "/tmp/preferred-dfs"); + EXPECT_EQ(preferred.fs_adapter_type, "hf3fs"); +} + +TEST_F(DistributedStorageConfigTest, EmptyPreferredRootOverridesAlias) { + env.legacy_root_dir.Set("/tmp/legacy-dfs"); + env.root_dir.Set(""); + + EXPECT_THROW(DistributedStorageConfig::FromEnvironment(), + std::filesystem::filesystem_error); +} + +TEST_F(DistributedStorageConfigTest, EmptyPreferredAdapterOverridesAlias) { + env.legacy_fs_adapter.Set("posix"); + env.fs_adapter.Set(""); + + const auto config = DistributedStorageConfig::FromEnvironment(); + + EXPECT_TRUE(config.fs_adapter_type.empty()); + EXPECT_FALSE(config.Validate()); +} + +TEST_F(DistributedStorageConfigTest, ConvertsRelativeRootToAbsolutePath) { + env.root_dir.Set("relative-dfs-root"); + + const auto config = DistributedStorageConfig::FromEnvironment(); + + EXPECT_EQ(config.fsdir, + std::filesystem::absolute("relative-dfs-root").string()); +} + +TEST_F(DistributedStorageConfigTest, + InvalidValuesUseDefaultsAndPreserveDiagnostics) { + env.health_check.Set("invalid"); + env.shard_count.Set("invalid"); + env.shard_capacity.Set("-1"); + env.alignment.Set("18446744073709551616"); + env.single_tenant.Set("invalid"); + env.eviction_enabled.Set("invalid"); + env.eviction_high_watermark.Set("invalid"); + env.deferred_free_seconds.Set("invalid"); + env.eviction_check_interval.Set("invalid"); + + ::testing::internal::CaptureStderr(); + const auto config = DistributedStorageConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + ExpectDefaultConfig(config); + for (const char* name : { + "MOONCAKE_DISTRIBUTED_HEALTH_CHECK", + "MOONCAKE_DFS_SHARD_COUNT", + "MOONCAKE_DFS_SHARD_CAPACITY", + "MOONCAKE_DFS_ALIGNMENT", + "MOONCAKE_DFS_SINGLE_TENANT", + "MOONCAKE_DFS_EVICTION_ENABLED", + "MOONCAKE_DFS_EVICTION_HIGH_WATERMARK", + "MOONCAKE_DFS_DEFERRED_FREE_SECONDS", + "MOONCAKE_DFS_EVICTION_CHECK_INTERVAL", + }) { + EXPECT_NE(logs.find(name), std::string::npos) << name; + } +} + +TEST_F(DistributedStorageConfigTest, + EmptyWatermarksUseDefaultsWithoutDiagnostics) { + env.eviction_high_watermark.Set(""); + env.eviction_low_watermark.Set(""); + + ::testing::internal::CaptureStderr(); + const auto config = DistributedStorageConfig::FromEnvironment(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + EXPECT_DOUBLE_EQ(config.eviction_high_watermark, 0.9); + EXPECT_DOUBLE_EQ(config.eviction_low_watermark, 0.7); + EXPECT_TRUE(logs.empty()); +} + +TEST(DistributedStorageConfigValidationTest, RejectsInvalidBaseSettings) { + using Mutation = + std::pair>; + const std::vector mutations{ + {"empty root", [](auto& config) { config.fsdir.clear(); }}, + {"relative root", [](auto& config) { config.fsdir = "relative/path"; }}, + {"unsupported adapter", + [](auto& config) { config.fs_adapter_type = "unsupported"; }}, + {"zero shards", [](auto& config) { config.shard_count = 0; }}, + {"zero capacity", [](auto& config) { config.shard_capacity = 0; }}, + {"zero alignment", [](auto& config) { config.alignment = 0; }}, + {"non-power-of-two alignment", + [](auto& config) { config.alignment = 3; }}, + {"unaligned capacity", + [](auto& config) { config.shard_capacity += 1; }}, + {"multi tenant", [](auto& config) { config.single_tenant = false; }}, + }; + + for (const auto& [name, mutate] : mutations) { + SCOPED_TRACE(name); + auto config = ValidConfig(); + mutate(config); + EXPECT_FALSE(config.Validate()); + } +} + +TEST(DistributedStorageConfigValidationTest, RejectsInvalidAllocatorSettings) { + using Mutation = + std::pair>; + const std::vector mutations{ + {"negative low watermark", + [](auto& config) { config.eviction_low_watermark = -0.1; }}, + {"high watermark above one", + [](auto& config) { config.eviction_high_watermark = 1.1; }}, + {"unordered watermarks", + [](auto& config) { config.eviction_low_watermark = 0.85; }}, + {"negative deferred free", + [](auto& config) { + config.deferred_free_duration = std::chrono::seconds(-1); + }}, + {"zero eviction interval", + [](auto& config) { + config.eviction_check_interval = std::chrono::seconds(0); + }}, + }; + + for (const auto& [name, mutate] : mutations) { + SCOPED_TRACE(name); + auto config = ValidConfig(); + mutate(config); + EXPECT_FALSE(config.ValidateForAllocator()); + } + + auto eviction_disabled = ValidConfig(); + eviction_disabled.eviction_enabled = false; + eviction_disabled.eviction_check_interval = std::chrono::seconds(0); + EXPECT_TRUE(eviction_disabled.ValidateForAllocator()); +} + +} // namespace + +} // namespace mooncake diff --git a/mooncake-store/tests/ha/master_service_ha_test.cpp b/mooncake-store/tests/ha/master_service_ha_test.cpp index cc2f0ae387..3eaffb4bb1 100644 --- a/mooncake-store/tests/ha/master_service_ha_test.cpp +++ b/mooncake-store/tests/ha/master_service_ha_test.cpp @@ -1487,6 +1487,9 @@ TEST_F(MasterServiceHATest, RemountRestoresCachelibMemoryReplica) { object.metadata.replicas.front() .get_memory_descriptor() .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.protocol_ = "rdma"; ASSERT_TRUE(service .RestoreFromStandbySnapshot( {object}, 7, {MakeStandbyMemorySegment(endpoint)}) @@ -1510,6 +1513,11 @@ TEST_F(MasterServiceHATest, RemountRestoresCachelibMemoryReplica) { ASSERT_TRUE(batch_after[0].has_value()); EXPECT_FALSE( HasInvalidMemoryHandleForTesting(service, kDefaultTenant, key)); + EXPECT_EQ(batch_after[0] + ->replicas.front() + .get_memory_descriptor() + .buffer_descriptor.protocol_, + "rdma"); EXPECT_EQ(SegmentAllocatedSizeForTesting(service, endpoint), 64); EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - metric_before, diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/coordinator_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/coordinator_test.cpp new file mode 100644 index 0000000000..2bcdcab79d --- /dev/null +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/coordinator_test.cpp @@ -0,0 +1,306 @@ +#include "ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_types.h" +#include "ha/snapshot/batch_oplog/metadata.h" +#include "ha/snapshot/snapshot_maintenance_lease.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "hot_standby_service.h" + +namespace mooncake::test { +namespace { + +class EmptyBackend final : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + auto it = values.find(std::string(key)); + if (it == values.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + + ErrorCode Put(std::string_view key, std::string_view value) override { + values[std::string(key)] = std::string(value); + return ErrorCode::OK; + } + + ErrorCode Range(std::string_view, std::string_view, size_t, + std::vector& output) override { + output.clear(); + return ErrorCode::OK; + } + + bool SupportsTxn() const override { return true; } + ErrorCode Txn(const KvTxn&) override { return ErrorCode::OK; } + + std::map values; +}; + +class UnusedObjectStore final : public SnapshotObjectStore { + public: + tl::expected UploadBuffer( + const std::string&, const std::vector&) override { + return {}; + } + tl::expected DownloadBuffer( + const std::string&, std::vector&) override { + return tl::make_unexpected("unused"); + } + tl::expected UploadString(const std::string&, + const std::string&) override { + return {}; + } + tl::expected DownloadString(const std::string&, + std::string&) override { + return tl::make_unexpected("unused"); + } + tl::expected DeleteObjectsWithPrefix( + const std::string&) override { + return {}; + } + tl::expected ListObjectsWithPrefix( + const std::string&, std::vector& output) override { + output.clear(); + return {}; + } + std::string GetConnectionInfo() const override { return "unused"; } +}; + +class RecordingBackend final : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + std::lock_guard lock(mutex_); + auto it = values_.find(std::string(key)); + if (it == values_.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + + ErrorCode Put(std::string_view key, std::string_view value) override { + std::lock_guard lock(mutex_); + const std::string owned_key(key); + values_[owned_key] = std::string(value); + create_revisions_.try_emplace(owned_key, next_revision_++); + return ErrorCode::OK; + } + + ErrorCode Range(std::string_view begin, std::string_view end, size_t limit, + std::vector& output) override { + std::lock_guard lock(mutex_); + output.clear(); + for (const auto& [key, value] : values_) { + if (key >= begin && key < end && + (limit == 0 || output.size() < limit)) { + output.push_back({key, value}); + } + } + return ErrorCode::OK; + } + + bool SupportsTxn() const override { return true; } + + ErrorCode Txn(const KvTxn& txn) override { + std::lock_guard lock(mutex_); + for (const auto& compare : txn.compares) { + auto it = values_.find(compare.key); + if (compare.kind == KvCompareKind::kKeyNotExists) { + if (it != values_.end()) + return ErrorCode::ETCD_TRANSACTION_FAIL; + } else if (compare.kind == KvCompareKind::kCreateRevisionEquals) { + auto revision = create_revisions_.find(compare.key); + if (revision == create_revisions_.end() || + revision->second != compare.expected_revision) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } else if (it == values_.end() || + it->second != compare.expected_value) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } + for (const auto& put : txn.puts) { + values_[put.key] = put.value; + create_revisions_.try_emplace(put.key, next_revision_++); + } + return ErrorCode::OK; + } + + bool Contains(std::string_view key) const { + std::lock_guard lock(mutex_); + return values_.contains(std::string(key)); + } + + private: + mutable std::mutex mutex_; + std::map values_; + std::map create_revisions_; + EtcdRevisionId next_revision_{1}; +}; + +class RecordingObjectStore final : public SnapshotObjectStore { + public: + tl::expected UploadBuffer( + const std::string& key, const std::vector& buffer) override { + objects_[key] = buffer; + return {}; + } + tl::expected DownloadBuffer( + const std::string& key, std::vector& buffer) override { + auto it = objects_.find(key); + if (it == objects_.end()) return tl::make_unexpected("not found"); + buffer = it->second; + return {}; + } + tl::expected UploadString( + const std::string& key, const std::string& value) override { + objects_[key] = std::vector(value.begin(), value.end()); + return {}; + } + tl::expected DownloadString( + const std::string& key, std::string& value) override { + std::vector bytes; + auto result = DownloadBuffer(key, bytes); + if (!result) return result; + value.assign(bytes.begin(), bytes.end()); + return {}; + } + tl::expected DeleteObjectsWithPrefix( + const std::string& prefix) override { + for (auto it = objects_.begin(); it != objects_.end();) { + if (it->first.starts_with(prefix)) + it = objects_.erase(it); + else + ++it; + } + return {}; + } + tl::expected ListObjectsWithPrefix( + const std::string& prefix, std::vector& output) override { + output.clear(); + for (const auto& [key, value] : objects_) { + (void)value; + if (key.starts_with(prefix)) output.push_back(key); + } + return {}; + } + tl::expected InspectObject( + const std::string& key) override { + auto it = objects_.find(key); + if (it == objects_.end()) return tl::make_unexpected("not found"); + return SnapshotObjectInspection{.stored_size = it->second.size(), + .crc32c = std::nullopt}; + } + std::string GetConnectionInfo() const override { return "recording"; } + + private: + std::map> objects_; +}; + +OpLogBatchRecord MakeBatch() { + OpLogEntry entry; + entry.sequence_id = 1; + entry.op_type = OpType::REMOVE; + entry.tenant_id = "tenant"; + entry.object_key = "key"; + entry.checksum = ComputeOpLogChecksum(entry.payload); + OpLogBatchRecord batch; + batch.batch_id = 1; + batch.first_seq = 1; + batch.last_seq = 1; + batch.entries.push_back(std::move(entry)); + return batch; +} + +} // namespace + +TEST(BatchOpLogSnapshotCoordinatorTest, EmptyStandbySkipsWithoutLease) { + HotStandbyConfig standby_config; + standby_config.enable_verification = false; + HotStandbyService standby(standby_config); + EmptyBackend backend; + UnusedObjectStore object_store; + size_t lease_factory_calls = 0; + BatchOpLogSnapshotCoordinatorConfig config; + config.snapshot_root = "snapshots"; + config.clock = [] { return std::chrono::steady_clock::now(); }; + BatchOpLogSnapshotCoordinator coordinator( + standby, backend, object_store, "cluster", std::move(config), [&] { + ++lease_factory_calls; + return std::unique_ptr(); + }); + + EXPECT_EQ(ErrorCode::OK, coordinator.RunOnce()); + EXPECT_EQ(0u, lease_factory_calls); + EXPECT_FALSE(coordinator.IsAttemptInFlight()); + EXPECT_EQ(0u, coordinator.GetStatus().attempts); + + coordinator.Start(); + EXPECT_TRUE(coordinator.IsRunning()); + coordinator.Stop(); + EXPECT_FALSE(coordinator.IsRunning()); +} + +TEST(BatchOpLogSnapshotCoordinatorTest, PublishesAfterCaptureAndResumesApply) { + auto backend = std::make_shared(); + ASSERT_EQ(ErrorCode::OK, backend->Put(BuildBatchRecordKey("cluster", 1), + EncodeOpLogBatchRecord(MakeBatch()))); + ASSERT_EQ( + ErrorCode::OK, + backend->Put(BuildDurablePrefixKey("cluster"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildProducerViewKey("cluster"), "7")); + const auto maintenance_key = + ha::BuildBatchOpLogSnapshotMaintenanceKey("cluster"); + ASSERT_EQ(ErrorCode::OK, backend->Put(maintenance_key, "101")); + + HotStandbyConfig standby_config; + standby_config.enable_verification = false; + standby_config.oplog_poll_interval_ms = 1; + HotStandbyService standby(standby_config); + standby.SetCatchUpBatchKvBackendForTesting(backend); + ASSERT_EQ(ErrorCode::OK, standby.Start("", "", "cluster")); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (std::chrono::steady_clock::now() < deadline) { + auto prefix = standby.GetLastAppliedBatchOpLogSnapshotPrefix(); + if (prefix && prefix->batch_id == 1) break; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + RecordingObjectStore object_store; + BatchOpLogSnapshotCoordinatorConfig config; + config.snapshot_root = "snapshots"; + config.snapshot_interval_seconds = 0; + BatchOpLogSnapshotCoordinator coordinator( + standby, *backend, object_store, "cluster", std::move(config), [] { + return SnapshotMaintenanceLease::MakeForTesting("cluster", "101", + 4); + }); + + EXPECT_EQ(ErrorCode::OK, coordinator.RunOnce()); + EXPECT_TRUE( + backend->Contains(ha::BuildBatchOpLogSnapshotLatestKey("cluster"))); + EXPECT_EQ(1u, coordinator.GetStatus().attempts); + EXPECT_TRUE(coordinator.GetStatus().catch_up_target.has_value()); + standby.Stop(); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h index 48659096c4..d0bd912094 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h +++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -165,6 +166,8 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { // private members static tl::expected CallPersistState( MasterService* service, const std::string& snapshot_id) { + std::unique_lock snapshot_lock( + service->snapshot_mutex_); // If snapshot_manager_ exists, use it; otherwise create a temporary one if (service->snapshot_manager_) { return service->snapshot_manager_->PersistState(snapshot_id); @@ -827,6 +830,15 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { << "Use 'service_.reset(new MasterService(...))' instead of " "'std::unique_ptr service_(...)'"; + // Freeze the original service before snapshot/restore validation. The + // eviction worker can otherwise mutate replica metadata while the + // snapshot is being serialized, making the post-restore comparison + // compare two different points in time. + service_->eviction_running_ = false; + if (service_->eviction_thread_.joinable()) { + service_->eviction_thread_.join(); + } + // Some test configs may not enable snapshot/restore, so the backend // is not created in the constructor. We create it here for TearDown // validation. diff --git a/mooncake-store/tests/local_hot_cache_config_test.cpp b/mooncake-store/tests/local_hot_cache_config_test.cpp new file mode 100644 index 0000000000..354dfd4de0 --- /dev/null +++ b/mooncake-store/tests/local_hot_cache_config_test.cpp @@ -0,0 +1,150 @@ +#include +#include + +#include +#include +#include + +#include "local_hot_cache.h" + +namespace mooncake { +namespace { + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + if (const char* value = std::getenv(name)) { + original_ = value; + } + unsetenv(name); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +struct LocalHotCacheEnvironment { + ScopedEnvVar total_size{"MC_STORE_LOCAL_HOT_CACHE_SIZE"}; + ScopedEnvVar block_size{"MC_STORE_LOCAL_HOT_BLOCK_SIZE"}; + ScopedEnvVar use_shm{"MC_STORE_LOCAL_HOT_CACHE_USE_SHM"}; + ScopedEnvVar admission_threshold{"MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD"}; +}; + +class LocalHotCacheConfigTest : public ::testing::Test { + protected: + LocalHotCacheEnvironment env; + + void SetUp() override { + google::InitGoogleLogging("LocalHotCacheConfigTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } +}; + +TEST_F(LocalHotCacheConfigTest, UsesExistingDefaultsWhenEnvironmentIsUnset) { + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 0); + EXPECT_EQ(config.block_size_bytes, 16 * 1024 * 1024); + EXPECT_FALSE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 2); +} + +TEST_F(LocalHotCacheConfigTest, ReadsValidValues) { + env.total_size.Set("33554432"); + env.block_size.Set("4194304"); + env.use_shm.Set("1"); + env.admission_threshold.Set("5"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 32 * 1024 * 1024); + EXPECT_EQ(config.block_size_bytes, 4 * 1024 * 1024); + EXPECT_TRUE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 5); +} + +TEST_F(LocalHotCacheConfigTest, DisabledCacheDoesNotReadDependentSettings) { + env.total_size.Set("0"); + env.block_size.Set("4194304"); + env.use_shm.Set("1"); + env.admission_threshold.Set("5"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 0); + EXPECT_EQ(config.block_size_bytes, 16 * 1024 * 1024); + EXPECT_FALSE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 2); +} + +TEST_F(LocalHotCacheConfigTest, InvalidCacheSizesDisableCache) { + for (const char* value : + {"", "0", "-1", "invalid", "18446744073709551616"}) { + env.total_size.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().total_size_bytes, 0) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, InvalidBlockSizesUseDefault) { + env.total_size.Set("33554432"); + + for (const char* value : + {"", "0", "-1", "invalid", "18446744073709551616"}) { + env.block_size.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().block_size_bytes, + 16 * 1024 * 1024) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, InvalidAdmissionThresholdsUseDefault) { + env.total_size.Set("33554432"); + + for (const char* value : + {"", "0", "-1", "256", "invalid", "18446744073709551616"}) { + env.admission_threshold.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().admission_threshold, 2) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, PreservesLegacyNumericPrefixParsing) { + env.total_size.Set("33554432suffix"); + env.block_size.Set("4194304suffix"); + env.admission_threshold.Set("5suffix"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 32 * 1024 * 1024); + EXPECT_EQ(config.block_size_bytes, 4 * 1024 * 1024); + EXPECT_EQ(config.admission_threshold, 5); +} + +TEST_F(LocalHotCacheConfigTest, SharedMemoryRequiresExactOne) { + env.total_size.Set("33554432"); + + for (const char* value : {"", "0", "true", "01", " 1"}) { + env.use_shm.Set(value); + EXPECT_FALSE(LocalHotCacheConfig::FromEnvironment().use_shm) << value; + } + + env.use_shm.Set("1"); + EXPECT_TRUE(LocalHotCacheConfig::FromEnvironment().use_shm); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 383ed6ae91..976d838770 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -1158,18 +1158,57 @@ TEST_F(MasterServiceTest, PutWithPreferredSegments) { TEST_F(MasterServiceTest, ResolveMooncakeHostIdUsesLocalHostnameAndRejectsLoopback) { + ScopedEnvVar host_id("MOONCAKE_HOST_ID"); + + EXPECT_EQ(ResolveMooncakeHostId("hostB:5000"), "hostB"); + EXPECT_EQ(ResolveMooncakeHostId("hostB:5001"), "hostB"); + EXPECT_EQ(ResolveMooncakeHostId("[2001:db8::1]:5000"), "2001:db8::1"); + EXPECT_TRUE(ResolveMooncakeHostId("localhost:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("127.0.0.1:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("0.0.0.0:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("::1").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::1]:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("::").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::]").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::]:5000").empty()); +} + +TEST_F(MasterServiceTest, ResolveMooncakeHostIdPrefersDeploymentOverride) { + ScopedEnvVar host_id("MOONCAKE_HOST_ID", " kubernetes-node-a "); + + EXPECT_EQ(ResolveMooncakeHostId("10.244.1.17:5000"), "kubernetes-node-a"); +} + +TEST_F(MasterServiceTest, ResolveMooncakeHostIdNormalizesEndpointOverride) { + ScopedEnvVar host_id("MOONCAKE_HOST_ID", " kubernetes-node-a:5000 "); + + EXPECT_EQ(ResolveMooncakeHostId("10.244.1.17:5000"), "kubernetes-node-a"); +} + +TEST_F(MasterServiceTest, ResolveMooncakeHostIdFallsBackForEmptyOverride) { { + ScopedEnvVar host_id("MOONCAKE_HOST_ID", ""); EXPECT_EQ(ResolveMooncakeHostId("hostB:5000"), "hostB"); - EXPECT_EQ(ResolveMooncakeHostId("hostB:5001"), "hostB"); - EXPECT_EQ(ResolveMooncakeHostId("[2001:db8::1]:5000"), "2001:db8::1"); - EXPECT_TRUE(ResolveMooncakeHostId("localhost:5000").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("127.0.0.1:5000").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("0.0.0.0:5000").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("::1").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("[::1]:5000").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("::").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("[::]").empty()); - EXPECT_TRUE(ResolveMooncakeHostId("[::]:5000").empty()); + } + + { + ScopedEnvVar host_id("MOONCAKE_HOST_ID", " \t "); + EXPECT_EQ(ResolveMooncakeHostId("hostB:5000"), "hostB"); + } +} + +TEST_F(MasterServiceTest, ResolveMooncakeHostIdRejectsInvalidOverride) { + const std::vector invalid_host_ids = { + "localhost", "localhost:5000", + "127.0.0.1", "127.0.0.1:5000", + "0.0.0.0", "0.0.0.0:5000", + "::1", "[::1]", + "[::1]:5000", "::", + "[::]", "[::]:5000"}; + for (const char* invalid_host_id : invalid_host_ids) { + ScopedEnvVar host_id("MOONCAKE_HOST_ID", invalid_host_id); + EXPECT_TRUE(ResolveMooncakeHostId("hostB:5000").empty()) + << invalid_host_id; } } diff --git a/mooncake-store/tests/region_driver_test.cpp b/mooncake-store/tests/region_driver_test.cpp new file mode 100644 index 0000000000..19c01a5a79 --- /dev/null +++ b/mooncake-store/tests/region_driver_test.cpp @@ -0,0 +1,181 @@ +#include "segment/region_driver.h" + +#include + +#include +#include + +#include "master_metric_manager.h" + +namespace mooncake { +namespace { + +constexpr size_t kRegionSize = 16U * 1024 * 1024; + +RegionResourceSpec MakeSpec(uintptr_t base = 0x100000000ULL) { + return {generate_uuid(), "memory", base, kRegionSize, "memory-endpoint"}; +} + +std::unique_ptr CreateTestDriver(RegionKind kind) { + RegionDriverConfig config; + config.memory_allocator = BufferAllocatorType::OFFSET; + if (kind == RegionKind::CXL) { + config.cxl = CxlRegionDriverConfig{"cxl-test", kRegionSize}; + } + auto drivers = CreateRegionDrivers(config); + if (!drivers) { + return nullptr; + } + auto driver = drivers->extract(kind); + return driver.empty() ? nullptr : std::move(driver.mapped()); +} + +TEST(RegionDriverTest, PreparedResourceRollsBackUntilCommitted) { + auto driver = CreateTestDriver(RegionKind::HOST_MEMORY); + ASSERT_NE(driver, nullptr); + const auto spec = MakeSpec(); + + { + auto prepared = driver->PrepareOpen(spec, {}); + ASSERT_TRUE(prepared.has_value()); + EXPECT_EQ(driver->GetResource(spec.id), nullptr); + } + EXPECT_EQ(driver->GetResource(spec.id), nullptr); + + auto prepared = driver->PrepareOpen(spec, {}); + ASSERT_TRUE(prepared.has_value()); + prepared->Commit(); + auto* resource = driver->GetResource(spec.id); + ASSERT_NE(resource, nullptr); + EXPECT_TRUE(resource->active); +} + +TEST(RegionDriverTest, ReplacementRollbackKeepsCommittedResource) { + auto driver = CreateTestDriver(RegionKind::HOST_MEMORY); + ASSERT_NE(driver, nullptr); + const auto spec = MakeSpec(); + auto first = driver->PrepareOpen(spec, {}); + ASSERT_TRUE(first.has_value()); + first->Commit(); + auto* committed = driver->GetResource(spec.id); + ASSERT_NE(committed, nullptr); + + { + auto replacement = driver->PrepareOpen(spec, {}); + ASSERT_TRUE(replacement.has_value()); + EXPECT_NE(&replacement->resource(), committed); + } + EXPECT_EQ(driver->GetResource(spec.id), committed); +} + +TEST(RegionDriverTest, RestoreInputValidatesEndpointBoundsAndPreservesOrder) { + const auto spec = MakeSpec(0x200000000ULL); + std::vector descriptors{ + {4096, spec.base + 8192, "tcp", spec.transport_endpoint}, + {4096, spec.base, "tcp", spec.transport_endpoint}}; + auto allocations = BuildRegionLiveAllocations(spec, descriptors); + ASSERT_TRUE(allocations.has_value()); + ASSERT_EQ(allocations->size(), 2U); + EXPECT_EQ((*allocations)[0].offset_from_base, 8192U); + EXPECT_EQ((*allocations)[1].offset_from_base, 0U); + + auto bad_endpoint = descriptors; + bad_endpoint[0].transport_endpoint_ = "other"; + EXPECT_EQ(BuildRegionLiveAllocations(spec, bad_endpoint).error(), + ErrorCode::INVALID_PARAMS); + auto segment_name_alias = descriptors; + segment_name_alias[0].transport_endpoint_ = spec.name; + EXPECT_EQ(BuildRegionLiveAllocations(spec, segment_name_alias).error(), + ErrorCode::INVALID_PARAMS); + auto out_of_bounds = descriptors; + out_of_bounds[0].buffer_address_ = spec.base + spec.size - 1024; + EXPECT_EQ(BuildRegionLiveAllocations(spec, out_of_bounds).error(), + ErrorCode::INVALID_PARAMS); +} + +TEST(RegionDriverTest, OffsetImportPreservesInputOrder) { + auto driver = CreateTestDriver(RegionKind::HOST_MEMORY); + ASSERT_NE(driver, nullptr); + const auto spec = MakeSpec(0x300000000ULL); + std::vector allocations{{8192, 4096}, {0, 4096}}; + auto prepared = driver->PrepareOpen(spec, allocations); + ASSERT_TRUE(prepared.has_value()); + ASSERT_EQ(prepared->imported_buffers().size(), 2U); + EXPECT_EQ( + reinterpret_cast(prepared->imported_buffers()[0]->data()), + spec.base + 8192); + EXPECT_EQ( + reinterpret_cast(prepared->imported_buffers()[1]->data()), + spec.base); +} + +TEST(RegionDriverTest, FailedRestoreDoesNotPublishResource) { + auto driver = CreateTestDriver(RegionKind::HOST_MEMORY); + ASSERT_NE(driver, nullptr); + const auto spec = MakeSpec(0x400000000ULL); + auto prepared = driver->PrepareOpen(spec, {{{0, 4096}, {0, 4096}}}); + EXPECT_FALSE(prepared.has_value()); + EXPECT_EQ(prepared.error(), ErrorCode::INVALID_PARAMS); + EXPECT_EQ(driver->GetResource(spec.id), nullptr); +} + +TEST(RegionDriverTest, CxlRejectsLiveRestoreInput) { + auto driver = CreateTestDriver(RegionKind::CXL); + ASSERT_NE(driver, nullptr); + RegionResourceSpec spec{generate_uuid(), "binding", 0, kRegionSize, + "transport"}; + auto prepared = driver->PrepareOpen(spec, {{{0, 4096}}}); + EXPECT_FALSE(prepared.has_value()); + EXPECT_EQ(prepared.error(), ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + EXPECT_EQ(driver->GetResource(spec.id), nullptr); +} + +TEST(RegionDriverTest, CxlTargetProducesCxlDescriptors) { + auto driver = CreateTestDriver(RegionKind::CXL); + ASSERT_NE(driver, nullptr); + RegionResourceSpec spec{generate_uuid(), "binding", 0, kRegionSize, + "transport"}; + auto prepared = driver->PrepareOpen(spec, {}); + ASSERT_TRUE(prepared.has_value()); + + auto buffer = prepared->resource().target->Allocate(4096); + ASSERT_NE(buffer, nullptr); + const auto descriptor = buffer->get_descriptor(); + EXPECT_EQ(descriptor.protocol_, "cxl"); + EXPECT_EQ(descriptor.transport_endpoint_, spec.name); +} + +TEST(RegionDriverTest, CxlDriverOwnsCapacityMetricLifetime) { + constexpr char kCxlPath[] = "region-driver-cxl-metric"; + auto& metrics = MasterMetricManager::instance(); + const int64_t total_before = metrics.get_total_mem_capacity(); + const int64_t segment_before = + metrics.get_segment_total_mem_capacity(kCxlPath); + + { + RegionDriverConfig config; + config.cxl = CxlRegionDriverConfig{kCxlPath, kRegionSize}; + auto drivers = CreateRegionDrivers(config); + ASSERT_TRUE(drivers.has_value()); + EXPECT_EQ(metrics.get_total_mem_capacity(), + total_before + static_cast(kRegionSize)); + EXPECT_EQ(metrics.get_segment_total_mem_capacity(kCxlPath), + segment_before + static_cast(kRegionSize)); + } + + EXPECT_EQ(metrics.get_total_mem_capacity(), total_before); + EXPECT_EQ(metrics.get_segment_total_mem_capacity(kCxlPath), segment_before); +} + +TEST(RegionDriverTest, InvalidCxlConfigIsReturnedExplicitly) { + RegionDriverConfig config; + config.cxl = + CxlRegionDriverConfig{"cxl-test", facebook::cachelib::Slab::kSize + 1}; + + auto drivers = CreateRegionDrivers(config); + ASSERT_FALSE(drivers.has_value()); + EXPECT_EQ(drivers.error(), ErrorCode::INVALID_PARAMS); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp index 04219025e3..3e52fda240 100644 --- a/mooncake-store/tests/registered_pinned_memory_test.cpp +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -1,7 +1,10 @@ #define MOONCAKE_STORE_TEST #include "../src/registered_pinned_memory.h" +#include "../src/config/registered_pinned_memory_config.h" #include +#include +#include #include @@ -11,6 +14,70 @@ namespace { using Manager = RegisteredPinnedMemoryManager; using UnregisterResult = Manager::UnregisterResult; +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + if (const char* value = std::getenv(name)) { + original_ = value; + } + unsetenv(name_.c_str()); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + ScopedEnvVar(const ScopedEnvVar&) = delete; + ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +class RegisteredPinnedMemoryConfigTest : public ::testing::Test { + protected: + ScopedEnvVar max_bytes{"MC_STORE_PIN_MEMORY_MAX_BYTES"}; +}; + +TEST_F(RegisteredPinnedMemoryConfigTest, UnsetAndEmptyDisableWithoutWarning) { + ::testing::internal::CaptureStderr(); + EXPECT_EQ(RegisteredPinnedMemoryConfig::FromEnvironment().max_bytes, 0); + EXPECT_TRUE(::testing::internal::GetCapturedStderr().empty()); + + max_bytes.Set(""); + ::testing::internal::CaptureStderr(); + EXPECT_EQ(RegisteredPinnedMemoryConfig::FromEnvironment().max_bytes, 0); + EXPECT_TRUE(::testing::internal::GetCapturedStderr().empty()); +} + +TEST_F(RegisteredPinnedMemoryConfigTest, ReadsZeroAndPositiveByteLimits) { + max_bytes.Set("0"); + EXPECT_EQ(RegisteredPinnedMemoryConfig::FromEnvironment().max_bytes, 0); + + max_bytes.Set(" 4096\t"); + EXPECT_EQ(RegisteredPinnedMemoryConfig::FromEnvironment().max_bytes, 4096); +} + +TEST_F(RegisteredPinnedMemoryConfigTest, InvalidValuesWarnAndDisable) { + for (const char* value : + {"-1", "+1", "1suffix", " ", "18446744073709551616"}) { + max_bytes.Set(value); + ::testing::internal::CaptureStderr(); + EXPECT_EQ(RegisteredPinnedMemoryConfig::FromEnvironment().max_bytes, 0) + << value; + const std::string logs = ::testing::internal::GetCapturedStderr(); + EXPECT_NE(logs.find("MC_STORE_PIN_MEMORY_MAX_BYTES"), std::string::npos) + << value; + } +} + struct FakePinState { bool register_succeeds = true; UnregisterResult unregister_result = UnregisterResult::kSuccess; @@ -44,7 +111,7 @@ class RegisteredPinnedMemoryManagerTest : public ::testing::Test { void SetUp() override { State() = FakePinState(); } Manager MakeManager(size_t limit) { - return Manager({true, limit}, {FakeRegister, FakeUnregister}); + return Manager({.max_bytes = limit}, {FakeRegister, FakeUnregister}); } std::shared_ptr Pin(Manager& manager, size_t offset, diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 9cb83588cf..ddb4312aae 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -3,8 +3,11 @@ #include #include +#include +#include #include #include +#include #include #include #include @@ -19,6 +22,8 @@ #include #include #include +#include +#include #include #include @@ -262,6 +267,123 @@ TEST_F(StorageBackendTest, CreateAcceptsValidConfig) { EXPECT_NE(result.value(), nullptr); } +TEST_F(StorageBackendTest, RemoveFileWaitsForStoreInWriteCriticalSection) { + StorageBackend backend(data_path, "unused", false); + ASSERT_TRUE(backend.Init(0)); + + const std::string path = data_path + "/concurrent_remove_fifo"; + std::error_code cleanup_ec; + fs::remove(path, cleanup_ec); + ASSERT_FALSE(cleanup_ec); + ASSERT_EQ(mkfifo(path.c_str(), 0600), 0) << strerror(errno); + + const int reader_fd = open(path.c_str(), O_RDONLY | O_NONBLOCK); + if (reader_fd < 0) { + fs::remove(path); + FAIL() << "Failed to open FIFO reader: " << strerror(errno); + } + const int pipe_capacity = fcntl(reader_fd, F_GETPIPE_SZ); + if (pipe_capacity <= 0) { + close(reader_fd); + fs::remove(path); + FAIL() << "Failed to query FIFO capacity: " << strerror(errno); + } + + // With no reader draining the FIFO, this write fills the pipe and blocks + // after StoreObject has acquired the path mutex. + const std::string value(static_cast(pipe_capacity) * 2, 'x'); + std::atomic store_done{false}; + auto store_future = std::async(std::launch::async, [&]() { + auto result = backend.StoreObject(path, value); + store_done.store(true, std::memory_order_release); + return result; + }); + + bool queue_probe_ok = true; + bool writer_blocked = false; + const auto write_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < write_deadline) { + int queued_bytes = 0; + if (ioctl(reader_fd, FIONREAD, &queued_bytes) != 0) { + queue_probe_ok = false; + break; + } + const auto status = store_future.wait_for(std::chrono::milliseconds(0)); + if (queued_bytes > 0 && status == std::future_status::timeout) { + writer_blocked = true; + break; + } + if (status == std::future_status::ready) { + break; + } + std::this_thread::yield(); + } + + bool remove_blocked = false; + std::optional> remove_future; + if (writer_blocked) { + std::promise remove_started_promise; + auto remove_started = remove_started_promise.get_future(); + remove_future.emplace(std::async(std::launch::async, [&]() { + remove_started_promise.set_value(); + backend.RemoveFile(path); + })); + remove_started.wait(); + remove_blocked = + remove_future->wait_for(std::chrono::milliseconds(200)) == + std::future_status::timeout; + } + + // Drain the FIFO only after checking that RemoveFile is blocked. This + // lets StoreObject finish and release the path mutex. + auto drain_future = std::async(std::launch::async, [&]() { + std::vector buffer(64 * 1024); + size_t drained_bytes = 0; + for (;;) { + const ssize_t n = read(reader_fd, buffer.data(), buffer.size()); + if (n > 0) { + drained_bytes += static_cast(n); + continue; + } + if (n == 0) { + if (store_done.load(std::memory_order_acquire)) { + return std::optional{drained_bytes}; + } + std::this_thread::yield(); + continue; + } + if (errno == EINTR || errno == EAGAIN) { + std::this_thread::yield(); + continue; + } + return std::optional{}; + } + }); + + auto store_result = store_future.get(); + auto drain_result = drain_future.get(); + if (remove_future.has_value()) { + remove_future->get(); + } else { + backend.RemoveFile(path); + } + close(reader_fd); + + const bool path_exists = fs::exists(path); + if (path_exists) { + fs::remove(path); + } + + EXPECT_TRUE(queue_probe_ok); + EXPECT_TRUE(writer_blocked); + EXPECT_TRUE(remove_blocked); + ASSERT_TRUE(store_result.has_value()); + ASSERT_TRUE(drain_result.has_value()); + EXPECT_EQ(drain_result.value(), value.size()); + EXPECT_FALSE(path_exists); +} + class OffsetAllocatorEnvironmentTest : public StorageBackendTest { protected: OffsetAllocatorEnvironment env; diff --git a/mooncake-store/tests/utils_test.cpp b/mooncake-store/tests/utils_test.cpp index ef224eb203..7bf4c1cc1f 100644 --- a/mooncake-store/tests/utils_test.cpp +++ b/mooncake-store/tests/utils_test.cpp @@ -2,6 +2,7 @@ #include "random.h" #include +#include #include #include #include @@ -268,3 +269,33 @@ TEST(UtilsTest, AutoPortBinderCustomRange) { EXPECT_GE(port, 50000); EXPECT_LE(port, 50100); } + +TEST(HugepageSizeEnvTest, Accepts512Mb) { + setenv("MC_STORE_USE_HUGEPAGE", "1", 1); + setenv("MC_STORE_HUGEPAGE_SIZE", "512MB", 1); + + unsigned int flags = 0; + EXPECT_EQ(get_hugepage_size_from_env(&flags), SZ_512MB); + EXPECT_TRUE(flags & MAP_HUGETLB); + EXPECT_TRUE(flags & MAP_HUGE_512MB); + + flags = 0; + EXPECT_EQ(get_hugepage_size_from_env(&flags, /*use_memfd=*/true), SZ_512MB); + EXPECT_TRUE(flags & MFD_HUGETLB); + EXPECT_TRUE(flags & MFD_HUGE_512MB); + + unsetenv("MC_STORE_HUGEPAGE_SIZE"); + unsetenv("MC_STORE_USE_HUGEPAGE"); +} + +TEST(HugepageSizeEnvTest, FallsBackTo2MbOnInvalidSize) { + setenv("MC_STORE_USE_HUGEPAGE", "1", 1); + setenv("MC_STORE_HUGEPAGE_SIZE", "256MB", 1); + + unsigned int flags = 0; + EXPECT_EQ(get_hugepage_size_from_env(&flags), SZ_2MB); + EXPECT_TRUE(flags & MAP_HUGE_2MB); + + unsetenv("MC_STORE_HUGEPAGE_SIZE"); + unsetenv("MC_STORE_USE_HUGEPAGE"); +} diff --git a/mooncake-transfer-engine/benchmark/CMakeLists.txt b/mooncake-transfer-engine/benchmark/CMakeLists.txt index e13773ec5b..8c9bdb2c3a 100644 --- a/mooncake-transfer-engine/benchmark/CMakeLists.txt +++ b/mooncake-transfer-engine/benchmark/CMakeLists.txt @@ -29,7 +29,8 @@ file(GLOB TEBENCH_SOURCES "*.cpp") # The TENT backend is only available when USE_TENT is enabled; drop its # translation unit (which pulls in tent/ headers) from non-TENT builds. if(NOT USE_TENT) - list(REMOVE_ITEM TEBENCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/tent_backend.cpp") + list(REMOVE_ITEM TEBENCH_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/tent_backend.cpp") list(APPEND TEBENCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/../tent/src/common/qos_metrics.cpp") endif() @@ -41,10 +42,11 @@ if(USE_TENT) target_link_libraries(tebench PUBLIC tent_link_group) else() # The classic backend still uses a couple of header-only helpers that live - # under tent/include (SimpleRandom in utils.h, bindToSocket in te_backend.cpp). - # Expose just the header path so tebench builds without the TENT library. - target_include_directories(tebench PRIVATE - "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include") + # under tent/include (SimpleRandom in utils.h, bindToSocket in + # te_backend.cpp). Expose just the header path so tebench builds without the + # TENT library. + target_include_directories( + tebench PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include") endif() if(USE_CUDA) target_link_libraries(tebench PUBLIC CUDA::cudart) @@ -63,19 +65,21 @@ else() set(TANGRT_RPATH "") endif() set_target_properties( - tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE - INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common${TANGRT_RPATH}") + tebench + PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH + "$ORIGIN/../lib:$ORIGIN/../../mooncake-common${TANGRT_RPATH}") if(BUILD_UNIT_TESTS) - add_executable(tebench_qos_metrics_test tests/qos_metrics_test.cpp - qos_metrics_adapter.cpp - workload_config.cpp utils.cpp) + add_executable( + tebench_qos_metrics_test tests/qos_metrics_test.cpp qos_metrics_adapter.cpp + workload_config.cpp utils.cpp) if(NOT USE_TENT) - target_sources(tebench_qos_metrics_test PRIVATE - ../tent/src/common/qos_metrics.cpp) + target_sources(tebench_qos_metrics_test + PRIVATE ../tent/src/common/qos_metrics.cpp) endif() - target_link_libraries(tebench_qos_metrics_test - PRIVATE transfer_engine gtest gtest_main) + target_link_libraries(tebench_qos_metrics_test PRIVATE transfer_engine gtest + gtest_main) if(USE_TENT) target_link_libraries(tebench_qos_metrics_test PRIVATE tent_common) endif() @@ -84,4 +88,17 @@ if(BUILD_UNIT_TESTS) PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include") add_test(NAME tebench_qos_metrics_test COMMAND tebench_qos_metrics_test) + + add_executable(tebench_target_metrics_test tests/target_metrics_test.cpp + target_metrics.cpp utils.cpp) + target_link_libraries(tebench_target_metrics_test PRIVATE transfer_engine + gtest gtest_main) + if(USE_TENT) + target_link_libraries(tebench_target_metrics_test PRIVATE tent_common) + endif() + target_include_directories( + tebench_target_metrics_test + PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include") + add_test(NAME tebench_target_metrics_test COMMAND tebench_target_metrics_test) endif() diff --git a/mooncake-transfer-engine/benchmark/bench_runner.h b/mooncake-transfer-engine/benchmark/bench_runner.h index 6accc9220f..2e66ef3210 100644 --- a/mooncake-transfer-engine/benchmark/bench_runner.h +++ b/mooncake-transfer-engine/benchmark/bench_runner.h @@ -56,11 +56,16 @@ class BenchRunner { virtual size_t getTargetCount() const = 0; + virtual size_t getTargetIndex(int thread_id) const = 0; + virtual uint64_t getTargetSegmentId(int thread_id) const = 0; virtual uint64_t getTargetBufferBase(int thread_id, uint64_t block_size, uint64_t batch_size) const = 0; + // Returns elapsed microseconds on success, or a negative value if the + // transfer failed. Implementations must not call exit() from worker + // threads. virtual double runSingleTransfer(uint64_t local_addr, uint64_t target_id, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, OpCode opcode, diff --git a/mooncake-transfer-engine/benchmark/main.cpp b/mooncake-transfer-engine/benchmark/main.cpp index 19f9bbf37a..b7559006d4 100644 --- a/mooncake-transfer-engine/benchmark/main.cpp +++ b/mooncake-transfer-engine/benchmark/main.cpp @@ -16,6 +16,7 @@ #include "bench_runner.h" #include "qos_metrics_adapter.h" +#include "target_metrics.h" #include "te_backend.h" #include "workload_config.h" #ifdef USE_TENT @@ -73,6 +74,12 @@ int processBatchSizes( XferBenchStats stats; std::vector qos_stats(qos_classes.size()); + std::vector target_stats(runner.getTargetCount()); + const auto target_names = + splitCommaSeparated(XferBenchConfig::target_seg_name); + LOG_ASSERT(target_names.size() == target_stats.size()); + for (size_t i = 0; i < target_stats.size(); ++i) + target_stats[i].segment_name = target_names[i]; XferBenchStats tight_stats; XferBenchStats loose_stats; std::mutex mutex; @@ -110,6 +117,7 @@ int processBatchSizes( uint64_t target_addr = runner.getTargetBufferBase( target_thread_id, address_stride_bytes, 1); uint64_t target_id = runner.getTargetSegmentId(target_thread_id); + const size_t target_index = runner.getTargetIndex(target_thread_id); const bool qos_enabled = !qos_classes.empty(); const size_t qos_class = qos_enabled ? qosClassForThread(qos_classes, thread_id) : 0; @@ -135,11 +143,18 @@ int processBatchSizes( deadline_us * 1000ull; }; + auto failTask = [&]() { + measurement_started.store(true, std::memory_order_release); + return -1; + }; + XferBenchTimer timer; while (timer.lap_us(false) < 1000000ull) { - runner.runSingleTransfer(local_addr, target_id, target_addr, - thread_block_size, thread_batch_size, - opcode, deadlineNs(), intent_type); + if (runner.runSingleTransfer( + local_addr, target_id, target_addr, thread_block_size, + thread_batch_size, opcode, deadlineNs(), intent_type) < 0) { + return failTask(); + } } if (measurement_ready.fetch_add(1, std::memory_order_acq_rel) + 1 == num_threads) { @@ -166,6 +181,7 @@ int processBatchSizes( auto val = runner.runSingleTransfer( local_addr, target_id, target_addr, thread_block_size, thread_batch_size, WRITE, deadlineNs(), intent_type); + if (val < 0) return failTask(); thread_instant_bandwidth.push_back( gbPerSecond(batch_bytes, val)); transfer_duration.push_back(val); @@ -174,6 +190,7 @@ int processBatchSizes( val = runner.runSingleTransfer( local_addr, target_id, target_addr, thread_block_size, thread_batch_size, READ, deadlineNs(), intent_type); + if (val < 0) return failTask(); thread_instant_bandwidth.push_back( gbPerSecond(batch_bytes, val)); if (XferBenchConfig::check_consistency) @@ -198,6 +215,7 @@ int processBatchSizes( auto val = runner.runSingleTransfer( local_addr, target_id, target_addr, thread_block_size, thread_batch_size, opcode, deadlineNs(), intent_type); + if (val < 0) return failTask(); thread_instant_bandwidth.push_back( gbPerSecond(batch_bytes, val)); if (read_verify) { @@ -208,10 +226,23 @@ int processBatchSizes( } } auto total_duration = timer.lap_us(); + const uint64_t bytes_per_operation = checkedMul( + thread_block_size, thread_batch_size, "operation payload size"); + const uint64_t transferred_bytes = + checkedMul(bytes_per_operation, transfer_duration.size(), + "thread transferred bytes"); std::lock_guard lock(mutex); stats.total_duration.add(total_duration); stats.transfer_duration.add(transfer_duration); stats.instant_bandwidth.add(thread_instant_bandwidth); + auto& target = target_stats[target_index]; + ++target.threads; + target.transferred_bytes = + checkedAdd(target.transferred_bytes, transferred_bytes, + "target transferred bytes"); + target.stats.total_duration.add(total_duration); + target.stats.transfer_duration.add(transfer_duration); + target.stats.instant_bandwidth.add(thread_instant_bandwidth); if (qos_enabled) { qos_stats[qos_class].total_duration.add(total_duration); qos_stats[qos_class].transfer_duration.add(transfer_duration); @@ -227,6 +258,18 @@ int processBatchSizes( if (rc != 0) return -1; if (workload_classes.empty()) printStats(block_size, batch_size, stats, num_threads); + auto target_report = calculateTargetMetrics( + block_size, batch_size, num_threads, XferBenchConfig::backend, + XferBenchConfig::op_type, &target_stats); + if (runner.getTargetCount() > 1) printTargetMetrics(target_report); + if (!XferBenchConfig::result_output_jsonl.empty()) { + std::string error; + if (!appendTargetMetricsJsonl(XferBenchConfig::result_output_jsonl, + target_report, &error)) { + LOG(ERROR) << error; + return -1; + } + } if (!qos_classes.empty()) { std::vector bytes_per_operation; for (const auto& config : workload_classes) { @@ -426,5 +469,5 @@ int main(int argc, char* argv[]) { } runner->stopInitiator(); } - return 0; + return interrupted ? EXIT_FAILURE : EXIT_SUCCESS; } diff --git a/mooncake-transfer-engine/benchmark/target_metrics.cpp b/mooncake-transfer-engine/benchmark/target_metrics.cpp new file mode 100644 index 0000000000..73cb5e1d70 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/target_metrics.cpp @@ -0,0 +1,144 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "target_metrics.h" + +#include +#include +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { + +TargetMetricsReport calculateTargetMetrics( + size_t block_size, size_t batch_size, int num_threads, + const std::string& backend, const std::string& op_type, + std::vector* stats) { + TargetMetricsReport report; + report.block_size = block_size; + report.batch_size = batch_size; + report.num_threads = num_threads; + report.backend = backend; + report.op_type = op_type; + report.targets.reserve(stats->size()); + double aggregate_duration_sum_us = 0.0; + int assigned_threads = 0; + + for (size_t i = 0; i < stats->size(); ++i) { + auto& input = (*stats)[i]; + TargetMetrics metrics; + metrics.index = i; + metrics.segment_name = input.segment_name; + metrics.threads = input.threads; + metrics.operations = input.stats.transfer_duration.count(); + metrics.transferred_bytes = input.transferred_bytes; + metrics.total_duration_us = input.stats.total_duration.avg(); + metrics.avg_transfer_us = input.stats.transfer_duration.avg(); + metrics.p99_us = input.stats.transfer_duration.p99(); + metrics.p999_us = input.stats.transfer_duration.p999(); + metrics.avg_instant_gbps = input.stats.instant_bandwidth.avg(); + if (metrics.operations != 0) { + metrics.avg_latency_us = metrics.total_duration_us * + metrics.threads / metrics.operations; + } + if (metrics.total_duration_us > 0.0) { + metrics.throughput_gbps = + static_cast(metrics.transferred_bytes) / 1000.0 / + metrics.total_duration_us; + } + report.aggregate_operations = + checkedAdd(report.aggregate_operations, metrics.operations, + "aggregate operations"); + report.aggregate_transferred_bytes = checkedAdd( + report.aggregate_transferred_bytes, metrics.transferred_bytes, + "aggregate transferred bytes"); + aggregate_duration_sum_us += + metrics.total_duration_us * metrics.threads; + assigned_threads += metrics.threads; + report.targets.push_back(std::move(metrics)); + } + if (assigned_threads > 0 && aggregate_duration_sum_us > 0.0) { + const double aggregate_duration_us = + aggregate_duration_sum_us / assigned_threads; + report.aggregate_throughput_gbps = + static_cast(report.aggregate_transferred_bytes) / 1000.0 / + aggregate_duration_us; + } + return report; +} + +void printTargetMetrics(const TargetMetricsReport& report) { + for (const auto& metrics : report.targets) { + std::cout << " [target-summary] index=" << metrics.index + << " name=" << metrics.segment_name + << " threads=" << metrics.threads + << " operations=" << metrics.operations + << " transferred_bytes=" << metrics.transferred_bytes + << " throughput=" << std::fixed << std::setprecision(6) + << metrics.throughput_gbps + << " GB/s p99_us=" << std::setprecision(1) << metrics.p99_us + << std::endl; + } +} + +bool appendTargetMetricsJsonl(const std::string& path, + const TargetMetricsReport& report, + std::string* error) { + nlohmann::json root = { + {"schema_version", 1}, + {"record_type", "target_metrics"}, + {"backend", report.backend}, + {"op_type", report.op_type}, + {"block_size", report.block_size}, + {"batch_size", report.batch_size}, + {"num_threads", report.num_threads}, + {"aggregate_operations", report.aggregate_operations}, + {"aggregate_transferred_bytes", report.aggregate_transferred_bytes}, + {"aggregate_throughput_gbps", report.aggregate_throughput_gbps}, + {"targets", nlohmann::json::array()}, + }; + for (const auto& metrics : report.targets) { + root["targets"].push_back({ + {"index", metrics.index}, + {"segment_name", metrics.segment_name}, + {"threads", metrics.threads}, + {"operations", metrics.operations}, + {"transferred_bytes", metrics.transferred_bytes}, + {"total_duration_us", metrics.total_duration_us}, + {"throughput_gbps", metrics.throughput_gbps}, + {"avg_latency_us", metrics.avg_latency_us}, + {"avg_transfer_us", metrics.avg_transfer_us}, + {"p99_us", metrics.p99_us}, + {"p999_us", metrics.p999_us}, + {"avg_instant_gbps", metrics.avg_instant_gbps}, + }); + } + + std::ofstream output(path, std::ios::app); + if (!output) { + *error = "failed to open target JSONL output: " + path; + return false; + } + output << root.dump() << '\n'; + if (!output) { + *error = "failed to write target JSONL output: " + path; + return false; + } + return true; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/target_metrics.h b/mooncake-transfer-engine/benchmark/target_metrics.h new file mode 100644 index 0000000000..333367edf6 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/target_metrics.h @@ -0,0 +1,75 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TEBENCH_TARGET_METRICS_H +#define TEBENCH_TARGET_METRICS_H + +#include +#include +#include + +#include "utils.h" + +namespace mooncake { +namespace tent { + +struct TargetBenchStats { + std::string segment_name; + int threads = 0; + uint64_t transferred_bytes = 0; + XferBenchStats stats; +}; + +struct TargetMetrics { + size_t index = 0; + std::string segment_name; + int threads = 0; + uint64_t operations = 0; + uint64_t transferred_bytes = 0; + double total_duration_us = 0.0; + double throughput_gbps = 0.0; + double avg_latency_us = 0.0; + double avg_transfer_us = 0.0; + double p99_us = 0.0; + double p999_us = 0.0; + double avg_instant_gbps = 0.0; +}; + +struct TargetMetricsReport { + size_t block_size = 0; + size_t batch_size = 0; + int num_threads = 0; + std::string backend; + std::string op_type; + uint64_t aggregate_operations = 0; + uint64_t aggregate_transferred_bytes = 0; + double aggregate_throughput_gbps = 0.0; + std::vector targets; +}; + +TargetMetricsReport calculateTargetMetrics( + size_t block_size, size_t batch_size, int num_threads, + const std::string& backend, const std::string& op_type, + std::vector* stats); + +void printTargetMetrics(const TargetMetricsReport& report); + +bool appendTargetMetricsJsonl(const std::string& path, + const TargetMetricsReport& report, + std::string* error); + +} // namespace tent +} // namespace mooncake + +#endif // TEBENCH_TARGET_METRICS_H diff --git a/mooncake-transfer-engine/benchmark/te_backend.h b/mooncake-transfer-engine/benchmark/te_backend.h index 05896b59f2..29d3ef239a 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.h +++ b/mooncake-transfer-engine/benchmark/te_backend.h @@ -63,6 +63,10 @@ class TEBenchRunner : public BenchRunner { size_t getTargetCount() const; + size_t getTargetIndex(int thread_id) const { + return targetIndex(thread_id); + } + uint64_t getTargetSegmentId(int thread_id) const; uint64_t getTargetBufferBase(int thread_id, uint64_t block_size, diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 8175289bdb..14be065ae1 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.cpp +++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp @@ -15,6 +15,8 @@ #include "tent_backend.h" #include "utils.h" #include "char_util.h" + +#include #include "tent/common/types.h" #include "tent/runtime/platform.h" #include "tent/runtime/topology.h" @@ -31,18 +33,18 @@ namespace mooncake { namespace tent { -volatile bool g_tent_running = true; -volatile bool g_tent_triggered_sig = false; +std::atomic g_tent_running{true}; +std::atomic g_tent_triggered_sig{false}; void signalHandlerV1(int signum) { - if (g_tent_triggered_sig) { + if (g_tent_triggered_sig.load()) { LOG(ERROR) << "Received signal " << signum << " again, forcefully terminating..."; std::exit(EXIT_FAILURE); } LOG(INFO) << "Received signal " << signum << ", stopping target server..."; - g_tent_running = false; - g_tent_triggered_sig = true; + g_tent_running.store(false); + g_tent_triggered_sig.store(true); } std::shared_ptr loadConfig() { @@ -60,11 +62,13 @@ std::shared_ptr loadConfig() { std::unordered_map transport_map = { {"rdma", "rdma"}, {"tcp", "tcp"}, + {"hp_tcp", "hp_tcp"}, {"shm", "shm"}, {"iouring", "io_uring"}, // Note: iouring -> io_uring {"gds", "gds"}, {"mnnvl", "mnnvl"}, {"nvlink", "nvlink"}, + {"ub", "ub"}, {"sunrise_link", "sunrise_link"}, {"mpcomm", "mpcomm"}}; @@ -91,8 +95,10 @@ static TransportType getTransportType(const std::string& xport_type) { if (xport_type == "nvlink") return NVLINK; if (xport_type == "tcp") return TCP; if (xport_type == "iouring") return IOURING; + if (xport_type == "ub") return UB; if (xport_type == "sunrise_link") return SUNRISE_LINK; if (xport_type == "mpcomm") return MPCOMM; + if (xport_type == "hp_tcp") return HP_TCP; return UNSPEC; } @@ -302,7 +308,7 @@ TENTBenchRunner::TENTBenchRunner() { TENTBenchRunner::~TENTBenchRunner() { freeBuffers(); } int TENTBenchRunner::runTarget() { - while (g_tent_running) sleep(1); + while (g_tent_running.load()) sleep(1); return 0; } @@ -350,7 +356,7 @@ int TENTBenchRunner::startInitiator(int num_threads) { LOG(INFO) << "Opened " << target_handles_.size() << " target segments"; threads_.resize(num_threads); current_task_.resize(threads_.size()); - g_tent_running = true; + g_tent_running.store(true); for (size_t i = 0; i < threads_.size(); ++i) threads_[i] = std::thread(&TENTBenchRunner::runner, this, i); return 0; @@ -359,7 +365,7 @@ int TENTBenchRunner::startInitiator(int num_threads) { int TENTBenchRunner::stopInitiator() { { std::unique_lock lk(mtx_); - g_tent_running = false; + g_tent_running.store(false); cv_task_.notify_all(); cv_done_.notify_all(); } @@ -434,14 +440,14 @@ void TENTBenchRunner::pinThread(int thread_id) { } int TENTBenchRunner::runner(int thread_id) { - while (g_tent_running) { + while (g_tent_running.load()) { std::function task; { std::unique_lock lk(mtx_); cv_task_.wait(lk, [&] { - return !g_tent_running || current_task_[thread_id]; + return !g_tent_running.load() || current_task_[thread_id]; }); - if (!g_tent_running) break; + if (!g_tent_running.load()) break; std::swap(task, current_task_[thread_id]); } if (task) task(thread_id); @@ -460,15 +466,37 @@ int TENTBenchRunner::runInitiatorTasks( current_task_[id] = func; pending_ = (int)threads_.size(); cv_task_.notify_all(); - cv_done_.wait(lk, [&] { return !g_tent_running || pending_ == 0; }); - return g_tent_running ? 0 : -1; + cv_done_.wait(lk, [&] { return !g_tent_running.load() || pending_ == 0; }); + return g_tent_running.load() ? 0 : -1; +} + +void TENTBenchRunner::noteTransferFailed() { + std::lock_guard lk(mtx_); + g_tent_running.store(false); + cv_task_.notify_all(); + cv_done_.notify_all(); } double TENTBenchRunner::runSingleTransfer( uint64_t local_addr, uint64_t target_id, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, OpCode opcode, uint64_t deadline_ns, IntentType intent_type) { + if (!g_tent_running.load()) return -1.0; + + auto abortTransfer = [&](const char* why) { + LOG(ERROR) << why; + noteTransferFailed(); + return -1.0; + }; + auto batch_id = engine_->allocateBatch(batch_size); + if (!batch_id) return abortTransfer("Failed to allocate transfer batch"); + + auto finish = [&](double duration) { + (void)engine_->freeBatch(batch_id); + return duration; + }; + std::vector requests; for (uint64_t i = 0; i < batch_size; ++i) { Request entry; @@ -485,26 +513,39 @@ double TENTBenchRunner::runSingleTransfer( requests.emplace_back(entry); } XferBenchTimer timer; + Status submitted; if (XferBenchConfig::notifi) { - // Use target_addr as msg for verification by peer Notification notifi{"benchmark", std::to_string(target_addr)}; - CHECK_FAIL(engine_->submitTransfer(batch_id, requests, notifi)); + submitted = engine_->submitTransfer(batch_id, requests, notifi); } else { - CHECK_FAIL(engine_->submitTransfer(batch_id, requests)); + submitted = engine_->submitTransfer(batch_id, requests); } - while (true) { + if (!submitted.ok()) { + LOG(ERROR) << "Failed to submit transfer: " << submitted.ToString(); + noteTransferFailed(); + return finish(-1.0); + } + while (g_tent_running.load()) { TransferStatus overall_status; - CHECK_FAIL(engine_->getTransferStatus(batch_id, overall_status)); + auto polled = engine_->getTransferStatus(batch_id, overall_status); + if (!polled.ok()) { + LOG(ERROR) << "Failed to poll transfer: " << polled.ToString(); + noteTransferFailed(); + return finish(-1.0); + } if (overall_status.s == TransferStatusEnum::COMPLETED) { - break; - } else if (overall_status.s == TransferStatusEnum::FAILED) { + return finish(timer.lap_us()); + } + if (overall_status.s == TransferStatusEnum::FAILED || + overall_status.s == TransferStatusEnum::TIMEOUT || + overall_status.s == TransferStatusEnum::CANCELED || + overall_status.s == TransferStatusEnum::INVALID) { LOG(ERROR) << "Failed transfer detected"; - exit(EXIT_FAILURE); + noteTransferFailed(); + return finish(-1.0); } } - auto duration = timer.lap_us(); - CHECK_FAIL(engine_->freeBatch(batch_id)); - return duration; + return finish(-1.0); } } // namespace tent diff --git a/mooncake-transfer-engine/benchmark/tent_backend.h b/mooncake-transfer-engine/benchmark/tent_backend.h index 4f6cba5ffc..15e3202655 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.h +++ b/mooncake-transfer-engine/benchmark/tent_backend.h @@ -89,6 +89,10 @@ class TENTBenchRunner : public BenchRunner { size_t getTargetCount() const; + size_t getTargetIndex(int thread_id) const { + return targetIndex(thread_id); + } + uint64_t getTargetSegmentId(int thread_id) const; uint64_t getTargetBufferBase(int thread_id, uint64_t block_size, @@ -166,6 +170,9 @@ class TENTBenchRunner : public BenchRunner { "target address"); } + // Returns elapsed microseconds on success, or a negative value if the + // transfer failed. Never calls exit() — workers set g_tent_running so + // main can join threads and stop the engine first. double runSingleTransfer(uint64_t local_addr, uint64_t target_id, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, OpCode opcode, @@ -178,6 +185,8 @@ class TENTBenchRunner : public BenchRunner { int runner(int thread_id); + void noteTransferFailed(); + size_t targetIndex(int thread_id) const; int localTargetThreadId(int thread_id) const; diff --git a/mooncake-transfer-engine/benchmark/tests/target_metrics_test.cpp b/mooncake-transfer-engine/benchmark/tests/target_metrics_test.cpp new file mode 100644 index 0000000000..0c43c426a3 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/tests/target_metrics_test.cpp @@ -0,0 +1,66 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "target_metrics.h" + +#include +#include + +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { +namespace { + +TEST(TargetMetricsTest, ReportsEachTargetAndWritesJsonl) { + std::vector stats(2); + stats[0].segment_name = "target-a"; + stats[0].threads = 2; + stats[0].transferred_bytes = 6000; + stats[0].stats.total_duration.add(1000.0); + stats[0].stats.total_duration.add(1000.0); + stats[0].stats.transfer_duration.add({10.0, 20.0, 30.0}); + stats[0].stats.instant_bandwidth.add({0.1, 0.2, 0.3}); + stats[1].segment_name = "target-b"; + + const auto report = + calculateTargetMetrics(1000, 2, 2, "tent", "read", &stats); + ASSERT_EQ(report.targets.size(), 2u); + EXPECT_EQ(report.aggregate_operations, 3u); + EXPECT_EQ(report.aggregate_transferred_bytes, 6000u); + EXPECT_NEAR(report.aggregate_throughput_gbps, 0.006, 1e-12); + EXPECT_EQ(report.targets[0].threads, 2); + EXPECT_NEAR(report.targets[0].avg_latency_us, 2000.0 / 3.0, 1e-12); + EXPECT_DOUBLE_EQ(report.targets[1].throughput_gbps, 0.0); + + const std::string path = "tebench_target_metrics_test.jsonl"; + std::remove(path.c_str()); + std::string error; + ASSERT_TRUE(appendTargetMetricsJsonl(path, report, &error)) << error; + std::ifstream input(path); + nlohmann::json record; + ASSERT_NO_THROW(input >> record); + EXPECT_EQ(record["schema_version"], 1); + EXPECT_EQ(record["record_type"], "target_metrics"); + ASSERT_EQ(record["targets"].size(), 2u); + EXPECT_EQ(record["targets"][0]["segment_name"], "target-a"); + EXPECT_EQ(record["targets"][1]["operations"], 0); + std::remove(path.c_str()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp index 25cbb3d417..05be1d488e 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/mooncake-transfer-engine/benchmark/utils.cpp @@ -71,6 +71,9 @@ DEFINE_double(qos_link_capacity_gbps, 0.0, "Link capacity in GB/s for total utilization (0 reports N/A)."); DEFINE_string(qos_output_jsonl, "", "Append versioned QoS metric records to this JSONL file."); +DEFINE_string(result_output_jsonl, "", + "Append versioned benchmark result records, including " + "per-target metrics, to this JSONL file."); DEFINE_uint64(request_interval_us, 0, "Per-thread delay before issuing each transfer batch, in " "microseconds. 0 disables pacing."); @@ -95,14 +98,15 @@ DEFINE_int32( "RPC server port used for p2p metadata service (0 = auto-select)."); DEFINE_string(xport_type, "", "Transport type: " - "rdma|shm|mnnvl|gds|iouring|sunrise_link|mpcomm|flagcx"); + "rdma|tcp|hp_tcp|shm|mnnvl|nvlink|gds|iouring|ub|sunrise_link|" + "mpcomm|flagcx"); DEFINE_string(backend, "tent", "Transport backend: classic|tent"); DEFINE_bool(notifi, false, "Enable RDMA notification for performance measurement."); -DEFINE_string( - tent_transport_hint, "unspec", - "tent only: per-request transport_hint. " - "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|sunrise_link|mpcomm"); +DEFINE_string(tent_transport_hint, "unspec", + "tent only: per-request transport_hint. " + "unspec|rdma|tcp|hp_tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|" + "ub|sunrise_link|mpcomm"); DEFINE_string(tent_intent_type, "unspec", "tent only: intent_type attached to every benchmark request. " "unspec|foreground_get|background_prefetch|migration|checkpoint|" @@ -132,6 +136,7 @@ std::string XferBenchConfig::qos_classes_json; std::string XferBenchConfig::workload_classes_json; double XferBenchConfig::qos_link_capacity_gbps = 0.0; std::string XferBenchConfig::qos_output_jsonl; +std::string XferBenchConfig::result_output_jsonl; uint64_t XferBenchConfig::request_interval_us = 0; uint64_t XferBenchConfig::deadline_us = 0; int XferBenchConfig::deadline_tight_threads = 0; @@ -171,6 +176,7 @@ void XferBenchConfig::loadFromFlags() { workload_classes_json = FLAGS_workload_classes_json; qos_link_capacity_gbps = FLAGS_qos_link_capacity_gbps; qos_output_jsonl = FLAGS_qos_output_jsonl; + result_output_jsonl = FLAGS_result_output_jsonl; request_interval_us = FLAGS_request_interval_us; deadline_us = FLAGS_deadline_us; deadline_tight_threads = FLAGS_deadline_tight_threads; diff --git a/mooncake-transfer-engine/benchmark/utils.h b/mooncake-transfer-engine/benchmark/utils.h index 1d0deb9c91..4c5ef70b72 100644 --- a/mooncake-transfer-engine/benchmark/utils.h +++ b/mooncake-transfer-engine/benchmark/utils.h @@ -79,6 +79,7 @@ struct XferBenchConfig { static std::string workload_classes_json; static double qos_link_capacity_gbps; static std::string qos_output_jsonl; + static std::string result_output_jsonl; static uint64_t request_interval_us; static uint64_t deadline_us; static int deadline_tight_threads; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h index f3bd20dd01..2ea3b754b0 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h @@ -41,7 +41,7 @@ class EndpointStore { virtual std::shared_ptr getEndpointByPtr( const RdmaEndPoint *endpoint_ptr) = 0; virtual std::shared_ptr insertEndpoint( - const std::string &peer_nic_path, RdmaContext *context) = 0; + const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) = 0; virtual int deleteEndpoint(const std::string &peer_nic_path) = 0; // Deletes the endpoint matching endpoint_ptr (by pointer identity, under // the store lock -- the pointer is never dereferenced, so a stale/freed @@ -83,7 +83,8 @@ class FIFOEndpointStore : public EndpointStore { std::shared_ptr getEndpointByPtr( const RdmaEndPoint *endpoint_ptr) override; std::shared_ptr insertEndpoint( - const std::string &peer_nic_path, RdmaContext *context) override; + const std::string &peer_nic_path, RdmaContext *context, + ibv_cq *cq) override; int deleteEndpoint(const std::string &peer_nic_path) override; int deleteEndpointByPtr( const RdmaEndPoint *endpoint_ptr, @@ -125,7 +126,8 @@ class SIEVEEndpointStore : public EndpointStore { std::shared_ptr getEndpointByPtr( const RdmaEndPoint *endpoint_ptr) override; std::shared_ptr insertEndpoint( - const std::string &peer_nic_path, RdmaContext *context) override; + const std::string &peer_nic_path, RdmaContext *context, + ibv_cq *cq) override; int deleteEndpoint(const std::string &peer_nic_path) override; int deleteEndpointByPtr( const RdmaEndPoint *endpoint_ptr, @@ -143,6 +145,11 @@ class SIEVEEndpointStore : public EndpointStore { } void testOnlyInsertWaiting(std::shared_ptr ep) override; + // Test-only: push a pre-constructed endpoint into the active map so + // pointer-identity lookup/delete paths can be exercised without standing up + // an RDMA device. + void testOnlyInsertEndpoint(const std::string &peer_nic_path, + std::shared_ptr ep); private: RWSpinlock endpoint_map_lock_; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h index 370d6c2660..c10a47c6d3 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h @@ -173,6 +173,10 @@ class RdmaContext { public: // EndPoint Management std::shared_ptr endpoint(const std::string &peer_nic_path); + std::shared_ptr endpoint(const std::string &peer_nic_path, + int cq_index); + std::shared_ptr findEndpoint( + const std::string &peer_nic_path); std::shared_ptr getEndpointByPtr( const RdmaEndPoint *endpoint_ptr); @@ -260,13 +264,20 @@ class RdmaContext { int eventFd() const { return event_fd_; } - ibv_cq *cq(); + ibv_cq *cq(int cq_index); std::atomic *cqOutstandingCount(int cq_index) { return &cq_list_[cq_index].outstanding; } int cqCount() const { return cq_list_.size(); } + int postingThreadForPeer(const std::string &peer_nic_path) const; + int cqIndexForPostingThread(int thread_id) const; + int cqIndexForPeer(const std::string &peer_nic_path) const; + int transferWorkerCount() const { return transfer_worker_count_; } + std::unique_lock lockEndpointLifecycle( + const std::string &peer_nic_path) const; + std::vector> lockAllEndpointLifecycles() const; int poll(int num_entries, ibv_wc *wc, int cq_index = 0); @@ -329,7 +340,9 @@ class RdmaContext { std::atomic next_comp_channel_index_; std::atomic next_comp_vector_index_; - std::atomic next_cq_list_index_; + + int transfer_worker_count_ = 0; + std::vector> endpoint_lifecycle_locks_; std::shared_ptr worker_pool_; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h index c9dd304b8c..8011ba114d 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h @@ -71,6 +71,7 @@ class RdmaEndPoint { public: void setPeerNicPath(const std::string &peer_nic_path); + std::string peerNicPath() const; int setupConnectionsByActive(); @@ -197,9 +198,10 @@ class RdmaEndPoint { static constexpr uint32_t kWaitExistingHandshakeInitialSleepUs = 50; static constexpr uint32_t kWaitExistingHandshakeMaxSleepUs = 2000; - // Maximum time (in seconds) to wait for outstanding WRs to drain in - // finishDestroy before forcing QP destruction. This guards against - // ibv_modify_qp-to-ERR failures that prevent WR flushing. + // Maximum time (in seconds) to wait for outstanding WRs to drain before + // treating the endpoint as leaked. Timed-out endpoints stay in the + // EndpointStore waiting list so stale in-flight references cannot turn + // into UAF. static constexpr double kFinishDestroyTimeoutSec = 30.0; // Maximum number of deconstructLocked retries in finishDestroy before @@ -210,7 +212,7 @@ class RdmaEndPoint { RdmaContext &context_; std::atomic status_; - RWSpinlock lock_; + mutable RWSpinlock lock_; std::vector qp_list_; uint64_t qp_generation_; @@ -225,8 +227,10 @@ class RdmaEndPoint { size_t max_inline_bytes_; std::atomic active_; + ibv_cq *cq_; std::atomic *cq_outstanding_; std::atomic inactive_time_; + bool finish_destroy_timeout_logged_ = false; int finish_destroy_retries_ = 0; }; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h index 5b1a8e0968..d65633ee69 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -114,9 +114,9 @@ class RdmaTransport : public Transport { virtual int onSetupRdmaConnections(const HandShakeDesc &peer_desc, HandShakeDesc &local_desc); - int sendHandshake(const std::string &peer_server_name, - const HandShakeDesc &local_desc, - HandShakeDesc &peer_desc) { + virtual int sendHandshake(const std::string &peer_server_name, + const HandShakeDesc &local_desc, + HandShakeDesc &peer_desc) { return metadata_->sendHandshake(peer_server_name, local_desc, peer_desc); } diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h index e7aad9fbd9..60266fb760 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h @@ -15,8 +15,13 @@ #ifndef WORKER_H #define WORKER_H -#include +#include +#include +#include +#include +#include #include +#include #include "config.h" #include "rdma_context.h" @@ -41,7 +46,6 @@ class WorkerPool { private: using SliceList = std::vector; - const static int kShardCount = 8; // Enqueue slices that were prepared by another WorkerPool. Used for // local-NIC failure handoff: the original worker keeps the remote path @@ -49,12 +53,16 @@ class WorkerPool { // worker queue. int submitPreparedPostSend( const std::vector &slice_list); - void enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount], + void enqueuePreparedSlices(const SliceList &slice_list, uint64_t submitted_slice_count); + void enqueueSliceToOwner(Transport::Slice *slice); + int postingThreadForPeer(const std::string &peer_nic_path) const; + int cqIndexForPostingThread(int thread_id) const; void performPostSend(int thread_id); void performPollCq(int thread_id); + void processCompletions(int thread_id, const std::vector &wc_list); void redispatch(std::vector &slice_list, int thread_id, bool handoff_to_local_worker = false); @@ -130,6 +138,7 @@ class WorkerPool { private: RdmaContext &context_; const int numa_socket_id_; + const int worker_count_; std::vector worker_thread_; std::atomic workers_running_; @@ -151,12 +160,10 @@ class WorkerPool { std::mutex cond_mutex_; std::condition_variable cond_var_; - std::unordered_map slice_queue_[kShardCount]; - std::atomic slice_queue_count_[kShardCount]; - TicketLock slice_queue_lock_[kShardCount]; - std::vector> collective_slice_queue_; + std::vector> worker_slice_queue_; + std::vector worker_slice_queue_lock_; std::atomic submitted_slice_count_, processed_slice_count_; std::atomic recovery_activate_after_ns_{0}; diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp index 54db6de43a..4385983c1f 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp @@ -47,22 +47,26 @@ std::shared_ptr FIFOEndpointStore::getEndpointByPtr( } std::shared_ptr FIFOEndpointStore::insertEndpoint( - const std::string &peer_nic_path, RdmaContext *context) { + const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) { RWSpinlock::WriteGuard guard(endpoint_map_lock_); if (endpoint_map_.find(peer_nic_path) != endpoint_map_.end()) { LOG(INFO) << "Endpoint " << peer_nic_path << " already exists in FIFOEndpointStore"; return endpoint_map_[peer_nic_path]; } + if (!cq) { + LOG(ERROR) << "Cannot insert endpoint " << peer_nic_path + << ": completion queue is null"; + return nullptr; + } auto endpoint = std::make_shared(*context); if (!endpoint) { LOG(ERROR) << "Failed to allocate memory for RdmaEndPoint"; return nullptr; } auto &config = globalConfig(); - int ret = - endpoint->construct(context->cq(), config.num_qp_per_ep, config.max_sge, - config.max_wr, config.max_inline); + int ret = endpoint->construct(cq, config.num_qp_per_ep, config.max_sge, + config.max_wr, config.max_inline); if (ret) return nullptr; while (this->getSize() >= max_size_) evictEndpoint(); @@ -80,8 +84,9 @@ int FIFOEndpointStore::deleteEndpoint(const std::string &peer_nic_path) { auto iter = endpoint_map_.find(peer_nic_path); // Begin two-phase destruction: mark endpoint as destroying and move QPs // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved - // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once - // all outstanding WRs have been drained. + // to waiting_list_ and will be fully destroyed by reclaimEndpoint() only + // after all outstanding WRs have been drained. Timed-out endpoints remain + // retired in waiting_list_ rather than being force-freed. if (iter != endpoint_map_.end()) { waiting_list_len_++; iter->second->beginDestroy(); @@ -220,22 +225,26 @@ std::shared_ptr SIEVEEndpointStore::getEndpointByPtr( } std::shared_ptr SIEVEEndpointStore::insertEndpoint( - const std::string &peer_nic_path, RdmaContext *context) { + const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) { RWSpinlock::WriteGuard guard(endpoint_map_lock_); if (endpoint_map_.find(peer_nic_path) != endpoint_map_.end()) { LOG(INFO) << "Endpoint " << peer_nic_path << " already exists in SIEVEEndpointStore"; return endpoint_map_[peer_nic_path].first; } + if (!cq) { + LOG(ERROR) << "Cannot insert endpoint " << peer_nic_path + << ": completion queue is null"; + return nullptr; + } auto endpoint = std::make_shared(*context); if (!endpoint) { LOG(ERROR) << "Failed to allocate memory for RdmaEndPoint"; return nullptr; } auto &config = globalConfig(); - int ret = - endpoint->construct(context->cq(), config.num_qp_per_ep, config.max_sge, - config.max_wr, config.max_inline); + int ret = endpoint->construct(cq, config.num_qp_per_ep, config.max_sge, + config.max_wr, config.max_inline); if (ret) return nullptr; while (this->getSize() >= max_size_) evictEndpoint(); @@ -252,8 +261,9 @@ int SIEVEEndpointStore::deleteEndpoint(const std::string &peer_nic_path) { auto iter = endpoint_map_.find(peer_nic_path); // Begin two-phase destruction: mark endpoint as destroying and move QPs // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved - // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once - // all outstanding WRs have been drained. + // to waiting_list_ and will be fully destroyed by reclaimEndpoint() only + // after all outstanding WRs have been drained. Timed-out endpoints remain + // retired in waiting_list_ rather than being force-freed. if (iter != endpoint_map_.end()) { iter->second.first->beginDestroy(); waiting_list_len_++; @@ -377,6 +387,14 @@ void SIEVEEndpointStore::testOnlyInsertWaiting( waiting_list_len_++; } +void SIEVEEndpointStore::testOnlyInsertEndpoint( + const std::string &peer_nic_path, std::shared_ptr ep) { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + endpoint_map_[peer_nic_path] = std::make_pair(ep, true); + fifo_list_.push_front(peer_nic_path); + fifo_map_[peer_nic_path] = fifo_list_.begin(); +} + size_t SIEVEEndpointStore::getTotalQPNumber() { RWSpinlock::ReadGuard guard(endpoint_map_lock_); size_t total_qps = 0; diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp index 6c38dfbfa0..ad32b56f03 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -166,7 +167,6 @@ RdmaContext::RdmaContext(RdmaTransport &engine, const std::string &device_name) [] { return static_cast(getCurrentTimeInNano()); }), next_comp_channel_index_(0), next_comp_vector_index_(0), - next_cq_list_index_(0), worker_pool_(nullptr), active_(true) { static std::once_flag g_once_flag; @@ -195,6 +195,23 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels, // Create endpoint store based on configuration auto &config = globalConfig(); + if (config.workers_per_ctx <= 0) { + LOG(ERROR) << "Invalid workers_per_ctx=" << config.workers_per_ctx + << " for device " << device_name_; + return ERR_INVALID_ARGUMENT; + } + transfer_worker_count_ = config.workers_per_ctx; + if (num_cq_list < static_cast(transfer_worker_count_)) { + LOG(INFO) << "Increasing RDMA CQ count for " << device_name_ << " from " + << num_cq_list << " to " << transfer_worker_count_ + << " to keep each transfer worker on a dedicated CQ"; + num_cq_list = static_cast(transfer_worker_count_); + } + endpoint_lifecycle_locks_.clear(); + endpoint_lifecycle_locks_.reserve( + static_cast(transfer_worker_count_)); + for (int i = 0; i < transfer_worker_count_; ++i) + endpoint_lifecycle_locks_.push_back(std::make_unique()); switch (config.endpoint_store_type) { case EndpointStoreType::FIFO: endpoint_store_ = @@ -724,6 +741,18 @@ RdmaContext::findMemoryRegionContaining(uintptr_t addr) const { std::shared_ptr RdmaContext::endpoint( const std::string &peer_nic_path) { + int cq_index = cqIndexForPeer(peer_nic_path); + if (cq_index < 0) return nullptr; + return endpoint(peer_nic_path, cq_index); +} + +std::shared_ptr RdmaContext::endpoint( + const std::string &peer_nic_path, int cq_index) { + if (cq_list_.empty()) { + LOG(ERROR) << "No CQ available for endpoint on " << deviceName(); + return nullptr; + } + if (!active_.load(std::memory_order_acquire)) { LOG(ERROR) << "Context is not active: " << deviceName(); return nullptr; @@ -739,11 +768,18 @@ std::shared_ptr RdmaContext::endpoint( return endpoint; } - endpoint = endpoint_store_->insertEndpoint(peer_nic_path, this); + endpoint = + endpoint_store_->insertEndpoint(peer_nic_path, this, cq(cq_index)); endpoint_store_->reclaimEndpoint(); return endpoint; } +std::shared_ptr RdmaContext::findEndpoint( + const std::string &peer_nic_path) { + if (!endpoint_store_) return nullptr; + return endpoint_store_->getEndpoint(peer_nic_path); +} + std::shared_ptr RdmaContext::getEndpointByPtr( const RdmaEndPoint *endpoint_ptr) { return endpoint_store_->getEndpointByPtr(endpoint_ptr); @@ -822,9 +858,53 @@ int RdmaContext::gidIndex() const { return gid_index_; } -ibv_cq *RdmaContext::cq() { - int index = (next_cq_list_index_++) % cq_list_.size(); - return cq_list_[index].native; +ibv_cq *RdmaContext::cq(int cq_index) { + if (cq_list_.empty()) return nullptr; + if (cq_index < 0) return nullptr; + return cq_list_[static_cast(cq_index) % cq_list_.size()].native; +} + +int RdmaContext::postingThreadForPeer(const std::string &peer_nic_path) const { + if (transfer_worker_count_ <= 0) { + LOG(ERROR) << "Invalid transfer_worker_count_=" + << transfer_worker_count_ << " for endpoint on " + << deviceName(); + return -1; + } + return static_cast(std::hash{}(peer_nic_path) % + static_cast(transfer_worker_count_)); +} + +int RdmaContext::cqIndexForPostingThread(int thread_id) const { + const int cq_count = cqCount(); + if (cq_count <= 0 || thread_id < 0) return -1; + if (thread_id < cq_count) return thread_id; + return thread_id % cq_count; +} + +int RdmaContext::cqIndexForPeer(const std::string &peer_nic_path) const { + return cqIndexForPostingThread(postingThreadForPeer(peer_nic_path)); +} + +std::unique_lock RdmaContext::lockEndpointLifecycle( + const std::string &peer_nic_path) const { + const int owner_thread = postingThreadForPeer(peer_nic_path); + if (owner_thread < 0 || + static_cast(owner_thread) >= endpoint_lifecycle_locks_.size()) { + return std::unique_lock(); + } + return std::unique_lock( + *endpoint_lifecycle_locks_[owner_thread]); +} + +std::vector> +RdmaContext::lockAllEndpointLifecycles() const { + std::vector> locks; + locks.reserve(endpoint_lifecycle_locks_.size()); + for (auto &lifecycle_lock : endpoint_lifecycle_locks_) { + locks.emplace_back(*lifecycle_lock); + } + return locks; } ibv_comp_channel *RdmaContext::compChannel() { diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp index 3da94a4b23..72c2cea291 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp @@ -97,6 +97,7 @@ RdmaEndPoint::RdmaEndPoint(RdmaContext &context) ready_wait_start_ts_(0), wr_depth_list_(nullptr), active_(true), + cq_(nullptr), cq_outstanding_(nullptr) {} RdmaEndPoint::~RdmaEndPoint() { @@ -118,6 +119,7 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list, } qp_list_.resize(num_qp_list); + cq_ = cq; cq_outstanding_ = static_cast *>(cq->cq_context); max_wr_depth_ = (int)max_wr_depth; @@ -172,9 +174,7 @@ int RdmaEndPoint::reconstruct() { return ret; } - // Get CQ from context for reconstruction - ibv_cq *cq = context_.cq(); - if (!cq) { + if (!cq_) { LOG(ERROR) << "No CQ available for endpoint reconstruction"; return ERR_ENDPOINT; } @@ -184,7 +184,7 @@ int RdmaEndPoint::reconstruct() { ready_wait_start_ts_.store(0, std::memory_order_relaxed); active_.store(true, std::memory_order_release); - return construct(cq, num_qp, max_sge_per_wr, max_wr_depth, + return construct(cq_, num_qp, max_sge_per_wr, max_wr_depth, max_inline_bytes); } @@ -218,7 +218,8 @@ int RdmaEndPoint::deconstructLocked() { if (!qp_list_[i]) continue; // already destroyed in a previous call int ret = ibv_destroy_qp(qp_list_[i]); if (ret) { - LOG(ERROR) << "Failed to destroy QP[" << i << "]: " << strerror(ret); + LOG(ERROR) << "Failed to destroy QP[" << i + << "]: " << strerror(ret); result = ERR_ENDPOINT; } else { qp_list_[i] = nullptr; @@ -299,8 +300,9 @@ bool RdmaEndPoint::finishDestroy() { // Fall through to the unified destroy path. } else { // Gate 3: two-phase path. Wait for inflight WRs to drain via CQ - // polling. If ibv_modify_qp-to-ERR failed in beginDestroy, WRs may - // never be flushed; enforce a timeout to avoid leaking forever. + // polling. If they never drain, keep the retired endpoint alive in + // the waiting list. Leaking this object is safer than forcing QP + // destruction while stale slice references may still exist. bool has_outstanding = false; for (size_t i = 0; i < qp_list_.size(); ++i) { if (wr_depth_list_[i].load(std::memory_order_relaxed) != 0) { @@ -313,8 +315,13 @@ bool RdmaEndPoint::finishDestroy() { inactive_time_.load(std::memory_order_relaxed)) / 1e9; if (elapsed < kFinishDestroyTimeoutSec) return false; - LOG(WARNING) << "finishDestroy timed out after " << elapsed - << "s with outstanding WRs, forcing destruction"; + if (!finish_destroy_timeout_logged_) { + LOG(ERROR) << "finishDestroy timed out after " << elapsed + << "s with outstanding WRs; keeping retired " + "endpoint alive to avoid UAF"; + finish_destroy_timeout_logged_ = true; + } + return false; } } @@ -347,6 +354,11 @@ void RdmaEndPoint::setPeerNicPath(const std::string &peer_nic_path) { peer_nic_path_ = peer_nic_path; } +std::string RdmaEndPoint::peerNicPath() const { + RWSpinlock::ReadGuard guard(lock_); + return peer_nic_path_; +} + int RdmaEndPoint::setupConnectionsByActive() { HandShakeDesc local_desc, peer_desc; std::string peer_server_name, peer_nic_name; @@ -1245,8 +1257,7 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, LOG(ERROR) << "[Handshake] " << message << ": local=" << context_.nicPath() << ", peer=" << peer_nic_path_ << ", qp_index=" << qp_index - << ", local_qp=" << qp->qp_num - << ", peer_qp=" << peer_qp_num + << ", local_qp=" << qp->qp_num << ", peer_qp=" << peer_qp_num << ", local_gid=" << context_.gid() << ", local_gid_index=" << local_gid_index << ", peer_gid=" << gidToString(peer_gid) diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index ebbd5a0e6c..e2850d9ae2 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -1069,6 +1069,8 @@ int RdmaTransport::onSetupRdmaConnections(const HandShakeDesc &peer_desc, } // Use existing endpoint or create new one. + auto endpoint_lifecycle_lock = + context->lockEndpointLifecycle(peer_desc.local_nic_path); auto endpoint = context->endpoint(peer_desc.local_nic_path); if (!endpoint) { local_desc.reply_msg = "Local RDMA endpoint unavailable for " + @@ -1106,10 +1108,16 @@ int RdmaTransport::initializeRdmaResources() { for (auto &device_name : hca_list) { auto context = std::make_shared(*this, device_name); auto &config = globalConfig(); - int ret = context->construct(config.num_cq_per_ctx, - config.num_comp_channels_per_ctx, - config.port, config.gid_index, - config.max_cqe, config.max_ep_per_ctx); + size_t cq_per_ctx = config.num_cq_per_ctx; + if (cq_per_ctx < static_cast(config.workers_per_ctx)) { + cq_per_ctx = static_cast(config.workers_per_ctx); + LOG(INFO) << "Increasing RDMA CQ count for " << device_name + << " to match workers_per_ctx=" << config.workers_per_ctx + << " for worker-owned endpoint polling"; + } + int ret = context->construct( + cq_per_ctx, config.num_comp_channels_per_ctx, config.port, + config.gid_index, config.max_cqe, config.max_ep_per_ctx); if (ret) { local_topology_->disableDevice(device_name); LOG(WARNING) << "Disable device " << device_name; diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp index a70434642c..8d17f211bb 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp @@ -31,8 +31,6 @@ namespace mooncake { -const static int kTransferWorkerCount = globalConfig().workers_per_ctx; - static std::string resolveBufferLocation( const TransferMetadata::BufferDesc &buffer, uint64_t offset) { std::string location = buffer.name; @@ -49,6 +47,26 @@ static const std::string &sourceLocationOrUnknown(Transport::Slice *slice) { return slice->source_location.empty() ? kUnknown : slice->source_location; } +struct ActiveEndpointSetupResult { + int ret = 0; + bool endpoint_current = false; +}; + +static ActiveEndpointSetupResult setupEndpointByActiveOutsideLifecycleGate( + RdmaContext &context, const std::string &peer_nic_path, + const std::shared_ptr &endpoint, + std::unique_lock &endpoint_lifecycle_lock) { + const bool had_lifecycle_gate = endpoint_lifecycle_lock.owns_lock(); + if (had_lifecycle_gate) endpoint_lifecycle_lock.unlock(); + + int ret = endpoint->setupConnectionsByActive(); + + if (had_lifecycle_gate) endpoint_lifecycle_lock.lock(); + auto current_endpoint = context.findEndpoint(peer_nic_path); + return {ret, current_endpoint.get() == endpoint.get() && + endpoint->active() && !endpoint->retired()}; +} + static int selectPeerDevice(RdmaTransport::SegmentDesc *peer_segment_desc, uint64_t offset, size_t length, const std::string &local_hca, int &buffer_id, @@ -95,44 +113,50 @@ static int selectPeerDevice(RdmaTransport::SegmentDesc *peer_segment_desc, return 0; } -static bool workerCanPost(int thread_id) { - return kTransferWorkerCount == 1 || thread_id != 0; -} - -static bool workerCanPoll(int thread_id) { - return kTransferWorkerCount == 1 || thread_id == 0; -} - -static void getPostingShardAssignment(int thread_id, int &post_tid, - int &post_count) { - assert(workerCanPost(thread_id)); - if (kTransferWorkerCount > 1) { - post_tid = thread_id - 1; - post_count = kTransferWorkerCount - 1; - } else { - post_tid = thread_id; - post_count = kTransferWorkerCount; - } -} - WorkerPool::WorkerPool(RdmaContext &context, int numa_socket_id) : context_(context), numa_socket_id_(numa_socket_id), + worker_count_(context.transferWorkerCount()), workers_running_(true), parked_worker_count_(0), redispatch_counter_(0), + worker_slice_queue_(worker_count_), + worker_slice_queue_lock_(worker_count_), submitted_slice_count_(0), processed_slice_count_(0) { - for (int i = 0; i < kShardCount; ++i) - slice_queue_count_[i].store(0, std::memory_order_relaxed); - collective_slice_queue_.resize(kTransferWorkerCount); - for (int i = 0; i < kTransferWorkerCount; ++i) + collective_slice_queue_.resize(worker_count_); + for (int i = 0; i < worker_count_; ++i) worker_thread_.emplace_back( std::thread(std::bind(&WorkerPool::transferWorker, this, i))); worker_thread_.emplace_back( std::thread(std::bind(&WorkerPool::monitorWorker, this))); } +int WorkerPool::postingThreadForPeer(const std::string &peer_nic_path) const { + return context_.postingThreadForPeer(peer_nic_path); +} + +int WorkerPool::cqIndexForPostingThread(int thread_id) const { + return context_.cqIndexForPostingThread(thread_id); +} + +void WorkerPool::enqueueSliceToOwner(Transport::Slice *slice) { + const int owner_thread = postingThreadForPeer(slice->peer_nic_path); + if (owner_thread < 0 || owner_thread >= worker_count_) { + LOG(ERROR) << "Invalid RDMA worker owner " << owner_thread + << " for peer " << slice->peer_nic_path; + slice->markFailed(); + processed_slice_count_.fetch_add(1); + return; + } + { + std::lock_guard lock( + worker_slice_queue_lock_[owner_thread]); + worker_slice_queue_[owner_thread][slice->peer_nic_path].push_back( + slice); + } +} + WorkerPool::~WorkerPool() { if (workers_running_) { cond_var_.notify_all(); @@ -179,7 +203,7 @@ int WorkerPool::submitPostSend( } #endif // CONFIG_CACHE_SEGMENT_DESC - SliceList slice_list_map[kShardCount]; + SliceList prepared_slice_list; uint64_t submitted_slice_count = 0; int all_rails_failed_count = 0; thread_local std::unordered_map failed_target_ids; @@ -272,12 +296,11 @@ int WorkerPool::submitPostSend( << reinterpret_cast(slice->rdma.dest_addr) << ", length=" << slice->length; } - int shard_id = (slice->target_id * 10007 + device_id) % kShardCount; - slice_list_map[shard_id].push_back(slice); + prepared_slice_list.push_back(slice); submitted_slice_count++; } - enqueuePreparedSlices(slice_list_map, submitted_slice_count); + enqueuePreparedSlices(prepared_slice_list, submitted_slice_count); // Context-level health tracking: if all slices failed due to no available // rails, increment the context failure counter. This detects catastrophic @@ -290,17 +313,9 @@ int WorkerPool::submitPostSend( return 0; } -void WorkerPool::enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount], +void WorkerPool::enqueuePreparedSlices(const SliceList &slice_list, uint64_t submitted_slice_count) { - for (int shard_id = 0; shard_id < kShardCount; ++shard_id) { - if (slice_list_map[shard_id].empty()) continue; - slice_queue_lock_[shard_id].lock(); - for (auto &slice : slice_list_map[shard_id]) - slice_queue_[shard_id][slice->peer_nic_path].push_back(slice); - slice_queue_count_[shard_id].fetch_add(slice_list_map[shard_id].size(), - std::memory_order_relaxed); - slice_queue_lock_[shard_id].unlock(); - } + for (auto &slice : slice_list) enqueueSliceToOwner(slice); submitted_slice_count_.fetch_add(submitted_slice_count); if (submitted_slice_count && @@ -315,7 +330,7 @@ int WorkerPool::submitPreparedPostSend( // Called by a different local RNIC's worker during local failover. The // slice already carries the chosen peer_nic_path and refreshed local lkey, // so enqueue it directly instead of running remote-path selection again. - SliceList slice_list_map[kShardCount]; + SliceList prepared_slice_list; uint64_t submitted_slice_count = 0; for (auto &slice : slice_list) { @@ -323,13 +338,11 @@ int WorkerPool::submitPreparedPostSend( slice->markFailed(); continue; } - auto shard_id = static_cast( - std::hash{}(slice->peer_nic_path) % kShardCount); - slice_list_map[shard_id].push_back(slice); + prepared_slice_list.push_back(slice); submitted_slice_count++; } - enqueuePreparedSlices(slice_list_map, submitted_slice_count); + enqueuePreparedSlices(prepared_slice_list, submitted_slice_count); return 0; } @@ -355,10 +368,17 @@ void WorkerPool::untrackPostedSlices( } void WorkerPool::performPostSend(int thread_id) { - int post_tid = 0; - int post_count = 0; - getPostingShardAssignment(thread_id, post_tid, post_count); auto &local_slice_queue = collective_slice_queue_[thread_id]; + { + std::lock_guard lock(worker_slice_queue_lock_[thread_id]); + for (auto &entry : worker_slice_queue_[thread_id]) { + if (entry.second.empty()) continue; + auto &local_entry = local_slice_queue[entry.first]; + local_entry.insert(local_entry.end(), entry.second.begin(), + entry.second.end()); + } + worker_slice_queue_[thread_id].clear(); + } // If this local RNIC is inactive/unhealthy, the remote rail is not the // problem. Move queued work to another local RNIC while preserving the @@ -368,38 +388,9 @@ void WorkerPool::performPostSend(int thread_id) { local_slice_queue.clear(); for (auto &entry : local_slice_queue_clone) redispatch(entry.second, thread_id, true); - - for (int shard_id = post_tid; shard_id < kShardCount; - shard_id += post_count) { - if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) == - 0) - continue; - slice_queue_lock_[shard_id].lock(); - auto slice_queue_clone = slice_queue_[shard_id]; - slice_queue_[shard_id].clear(); - slice_queue_count_[shard_id].store(0, std::memory_order_relaxed); - slice_queue_lock_[shard_id].unlock(); - for (auto &entry : slice_queue_clone) - redispatch(entry.second, thread_id, true); - } return; } - for (int shard_id = post_tid; shard_id < kShardCount; - shard_id += post_count) { - if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) == 0) - continue; - - slice_queue_lock_[shard_id].lock(); - for (auto &entry : slice_queue_[shard_id]) { - for (auto &slice : entry.second) - local_slice_queue[entry.first].push_back(slice); - entry.second.clear(); - } - slice_queue_count_[shard_id].store(0, std::memory_order_relaxed); - slice_queue_lock_[shard_id].unlock(); - } - // Redispatch slices to other endpoints, for temporary failures thread_local int tl_redispatch_counter = 0; if (tl_redispatch_counter < @@ -446,12 +437,16 @@ void WorkerPool::performPostSend(int thread_id) { entry.second.clear(); continue; } + auto endpoint_lifecycle_lock = + context_.lockEndpointLifecycle(entry.first); #ifdef CONFIG_CACHE_ENDPOINT auto &endpoint = endpoint_map[entry.first]; if (endpoint == nullptr || !endpoint->active()) - endpoint = context_.endpoint(entry.first); + endpoint = context_.endpoint(entry.first, + cqIndexForPostingThread(thread_id)); #else - auto endpoint = context_.endpoint(entry.first); + auto endpoint = + context_.endpoint(entry.first, cqIndexForPostingThread(thread_id)); #endif if (!endpoint) { for (auto &slice : entry.second) failed_slice_list.push_back(slice); @@ -459,7 +454,22 @@ void WorkerPool::performPostSend(int thread_id) { continue; } if (!endpoint->connected()) { - int setup_ret = endpoint->setupConnectionsByActive(); + auto setup_result = setupEndpointByActiveOutsideLifecycleGate( + context_, entry.first, endpoint, endpoint_lifecycle_lock); + if (!setup_result.endpoint_current) { + LOG(WARNING) + << "Worker: Endpoint changed while active handshake was " + "outstanding: " + << entry.first << ", retrying queued slices"; +#ifdef CONFIG_CACHE_ENDPOINT + endpoint.reset(); +#endif + for (auto &slice : entry.second) + failed_slice_list.push_back(slice); + entry.second.clear(); + continue; + } + int setup_ret = setup_result.ret; if (setup_ret) { // Active handshake setup failures are ambiguous: the failed // side may be the peer rail, or this local RNIC may have just @@ -550,6 +560,9 @@ void WorkerPool::performPostSend(int thread_id) { } void WorkerPool::performPollCq(int thread_id) { + if (context_.cqCount() <= 0) return; + if (thread_id < 0 || thread_id >= worker_count_) return; + const uint64_t poll_ts = getCurrentTimeInNano(); const uint64_t previous_poll_ts = last_poll_ts_ns_.exchange(poll_ts, std::memory_order_relaxed); @@ -564,7 +577,49 @@ void WorkerPool::performPollCq(int thread_id) { } } - // Slices this loop drove to a terminal state, successes and + const static size_t kPollCount = 64; + std::unordered_map *, int> qp_depth_set; + std::vector wc_list; + const int cq_index = cqIndexForPostingThread(thread_id); + if (cq_index < 0) return; + ibv_wc wc[kPollCount]; + int nr_poll = context_.poll(kPollCount, wc, cq_index); + if (nr_poll < 0) { + LOG(ERROR) << "Worker: Failed to poll completion queues"; + return; + } + + if (nr_poll > 0 && globalConfig().track_rdma_posted_slices) { + std::lock_guard lock(posted_slices_mutex_); + for (int i = 0; i < nr_poll; ++i) { + auto *slice = reinterpret_cast(wc[i].wr_id); + posted_slices_.erase(slice); + } + } + + for (int i = 0; i < nr_poll; ++i) { + Transport::Slice *slice = (Transport::Slice *)wc[i].wr_id; + assert(slice); + assert(postingThreadForPeer(slice->peer_nic_path) == thread_id); + if (qp_depth_set.count(slice->rdma.qp_depth)) + qp_depth_set[slice->rdma.qp_depth]++; + else + qp_depth_set[slice->rdma.qp_depth] = 1; + wc_list.push_back(wc[i]); + } + if (nr_poll) + context_.cqOutstandingCount(cq_index)->fetch_sub( + nr_poll, std::memory_order_acq_rel); + + for (auto &entry : qp_depth_set) + entry.first->fetch_sub(entry.second, std::memory_order_acq_rel); + + if (!wc_list.empty()) processCompletions(thread_id, wc_list); +} + +void WorkerPool::processCompletions(int thread_id, + const std::vector &wc_list) { + // Slices this call drove to a terminal state, successes and // retry-exhausted failures alike; folded into processed_slice_count_, // which gates worker parking. Every terminal outcome below goes through // finalize_slice() so it is counted exactly once. Slices handed to @@ -572,8 +627,8 @@ void WorkerPool::performPollCq(int thread_id) { int processed_slice_count = 0; // Successful completions only, kept apart because it clears the context // health counter. A completion error is no evidence that this RNIC can - // still move data -- including IBV_WC_WR_FLUSH_ERR, which only reports - // WRs the hardware discarded after the QP had already entered ERR. + // still move data, including IBV_WC_WR_FLUSH_ERR, which only reports WRs + // the hardware discarded after the QP had already entered ERR. int succeeded_slice_count = 0; auto finalize_slice = [&](Transport::Slice *slice, bool success) { if (success) { @@ -584,8 +639,6 @@ void WorkerPool::performPollCq(int thread_id) { } processed_slice_count++; }; - const static size_t kPollCount = 64; - std::unordered_map *, int> qp_depth_set; std::unordered_set local_failed_endpoints; // Peer NIC paths already charged a local-fault error this pass. A QP that // takes a local fault completes every WR it still holds, so one burst must @@ -594,133 +647,107 @@ void WorkerPool::performPollCq(int thread_id) { bool recorded_local_context_failure = false; SliceList failed_slice_list; SliceList local_failed_slice_list; - for (int cq_index = 0; cq_index < context_.cqCount(); cq_index++) { - ibv_wc wc[kPollCount]; - int nr_poll = context_.poll(kPollCount, wc, cq_index); - if (nr_poll < 0) { - LOG(ERROR) << "Worker: Failed to poll completion queues"; - continue; - } - if (nr_poll > 0 && globalConfig().track_rdma_posted_slices) { - std::lock_guard lock(posted_slices_mutex_); - for (int i = 0; i < nr_poll; ++i) { - auto *slice = reinterpret_cast(wc[i].wr_id); - posted_slices_.erase(slice); + for (const auto &wc : wc_list) { + Transport::Slice *slice = (Transport::Slice *)wc.wr_id; + assert(slice); + if (wc.status != IBV_WC_SUCCESS) { + // Flush errors are generated when QPs transition to ERR state + // during normal endpoint destruction (beginDestroy). They are not + // real network errors and should not trigger rail failure handling + // or endpoint deletion. + if (wc.status == IBV_WC_WR_FLUSH_ERR) { + if (!context_.active()) { + if (globalConfig().trace) + LOG(INFO) << "Worker: WR flush error on inactive " + << "local context " << context_.deviceName() + << " (peer_nic: " << slice->peer_nic_path + << "), handing off if retry allows"; + if (shouldRetrySlice(slice)) + local_failed_slice_list.push_back(slice); + else + finalize_slice(slice, false); + } else { + if (globalConfig().trace) + LOG(INFO) << "Worker: WR flush error (peer_nic: " + << slice->peer_nic_path + << "), redispatching if retry allows"; + if (shouldRetrySlice(slice)) + failed_slice_list.push_back(slice); + else + finalize_slice(slice, false); + } + continue; } - } - for (int i = 0; i < nr_poll; ++i) { - Transport::Slice *slice = (Transport::Slice *)wc[i].wr_id; - assert(slice); - if (qp_depth_set.count(slice->rdma.qp_depth)) - qp_depth_set[slice->rdma.qp_depth]++; - else - qp_depth_set[slice->rdma.qp_depth] = 1; - // __sync_fetch_and_sub(slice->rdma.qp_depth, 1); - if (wc[i].status != IBV_WC_SUCCESS) { - // Flush errors are generated when QPs transition to ERR state - // during normal endpoint destruction (beginDestroy). They are - // not real network errors and should not trigger rail failure - // handling or endpoint deletion. - if (wc[i].status == IBV_WC_WR_FLUSH_ERR) { - if (!context_.active()) { - if (globalConfig().trace) - LOG(INFO) - << "Worker: WR flush error on inactive " - << "local context " << context_.deviceName() - << " (peer_nic: " << slice->peer_nic_path - << "), handing off if retry allows"; - if (shouldRetrySlice(slice)) - local_failed_slice_list.push_back(slice); - else - finalize_slice(slice, false); - } else { - if (globalConfig().trace) - LOG(INFO) << "Worker: WR flush error (peer_nic: " - << slice->peer_nic_path - << "), redispatching if retry allows"; - if (shouldRetrySlice(slice)) - failed_slice_list.push_back(slice); - else - finalize_slice(slice, false); - } - continue; + auto endpoint_lifecycle_lock = + context_.lockEndpointLifecycle(slice->peer_nic_path); + + // Completion errors are split by local context health. Local faults + // hand off to another local RNIC; remote/default faults keep this + // local context and switch peer rails. + LOG(ERROR) << "Worker: Process failed for slice (opcode: " + << slice->opcode + << ", source_addr: " << slice->source_addr + << ", length: " << slice->length + << ", dest_addr: " << (void *)slice->rdma.dest_addr + << ", local_nic: " << context_.deviceName() + << ", peer_nic: " << slice->peer_nic_path + << ", dest_rkey: " << slice->rdma.dest_rkey + << ", retry_cnt: " << slice->rdma.retry_cnt + << ", max_retry_cnt: " << slice->rdma.max_retry_cnt + << "): " << ibv_wc_status_str(wc.status); + auto *retry_list = &failed_slice_list; + const bool local_wc_failure = isLocalWcFailure(wc); + if (!context_.active() || local_wc_failure) { + // A local completion fault retires the endpoint, and the slice + // is handed to another local RNIC that rebuilds its own + // endpoint to the same peer NIC. If the fault keeps recurring + // nothing throttles that cycle: the rail is deliberately not + // paused on the local path, and the context failure counter is + // cleared by any concurrent success, so both RNICs re-handshake + // the same peer as fast as the workers spin. Charge the path an + // error instead -- without an immediate pause, so a one-off + // fault still costs nothing -- and let kRailErrorThreshold stop + // the rebuild loop from this context. See issue #3299. + if (local_wc_failure && + local_failed_peer_paths.insert(slice->peer_nic_path) + .second) { + markRailFailed(slice->peer_nic_path); } - - // Completion errors are split by local context health. Local - // faults hand off to another local RNIC; remote/default faults - // keep this local context and switch peer rails. - LOG(ERROR) << "Worker: Process failed for slice (opcode: " - << slice->opcode - << ", source_addr: " << slice->source_addr - << ", length: " << slice->length - << ", dest_addr: " << (void *)slice->rdma.dest_addr - << ", local_nic: " << context_.deviceName() - << ", peer_nic: " << slice->peer_nic_path - << ", dest_rkey: " << slice->rdma.dest_rkey - << ", retry_cnt: " << slice->rdma.retry_cnt - << ", max_retry_cnt: " << slice->rdma.max_retry_cnt - << "): " << ibv_wc_status_str(wc[i].status); - auto *retry_list = &failed_slice_list; - const bool local_wc_failure = isLocalWcFailure(wc[i]); - if (!context_.active() || local_wc_failure) { - // A local completion fault retires the endpoint, and the - // slice is handed to another local RNIC that rebuilds its - // own endpoint to the same peer NIC. If the fault keeps - // recurring nothing throttles that cycle: the rail is - // deliberately not paused on the local path, and the - // context failure counter is cleared by any concurrent - // success, so both RNICs re-handshake the same peer as fast - // as the workers spin. Charge the path an error instead -- - // without an immediate pause, so a one-off fault still - // costs nothing -- and let kRailErrorThreshold stop the - // rebuild loop from this context. See issue #3299. - if (local_wc_failure && - local_failed_peer_paths.insert(slice->peer_nic_path) - .second) { - markRailFailed(slice->peer_nic_path); - } - if (!recorded_local_context_failure) { - handleLocalFailure(slice->peer_nic_path, - slice->rdma.endpoint); - recorded_local_context_failure = true; - if (slice->rdma.endpoint) - local_failed_endpoints.insert(slice->rdma.endpoint); - } else if (slice->rdma.endpoint && - !local_failed_endpoints.count( - slice->rdma.endpoint)) { - context_.deleteEndpointByPtr(slice->rdma.endpoint); + if (!recorded_local_context_failure) { + handleLocalFailure(slice->peer_nic_path, + slice->rdma.endpoint); + recorded_local_context_failure = true; + if (slice->rdma.endpoint) local_failed_endpoints.insert(slice->rdma.endpoint); - } - retry_list = &local_failed_slice_list; - } else { - if (hasAvailablePeerRailAlternative(slice, - slice->peer_nic_path)) { - markRailFailed(slice->peer_nic_path, true); - redispatch_counter_++; - } - if (slice->rdma.endpoint) { - context_.deleteEndpointByPtr(slice->rdma.endpoint); - } + } else if (slice->rdma.endpoint && + !local_failed_endpoints.count( + slice->rdma.endpoint)) { + context_.deleteEndpointByPtr(slice->rdma.endpoint); + local_failed_endpoints.insert(slice->rdma.endpoint); } - if (shouldRetrySlice(slice)) { - retry_list->push_back(slice); - } else { - finalize_slice(slice, false); + retry_list = &local_failed_slice_list; + } else { + if (hasAvailablePeerRailAlternative(slice, + slice->peer_nic_path)) { + markRailFailed(slice->peer_nic_path, true); + redispatch_counter_++; + } + if (slice->rdma.endpoint) { + context_.deleteEndpointByPtr(slice->rdma.endpoint); } + } + if (shouldRetrySlice(slice)) { + retry_list->push_back(slice); } else { - finalize_slice(slice, true); + finalize_slice(slice, false); } + } else { + finalize_slice(slice, true); } - if (nr_poll) - context_.cqOutstandingCount(cq_index)->fetch_sub( - nr_poll, std::memory_order_acq_rel); } - for (auto &entry : qp_depth_set) - entry.first->fetch_sub(entry.second, std::memory_order_acq_rel); - if (processed_slice_count) processed_slice_count_.fetch_add(processed_slice_count); // Clear the consecutive-failure counter only on proven data movement, so @@ -740,7 +767,6 @@ void WorkerPool::redispatch(std::vector &slice_list, int thread_id, bool handoff_to_local_worker) { std::unordered_map> segment_desc_map; - const bool use_local_queue = workerCanPost(thread_id); int shared_redispatch_count = 0; // Remote redispatch needs target metadata to choose a new peer RNIC. // Local handoff keeps the peer RNIC fixed and only switches source RNIC, so @@ -851,17 +877,12 @@ void WorkerPool::redispatch(std::vector &slice_list, << ", retry_cnt=" << slice->rdma.retry_cnt; } slice->ts = 0; - if (use_local_queue) { + const int owner_thread = postingThreadForPeer(peer_nic_path); + if (owner_thread == thread_id) { collective_slice_queue_[thread_id][peer_nic_path].push_back( slice); } else { - int shard_id = - (slice->target_id * 10007 + device_id) % kShardCount; - slice_queue_lock_[shard_id].lock(); - slice_queue_[shard_id][peer_nic_path].push_back(slice); - slice_queue_count_[shard_id].fetch_add( - 1, std::memory_order_relaxed); - slice_queue_lock_[shard_id].unlock(); + enqueueSliceToOwner(slice); shared_redispatch_count++; } } @@ -945,21 +966,17 @@ bool WorkerPool::tryHandoffToAnotherLocalWorker(Transport::Slice *slice) { } bool WorkerPool::hasOutstandingCq(int thread_id) { - if (!workerCanPoll(thread_id)) return false; - for (int cq_index = 0; cq_index < context_.cqCount(); ++cq_index) { - if (context_.cqOutstandingCount(cq_index)->load( - std::memory_order_relaxed) > 0) - return true; - } - return false; + if (context_.cqCount() <= 0) return false; + const int cq_index = cqIndexForPostingThread(thread_id); + if (cq_index < 0) return false; + return context_.cqOutstandingCount(cq_index)->load( + std::memory_order_relaxed) > 0; } void WorkerPool::transferWorker(int thread_id) { bindToSocket(numa_socket_id_); const static uint64_t kWaitPeriodInNano = 100000000; // 100ms uint64_t last_wait_ts = getCurrentTimeInNano(); - const bool can_post = workerCanPost(thread_id); - const bool can_poll = workerCanPoll(thread_id); while (workers_running_.load(std::memory_order_relaxed)) { auto processed_slice_count = processed_slice_count_.load(std::memory_order_relaxed); @@ -984,13 +1001,9 @@ void WorkerPool::transferWorker(int thread_id) { } continue; } - if (can_post) { - performPostSend(thread_id); - } + performPostSend(thread_id); #ifndef USE_FAKE_POST_SEND - if (can_poll) { - performPollCq(thread_id); - } + performPollCq(thread_id); #endif last_wait_ts = getCurrentTimeInNano(); } @@ -1016,6 +1029,7 @@ int WorkerPool::doProcessContextEvents() { << context_.deviceName(); if (event.event_type == IBV_EVENT_QP_FATAL) { auto endpoint_ptr = (RdmaEndPoint *)event.element.qp->qp_context; + auto endpoint = context_.getEndpointByPtr(endpoint_ptr); /** * There might be a deadlock if we call endpoint->set_active(false) @@ -1033,13 +1047,17 @@ int WorkerPool::doProcessContextEvents() { event_acked = true; /** - * After ack the event, the endpoint might be destroyed if it - * happened to be destroying event.element.qp. Therefore, we cannot - * just dereference endpoint_ptr. Instead, we need to get the - * shared_ptr of the endpoint from context_ and use that shared_ptr - * to access the endpoint. + * After ack the event, use the tracked endpoint's peer path to + * serialize deletion with the same lifecycle gate as post/send and + * passive setup. */ - context_.deleteEndpointByPtr(endpoint_ptr); + if (endpoint) { + auto endpoint_lifecycle_lock = + context_.lockEndpointLifecycle(endpoint->peerNicPath()); + context_.deleteEndpointByPtr(endpoint.get()); + } else { + LOG(WARNING) << "QP fatal event endpoint is no longer tracked"; + } } else if (handleContextEvent(event.event_type, false, &event)) { event_acked = true; } @@ -1083,6 +1101,7 @@ bool WorkerPool::handleContextEvent(ibv_event_type event_type, */ if (event != nullptr) ibv_ack_async_event(event); + auto endpoint_lifecycle_locks = context_.lockAllEndpointLifecycles(); context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() << " is now inactive due to " @@ -1094,6 +1113,8 @@ bool WorkerPool::handleContextEvent(ibv_event_type event_type, if (event != nullptr) ibv_ack_async_event(event); if (gid_refresh_result != GidRefreshResult::UNCHANGED) { + auto endpoint_lifecycle_locks = + context_.lockAllEndpointLifecycles(); context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() << (injected_for_test ? " injected GID refresh result=" @@ -1148,6 +1169,7 @@ void WorkerPool::maybeActivateRecoveredContext() { return; } if (gid_refresh_result == GidRefreshResult::CHANGED) { + auto endpoint_lifecycle_locks = context_.lockAllEndpointLifecycles(); context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() << " GID changed during recovery, disconnected all endpoints"; diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index 1618037f8b..25db90bddb 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -329,10 +329,10 @@ TcpTransport::TcpTransport() } constexpr size_t kDefaultLanesPerPeer = 4; - constexpr size_t kDefaultQueuedTransfersPerPeer = 1024; - constexpr size_t kMaxQueuedTransfersPerPeer = 65535; - constexpr size_t kDefaultPendingAdmissionsPerPeer = 1024; - constexpr size_t kMaxPendingAdmissionsPerPeer = 65535; + constexpr size_t kDefaultQueuedTransfersPerPeer = 65535; + constexpr size_t kMaxQueuedTransfersPerPeer = 1048576; + constexpr size_t kDefaultPendingAdmissionsPerPeer = 65535; + constexpr size_t kMaxPendingAdmissionsPerPeer = 1048576; constexpr size_t kDefaultAdmissionTimeoutMs = 1000; constexpr size_t kMaxAdmissionTimeoutMs = 600000; diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json index 632346afb5..c9de243100 100644 --- a/mooncake-transfer-engine/tent/config/transfer-engine.json +++ b/mooncake-transfer-engine/tent/config/transfer-engine.json @@ -82,6 +82,19 @@ "retry_max_delay_ms": 2000, "max_concurrent_tasks": 16 }, + "hp_tcp": { + "enable": false, + "bind_address": "", + "advertise_address": "", + "port": 0, + "worker_count": 16, + "connections_per_peer": 4, + "max_outstanding_tasks": 4096, + "max_outstanding_bytes": 4294967296, + "max_transfer_bytes": 1073741824, + "connect_timeout_ms": 2000, + "progress_timeout_ms": 30000 + }, "gds": { "enable" : true }, @@ -98,7 +111,7 @@ "name": "default_memory", "segment_type": "memory", "devices": ["mlx5_0", "mlx5_2"], - "transports": ["nvlink", "rdma", "shm"] + "transports": ["nvlink", "rdma", "shm", "hp_tcp"] }, { "name": "file_storage", diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index ea94411b19..0bb1b744e4 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -57,6 +57,7 @@ enum TransportType : int { TPU, UB, MPCOMM, + HP_TCP, // Sentinel: must remain the last enumerator. kNumTransportTypes, }; @@ -96,6 +97,8 @@ inline const char* transportTypeName(TransportType type) { return "ub"; case MPCOMM: return "mpcomm"; + case HP_TCP: + return "hp_tcp"; case kNumTransportTypes: return "unknown"; } @@ -116,6 +119,7 @@ inline TransportType parseTransportType(const std::string& str) { if (str == "tpu") return TPU; if (str == "ub") return UB; if (str == "mpcomm") return MPCOMM; + if (str == "hp_tcp") return HP_TCP; return UNSPEC; } @@ -162,6 +166,27 @@ enum TransferStatusEnum { FAILED }; +// Rank for aggregating batch status. Unknown values rank with FAILED so +// getBatchStatus never throws from unordered_map::at during teardown. +inline int transferStatusSeverity(TransferStatusEnum s) { + switch (s) { + case INITIAL: + case PENDING: + case COMPLETED: + return 0; + case INVALID: + return 1; + case CANCELED: + return 2; + case TIMEOUT: + return 3; + case FAILED: + return 4; + default: + return 4; + } +} + struct TransferStatus { TransferStatusEnum s; size_t transferred_bytes; diff --git a/mooncake-transfer-engine/tent/include/tent/platform/cuda.h b/mooncake-transfer-engine/tent/include/tent/platform/cuda.h index a808a2b6f0..9b5489cb46 100644 --- a/mooncake-transfer-engine/tent/include/tent/platform/cuda.h +++ b/mooncake-transfer-engine/tent/include/tent/platform/cuda.h @@ -126,6 +126,9 @@ class CudaPlatform : public Platform { int deviceId = CUDAStreamPool::kCurrentDevice); private: + // Device owning `addr`, or kCurrentDevice when `addr` is not device memory. + int getPointerDeviceId(void* addr); + std::shared_ptr conf; CUDAStreamPool stream_pool; }; diff --git a/mooncake-transfer-engine/tent/include/tent/platform/rocm.h b/mooncake-transfer-engine/tent/include/tent/platform/rocm.h index 3c5bd2ff7f..5ddde8baaa 100644 --- a/mooncake-transfer-engine/tent/include/tent/platform/rocm.h +++ b/mooncake-transfer-engine/tent/include/tent/platform/rocm.h @@ -121,6 +121,9 @@ class RocmPlatform : public Platform { int deviceId = HIPStreamPool::kCurrentDevice); private: + // Device owning `addr`, or kCurrentDevice when `addr` is not device memory. + int getPointerDeviceId(void* addr); + std::shared_ptr conf; HIPStreamPool stream_pool_; }; diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 23a96aa7f5..379704544b 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -47,6 +47,8 @@ enum RpcFuncID { Unpin, SubscribeSegmentUpdate, NotifySegmentUpdated, + // Appended to preserve the numeric values of the existing RPCs. + BootstrapUb, }; class ClientPool; 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 11e3da356a..46cc859f1c 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -48,7 +49,7 @@ struct BootstrapDesc { // RDMA address of local_nic_path. uint16_t local_lid = 0; std::string local_gid; - std::string reply_msg; // non-empty means the bootstrap callback failed + std::string reply_msg; // on error uint32_t notify_qp_num = 0; // Notification QP number (0 = not supported) public: @@ -57,6 +58,65 @@ struct BootstrapDesc { notify_qp_num); }; +// UB/URMA has Jetty, JFC and EID concepts that are not wire-compatible with +// RDMA QPs, CQs and GIDs. Keep a dedicated bootstrap envelope so neither +// transport has to smuggle native identifiers through the other's fields. +struct UbBootstrapDesc { + uint32_t protocol_version = 1; + std::string segment_name; + std::string local_nic_path; + std::string peer_nic_path; + std::string local_device_name; + int local_device_id = -1; + int local_eid_index = -1; + std::string local_eid; + std::vector jetty_ids; + // UASID is part of urma_jetty_id_t on providers that use nonzero address + // spaces. Kept parallel to jetty_ids for protocol-v1 compatibility. + std::vector jetty_uasids; + uint64_t endpoint_generation = 0; + uint64_t segment_generation = 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}}; +} + +inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { + desc.protocol_version = j.value("protocol_version", 0u); + if (desc.protocol_version != 1) { + throw std::invalid_argument("unsupported UB bootstrap version"); + } + desc.segment_name = j.value("segment_name", ""); + desc.local_nic_path = j.value("local_nic_path", ""); + desc.peer_nic_path = j.value("peer_nic_path", ""); + desc.local_device_name = j.value("local_device_name", ""); + desc.local_device_id = j.value("local_device_id", -1); + desc.local_eid_index = j.value("local_eid_index", -1); + desc.local_eid = j.value("local_eid", ""); + desc.jetty_ids = j.value("jetty_ids", std::vector{}); + 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.capabilities = j.value("capabilities", std::vector{}); + desc.reply_msg = j.value("reply_msg", ""); +} + struct XferDataDesc { uint64_t peer_mem_addr; size_t length; @@ -65,6 +125,9 @@ struct XferDataDesc { using OnReceiveBootstrap = std::function; +using OnReceiveUbBootstrap = std::function; + using OnNotify = std::function; class ControlClient { @@ -86,6 +149,10 @@ class ControlClient { static Status decodeBootstrapResponse(const std::string& response_raw, BootstrapDesc& response); + static Status bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response); + static Status sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length); @@ -141,6 +208,11 @@ class ControlService { void setBootstrapRdmaCallback(const OnReceiveBootstrap& callback); + void setBootstrapUbCallback(const OnReceiveUbBootstrap& callback) { + std::lock_guard lock(ub_bootstrap_callback_mutex_); + ub_bootstrap_callback_ = callback; + } + void setNotifyCallback(const OnNotify& callback); Status start(uint16_t& port, bool ipv6_ = false, size_t threads = 1); @@ -152,6 +224,8 @@ class ControlService { void onBootstrapRdma(const std::string_view& request, std::string& response); + void onBootstrapUb(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); @@ -189,6 +263,9 @@ class ControlService { OnReceiveBootstrap bootstrap_callback_; static thread_local const ControlService* active_bootstrap_service_; + std::mutex ub_bootstrap_callback_mutex_; + OnReceiveUbBootstrap ub_bootstrap_callback_; + std::mutex notify_cb_mutex_; std::condition_variable notify_cb_cv_; size_t notify_callbacks_in_flight_ = 0; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/hp_tcp_transport_config.h b/mooncake-transfer-engine/tent/include/tent/runtime/hp_tcp_transport_config.h new file mode 100644 index 0000000000..98fac09ed5 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/hp_tcp_transport_config.h @@ -0,0 +1,38 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_RUNTIME_HP_TCP_TRANSPORT_CONFIG_H_ +#define TENT_RUNTIME_HP_TCP_TRANSPORT_CONFIG_H_ + +#include +#include +#include + +#include "tent/common/config.h" + +namespace mooncake::tent { + +struct HighPerformanceTcpParams { + std::string bind_address; + std::string advertise_address; + uint16_t port{0}; + size_t worker_count{16}; + size_t connections_per_peer{4}; + uint64_t max_outstanding_tasks{4096}; + uint64_t max_outstanding_bytes{1ULL << 32}; + uint64_t max_transfer_bytes{1ULL << 30}; + uint64_t connect_timeout_ms{2000}; + uint64_t progress_timeout_ms{30000}; +}; + +struct HpTcpTransportConfig { + bool enabled{false}; + HighPerformanceTcpParams params; +}; + +// This is intentionally the sole parser for transports/hp_tcp. Config::get() +// silently substitutes defaults for type mismatches and is not suitable here. +Status ParseHpTcpTransportConfig(const Config& config, + HpTcpTransportConfig* out); + +} // namespace mooncake::tent + +#endif // TENT_RUNTIME_HP_TCP_TRANSPORT_CONFIG_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h index b1fbb9e89f..86256f47b5 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h @@ -80,12 +80,8 @@ class Topology { void clear(); - // Preserve the original one-argument symbol for source and binary - // compatibility with callers that do not opt into UB discovery. Status discover(const std::vector& platforms); - Status discover(const std::vector& platforms, bool discover_ub); - Status parse(const std::string& json_content); // Parse classic TE NIC priority matrix: @@ -116,6 +112,10 @@ class Topology { const MemEntry* getMemEntry(const std::string& name) const; + // True only when both NUMA ids are known and differ. Unknown (-1) is not + // treated as remote; rank is ignored because probes disagree on placement. + bool isCrossNuma(const MemEntry& mem, NicID nic_id) const; + NicID getNicId(const std::string& name) const; MemID getMemId(const std::string& name) const; 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 681e98a7a8..6cb0eed42d 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -33,6 +33,7 @@ #include "tent/runtime/admission_queue.h" #include "tent/runtime/transport.h" #include "tent/runtime/transport_selector.h" +#include "tent/runtime/hp_tcp_transport_config.h" namespace mooncake { namespace tent { @@ -63,11 +64,13 @@ void waitBeforeNextPoll(uint64_t poll_count); struct TaskInfo { TransportType type{UNSPEC}; int sub_task_id{-1}; - bool derived{false}; // merged by other tasks - int xport_priority{0}; // transport priority (for fallback) - int failover_count{0}; // number of failover attempts - uint64_t device_mask{~0ULL}; // Device mask for quota allocation - std::string qp_pool; // Named QP pool (RFC #2568 step 3), "" = none + bool derived{false}; // merged by other tasks + int xport_priority{0}; // transport priority (for fallback) + int failover_count{0}; // number of failover attempts + int metadata_refresh_retry_count{0}; // same-transport stale-cache retries + bool suppress_failover{false}; // permanent transport result + uint64_t device_mask{~0ULL}; // Device mask for quota allocation + std::string qp_pool; // Named QP pool (RFC #2568 step 3), "" = none Request request; bool staging{false}; bool cancel_requested{false}; @@ -101,6 +104,8 @@ struct TaskInfo { derived(other.derived), xport_priority(other.xport_priority), failover_count(other.failover_count), + metadata_refresh_retry_count(other.metadata_refresh_retry_count), + suppress_failover(other.suppress_failover), device_mask(other.device_mask), qp_pool(other.qp_pool), request(other.request), @@ -122,6 +127,8 @@ struct TaskInfo { derived(other.derived), xport_priority(other.xport_priority), failover_count(other.failover_count), + metadata_refresh_retry_count(other.metadata_refresh_retry_count), + suppress_failover(other.suppress_failover), device_mask(other.device_mask), qp_pool(std::move(other.qp_pool)), request(std::move(other.request)), @@ -144,6 +151,8 @@ struct TaskInfo { derived = other.derived; xport_priority = other.xport_priority; failover_count = other.failover_count; + metadata_refresh_retry_count = other.metadata_refresh_retry_count; + suppress_failover = other.suppress_failover; device_mask = other.device_mask; qp_pool = other.qp_pool; request = other.request; @@ -171,6 +180,8 @@ struct TaskInfo { derived = other.derived; xport_priority = other.xport_priority; failover_count = other.failover_count; + metadata_refresh_retry_count = other.metadata_refresh_retry_count; + suppress_failover = other.suppress_failover; device_mask = other.device_mask; qp_pool = std::move(other.qp_pool); request = std::move(other.request); @@ -503,6 +514,7 @@ class TransferEngineImpl { private: std::shared_ptr conf_; + HpTcpTransportConfig hp_tcp_transport_config_; std::shared_ptr metadata_; std::shared_ptr topology_; std::unique_ptr transport_selector_; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index 260ec49e11..0299475a7a 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -79,6 +79,11 @@ class Transport { virtual Status uninstall() { return Status::OK(); } + // Called before registered ranges and sub-batches are reclaimed. Most + // transports have no background work; transports with async I/O use this + // barrier to settle work while their buffer registry is still alive. + virtual Status quiesce() { return Status::OK(); } + virtual const Capabilities capabilities() const { return caps; } virtual Status allocateSubBatch(SubBatchRef& batch, size_t max_size) { @@ -111,6 +116,12 @@ class Transport { "getTransferStatus not implemented" LOC_MARK); } + virtual Status retryTransferTask(SubBatchRef batch, int task_id, + const Request& request) { + return Status::NotImplemented( + "retryTransferTask not implemented" LOC_MARK); + } + // Cancellation is best effort: implementations must prevent work that has // not reached the device from being submitted, but work already posted to // a device may still complete. Callers must continue polling until the @@ -155,6 +166,10 @@ class Transport { "removeMemoryBuffer not implemented" LOC_MARK); } + // Some transports keep private local-only registrations that must not be + // advertised through BufferDesc::transports. + virtual bool tracksLocalBuffer(const BufferDesc&) const { return false; } + virtual bool supportNotification() const { return false; } virtual Status sendNotification(SegmentID target_id, diff --git a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h index 0a168d85ba..0cc971756b 100644 --- a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h +++ b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h @@ -112,6 +112,7 @@ typedef struct tent_notifi_info tent_notifi_info; #define TRANSPORT_TPU (10) #define TRANSPORT_UB (11) #define TRANSPORT_MPCOMM (12) +#define TRANSPORT_HP_TCP (13) struct tent_memory_options { char location[64]; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_buffer_registry.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_buffer_registry.h new file mode 100644 index 0000000000..7c22b3b510 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_buffer_registry.h @@ -0,0 +1,90 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_BUFFER_REGISTRY_H_ +#define TENT_HP_TCP_BUFFER_REGISTRY_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/common/types.h" +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" + +namespace mooncake::tent { + +class HighPerformanceTcpBufferRegistry { + public: + HighPerformanceTcpBufferRegistry(); + + struct Entry { + uint64_t base{0}; + uint64_t length{0}; + uint64_t registration_id{0}; + Permission permission{kLocalReadWrite}; + + std::mutex mutex; + std::condition_variable drained; + uint64_t active_leases{0}; + }; + + class Lease { + public: + Lease() = default; + Lease(const Lease&) = delete; + Lease& operator=(const Lease&) = delete; + Lease(Lease&& other) noexcept; + Lease& operator=(Lease&& other) noexcept; + ~Lease(); + + void reset(); + void* data() const; + uint64_t base() const; + uint64_t length() const; + explicit operator bool() const { return entry_ != nullptr; } + + private: + friend class HighPerformanceTcpBufferRegistry; + explicit Lease(std::shared_ptr entry); + std::shared_ptr entry_; + }; + + Status add(uint64_t base, uint64_t length, Permission permission, + uint64_t* registration_id); + Status remove(uint64_t base, uint64_t length); + + // Prevent new registrations and leases during quiesce. Existing leases + // remain valid and drain through normal unregister/session completion. + void close(); + Status reopen(); + + // Local access checks lifetime/range only. MemoryOptions::perm is a remote + // authorization policy and must not reject the local side of a transfer. + Status acquireLocalLease(uint64_t addr, uint64_t length, Lease* lease); + + Status acquireRemoteLease(uint64_t addr, uint64_t length, + uint64_t registration_id, + HighPerformanceTcpOpcode opcode, Lease* lease, + HighPerformanceTcpStatus* failure = nullptr); + + bool tracks(uint64_t base, uint64_t length) const; + + private: + Status acquire(uint64_t addr, uint64_t length, uint64_t registration_id, + HighPerformanceTcpOpcode opcode, bool remote, Lease* lease, + HighPerformanceTcpStatus* failure); + + mutable std::mutex registry_mutex_; + std::map> entries_; + // Fence stale capabilities from an earlier registry incarnation while the + // sequence prevents reuse within this incarnation. + const uint64_t registration_namespace_; + uint64_t next_registration_sequence_{1}; + bool closing_{false}; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_BUFFER_REGISTRY_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_client.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_client.h new file mode 100644 index 0000000000..f3076c561e --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_client.h @@ -0,0 +1,122 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_CLIENT_H_ +#define TENT_HP_TCP_CLIENT_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/common/types.h" +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +namespace mooncake::tent { + +// Process-wide client runtime for the HP TCP transport. It owns no thread of +// its own: every lane is created, queued, connected and driven exclusively on +// the worker selected by AffinityKey. +class HighPerformanceTcpClient { + public: + struct Config { + uint64_t max_transfer_bytes{1ULL << 30}; + size_t chunk_size{1ULL << 20}; + uint64_t connect_timeout_ms{2000}; + uint64_t progress_timeout_ms{30000}; + size_t connections_per_peer{4}; + }; + + struct Operation { + SegmentID peer_id{0}; + std::string incarnation; + std::string host; + uint16_t port{0}; + uint32_t lane_id{0}; + uint64_t registration_id{0}; + uint64_t remote_addr{0}; + void* local_addr{nullptr}; + uint64_t length{0}; + HighPerformanceTcpOpcode opcode{HighPerformanceTcpOpcode::kRead}; + uint64_t request_id{0}; + std::function)> + complete; + }; + + HighPerformanceTcpClient(Config config, HighPerformanceTcpWorkers* workers); + ~HighPerformanceTcpClient(); + + HighPerformanceTcpClient(const HighPerformanceTcpClient&) = delete; + HighPerformanceTcpClient& operator=(const HighPerformanceTcpClient&) = + delete; + + // Must execute on owner_worker. The transport reaches this method through + // the owner's ASIO event queue after global admission succeeds. + void enqueueOnOwner(size_t owner_worker, Operation operation); + + // Non-worker quiesce barrier. Cancels every queued/connecting/in-flight + // operation and waits until all operation callbacks have retired. + Status cancelAll(TransferStatusEnum terminal = CANCELED); + + // Best-effort cancellation for one logical request. If the request is + // still in the transport dispatch queue, the adapter's cancel flag settles + // it; if it already reached a lane, this posts cancellation to that owner. + Status cancelRequest(size_t owner_worker, uint64_t request_id); + + uint64_t connectionsCreatedForTest() const { + return connections_created_.load(std::memory_order_acquire); + } + uint64_t activeOperations() const { + return active_operations_.load(std::memory_order_acquire); + } + + private: + struct LaneKey { + SegmentID peer_id{0}; + std::string incarnation; + std::string host; + uint16_t port{0}; + uint32_t lane_id{0}; + + bool operator==(const LaneKey& other) const; + }; + + struct LaneKeyHash { + size_t operator()(const LaneKey& key) const; + }; + + class Lane; + + struct WorkerState { + std::unordered_map, LaneKeyHash> lanes; + }; + + void cancelWorker(size_t worker_id, TransferStatusEnum terminal); + void cancelRequestOnWorker(size_t worker_id, uint64_t request_id); + void operationStarted(); + void operationFinished(); + + Config config_; + HighPerformanceTcpWorkers* workers_{nullptr}; + std::vector worker_states_; + std::atomic stopping_{false}; + + std::atomic connections_created_{0}; + std::atomic active_operations_{0}; + mutable std::mutex active_mutex_; + std::condition_variable active_cv_; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_CLIENT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_protocol.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_protocol.h new file mode 100644 index 0000000000..76a51d34f9 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_protocol.h @@ -0,0 +1,89 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_PROTOCOL_H_ +#define TENT_HP_TCP_PROTOCOL_H_ + +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/common/types.h" + +namespace mooncake::tent { + +constexpr uint32_t kHighPerformanceTcpMagic = 0x4d435450; // "MCTP" +constexpr uint16_t kHighPerformanceTcpVersion = 1; +constexpr size_t kHighPerformanceTcpRequestSize = 48; +constexpr size_t kHighPerformanceTcpResponseSize = 32; + +enum class HighPerformanceTcpOpcode : uint8_t { + kRead = 1, + kWrite = 2, +}; + +enum class HighPerformanceTcpStatus : uint16_t { + kOk = 0, + kBadVersion = 1, + kBadOpcode = 2, + kBadLength = 3, + kRangeRejected = 4, + kPermissionDenied = 5, + kStaleRegistration = 6, + kShuttingDown = 7, + kInternalError = 8, +}; + +struct HighPerformanceTcpRequestFrame { + HighPerformanceTcpOpcode opcode{HighPerformanceTcpOpcode::kRead}; + uint64_t request_id{0}; + uint64_t registration_id{0}; + uint64_t remote_addr{0}; + uint64_t length{0}; +}; + +struct HighPerformanceTcpResponseFrame { + HighPerformanceTcpStatus status{HighPerformanceTcpStatus::kOk}; + uint64_t request_id{0}; + uint64_t committed_bytes{0}; +}; + +std::array +EncodeHighPerformanceTcpRequest(const HighPerformanceTcpRequestFrame& frame); + +Status DecodeHighPerformanceTcpRequest( + const uint8_t* bytes, size_t size, HighPerformanceTcpRequestFrame* frame, + HighPerformanceTcpStatus* wire_error = nullptr); + +std::array +EncodeHighPerformanceTcpResponse(const HighPerformanceTcpResponseFrame& frame); + +Status DecodeHighPerformanceTcpResponse(const uint8_t* bytes, size_t size, + HighPerformanceTcpResponseFrame* frame); + +struct HighPerformanceTcpEndpointAttr { + std::string incarnation; + std::string host; + uint16_t port{0}; + uint64_t max_transfer_bytes{0}; +}; + +struct HighPerformanceTcpBufferAttr { + uint64_t registration_id{0}; + std::string permission; +}; + +Status EncodeHighPerformanceTcpEndpointAttr( + const HighPerformanceTcpEndpointAttr& attr, std::string* encoded); +Status DecodeHighPerformanceTcpEndpointAttr( + const std::string& encoded, HighPerformanceTcpEndpointAttr* attr); +Status EncodeHighPerformanceTcpBufferAttr( + const HighPerformanceTcpBufferAttr& attr, std::string* encoded); +Status DecodeHighPerformanceTcpBufferAttr(const std::string& encoded, + HighPerformanceTcpBufferAttr* attr); + +const char* HighPerformanceTcpPermissionName(Permission permission); + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_PROTOCOL_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_server.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_server.h new file mode 100644 index 0000000000..c1b65e3e94 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_server.h @@ -0,0 +1,87 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_SERVER_H_ +#define TENT_HP_TCP_SERVER_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +namespace mooncake::tent { + +class HighPerformanceTcpServer { + public: + struct Config { + std::string bind_address; + uint16_t port{0}; + uint64_t max_transfer_bytes{1ULL << 30}; + size_t chunk_size{1ULL << 20}; + uint64_t progress_timeout_ms{30000}; + size_t max_connections{4096}; + }; + + HighPerformanceTcpServer(Config config, + HighPerformanceTcpBufferRegistry* registry, + HighPerformanceTcpWorkers* workers); + ~HighPerformanceTcpServer(); + + HighPerformanceTcpServer(const HighPerformanceTcpServer&) = delete; + HighPerformanceTcpServer& operator=(const HighPerformanceTcpServer&) = + delete; + + Status start(uint16_t* bound_port); + Status stopAccepting(); + Status cancelAll(); + Status stop(); + + size_t activeSessionsForTest() const { + return active_sessions_.load(std::memory_order_acquire); + } + + private: + class Session; + + Status startAccept(); + void installAcceptedSocket(size_t worker_id, + std::shared_ptr socket); + bool reserveConnection(); + void onSessionClosed(size_t worker_id, + const std::shared_ptr& session); + void cancelWorkerSessions(size_t worker_id); + + Config config_; + HighPerformanceTcpBufferRegistry* registry_{nullptr}; + HighPerformanceTcpWorkers* workers_{nullptr}; + + asio::io_context accept_io_; + std::optional> + accept_guard_; + std::unique_ptr acceptor_; + std::thread accept_thread_; + std::atomic started_{false}; + std::atomic stopping_{false}; + std::atomic next_worker_{0}; + std::atomic active_sessions_{0}; + + // Each set is touched only on the corresponding worker context. + std::vector>> sessions_; + mutable std::mutex sessions_wait_mutex_; + std::condition_variable sessions_wait_cv_; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_SERVER_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_task.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_task.h new file mode 100644 index 0000000000..55f747d9c7 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_task.h @@ -0,0 +1,85 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_TASK_H_ +#define TENT_HP_TCP_TASK_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/types.h" +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +namespace mooncake::tent { + +// TENT-visible task shell. Socket callbacks never own the SubBatch; they hold a +// shared_ptr to this state instead. A task becomes admission-reserved only in +// the no-throw batch commit section. +class HighPerformanceTcpTaskState { + public: + HighPerformanceTcpTaskState( + uint64_t reserved_bytes, BatchID progress_batch_id, + std::function notify_progress, + HighPerformanceTcpBufferRegistry::Lease local_lease) + : progress_batch_id_(progress_batch_id), + notify_progress_(std::move(notify_progress)), + local_lease_(std::move(local_lease)), + reserved_bytes_(reserved_bytes) {} + + HighPerformanceTcpTaskState(const HighPerformanceTcpTaskState&) = delete; + HighPerformanceTcpTaskState& operator=(const HighPerformanceTcpTaskState&) = + delete; + + void activateReservation( + HighPerformanceTcpAdmissionController* admission) noexcept { + admission_ = admission; + reservation_active_.store(true, std::memory_order_release); + } + + bool completeOnce(TransferStatusEnum terminal, size_t bytes, + std::optional remote_status = + std::nullopt) noexcept; + TransferStatus snapshot() const noexcept; + std::optional remoteStatus() const noexcept; + + void setDispatchIdentity(size_t owner_worker, + uint64_t request_id) noexcept { + owner_worker_ = owner_worker; + request_id_ = request_id; + } + size_t ownerWorker() const noexcept { return owner_worker_; } + uint64_t requestId() const noexcept { return request_id_; } + + void requestCancel() noexcept { + cancel_requested_.store(true, std::memory_order_release); + } + bool cancelRequested() const noexcept { + return cancel_requested_.load(std::memory_order_acquire); + } + + private: + std::atomic status_{PENDING}; + std::atomic bytes_{0}; + std::atomic remote_status_{ + HighPerformanceTcpStatus::kOk}; + std::atomic completion_claimed_{false}; + std::atomic cancel_requested_{false}; + + BatchID progress_batch_id_{0}; + std::function notify_progress_; + HighPerformanceTcpBufferRegistry::Lease local_lease_; + + HighPerformanceTcpAdmissionController* admission_{nullptr}; + uint64_t reserved_bytes_{0}; + std::atomic reservation_active_{false}; + + size_t owner_worker_{0}; + uint64_t request_id_{0}; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_TASK_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_transport.h new file mode 100644 index 0000000000..03bd9fb91e --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_transport.h @@ -0,0 +1,104 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_TRANSPORT_H_ +#define TENT_HP_TCP_TRANSPORT_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/hp_tcp_transport_config.h" +#include "tent/runtime/transport.h" +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" +#include "tent/transport/hp_tcp/hp_tcp_client.h" +#include "tent/transport/hp_tcp/hp_tcp_server.h" +#include "tent/transport/hp_tcp/hp_tcp_task.h" +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +namespace mooncake::tent { + +struct HighPerformanceTcpSubBatch : Transport::SubBatch { + std::vector> tasks; + size_t max_size{0}; + + size_t size() const override { return tasks.size(); } +}; + +class HighPerformanceTcpTransport final : public Transport { + public: + HighPerformanceTcpTransport(); + explicit HighPerformanceTcpTransport(HighPerformanceTcpParams params); + ~HighPerformanceTcpTransport() override; + + Status install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr config) override; + Status uninstall() override; + Status quiesce() override; + + Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override; + Status freeSubBatch(SubBatchRef& batch) override; + Status submitTransferTasks(SubBatchRef batch, + const std::vector& requests) override; + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override; + Status retryTransferTask(SubBatchRef batch, int task_id, + const Request& request) override; + + bool supportsCancellation() const override { return true; } + Status cancelTransferTask(SubBatchRef batch, int task_id) override; + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) override; + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override; + Status removeMemoryBuffer(BufferDesc& desc) override; + bool tracksLocalBuffer(const BufferDesc& desc) const override { + return registry_.tracks(desc.addr, desc.length); + } + + bool supportNotification() const override { return true; } + Status sendNotification(SegmentID target_id, + const Notification& notification) override; + Status receiveNotification( + std::vector& notifications) override; + + const char* getName() const override { return "hp_tcp"; } + + private: + friend class HighPerformanceTcpTransportTestPeer; + + struct TaskPlan; + + Status validateParams() const; + Status planTask(const Request& request, HighPerformanceTcpSubBatch* batch, + TaskPlan* plan); + Status rollbackPublishedEndpoint( + const std::optional& previous_attr); + Status stopRuntime(); + std::string makeIncarnation() const; + + HighPerformanceTcpParams params_; + std::shared_ptr metadata_; + + std::unique_ptr admission_; + std::unique_ptr workers_; + std::unique_ptr client_; + std::unique_ptr server_; + HighPerformanceTcpBufferRegistry registry_; + + std::atomic next_request_id_{1}; + std::atomic installed_{false}; + std::atomic stopping_{false}; + mutable std::mutex lifecycle_mutex_; + + RWSpinlock notify_lock_; + std::vector notifications_; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_workers.h b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_workers.h new file mode 100644 index 0000000000..8ebf5f675a --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/hp_tcp/hp_tcp_workers.h @@ -0,0 +1,119 @@ +// Copyright 2026 KVCache.AI +#ifndef TENT_HP_TCP_WORKERS_H_ +#define TENT_HP_TCP_WORKERS_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" + +namespace mooncake::tent { + +// A single process-wide bound for accepted HP TCP work. Reservations remain +// active until the corresponding task reaches a terminal state. +class HighPerformanceTcpAdmissionController { + public: + HighPerformanceTcpAdmissionController(uint64_t max_tasks, + uint64_t max_bytes) + : max_tasks_(max_tasks), max_bytes_(max_bytes) {} + + Status tryReserve(uint64_t tasks, uint64_t bytes); + void release(uint64_t tasks, uint64_t bytes); + void close(); + // Returns an error rather than waiting forever if accounting was + // corrupted. The counters are intentionally left unchanged in that case. + Status waitForZero(); + bool failed() const; + uint64_t outstandingTasks() const; + uint64_t outstandingBytes() const; + + private: + const uint64_t max_tasks_; + const uint64_t max_bytes_; + mutable std::mutex mutex_; + std::condition_variable zero_cv_; + uint64_t tasks_{0}; + uint64_t bytes_{0}; + bool accepting_{true}; + std::atomic failed_{false}; +}; + +// Thin owner-thread pool for sockets. ASIO is the queue; this class adds +// deterministic affinity and lifecycle management. Accepted transfer work is +// bounded by HighPerformanceTcpAdmissionController. +class HighPerformanceTcpWorkers { + public: + struct Config { + size_t worker_count{16}; + }; + + using Task = std::function; + struct Command { + size_t worker_id{0}; + Task run; + std::function cancel; + }; + + HighPerformanceTcpWorkers(); + explicit HighPerformanceTcpWorkers(Config config); + ~HighPerformanceTcpWorkers(); + + HighPerformanceTcpWorkers(const HighPerformanceTcpWorkers&) = delete; + HighPerformanceTcpWorkers& operator=(const HighPerformanceTcpWorkers&) = + delete; + + Status start(); + Status stop(); + Status tryCommitBatch(std::vector& commands, + HighPerformanceTcpAdmissionController* admission, + uint64_t reserve_tasks, uint64_t reserve_bytes, + const std::function& on_commit); + + Status barrier(); + size_t affinityOwner(uint64_t peer, uint32_t lane) const; + bool running() const { return running_.load(std::memory_order_acquire); } + bool hasFailedWorker() const { + return failed_.load(std::memory_order_acquire); + } + // A stopped context is kept alive until client and server teardown has + // destroyed every socket and resolver that uses it, but it cannot run + // control work anymore. + bool controlContextAvailable() const { + return running() && !workers_.empty(); + } + size_t workerCount() const { return config_.worker_count; } + asio::io_context& ioContext(size_t worker_id); + bool onWorkerThread() const; + + private: + struct WorkerContext { + WorkerContext() : guard(asio::make_work_guard(io)) {} + asio::io_context io; + asio::executor_work_guard guard; + std::thread thread; + }; + + void runCommand(Command command); + void markWorkerFailed() noexcept; + + Config config_; + std::vector> workers_; + std::atomic running_{false}; + std::atomic failed_{false}; + bool started_{false}; + mutable std::mutex lifecycle_mutex_; + mutable std::mutex submit_mutex_; +}; + +} // namespace mooncake::tent + +#endif // TENT_HP_TCP_WORKERS_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h index 434694fdd1..3126812c6a 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h @@ -46,6 +46,7 @@ class RdmaTransport; class RdmaContext { friend class RdmaCQ; friend class RdmaEndPoint; + friend class RdmaContextTestPeer; public: RdmaContext(RdmaTransport &transport); @@ -113,7 +114,9 @@ class RdmaContext { // The one port this context opened; 0 for a slot that never constructed. uint8_t portNum() const { return params_ ? params_->device.port : 0; } - // Negotiated port speed in Gbps, 0 when it could not be determined. + // Port speed in Gbps: the effective speed from ibv_query_port_speed() + // where the library provides it (LAG-aware), otherwise the negotiated + // link rate. 0 when neither could be determined. double linkSpeedGbps() const; // Re-read the port's negotiated speed and width from the hardware, so @@ -123,6 +126,15 @@ class RdmaContext { // open or the query fails. int refreshPortAttributes(); + // Read the port's current state from the hardware. The monitor thread + // polls this for paused contexts so a link that came back without a + // usable IBV_EVENT_PORT_ACTIVE is still noticed. Deliberately leaves the + // cached speed/width alone: refreshLinkSpeed() diffs against them, so + // refreshing here would hide a renegotiated rate from the selector. + // Returns -1 and leaves *state untouched if no device is open or the + // query fails. + int queryPortState(ibv_port_state *state) const; + int eventFd() const { return event_fd_; } RdmaCQ *cq(int index); @@ -131,6 +143,18 @@ class RdmaContext { RdmaParams ¶ms() const { return *params_.get(); } + // True while linkSpeedGbps() is derived from ibv_query_port_speed() + // rather than the encoded speed x width. + bool effectiveSpeedKnown() const { + return effective_speed_mbps_.load(std::memory_order_relaxed) > 0; + } + + // ibv_query_port_speed() errors since the device was opened. The verb + // being absent, or reporting 0, is not an error. + uint64_t effectiveSpeedQueryFailures() const { + return effective_speed_query_failures_.load(std::memory_order_relaxed); + } + // PCIe Relaxed Ordering support bool isRelaxedOrderingEnabled() const { return relaxed_ordering_enabled_; } @@ -141,6 +165,12 @@ class RdmaContext { int openDevice(const std::string &device_name, uint8_t port); // Decode one ibv_query_port result into active_speed_/active_width_. void recordPortSpeed(const ibv_port_attr &port_attr); + // Ask ibv_query_port_speed() for the effective speed when the library + // has it; records 0 when the verb is absent or reports nothing, so + // linkSpeedGbps() falls back. A verb *error* keeps the last known value + // instead: on a degraded LAG, falling back would briefly restore the + // higher encoded rate. Failures are counted and logged on transition. + void queryEffectiveSpeed(); // Release every resource currently owned by this context. This is // intentionally state-independent so it can clean up a partially completed @@ -169,6 +199,11 @@ class RdmaContext { // atomic so a reader added elsewhere stays well-defined. std::atomic active_speed_{0}; std::atomic active_width_{0}; + // From ibv_query_port_speed(), converted to Mb/s; 0 = unavailable. + std::atomic effective_speed_mbps_{0}; + std::atomic effective_speed_query_failures_{0}; + // Whether the last query errored; drives the transition logging. + std::atomic effective_speed_query_failing_{false}; int gid_index_ = -1; ibv_gid gid_; @@ -184,7 +219,11 @@ class RdmaContext { // PCIe Relaxed Ordering support bool relaxed_ordering_enabled_ = false; - const IbvSymbols &verbs_; + // The context's own copy of the loader's verbs table (copied once at + // construction, read-only afterwards). A copy rather than a reference so + // tests can substitute individual entries and drive the port-attribute + // and event paths without an RNIC. + IbvSymbols verbs_; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h index 896c54dff6..76e3ada01d 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h @@ -196,6 +196,12 @@ class RdmaEndPoint : public std::enable_shared_from_this { void postNotifyRecv(size_t idx); void repostAllNotifyRecvs(); + static char* notifySlotPtr(char* base, size_t idx); + static bool encodeNotifyPayload(char* slot, const std::string& name, + const std::string& msg, uint32_t* out_len); + static bool decodeNotifyPayload(const char* data, size_t byte_len, + std::string* name, std::string* msg); + private: friend class EndpointTestAccess; @@ -226,13 +232,14 @@ class RdmaEndPoint : public std::enable_shared_from_this { // Notification QP (one per endpoint for control plane operations) ibv_qp* notify_qp_ = nullptr; - // Notification buffers + // Notification buffers. Send and recv each use one contiguous host buffer + // split into kNotifyMaxPendingSends slots, registered with a single MR. static constexpr size_t kNotifyBufferSize = 65536; // 64 KB static constexpr size_t kNotifyMaxPendingSends = 256; - std::vector> notify_recv_buffers_; - std::vector notify_recv_mrs_; // Memory regions for recv buffers - std::vector notify_send_buffer_; // Single contiguous send buffer - ibv_mr* notify_send_mr_ = nullptr; // Single MR for all send slots + std::vector notify_recv_buffer_; + ibv_mr* notify_recv_mr_ = nullptr; + std::vector notify_send_buffer_; + ibv_mr* notify_send_mr_ = nullptr; // Serializes notification buffer/QP access against deconstruction. std::mutex notify_resource_mutex_; std::mutex notify_send_mutex_; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/ibv_loader.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/ibv_loader.h index 4b0f68d719..a6c7196235 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/ibv_loader.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/ibv_loader.h @@ -31,6 +31,11 @@ struct IbvSymbols { int index, union ibv_gid* gid); int (*ibv_query_port_default)(ibv_context* context, uint8_t port_num, ibv_port_attr* port_attr); + // Optional (rdma-core >= 62): effective port speed in 100 Mb/s units, + // LAG-aware. nullptr on older libraries; callers fall back to the + // ibv_port_attr encodings. + int (*ibv_query_port_speed)(ibv_context* context, uint32_t port_num, + uint64_t* port_speed); const char* (*ibv_get_device_name)(struct ibv_device* device); ibv_pd* (*ibv_alloc_pd)(ibv_context* context); @@ -84,4 +89,4 @@ class IbvLoader { } // namespace tent } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h index 6e08b44c02..99e9beed55 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h @@ -182,6 +182,10 @@ class DeviceSelector { double numa_tier_weights[Topology::DevicePriorityRanks] = {1.0, 5.0, 10.0}; + // Hard-exclude known cross-NUMA NICs; unknown NUMA keeps + // numa_tier_weights. + bool strict_local_numa = false; + // EWMA bandwidth learning rate (0.0 = full adaptation, 1.0 = no // learning) double bandwidth_learning_rate = 0.01; @@ -222,6 +226,9 @@ class DeviceSelector { return sched_params_; } + // Startup warning if the flag excludes every NIC or classifies none. + void auditStrictLocalNuma() const; + private: std::shared_ptr local_topology_; std::unordered_map devices_; @@ -240,6 +247,19 @@ class DeviceSelector { it->second.available.load(std::memory_order_relaxed); } + const char *noEligibleDeviceReason() const { + return sched_params_.strict_local_numa + ? "no eligible devices (strict_local_numa excludes " + "cross-NUMA NICs)" + : "no eligible devices"; + } + + bool isNumaEligible(const Topology::MemEntry *entry, int dev_id) const { + if (!sched_params_.strict_local_numa) return true; + if (!entry || !local_topology_) return true; + return !local_topology_->isCrossNuma(*entry, dev_id); + } + Status buildCandidates(const Topology::MemEntry *entry, uint64_t slice_bytes, uint64_t device_mask, std::vector &candidates, diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h index 9868100007..2a354c6280 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -84,6 +84,9 @@ class Workers { // dev_id is the NicID (context_set_ index), which is also the // DeviceSelector id, so port events can flip that device's availability. + // Drains the whole async event queue: the async fd is edge-triggered and + // ibv_get_async_event() dequeues one record per call, so an event left + // behind waits for an unrelated later event to release it. int handleContextEvents(int dev_id, std::shared_ptr& context); // The decision half of handleContextEvents, kept apart from @@ -94,11 +97,24 @@ class Workers { void applyContextEvent(int dev_id, RdmaContext& context, const ibv_async_event& event); + // Everything a recovered port needs: resume the context, re-seed its + // bandwidth and make it selectable again. Shared by the + // IBV_EVENT_PORT_ACTIVE path and by resumePausedContexts(). + void activateContext(int dev_id, RdmaContext& context); + // Re-read the link speed after a port event and re-seed the selector if // it changed; a link that returns at the same speed keeps what it // learned. void refreshLinkSpeed(int dev_id, RdmaContext& context); + // 1 Hz heartbeat from monitorThread(): safety net for a lost + // IBV_EVENT_PORT_ACTIVE. Nothing else ever leaves DEVICE_PAUSED, so a + // context whose recovery event was dropped (edge-triggered fd, event + // queue overflow, ...) would fail every transfer on that NIC forever. + // Polls the port state of paused contexts and activates the ones the + // hardware reports as up. + void resumePausedContexts(); + Status generatePostPath(RdmaSlice* slice); private: @@ -124,6 +140,8 @@ class Workers { int getDeviceByFlatIndex(const RouteHint& hint, size_t flat_idx); + bool strictLocalNuma() const; + // True if the (sdev -> tdev) NIC pair is known-unable to GPUDirect-DMA to // the source/target GPU (learned from prior completion errors). Used to // steer selection away from dead rails before posting. diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h new file mode 100644 index 0000000000..ddd410f961 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h @@ -0,0 +1,61 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ +#define TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ + +#include +#include +#include +#include +#include + +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake { +namespace tent { +namespace ub { + +// Heuristic used when device_filter is empty: prefer UBAGG bonding devices +// over underlying physical ports (e.g. udmac*) that appear alongside them. +inline bool isBondingDeviceName(std::string_view name) { + std::string lower(name); + std::transform( + lower.begin(), lower.end(), lower.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + if (lower.rfind("bonding", 0) == 0) return true; + if (lower.find(":bonding") != std::string::npos) return true; + if (lower.find("_bond") != std::string::npos) return true; + if (lower.find("-bond") != std::string::npos) return true; + return false; +} + +inline bool isBondingDevice(const DeviceInfo& device) { + return isBondingDeviceName(device.native_device_name) || + isBondingDeviceName(device.topology_name); +} + +// When explicit_filter is true, returns devices unchanged (caller already +// applied device_filter). When false and at least one bonding device is +// present, returns only bonding devices; otherwise returns all devices. +inline std::vector preferBondingDevicesIfPresent( + const std::vector& devices, bool explicit_filter) { + if (explicit_filter || devices.empty()) return devices; + const bool has_bonding = + std::any_of(devices.begin(), devices.end(), + [](const DeviceInfo& d) { return isBondingDevice(d); }); + if (!has_bonding) return devices; + + std::vector selected; + selected.reserve(devices.size()); + for (const auto& device : devices) { + if (isBondingDevice(device)) selected.push_back(device); + } + return selected; +} + +} // namespace ub +} // namespace tent +} // namespace mooncake + +#endif // TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h new file mode 100644 index 0000000000..3dfe4533ba --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h @@ -0,0 +1,166 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_ENDPOINT_H_ +#define TENT_TRANSPORT_UB_ENDPOINT_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/common/types.h" +#include "tent/runtime/control_plane.h" +#include "tent/runtime/topology.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +// One cache entry represents a local-device to remote-device path. The peer +// NIC path is part of the identity because a remote topology ID may be reused +// after the peer republishes its topology. +struct UbEndpointKey { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + Topology::NicID remote_topology_id{-1}; + std::string peer_nic_path; + + bool operator==(const UbEndpointKey&) const = default; + + [[nodiscard]] bool valid() const noexcept { + return local_topology_id >= 0 && remote_topology_id >= 0 && + !peer_nic_path.empty(); + } +}; + +struct UbEndpointKeyHash { + size_t operator()(const UbEndpointKey& key) const noexcept; +}; + +// A UB endpoint is one immutable incarnation of a Jetty set. Its generation +// is allocated process-wide and is never reused. Lifecycle operations are +// serialized internally; a failed endpoint can only retire, never reconnect. +class UbEndpoint final : public std::enable_shared_from_this { + public: + enum class State : uint8_t { + kUninitialized, + kHandshaking, + kPrepared, + kBinding, + kReady, + kFailed, + kDestroying, + kDestroyed, + }; + + UbEndpoint(UbEndpointKey key, UbContextPtr context, + std::shared_ptr adapter, uint32_t jetty_count, + JettyOptions jetty_options = {}); + ~UbEndpoint(); + + UbEndpoint(const UbEndpoint&) = delete; + UbEndpoint& operator=(const UbEndpoint&) = delete; + + // Creates the local Jetty set. Concurrent calls share the same attempt. + // An error permanently moves this incarnation to kFailed. + Status prepare(); + + // Binds each local Jetty to the peer EID and the corresponding peer Jetty + // ID. The bootstrap must describe exactly one peer Jetty per local Jetty. + Status bind(const UbBootstrapDesc& peer); + + // Builds the local half of the native UB bootstrap after prepare(). + Status makeBootstrapDesc(const std::string& segment_name, + const std::string& local_nic_path, + const std::string& peer_nic_path, + uint64_t segment_generation, + UbBootstrapDesc& output) const; + + // Admission to the posting path is synchronized with retirement: once + // kDestroying is visible no new work can acquire the endpoint. Every + // successful acquire must have one matching release. + [[nodiscard]] bool tryAcquireOutstanding(uint64_t bytes = 0) noexcept; + void releaseOutstanding(uint64_t bytes = 0) noexcept; + + // Stops new posts immediately. With outstanding work, native resources + // remain intact until either completions drain naturally or quiesce() + // establishes an explicit no-more-DMA fence. Idempotent. + Status retire(); + + // Establishes a native drain fence for every Jetty, then resets/unbinds + // them. Returned completions still own the corresponding logical tokens + // and must be dispatched by Workers. A failed fence leaves all resources + // alive so shutdown can be retried safely. + Status quiesce(uint32_t timeout_ms, std::vector& completions); + + [[nodiscard]] const UbEndpointKey& key() const noexcept { return key_; } + [[nodiscard]] const UbContextPtr& context() const noexcept { + return context_; + } + [[nodiscard]] uint64_t generation() const noexcept { return generation_; } + [[nodiscard]] State state() const noexcept { + return state_.load(std::memory_order_acquire); + } + [[nodiscard]] bool ready() const noexcept { + return state() == State::kReady; + } + [[nodiscard]] bool failed() const noexcept { + return state() == State::kFailed; + } + [[nodiscard]] bool reusable() const noexcept; + [[nodiscard]] uint64_t peerGeneration() const noexcept { + return peer_generation_.load(std::memory_order_acquire); + } + [[nodiscard]] uint64_t outstandingWrs() const noexcept { + return outstanding_wrs_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t outstandingBytes() const noexcept { + return outstanding_bytes_.load(std::memory_order_relaxed); + } + [[nodiscard]] size_t jettyCount() const; + [[nodiscard]] JettyPtr jetty(size_t index) const; + [[nodiscard]] size_t jfcIndex(size_t jetty_index) const; + [[nodiscard]] std::vector jetties() const; + [[nodiscard]] Status lifecycleStatus() const; + + private: + static uint64_t allocateGeneration() noexcept; + + Status failLocked(Status status); + Status resetAndUnbindLocked(); + Status deleteJettysLocked(); + Status finishRetireLocked(); + static void rememberFirstError(const Status& candidate, Status& first); + + const UbEndpointKey key_; + const UbContextPtr context_; + const std::shared_ptr adapter_; + const uint32_t jetty_count_; + const JettyOptions jetty_options_; + const uint64_t generation_; + + mutable std::mutex lifecycle_mutex_; + std::vector jetties_; + std::vector jfc_indices_; + UbBootstrapDesc peer_; + Status lifecycle_status_; + Status retire_status_; + bool native_quiesced_{false}; + std::atomic state_{State::kUninitialized}; + std::atomic peer_generation_{0}; + std::atomic outstanding_wrs_{0}; + std::atomic outstanding_bytes_{0}; +}; + +using UbEndpointPtr = std::shared_ptr; +// Keep the spelling used by the design document available to callers. +using UbEndPoint = UbEndpoint; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_ENDPOINT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h new file mode 100644 index 0000000000..06d394c1d9 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h @@ -0,0 +1,67 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ +#define TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/endpoint.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +// Generation-aware endpoint cache. Failed and retired incarnations are +// unpublished before cleanup, so a subsequent lookup always allocates a new +// generation and can never resurrect a failed Jetty set. +class EndpointStore final { + public: + EndpointStore(std::shared_ptr adapter, size_t max_size, + uint32_t jetty_count, JettyOptions jetty_options = {}); + ~EndpointStore(); + + EndpointStore(const EndpointStore&) = delete; + EndpointStore& operator=(const EndpointStore&) = delete; + + std::shared_ptr get(const UbEndpointKey& key); + Status getOrCreate(const UbEndpointKey& key, const UbContextPtr& context, + std::shared_ptr& endpoint); + + // Only the exact incarnation is removed. A late timeout from generation N + // therefore cannot evict a replacement generation N+1. + bool retire(const UbEndpointKey& key, uint64_t generation); + bool retire(const std::shared_ptr& endpoint); + Status clear(); + [[nodiscard]] size_t size() const; + + private: + struct Entry { + std::shared_ptr endpoint; + uint64_t insertion_order{0}; + }; + + std::shared_ptr adapter_; + const size_t max_size_; + const uint32_t jetty_count_; + const JettyOptions jetty_options_; + mutable std::mutex mutex_; + std::unordered_map endpoints_; + // Unpublished endpoints whose native cleanup failed remain owned until a + // later clear()/shutdown retry. They must never fall through a destructor + // while an ERROR Jetty still lacks its flush fence. + std::vector> quarantined_; + uint64_t next_insertion_order_{1}; +}; + +using UbEndpointStore = EndpointStore; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h index 175dc8739c..d34074f7ed 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h @@ -40,6 +40,7 @@ struct UbParams { uint32_t jetty_per_endpoint = 6; uint32_t max_endpoints = 65536; size_t slice_size = 64 * 1024; + size_t max_slices_per_task = 64 * 1024; uint32_t max_retries = 8; uint32_t slice_timeout_ms = 5000; uint32_t endpoint_cooldown_ms = 1000; @@ -78,6 +79,15 @@ struct UbParams { } parsed.slice_size = static_cast(slice_size); + uint64_t max_slices_per_task = parsed.max_slices_per_task; + CHECK_STATUS(readPositive(config, "transports/ub/max_slices_per_task", + max_slices_per_task, max_slices_per_task)); + if (max_slices_per_task > std::numeric_limits::max()) { + return Status::InvalidArgument( + "transports/ub/max_slices_per_task exceeds size_t"); + } + parsed.max_slices_per_task = static_cast(max_slices_per_task); + CHECK_STATUS(readNonNegative(config, "transports/ub/max_retries", parsed.max_retries, parsed.max_retries)); CHECK_STATUS(readPositive(config, "transports/ub/slice_timeout_ms", diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h new file mode 100644 index 0000000000..c135899cdc --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h @@ -0,0 +1,195 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/transport/ub/rail_monitor.h" + +namespace mooncake::tent::ub { + +struct QuotaLimits { + uint64_t max_inflight_bytes{std::numeric_limits::max()}; + uint64_t max_outstanding_wrs{std::numeric_limits::max()}; + + bool operator==(const QuotaLimits&) const = default; +}; + +struct QuotaUsage { + uint64_t inflight_bytes{0}; + uint64_t outstanding_wrs{0}; + + bool operator==(const QuotaUsage&) const = default; +}; + +struct QuotaStats { + QuotaLimits limits{}; + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; +}; + +// A reservation is a copyable release token. QuotaManager retains the +// authoritative path and charge, so changing those descriptive fields cannot +// decrement the wrong counters, and releasing a copied token cannot do so +// twice. +struct QuotaReservation { + uint64_t id{0}; + UbPostPath path{}; + uint64_t bytes{0}; + uint64_t wrs{0}; + + [[nodiscard]] bool valid() const { return id != 0; } +}; + +struct DeviceQuotaStats : QuotaStats { + Topology::NicID local_topology_id{-1}; +}; + +struct PathQuotaStats : QuotaStats { + UbPostPath path{}; +}; + +struct AggregateQuotaStats { + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + size_t active_reservations{0}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; + uint64_t duplicate_release_attempts{0}; +}; + +// A lock-consistent capacity snapshot used by path selection. Pressure is the +// projected utilization after charging the requested work and is normalized +// to [0, 1]. Unlimited dimensions contribute no pressure. +struct QuotaAvailability { + bool can_acquire{false}; + double normalized_inflight{1.0}; + double normalized_outstanding_wrs{1.0}; +}; + +// Atomically enforces both physical-device and posting-path capacity. This is +// an internal sender-side limit and is intentionally independent of TENT's +// receiver-credit protocol. +class QuotaManager { + public: + explicit QuotaManager(QuotaLimits default_device_limits = {}, + QuotaLimits default_path_limits = {}); + + QuotaManager(const QuotaManager&) = delete; + QuotaManager& operator=(const QuotaManager&) = delete; + + void setDefaultDeviceLimits(const QuotaLimits& limits); + void setDefaultPathLimits(const QuotaLimits& limits); + [[nodiscard]] QuotaLimits defaultDeviceLimits() const; + [[nodiscard]] QuotaLimits defaultPathLimits() const; + + bool setDeviceLimits(Topology::NicID local_topology_id, + const QuotaLimits& limits); + bool clearDeviceLimits(Topology::NicID local_topology_id); + bool setPathLimits(const UbPostPath& path, const QuotaLimits& limits); + bool clearPathLimits(const UbPostPath& path); + + // Acquires device and path charges as one transaction. A zero-byte work + // request is supported, but wrs must be nonzero. + [[nodiscard]] std::optional tryAcquire( + const UbPostPath& path, uint64_t bytes, uint64_t wrs = 1); + + // Tries paths in caller-provided preference order under one lock. This is + // the commit point for multi-rail selection: if a preflight snapshot races + // with another posting worker, later rails are considered before the + // request is deferred. + [[nodiscard]] std::optional tryAcquireFirst( + const std::vector& paths, uint64_t bytes, uint64_t wrs = 1); + + // Returns projected device/path pressure without reserving capacity. + [[nodiscard]] QuotaAvailability availability(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs = 1) const; + + // The first release returns true. Releasing the same token again is a + // harmless no-op and returns false; usage can never underflow. + bool release(const QuotaReservation& reservation); + + [[nodiscard]] DeviceQuotaStats deviceStats( + Topology::NicID local_topology_id) const; + [[nodiscard]] PathQuotaStats pathStats(const UbPostPath& path) const; + [[nodiscard]] std::vector allDeviceStats() const; + [[nodiscard]] std::vector allPathStats() const; + [[nodiscard]] AggregateQuotaStats aggregateStats() const; + [[nodiscard]] size_t activeReservationCount() const; + + private: + struct QuotaRecord { + std::optional override_limits; + // Quota is charged to a physical rail, while this preserves the most + // recent endpoint incarnation for diagnostics. + UbPostPath latest_path{}; + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; + }; + + struct ActiveReservation { + UbPostPath path{}; + uint64_t bytes{0}; + uint64_t wrs{0}; + }; + + static bool fits(uint64_t current, uint64_t charge, uint64_t limit); + static double normalizedUsage(uint64_t current, uint64_t charge, + uint64_t limit); + static uint64_t saturatingAdd(uint64_t lhs, uint64_t rhs); + static void addUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs); + static void releaseUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs); + static void updatePeak(const QuotaUsage& usage, QuotaUsage& peak); + static QuotaLimits effectiveLimits(const QuotaRecord& record, + const QuotaLimits& defaults); + static QuotaStats makeStats(const QuotaRecord& record, + const QuotaLimits& defaults); + std::optional tryAcquireLocked( + const UbPostPath& path, uint64_t bytes, uint64_t wrs, + bool count_aggregate_reject); + QuotaAvailability availabilityLocked(const UbPostPath& path, uint64_t bytes, + uint64_t wrs) const; + uint64_t nextReservationIdLocked(); + + mutable std::mutex mutex_; + QuotaLimits default_device_limits_; + QuotaLimits default_path_limits_; + std::unordered_map devices_; + std::unordered_map paths_; + std::unordered_map active_reservations_; + AggregateQuotaStats aggregate_stats_{}; + uint64_t next_reservation_id_{1}; +}; + +using UbQuotaManager = QuotaManager; + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h new file mode 100644 index 0000000000..1dbf0fd6bf --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h @@ -0,0 +1,165 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/transport/ub/slice.h" + +namespace mooncake::tent::ub { + +// Health and learned bandwidth belong to a physical rail, rather than to one +// endpoint incarnation. Endpoint generation is therefore reported in the +// statistics but intentionally excluded from this key. +struct UbRailKey { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + int remote_device_id{-1}; + + bool operator==(const UbRailKey&) const = default; + + [[nodiscard]] bool valid() const { + return local_topology_id >= 0 && remote_device_id >= 0; + } + + static UbRailKey fromPath(const UbPostPath& path) { + return {path.local_topology_id, path.remote_segment_id, + path.remote_device_id}; + } +}; + +struct UbRailKeyHash { + size_t operator()(const UbRailKey& key) const noexcept { + size_t seed = std::hash{}(key.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(key.remote_segment_id)); + combine(std::hash{}(key.remote_device_id)); + return seed; + } +}; + +struct RailMonitorConfig { + uint32_t error_threshold{3}; + uint64_t error_window_ns{10'000'000'000ULL}; + uint64_t cooldown_ns{30'000'000'000ULL}; + // Weight assigned to the newest completion sample. + double ewma_alpha{0.2}; + + [[nodiscard]] bool valid() const { + return error_threshold != 0 && error_window_ns != 0 && + cooldown_ns != 0 && ewma_alpha > 0.0 && ewma_alpha <= 1.0; + } +}; + +struct RailStats { + UbRailKey key{}; + bool paused{false}; + uint32_t errors_in_window{0}; + uint64_t successful_completions{0}; + uint64_t completed_bytes{0}; + uint64_t completion_errors{0}; + uint64_t timeouts{0}; + uint64_t recoveries{0}; + uint64_t endpoint_rebuilds{0}; + uint64_t pauses{0}; + // Bytes/second and nanoseconds respectively. -1 means no valid sample. + double ewma_bandwidth_bytes_per_second{-1.0}; + double ewma_latency_ns{-1.0}; + uint64_t last_success_ns{0}; + uint64_t last_error_ns{0}; + uint64_t pause_started_ns{0}; + uint64_t cooldown_until_ns{0}; + uint64_t latest_endpoint_generation{0}; +}; + +// Thread-safe rolling health and telemetry for UB posting paths. +class RailMonitor { + public: + explicit RailMonitor(RailMonitorConfig config = {}); + + RailMonitor(const RailMonitor&) = delete; + RailMonitor& operator=(const RailMonitor&) = delete; + + // Rejects invalid configurations without changing the active one. + bool configure(const RailMonitorConfig& config); + [[nodiscard]] RailMonitorConfig config() const; + + // Registration is optional; all record operations create the rail lazily. + bool registerPath(const UbPostPath& path); + [[nodiscard]] bool available(const UbPostPath& path, uint64_t now_ns = 0); + + void recordSuccess(const UbPostPath& path, uint64_t bytes, + uint64_t latency_ns, uint64_t now_ns = 0); + void recordError(const UbPostPath& path, uint64_t now_ns = 0); + void recordTimeout(const UbPostPath& path, uint64_t now_ns = 0); + // Records at most one rebuild for each endpoint generation on a physical + // rail. Returns true when telemetry advanced, allowing EndpointStore to + // call this safely from converging/retried rebuild paths. + bool recordEndpointRebuild(const UbPostPath& path, uint64_t now_ns = 0); + + [[nodiscard]] RailStats stats(const UbPostPath& path, uint64_t now_ns = 0); + [[nodiscard]] std::vector allStats(uint64_t now_ns = 0); + + // Adds the best usable remote path sample for each local device. Returns + // -1 until at least one valid completion sample has been observed. + [[nodiscard]] double aggregateBandwidth(uint64_t now_ns = 0); + [[nodiscard]] size_t pathCount() const; + + private: + struct RailState { + RailStats stats{}; + std::deque recent_errors; + // Event timestamps may arrive out of order from different pollers. + // Health decisions never move behind this per-rail watermark. + uint64_t observed_through_ns{0}; + // Late errors at or before a completed recovery epoch must not + // resurrect an already-expired pause. + uint64_t ignore_errors_through_ns{0}; + // Endpoint generations are process-wide monotonic. Keep a separate + // watermark from latest_endpoint_generation because path registration + // may observe the replacement before rebuild telemetry is emitted. + uint64_t recorded_rebuild_generation{0}; + }; + + using RailMap = std::unordered_map; + + static uint64_t normalizedNow(uint64_t now_ns); + static uint64_t deadlineAfter(uint64_t now_ns, uint64_t duration_ns); + static uint64_t observeTimeLocked(RailState& state, uint64_t event_ns); + RailState& getOrCreateLocked(const UbPostPath& path); + static void insertErrorLocked(RailState& state, uint64_t event_ns); + void pruneErrorsLocked(RailState& state, uint64_t now_ns); + void refreshCooldownLocked(RailState& state, uint64_t now_ns); + void recordFailureLocked(const UbPostPath& path, uint64_t now_ns, + bool timeout); + + mutable std::mutex mutex_; + RailMonitorConfig config_; + RailMap rails_; +}; + +using UbRailMonitor = RailMonitor; + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h new file mode 100644 index 0000000000..3b70c3dc21 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h @@ -0,0 +1,707 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/types.h" +#include "tent/runtime/topology.h" + +namespace mooncake::tent::ub { + +inline uint64_t steadyNowNs() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +// Identifies one local-device to remote-device posting path. Endpoint +// generation is part of the identity so a completion from a retired endpoint +// cannot be mistaken for work posted on its replacement. +struct UbPostPath { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + int remote_device_id{-1}; + uint64_t endpoint_generation{0}; + + bool operator==(const UbPostPath&) const = default; + + [[nodiscard]] bool valid() const { + return local_topology_id >= 0 && remote_device_id >= 0 && + endpoint_generation != 0; + } +}; + +struct UbPostPathHash { + size_t operator()(const UbPostPath& path) const noexcept { + size_t seed = std::hash{}(path.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(path.remote_segment_id)); + combine(std::hash{}(path.remote_device_id)); + combine(std::hash{}(path.endpoint_generation)); + return seed; + } +}; + +enum class UbSliceState : uint8_t { + kInitial, + kQueued, + kPosting, + kPosted, + kRetryPending, + kCompleted, + kFailed, + kTimedOut, + kCanceled, + kInvalid, +}; + +inline bool isTerminal(UbSliceState state) { + switch (state) { + case UbSliceState::kCompleted: + case UbSliceState::kFailed: + case UbSliceState::kTimedOut: + case UbSliceState::kCanceled: + case UbSliceState::kInvalid: + return true; + default: + return false; + } +} + +inline TransferStatusEnum transferStatus(UbSliceState state) { + switch (state) { + case UbSliceState::kInitial: + return INITIAL; + case UbSliceState::kCompleted: + return COMPLETED; + case UbSliceState::kFailed: + return FAILED; + case UbSliceState::kTimedOut: + return TIMEOUT; + case UbSliceState::kCanceled: + return CANCELED; + case UbSliceState::kInvalid: + return INVALID; + default: + return PENDING; + } +} + +inline bool isTerminal(TransferStatusEnum status) { + return status == COMPLETED || status == FAILED || status == TIMEOUT || + status == CANCELED || status == INVALID; +} + +inline UbSliceState sliceState(TransferStatusEnum status) { + switch (status) { + case COMPLETED: + return UbSliceState::kCompleted; + case FAILED: + return UbSliceState::kFailed; + case TIMEOUT: + return UbSliceState::kTimedOut; + case CANCELED: + return UbSliceState::kCanceled; + case INVALID: + return UbSliceState::kInvalid; + case INITIAL: + return UbSliceState::kInitial; + case PENDING: + default: + return UbSliceState::kQueued; + } +} + +struct UbSliceSpec { + void* local_address{nullptr}; + uint64_t remote_address{0}; + size_t length{0}; + size_t request_offset{0}; + uint32_t max_retries{0}; +}; + +struct UbAttemptToken { + uint64_t slice_id{0}; + uint32_t attempt{0}; + UbPostPath path{}; + + bool operator==(const UbAttemptToken&) const = default; + + [[nodiscard]] bool valid() const { + return slice_id != 0 && attempt != 0 && path.valid(); + } +}; + +enum class UbAttemptResolution : uint8_t { + kIgnored, + kRetryScheduled, + kTerminal, +}; + +struct UbSliceSnapshot { + uint64_t id{0}; + UbSliceState state{UbSliceState::kInitial}; + UbPostPath path{}; + uint32_t attempt{0}; + uint32_t retry_count{0}; + uint32_t max_retries{0}; + bool cancel_requested{false}; + size_t transferred_bytes{0}; + uint64_t created_ns{0}; + uint64_t queued_ns{0}; + uint64_t attempt_started_ns{0}; + uint64_t posted_ns{0}; + uint64_t terminal_ns{0}; +}; + +struct UbTask; +class UbSlice; + +// Safe user-context payload for a posted operation. Holding this token keeps +// the slice alive until the device completion has been dispatched, while the +// slice itself keeps only a weak reference to its parent task. +struct UbCompletionToken { + std::shared_ptr slice; + UbAttemptToken attempt{}; + + [[nodiscard]] bool valid() const; + bool markPosted(uint64_t now_ns = 0) const; + UbAttemptResolution resolve(TransferStatusEnum outcome, + size_t transferred_bytes, bool retryable, + uint64_t now_ns = 0) const; +}; + +// A logical slice is owned by its UbTask and by any in-flight completion +// token. It keeps only a weak reference back to the task, avoiding a cycle and +// allowing a late completion to be discarded safely after batch teardown. +class UbSlice : public std::enable_shared_from_this { + public: + using Ptr = std::shared_ptr; + + UbSlice(const UbSlice&) = delete; + UbSlice& operator=(const UbSlice&) = delete; + + [[nodiscard]] uint64_t id() const { return id_; } + [[nodiscard]] const UbSliceSpec& spec() const { return spec_; } + + bool markQueued(uint64_t now_ns = 0) { + std::lock_guard lock(mutex_); + if (cancel_requested_.load(std::memory_order_acquire) || + isTerminal(state_)) { + return false; + } + if (state_ != UbSliceState::kInitial && + state_ != UbSliceState::kRetryPending) { + return false; + } + state_ = UbSliceState::kQueued; + queued_ns_ = normalizedNow(now_ns); + return true; + } + + // Claims this slice for one posting attempt. A cancellation racing after + // this point is best-effort: the posting thread owns the device boundary + // and must resolve the attempt instead of reporting an immediate cancel. + std::optional beginAttempt(const UbPostPath& path, + uint64_t now_ns = 0) { + if (!path.valid()) return std::nullopt; + bool canceled = false; + TransferStatusEnum terminal_status = PENDING; + size_t terminal_bytes = 0; + std::optional token; + { + std::lock_guard lock(mutex_); + if (isTerminal(state_)) return std::nullopt; + if (cancel_requested_.load(std::memory_order_acquire)) { + if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + canceled = true; + terminal_status = CANCELED; + } + } else if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + state_ = UbSliceState::kPosting; + path_ = path; + ++attempt_; + attempt_started_ns_ = normalizedNow(now_ns); + posted_ns_ = 0; + token = UbAttemptToken{id_, attempt_, path_}; + } + } + if (canceled) notifyTaskTerminal(terminal_status, terminal_bytes); + return token; + } + + [[nodiscard]] std::optional completionToken( + const UbAttemptToken& token) { + auto self = weak_from_this().lock(); + if (!self) return std::nullopt; + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + (state_ != UbSliceState::kPosting && + state_ != UbSliceState::kPosted)) { + return std::nullopt; + } + return UbCompletionToken{std::move(self), token}; + } + + bool markPosted(const UbAttemptToken& token, uint64_t now_ns = 0) { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + state_ != UbSliceState::kPosting) { + return false; + } + state_ = UbSliceState::kPosted; + posted_ns_ = normalizedNow(now_ns); + return true; + } + + // Linearization point immediately before crossing the native post + // boundary. Cancellation that wins this mutex prevents the WR from being + // posted and terminalizes the slice; cancellation after this point is + // best-effort and waits for completion/drain. + bool tryCommitPost(const UbAttemptToken& token, uint64_t now_ns = 0) { + bool canceled = false; + { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + state_ != UbSliceState::kPosting) { + return false; + } + if (cancel_requested_.load(std::memory_order_acquire)) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + canceled = true; + } else { + state_ = UbSliceState::kPosted; + posted_ns_ = normalizedNow(now_ns); + } + } + if (canceled) notifyTaskTerminal(CANCELED, 0); + return !canceled; + } + + // Resolves exactly one posting attempt. Retriable errors move the logical + // slice back to kRetryPending without notifying the task. Duplicate or + // stale completions are ignored by matching attempt and endpoint + // generation. Only the final resolution contributes task bytes/status. + UbAttemptResolution resolveAttempt(const UbAttemptToken& token, + TransferStatusEnum outcome, + size_t transferred_bytes, bool retryable, + uint64_t now_ns = 0) { + bool terminal = false; + TransferStatusEnum terminal_status = PENDING; + size_t terminal_bytes = 0; + UbAttemptResolution resolution = UbAttemptResolution::kIgnored; + { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + (state_ != UbSliceState::kPosting && + state_ != UbSliceState::kPosted) || + !isTerminal(outcome)) { + return UbAttemptResolution::kIgnored; + } + + const bool canceled = + cancel_requested_.load(std::memory_order_acquire); + if (outcome != COMPLETED && retryable && !canceled && + retry_count_ < spec_.max_retries) { + ++retry_count_; + state_ = UbSliceState::kRetryPending; + queued_ns_ = normalizedNow(now_ns); + resolution = UbAttemptResolution::kRetryScheduled; + } else { + terminal_status = + (canceled && outcome != COMPLETED) ? CANCELED : outcome; + setTerminalLocked(terminal_status, transferred_bytes, + normalizedNow(now_ns)); + terminal_bytes = transferred_bytes_; + terminal = true; + resolution = UbAttemptResolution::kTerminal; + } + } + if (terminal) notifyTaskTerminal(terminal_status, terminal_bytes); + return resolution; + } + + // Resolves work that failed (or completed locally) before any adapter post. + // It deliberately refuses to terminalize kPosting/kPosted work; those must + // drain through resolveAttempt(). + UbAttemptResolution resolveBeforePost(TransferStatusEnum outcome, + size_t transferred_bytes, + bool retryable, uint64_t now_ns = 0) { + bool terminal = false; + TransferStatusEnum terminal_status = outcome; + size_t terminal_bytes = 0; + UbAttemptResolution resolution = UbAttemptResolution::kIgnored; + { + std::lock_guard lock(mutex_); + if (!isTerminal(outcome) || isTerminal(state_) || + state_ == UbSliceState::kPosting || + state_ == UbSliceState::kPosted) { + return UbAttemptResolution::kIgnored; + } + const bool canceled = + cancel_requested_.load(std::memory_order_acquire); + if (outcome != COMPLETED && retryable && !canceled && + retry_count_ < spec_.max_retries) { + ++retry_count_; + state_ = UbSliceState::kRetryPending; + queued_ns_ = normalizedNow(now_ns); + resolution = UbAttemptResolution::kRetryScheduled; + } else { + if (canceled && outcome != COMPLETED) { + terminal_status = CANCELED; + } + setTerminalLocked(terminal_status, transferred_bytes, + normalizedNow(now_ns)); + terminal_bytes = transferred_bytes_; + terminal = true; + resolution = UbAttemptResolution::kTerminal; + } + } + if (terminal) notifyTaskTerminal(terminal_status, terminal_bytes); + return resolution; + } + + bool tryResolveBeforePost(TransferStatusEnum outcome, + size_t transferred_bytes = 0, + uint64_t now_ns = 0) { + return resolveBeforePost(outcome, transferred_bytes, false, now_ns) == + UbAttemptResolution::kTerminal; + } + + // Best-effort cancellation. Unclaimed work becomes terminal immediately; + // posting or posted work only observes the flag and must still be resolved + // by its device completion. + bool requestCancellation(uint64_t now_ns = 0) { + cancel_requested_.store(true, std::memory_order_release); + bool terminal = false; + { + std::lock_guard lock(mutex_); + if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + terminal = true; + } + } + if (terminal) notifyTaskTerminal(CANCELED, 0); + return terminal; + } + + [[nodiscard]] bool cancellationRequested() const { + return cancel_requested_.load(std::memory_order_acquire); + } + + [[nodiscard]] UbSliceSnapshot snapshot() const { + std::lock_guard lock(mutex_); + return UbSliceSnapshot{ + id_, + state_, + path_, + attempt_, + retry_count_, + spec_.max_retries, + cancel_requested_.load(std::memory_order_acquire), + transferred_bytes_, + created_ns_, + queued_ns_, + attempt_started_ns_, + posted_ns_, + terminal_ns_}; + } + + private: + friend struct UbTask; + + UbSlice(uint64_t id, UbSliceSpec spec, std::weak_ptr task, + uint64_t created_ns) + : id_(id), + spec_(std::move(spec)), + task_(std::move(task)), + created_ns_(normalizedNow(created_ns)) {} + + static uint64_t normalizedNow(uint64_t now_ns) { + return now_ns == 0 ? steadyNowNs() : now_ns; + } + + bool matchesActiveAttemptLocked(const UbAttemptToken& token) const { + return token.slice_id == id_ && token.attempt == attempt_ && + token.path == path_; + } + + void setTerminalLocked(TransferStatusEnum status, size_t bytes, + uint64_t now_ns) { + state_ = sliceState(status); + transferred_bytes_ = std::min(bytes, spec_.length); + terminal_ns_ = now_ns; + } + + void notifyTaskTerminal(TransferStatusEnum status, size_t bytes); + + const uint64_t id_; + const UbSliceSpec spec_; + std::weak_ptr task_; + + mutable std::mutex mutex_; + UbSliceState state_{UbSliceState::kInitial}; + UbPostPath path_{}; + uint32_t attempt_{0}; + uint32_t retry_count_{0}; + std::atomic cancel_requested_{false}; + size_t transferred_bytes_{0}; + uint64_t created_ns_{0}; + uint64_t queued_ns_{0}; + uint64_t attempt_started_ns_{0}; + uint64_t posted_ns_{0}; + uint64_t terminal_ns_{0}; +}; + +struct UbTaskSnapshot { + TransferStatus status{PENDING, 0}; + size_t total_bytes{0}; + size_t total_slices{0}; + size_t resolved_slices{0}; + size_t remaining_slices{0}; + size_t successful_slices{0}; + bool sealed{false}; + bool cancel_requested{false}; + uint64_t created_ns{0}; + uint64_t deadline_ns{0}; + uint64_t terminal_ns{0}; +}; + +struct UbTask : public std::enable_shared_from_this { + public: + using Ptr = std::shared_ptr; + using TerminalCallback = std::function; + + static Ptr create(Request request, TerminalCallback terminal_callback = {}, + uint64_t created_ns = 0) { + return Ptr(new UbTask(std::move(request), std::move(terminal_callback), + created_ns)); + } + + UbTask(const UbTask&) = delete; + UbTask& operator=(const UbTask&) = delete; + + [[nodiscard]] const Request& request() const { return request_; } + + // Slices must be added before seal(). Workers may begin processing only + // after sealing, which prevents an early completion from finalizing a task + // while its remaining slices are still being constructed. + UbSlice::Ptr addSlice(UbSliceSpec spec, uint64_t created_ns = 0) { + UbSlice::Ptr slice; + bool cancel = false; + { + std::lock_guard lock(mutex_); + if (sealed_) return nullptr; + slice = UbSlice::Ptr(new UbSlice(next_slice_id_++, std::move(spec), + weak_from_this(), created_ns)); + slices_.push_back(slice); + cancel = cancel_requested_.load(std::memory_order_acquire); + } + if (cancel) slice->requestCancellation(created_ns); + return slice; + } + + bool seal() { + TerminalCallback callback; + TransferStatus final_status{}; + bool notify = false; + { + std::lock_guard lock(mutex_); + if (sealed_) return false; + sealed_ = true; + notify = maybeFinalizeLocked(final_status, callback); + } + if (notify && callback) callback(final_status); + return true; + } + + // Returns the number of slices canceled before reaching the adapter. Any + // posting/posted slices remain pending until their attempts resolve. + size_t requestCancellation(uint64_t now_ns = 0) { + cancel_requested_.store(true, std::memory_order_release); + std::vector slices; + { + std::lock_guard lock(mutex_); + if (isTerminal(status_.s)) return 0; + slices = slices_; + } + size_t canceled = 0; + for (const auto& slice : slices) { + if (slice->requestCancellation(now_ns)) ++canceled; + } + return canceled; + } + + [[nodiscard]] bool cancellationRequested() const { + return cancel_requested_.load(std::memory_order_acquire); + } + + [[nodiscard]] TransferStatus transferStatus() const { + std::lock_guard lock(mutex_); + return status_; + } + + [[nodiscard]] UbTaskSnapshot snapshot() const { + std::lock_guard lock(mutex_); + return UbTaskSnapshot{status_, + request_.length, + slices_.size(), + resolved_slice_ids_.size(), + slices_.size() - resolved_slice_ids_.size(), + successful_slices_, + sealed_, + cancel_requested_.load(std::memory_order_acquire), + created_ns_, + request_.deadline_ns, + terminal_ns_}; + } + + [[nodiscard]] std::vector slices() const { + std::lock_guard lock(mutex_); + return slices_; + } + + private: + friend class UbSlice; + + UbTask(Request request, TerminalCallback terminal_callback, + uint64_t created_ns) + : request_(std::move(request)), + created_ns_(created_ns == 0 ? steadyNowNs() : created_ns), + terminal_callback_(std::move(terminal_callback)) {} + + static int terminalSeverity(TransferStatusEnum status) { + switch (status) { + case INVALID: + return 4; + case FAILED: + return 3; + case TIMEOUT: + return 2; + case CANCELED: + return 1; + default: + return 0; + } + } + + void onSliceTerminal(uint64_t slice_id, TransferStatusEnum status, + size_t bytes) { + TerminalCallback callback; + TransferStatus final_status{}; + bool notify = false; + { + std::lock_guard lock(mutex_); + if (!resolved_slice_ids_.insert(slice_id).second) return; + if (std::numeric_limits::max() - status_.transferred_bytes < + bytes) { + status_.transferred_bytes = std::numeric_limits::max(); + } else { + status_.transferred_bytes += bytes; + } + if (status == COMPLETED) { + ++successful_slices_; + } else if (terminalSeverity(status) > + terminalSeverity(aggregate_error_)) { + aggregate_error_ = status; + } + notify = maybeFinalizeLocked(final_status, callback); + } + if (notify && callback) callback(final_status); + } + + bool maybeFinalizeLocked(TransferStatus& final_status, + TerminalCallback& callback) { + if (!sealed_ || terminal_notified_ || + resolved_slice_ids_.size() != slices_.size()) { + return false; + } + + if (slices_.empty() || successful_slices_ == slices_.size()) { + status_.s = COMPLETED; + } else { + status_.s = aggregate_error_ == PENDING ? FAILED : aggregate_error_; + } + terminal_notified_ = true; + terminal_ns_ = steadyNowNs(); + final_status = status_; + callback = std::move(terminal_callback_); + return true; + } + + const Request request_; + const uint64_t created_ns_; + mutable std::mutex mutex_; + std::vector slices_; + std::unordered_set resolved_slice_ids_; + uint64_t next_slice_id_{1}; + size_t successful_slices_{0}; + TransferStatusEnum aggregate_error_{PENDING}; + TransferStatus status_{PENDING, 0}; + std::atomic cancel_requested_{false}; + bool sealed_{false}; + bool terminal_notified_{false}; + uint64_t terminal_ns_{0}; + TerminalCallback terminal_callback_; +}; + +inline bool UbCompletionToken::valid() const { + return slice != nullptr && attempt.valid(); +} + +inline bool UbCompletionToken::markPosted(uint64_t now_ns) const { + return valid() && slice->markPosted(attempt, now_ns); +} + +inline UbAttemptResolution UbCompletionToken::resolve( + TransferStatusEnum outcome, size_t transferred_bytes, bool retryable, + uint64_t now_ns) const { + if (!valid()) return UbAttemptResolution::kIgnored; + return slice->resolveAttempt(attempt, outcome, transferred_bytes, retryable, + now_ns); +} + +inline void UbSlice::notifyTaskTerminal(TransferStatusEnum status, + size_t bytes) { + if (auto task = task_.lock()) task->onSliceTerminal(id_, status, bytes); +} + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h new file mode 100644 index 0000000000..ea12b9110e --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h @@ -0,0 +1,72 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_UB_TRANSPORT_H_ +#define TENT_TRANSPORT_UB_UB_TRANSPORT_H_ + +#include +#include + +#include "tent/runtime/control_plane.h" +#include "tent/runtime/transport.h" + +namespace mooncake::tent { + +namespace ub { +class UrmaAdapter; +struct UbTask; +} // namespace ub + +struct UbSubBatch final : public Transport::SubBatch { + std::vector> task_list; + size_t max_size{0}; + + size_t size() const override { return task_list.size(); } +}; + +// TENT-native UB transport. The implementation owns its task/slice scheduler, +// endpoint store and URMA resources; it never converts requests to Classic TE +// types or calls the Classic UbTransport/UbWorkerPool data path. +class UbTransport final : public Transport { + public: + explicit UbTransport(std::shared_ptr adapter = nullptr); + ~UbTransport() override; + + UbTransport(const UbTransport&) = delete; + UbTransport& operator=(const UbTransport&) = delete; + + Status install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr) override; + Status uninstall() override; + + Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override; + Status freeSubBatch(SubBatchRef& batch) override; + Status submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) override; + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override; + + bool supportsCancellation() const override { return true; } + Status cancelTransferTask(SubBatchRef batch, int task_id) override; + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) override; + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override; + Status removeMemoryBuffer(BufferDesc& desc) override; + bool warmupMemory(void* addr, size_t length) override; + + const char* getName() const override { return "ub"; } + double getEstimatedBandwidth() const override; + bool supportNotification() const override { return false; } + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace mooncake::tent + +#endif // TENT_TRANSPORT_UB_UB_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h new file mode 100644 index 0000000000..774169f799 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -0,0 +1,156 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_WORKERS_H_ +#define TENT_TRANSPORT_UB_WORKERS_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/topology.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +class UbEndpoint; + +struct EndpointResolveRequest { + UbContextPtr local_context; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + const SegmentDesc* remote_segment{nullptr}; + Topology::NicID remote_topology_id{-1}; + uint64_t segment_generation{0}; +}; + +using EndpointResolver = std::function&)>; +using EndpointRetirer = std::function&)>; + +// Native UB scheduler. Posting lanes own request selection and URMA post; +// poller lanes own completion dispatch. A monotonic numeric token, never a raw +// UbSlice pointer, crosses the adapter boundary. +class UbWorkers final { + public: + UbWorkers(std::shared_ptr adapter, + std::vector contexts, + std::shared_ptr local_topology, + SegmentManager* segment_manager, UbBufferManager* buffers, + RailMonitor* rail_monitor, QuotaManager* quota, UbParams params, + EndpointResolver endpoint_resolver, + EndpointRetirer endpoint_retirer = {}); + ~UbWorkers(); + + UbWorkers(const UbWorkers&) = delete; + UbWorkers& operator=(const UbWorkers&) = delete; + + Status start(); + Status stop(); + Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL); + Status cancel(const UbTask::Ptr& task); + + [[nodiscard]] bool running() const noexcept { + return accepting_.load(std::memory_order_acquire); + } + [[nodiscard]] size_t queuedCount() const; + [[nodiscard]] size_t inflightCount() const; + + private: + struct PendingSlice { + // Keep the task alive until every queued/in-flight slice reaches a + // terminal state. Callers may release the sub-batch handle before the + // asynchronous data path has drained. + UbTask::Ptr task; + UbSlice::Ptr slice; + uint64_t device_mask{~0ULL}; + int priority{PRIO_HIGH}; + SegmentID target_id{LOCAL_SEGMENT_ID}; + Request::OpCode opcode{Request::READ}; + }; + struct Route; + struct Inflight; + + void postingLoop(size_t worker_index); + void pollingLoop(size_t poller_index); + bool popPending(PendingSlice& pending); + void enqueueRetry(const PendingSlice& pending); + void deferPending(const PendingSlice& pending); + void processPending(const PendingSlice& pending, size_t worker_index); + Status buildRoute(const PendingSlice& pending, Route& route); + Status chooseAndResolveEndpoint(const PendingSlice& pending, Route& route, + std::shared_ptr& endpoint, + UbPostPath& path); + void handleCompletion(const Completion& completion); + void scanTimeouts(); + void releaseInflight(const std::shared_ptr& inflight); + void resolveInflight(const std::shared_ptr& inflight, + TransferStatusEnum outcome, size_t bytes, + bool retryable); + void recordTimeoutOnce(const std::shared_ptr& inflight, + uint64_t now_ns); + void rememberEndpointDrain(const std::shared_ptr& endpoint); + void forgetEndpointDrain(const std::shared_ptr& endpoint); + void failUnposted(const PendingSlice& pending, + TransferStatusEnum outcome = FAILED); + uint64_t nextCompletionToken(); + std::vector orderedLocalDevices( + const PendingSlice& pending) const; + static std::vector orderedRemoteDevices( + const SegmentDesc& segment, const BufferDesc& buffer); + + std::shared_ptr adapter_; + std::vector contexts_; + std::unordered_map context_by_topology_id_; + std::shared_ptr local_topology_; + SegmentManager* segment_manager_; + UbBufferManager* buffers_; + RailMonitor* rail_monitor_; + QuotaManager* quota_; + const UbParams params_; + EndpointResolver endpoint_resolver_; + EndpointRetirer endpoint_retirer_; + + mutable std::mutex queue_mutex_; + std::condition_variable queue_cv_; + std::array, PRIO_LOW + 1> queues_; + + mutable std::mutex inflight_mutex_; + mutable std::mutex endpoint_drain_mutex_; + std::condition_variable inflight_cv_; + std::unordered_map> inflight_; + // Endpoints whose ERROR transition has not yet reached its native flush + // fence remain owned here even if every logical token completes naturally. + std::unordered_map> + draining_endpoints_; + + std::vector> all_jfcs_; + std::unordered_map context_by_jfc_; + std::vector posting_threads_; + std::vector polling_threads_; + std::atomic accepting_{false}; + std::atomic posting_{false}; + std::atomic polling_{false}; + std::atomic timeout_scans_enabled_{false}; + std::atomic next_token_{1}; +}; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_WORKERS_H_ diff --git a/mooncake-transfer-engine/tent/src/CMakeLists.txt b/mooncake-transfer-engine/tent/src/CMakeLists.txt index 594d64ee16..1619f457bf 100644 --- a/mooncake-transfer-engine/tent/src/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/CMakeLists.txt @@ -153,6 +153,8 @@ foreach( tent_xport_rdma tent_xport_shm tent_xport_tcp + tent_xport_hp_tcp + tent_xport_hp_tcp_core tent_xport_ub tent_xport_ascend_direct tent_xport_sunrise_link diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index cec776d2ac..576cf5b7b5 100644 --- a/mooncake-transfer-engine/tent/src/common/config.cpp +++ b/mooncake-transfer-engine/tent/src/common/config.cpp @@ -85,6 +85,16 @@ static inline void setConfig(Config& config, const std::string& env_key, if (val) config.setFromString(config_key, std::string(val)); } +// Bool env vars: setFromString() stores "1" as int, so get(..., false) misses +// it. +static inline void setBoolConfig(Config& config, const std::string& env_key, + const std::string& config_key) { + const char* val = std::getenv(env_key.c_str()); + if (!val) return; + config.set(config_key, + ConfigHelper::parseBool(val, config.get(config_key, false))); +} + // Like setConfig, but parses the env value as a comma-separated list and // stores it as a string array. Empty/whitespace-only items are dropped so a // trailing comma or spaces around names are tolerated (e.g. "mlx5_0, mlx5_1"). @@ -135,6 +145,7 @@ Status ConfigHelper::loadFromEnv(Config& config) { } // Legacy keys for backward compatibility (MC_* env vars) + setConfig(config, "MOONCAKE_LOCAL_HOSTNAME", "rpc_server_hostname"); setConfig(config, "MC_RDMA_BIND_ADDRESS", "transports/rdma/bind_address"); setConfig(config, "MC_NUM_CQ_PER_CTX", "transports/rdma/device/num_cq_list"); @@ -167,6 +178,8 @@ Status ConfigHelper::loadFromEnv(Config& config) { "transports/rdma/disable_gpu_direct_rdma"); setConfig(config, "MC_LOG_RDMA_SLICE_AFFINITY", "transports/rdma/log_slice_affinity"); + setBoolConfig(config, "MC_STRICT_LOCAL_NUMA", + "transports/rdma/strict_local_numa"); // Restrict which RDMA NICs the engine discovers/uses (comma-separated // device names). MC_TE_FILTERS is an allow-list — same name and semantics // as the legacy Transfer Engine's device whitelist, so a single env works diff --git a/mooncake-transfer-engine/tent/src/platform/ascend/ascend_allocator.cpp b/mooncake-transfer-engine/tent/src/platform/ascend/ascend_allocator.cpp index b7bd61d608..1f6645c0b8 100644 --- a/mooncake-transfer-engine/tent/src/platform/ascend/ascend_allocator.cpp +++ b/mooncake-transfer-engine/tent/src/platform/ascend/ascend_allocator.cpp @@ -51,6 +51,12 @@ Status AscendPlatform::free(void* ptr, size_t size) { } Status AscendPlatform::copy(void* dst, void* src, size_t length) { + // Unlike CUDA/ROCm, the copy is not routed to the device owning the + // buffer; aclrtMemcpy runs in the caller thread's ACL context. Routing it + // needs a driver-id -> user-id lookup ACL does not expose: location.id is + // a driver id, aclrtSetDevice takes an ASCEND_RT_VISIBLE_DEVICES user id. + // TODO: left to someone with Ascend expertise and NPU hardware to verify; + // getting that id mapping wrong picks the wrong device silently. CHECK_ASCEND(aclrtMemcpy(dst, length, src, length, ACL_MEMCPY_DEFAULT)); return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_allocator.cpp b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_allocator.cpp index d60ddffa2f..5f00fe01fe 100644 --- a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_allocator.cpp +++ b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_allocator.cpp @@ -60,8 +60,20 @@ Status CudaPlatform::copy(void* dst, void* src, size_t length) { // as the latter relies on the legacy default stream and can introduce // unintended synchronization or even deadlocks in downstream // components (e.g. mooncake-pg). + // + // cudaMemcpyAsync routes the copy through its stream's device context, so + // the stream must live on the device owning the device-side buffer. + // Control-plane RPC worker threads sit on cuda:0 while a registered buffer + // may live on cuda:R; taking the stream from the buffer's device routes the + // copy correctly without mutating the calling thread's current device. + // Host-only copies keep the current device. + int device_id = getPointerDeviceId(dst); + if (device_id == CUDAStreamPool::kCurrentDevice) { + device_id = getPointerDeviceId(src); + } + CUDAStreamHandle stream; - CHECK_STATUS(getStreamFromPool(stream)); + CHECK_STATUS(getStreamFromPool(stream, device_id)); CHECK_CUDA( cudaMemcpyAsync(dst, src, length, cudaMemcpyDefault, stream.get())); CHECK_CUDA(cudaStreamSynchronize(stream.get())); diff --git a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp index 6a52601c9a..2d4cda8d17 100644 --- a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp @@ -310,6 +310,26 @@ MemoryType CudaPlatform::getMemoryType(void* addr) { return MTYPE_CPU; } +int CudaPlatform::getPointerDeviceId(void* addr) { + // Same guards as getMemoryType(): the cudaPointerAttributes layout changes + // across CUDA majors, so the struct must not be read on a runtime whose ABI + // does not match what we built against. + if (!cudaDevicePresent() || !cudaAbiMatches()) { + return CUDAStreamPool::kCurrentDevice; + } + cudaPointerAttributes attributes{}; + if (cudaPointerGetAttributes(&attributes, addr) != cudaSuccess) { + // Clear the latched error so it cannot surface at an unrelated + // cudaGetLastError() call site. + cudaGetLastError(); + return CUDAStreamPool::kCurrentDevice; + } + if (attributes.type != cudaMemoryTypeDevice) { + return CUDAStreamPool::kCurrentDevice; + } + return attributes.device; +} + static inline uintptr_t alignPage(uintptr_t address) { const static size_t kPageSize = 4096; return address & ~(kPageSize - 1); diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp index bcc758f2e8..fc68c401cd 100644 --- a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp +++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_allocator.cpp @@ -56,8 +56,19 @@ Status RocmPlatform::free(void* ptr, size_t size) { } Status RocmPlatform::copy(void* dst, void* src, size_t length) { + // hipMemcpyAsync routes the copy through its stream's device context, so + // the stream must live on the device owning the device-side buffer. + // Control-plane RPC worker threads sit on device 0 while a registered + // buffer may live on device R; taking the stream from the buffer's device + // routes the copy correctly without mutating the calling thread's current + // device. Host-only copies keep the current device. + int device_id = getPointerDeviceId(dst); + if (device_id == HIPStreamPool::kCurrentDevice) { + device_id = getPointerDeviceId(src); + } + HIPStreamHandle stream; - CHECK_STATUS(getStreamFromPool(stream)); + CHECK_STATUS(getStreamFromPool(stream, device_id)); CHECK_HIP(hipMemcpyAsync(dst, src, length, hipMemcpyDefault, stream.get())); CHECK_HIP(hipStreamSynchronize(stream.get())); return Status::OK(); diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp index 28e472d02b..50f6de1bf1 100644 --- a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp @@ -263,6 +263,20 @@ MemoryType RocmPlatform::getMemoryType(void* addr) { return MTYPE_CPU; } +int RocmPlatform::getPointerDeviceId(void* addr) { + hipPointerAttribute_t attributes{}; + if (hipPointerGetAttributes(&attributes, addr) != hipSuccess) { + // Clear the latched error so it cannot surface at an unrelated + // hipGetLastError() call site. + hipGetLastError(); + return HIPStreamPool::kCurrentDevice; + } + if (attributes.type != hipMemoryTypeDevice) { + return HIPStreamPool::kCurrentDevice; + } + return attributes.device; +} + static inline uintptr_t alignPage(uintptr_t address) { const static size_t kPageSize = 4096; return address & ~(kPageSize - 1); diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index ab15831bad..a78c1bc6d7 100644 --- a/mooncake-transfer-engine/tent/src/python/pybind.cpp +++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp @@ -37,6 +37,8 @@ static_assert(static_cast(TransportType::UB) == TRANSPORT_UB, "UB wire value must match the C API macro"); static_assert(static_cast(TransportType::MPCOMM) == TRANSPORT_MPCOMM, "MPCOMM wire value must match the C API macro"); +static_assert(static_cast(TransportType::HP_TCP) == TRANSPORT_HP_TCP, + "HP_TCP wire value must match the C API macro"); // ============================================================================= // Custom Exception Hierarchy @@ -312,6 +314,7 @@ PYBIND11_MODULE(tent, m) { .value("TPU", TransportType::TPU) .value("UB", TransportType::UB) .value("MPCOMM", TransportType::MPCOMM) + .value("HP_TCP", TransportType::HP_TCP) .export_values(); py::enum_(m, "IntentType") diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 1abf18251e..c6d7bc2b4b 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -87,6 +87,24 @@ Status ControlClient::bootstrap(const std::string& server_addr, return decodeBootstrapResponse(response_raw, response); } +Status ControlClient::bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response) { + std::string request_raw, response_raw; + json j = request; + request_raw = j.dump(); + CHECK_STATUS( + tl_rpc_agent.call(server_addr, BootstrapUb, request_raw, response_raw)); + try { + response = json::parse(response_raw).get(); + } catch (const std::exception& e) { + return Status::MalformedJson( + std::string("Malformed UB bootstrap response: ") + e.what() + + LOC_MARK); + } + return Status::OK(); +} + Status ControlClient::sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length) { @@ -102,8 +120,11 @@ Status ControlClient::sendData(const std::string& server_addr, // 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); - loader.copy(request.data() + sizeof(desc), local_mem_addr, length); + CHECK_STATUS( + loader.copy(request.data() + sizeof(desc), local_mem_addr, length)); } auto status = tl_rpc_agent.callOwned(server_addr, SendData, std::move(request), response); @@ -123,9 +144,8 @@ Status ControlClient::recvData(const std::string& server_addr, if (!status.ok()) return status; if (response.size() != length) return Status::RpcServiceError( - "RecvData failed: target address not in registered buffer"); - Platform::getLoader().copy(local_mem_addr, response.data(), length); - return Status::OK(); + response.empty() ? "RecvData failed: empty response" : response); + return Platform::getLoader().copy(local_mem_addr, response.data(), length); } inline void to_json(nlohmann::json& j, const Notification& n) { @@ -271,6 +291,11 @@ ControlService::ControlService(const std::string& type, // io_context serializes every bulk transfer and stalls Probe/Bootstrap // on the same thread. Offload matches Delegate: the connection coroutine // suspends, copies run on the blocking executor, and other RPCs proceed. + rpc_server_->registerFunction( + BootstrapUb, + [this](const std::string_view& request, std::string& response) { + onBootstrapUb(request, response); + }); rpc_server_->registerFunction( SendData, [this](const std::string_view& request, std::string& response) { @@ -322,6 +347,10 @@ ControlService::ControlService(const std::string& type, ControlService::~ControlService() { // Stop RPC workers while callback state and synchronization primitives are // still alive. Member destruction would otherwise tear them down first. + { + std::lock_guard lock(ub_bootstrap_callback_mutex_); + ub_bootstrap_callback_ = {}; + } rpc_server_.reset(); } @@ -434,6 +463,37 @@ void ControlService::onBootstrapRdma(const std::string_view& request, response = j.dump(); } +void ControlService::onBootstrapUb(const std::string_view& request, + std::string& response) { + UbBootstrapDesc response_desc; + try { + auto request_desc = + json::parse(std::string(request)).get(); + int ret = -1; + { + // Callback replacement during uninstall is serialized with + // invocation, so an in-flight bootstrap cannot outlive the UB + // transport object it targets. + std::lock_guard lock(ub_bootstrap_callback_mutex_); + if (ub_bootstrap_callback_) { + ret = ub_bootstrap_callback_(request_desc, response_desc); + } else { + response_desc.reply_msg = + "UB bootstrap callback is not registered"; + } + } + if (ret != 0 && response_desc.reply_msg.empty()) { + response_desc.reply_msg = + "UB bootstrap callback failed, ret=" + std::to_string(ret); + } + } catch (const std::exception& e) { + response_desc.reply_msg = + std::string("Malformed UB bootstrap request: ") + e.what(); + } + json j = response_desc; + response = j.dump(); +} + void ControlService::onSendData(const std::string_view& request, std::string& response) { if (request.size() < sizeof(XferDataDesc)) { @@ -452,7 +512,15 @@ void ControlService::onSendData(const std::string_view& request, } if (local_desc->findBuffer(peer_mem_addr, length)) { - Platform::getLoader().copy((void*)peer_mem_addr, &desc[1], length); + 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(); + } } else { response = "SendData failed: target address not in registered buffer"; } @@ -469,10 +537,18 @@ void ControlService::onRecvData(const std::string_view& request, 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 constexpr size_t kMaxTransferSize = 1ULL << 30; // 1GB max per RPC if (length > kMaxTransferSize) { - response = "RecvData failed: length exceeds maximum allowed"; + fail("RecvData failed: length exceeds maximum allowed"); return; } @@ -484,10 +560,14 @@ void ControlService::onRecvData(const std::string_view& request, length); } else { response.resize(length); - loader.copy(response.data(), (void*)peer_mem_addr, length); + auto status = + loader.copy(response.data(), (void*)peer_mem_addr, length); + if (!status.ok()) { + fail("RecvData failed: copy: " + status.ToString()); + } } } else { - response = "RecvData failed: target address not in registered buffer"; + fail("RecvData failed: target address not in registered buffer"); } } diff --git a/mooncake-transfer-engine/tent/src/runtime/hp_tcp_transport_config.cpp b/mooncake-transfer-engine/tent/src/runtime/hp_tcp_transport_config.cpp new file mode 100644 index 0000000000..3819bf4843 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/hp_tcp_transport_config.cpp @@ -0,0 +1,124 @@ +// Copyright 2026 KVCache.AI +#include "tent/runtime/hp_tcp_transport_config.h" + +#include +#include + +namespace mooncake::tent { +namespace { + +Status Invalid(const std::string& path, const std::string& detail) { + return Status::InvalidArgument(path + " " + detail + LOC_MARK); +} + +template +Status ReadUnsigned(const json& object, std::string_view key, T* output) { + auto it = object.find(std::string(key)); + if (it == object.end()) return Status::OK(); + uint64_t value = 0; + if (it->is_number_unsigned()) { + value = it->get(); + } else if (it->is_number_integer()) { + const auto signed_value = it->get(); + if (signed_value < 0) { + return Invalid("transports/hp_tcp/" + std::string(key), + "must be non-negative"); + } + value = static_cast(signed_value); + } else { + return Invalid("transports/hp_tcp/" + std::string(key), + "must be an integer"); + } + if (value > std::numeric_limits::max()) { + return Invalid("transports/hp_tcp/" + std::string(key), + "is out of range"); + } + *output = static_cast(value); + return Status::OK(); +} + +Status ReadString(const json& object, std::string_view key, + std::string* output) { + auto it = object.find(std::string(key)); + if (it == object.end()) return Status::OK(); + if (!it->is_string()) { + return Invalid("transports/hp_tcp/" + std::string(key), + "must be a string"); + } + *output = it->get(); + return Status::OK(); +} + +} // namespace + +Status ParseHpTcpTransportConfig(const Config& config, + HpTcpTransportConfig* out) { + if (out == nullptr) { + return Status::InvalidArgument( + "HP TCP configuration output is null" LOC_MARK); + } + + HpTcpTransportConfig parsed; + std::string subtree; + if (!config.dumpSubtree("transports/hp_tcp", &subtree)) { + *out = std::move(parsed); + return Status::OK(); + } + + json hp_tcp; + try { + hp_tcp = json::parse(subtree); + } catch (const std::exception& error) { + return Status::MalformedJson( + std::string("Invalid transports/hp_tcp configuration: ") + + error.what() + LOC_MARK); + } + if (!hp_tcp.is_object()) { + return Invalid("transports/hp_tcp", "must be an object"); + } + + if (auto it = hp_tcp.find("enable"); it != hp_tcp.end()) { + if (!it->is_boolean()) + return Invalid("transports/hp_tcp/enable", "must be a boolean"); + parsed.enabled = it->get(); + } + CHECK_STATUS( + ReadString(hp_tcp, "bind_address", &parsed.params.bind_address)); + CHECK_STATUS(ReadString(hp_tcp, "advertise_address", + &parsed.params.advertise_address)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "port", &parsed.params.port)); + CHECK_STATUS( + ReadUnsigned(hp_tcp, "worker_count", &parsed.params.worker_count)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "connections_per_peer", + &parsed.params.connections_per_peer)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "max_outstanding_tasks", + &parsed.params.max_outstanding_tasks)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "max_outstanding_bytes", + &parsed.params.max_outstanding_bytes)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "max_transfer_bytes", + &parsed.params.max_transfer_bytes)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "connect_timeout_ms", + &parsed.params.connect_timeout_ms)); + CHECK_STATUS(ReadUnsigned(hp_tcp, "progress_timeout_ms", + &parsed.params.progress_timeout_ms)); + + const auto& hp = parsed.params; + if (parsed.enabled && + (hp.worker_count == 0 || hp.connections_per_peer == 0 || + hp.max_outstanding_tasks == 0 || hp.max_outstanding_bytes == 0 || + hp.max_transfer_bytes == 0 || hp.connect_timeout_ms == 0 || + hp.progress_timeout_ms == 0)) { + return Invalid("transports/hp_tcp", + "contains zero or inconsistent limits"); + } + + if (parsed.enabled && config.get("transports/tcp/enable", true)) { + return Invalid("transports", + "tcp and hp_tcp cannot be enabled together"); + } + + *out = std::move(parsed); + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/topology.cpp b/mooncake-transfer-engine/tent/src/runtime/topology.cpp index dcb9b41d79..2bf81b1664 100644 --- a/mooncake-transfer-engine/tent/src/runtime/topology.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/topology.cpp @@ -108,11 +108,6 @@ void Topology::print() const { } Status Topology::discover(const std::vector& platforms) { - return discover(platforms, false); -} - -Status Topology::discover(const std::vector& platforms, - bool discover_ub) { clear(); for (auto& entry : platforms) { CHECK_STATUS(entry->probe(nic_list_, mem_list_)); @@ -121,7 +116,7 @@ Status Topology::discover(const std::vector& platforms, // UB discovery is intentionally adapter-backed instead of inferring UB // devices from verbs/sysfs names. One topology NIC is emitted per EID and // carries both the globally serialized identity and the native URMA name. - auto adapter = discover_ub ? ub::createDefaultUrmaAdapter() : nullptr; + auto adapter = ub::createDefaultUrmaAdapter(); if (adapter && adapter->available()) { auto status = adapter->initialize(); if (status.ok()) { @@ -182,7 +177,6 @@ Status Topology::discover(const std::vector& platforms, } } #endif - (void)discover_ub; return Status::OK(); } @@ -419,6 +413,13 @@ const Topology::NicEntry* Topology::getNicEntry(NicID id) const { return &nic_list_[id]; } +bool Topology::isCrossNuma(const MemEntry& mem, NicID nic_id) const { + if (mem.numa_node < 0) return false; + const auto* nic = getNicEntry(nic_id); + if (!nic || nic->numa_node < 0) return false; + return nic->numa_node != mem.numa_node; +} + const Topology::MemEntry* Topology::getMemEntry(MemID id) const { if (id < 0 || id >= (int)mem_list_.size()) return nullptr; return &mem_list_[id]; 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 4e15b63d02..06f818b85e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -55,6 +55,7 @@ constexpr uint8_t kRedisDefaultDbIndex = 0; // pass or two and a healthy-but-inflight batch reports PENDING (which resets // the counter), so only a permanently failing poll or queue retire gets here. constexpr size_t kMaxReclaimAttempts = 4096; +constexpr int kMaxHpTcpMetadataRefreshRetries = 1; } // namespace struct Batch { @@ -336,6 +337,7 @@ Status TransferEngineImpl::setupLocalSegment() { } Status TransferEngineImpl::construct() { + CHECK_STATUS(ParseHpTcpTransportConfig(*conf_, &hp_tcp_transport_config_)); auto metadata_type = conf_->get("metadata_type", "p2p"); auto metadata_servers = conf_->get("metadata_servers", ""); @@ -426,11 +428,30 @@ Status TransferEngineImpl::construct() { CHECK_STATUS(loadTransports()); std::string transport_string; - for (auto& transport : transport_list_) { + for (size_t transport_index = 0; transport_index < transport_list_.size(); + ++transport_index) { + auto& transport = transport_list_[transport_index]; if (transport) { auto status = transport->install(local_segment_name_, metadata_, topology_, conf_); if (!status.ok()) { + if (hp_tcp_transport_config_.enabled && + transport_index == + static_cast(TransportType::HP_TCP)) { + // HP TCP is explicitly required, so a failed install is a + // construction failure rather than an optional-transport + // skip. Unwind the already-started control service and + // any preceding transports immediately; the destructor's + // later deconstruct() call is intentionally idempotent. + const Status cleanup = deconstruct(); + if (!cleanup.ok()) { + LOG(ERROR) + << "Failed to unwind TENT after required HP TCP " + "install failure: " + << cleanup.ToString(); + } + return status; + } LOG(WARNING) << "Transport " << transport->getName() << " skipped: " << status.ToString(); transport = nullptr; @@ -525,6 +546,16 @@ Status TransferEngineImpl::deconstruct() { progress_worker_->stop(); } + for (auto& transport : transport_list_) { + if (!transport) continue; + const Status status = transport->quiesce(); + if (!status.ok()) { + LOG(ERROR) << "Transport " << transport->getName() + << " quiesce failed during teardown: " + << status.ToString(); + } + } + // Destroy staging_proxy_ first: its destructor calls back into // unregisterLocalMemory/freeLocalMemory, which require // local_segment_tracker_ and metadata_ to be alive. @@ -710,15 +741,19 @@ Status TransferEngineImpl::allocateLocalMemory(void** addr, size_t size, options.type = SHM; else if (transport_list_[RDMA]) options.type = RDMA; - else + else if (transport_list_[TCP]) options.type = TCP; + else + options.type = HP_TCP; } else { if (transport_list_[MNNVL]) options.type = MNNVL; else if (transport_list_[RDMA]) options.type = RDMA; - else + else if (transport_list_[TCP]) options.type = TCP; + else + options.type = HP_TCP; } return allocateLocalMemory(addr, size, options); } @@ -730,6 +765,8 @@ Status TransferEngineImpl::allocateLocalMemory(void** addr, size_t size, options.type = RDMA; else if (transport_list_[TCP]) options.type = TCP; + else if (transport_list_[HP_TCP]) + options.type = HP_TCP; else return Status::InvalidArgument( "Not supported type in memory options" LOC_MARK); @@ -754,13 +791,23 @@ Status TransferEngineImpl::freeLocalMemory(void* addr) { ++it) { if (it->addr == addr) { auto status = it->transport->freeLocalMemory(addr, it->size); + if (!status.ok()) { + LOG(WARNING) + << "Failed to free local memory, addr=" << addr + << ", size=" << it->size << ": " << status.ToString(); + return status; + } allocated_memory_.erase(it); - return status; + return Status::OK(); } } return Status::InvalidArgument("Address region not registered" LOC_MARK); } +// Forward declaration: getTypeEnum() is defined below but is needed by the +// registerLocalMemory location-override validation. +static MemoryType getTypeEnum(const std::string& type); + Status TransferEngineImpl::registerLocalMemory(void* addr, size_t size, Permission permission) { MemoryOptions options; @@ -793,6 +840,7 @@ std::vector TransferEngineImpl::getSupportedTransports( if (transport_list_[AscendDirect]) result.push_back(AscendDirect); if (transport_list_[SHM]) result.push_back(SHM); if (transport_list_[TCP]) result.push_back(TCP); + if (transport_list_[HP_TCP]) result.push_back(HP_TCP); if (transport_list_[GDS]) result.push_back(GDS); if (transport_list_[MPCOMM]) result.push_back(MPCOMM); if (transport_list_[TPU]) result.push_back(TPU); @@ -841,17 +889,53 @@ Status TransferEngineImpl::registerLocalMemory(std::vector addr_list, desc.regions = coalesceRegions(entries); } desc.ref_count = 1; - if (options.location != kWildcardLocation) - desc.location = options.location; + // The probe is the source of truth for transport selection: it + // classifies the memory (cpu/cuda/...). A caller-supplied location + // may refine the probe within the SAME memory type (e.g. probe + // "cpu:0" -> caller "cpu:1"), but must not replace it with an + // incompatible or unknown type. Classic TE encodes NUMA-segmented + // host DRAM as "segments:4096:0,1"; that is a TE dialect TENT's + // type system does not understand, and blindly adopting it would + // make getTypeEnum() return MTYPE_UNKNOWN and break transport + // selection. Validate before overriding: keep the probe when the + // caller is a wildcard, unknown, or a different type. + if (options.location != kWildcardLocation && + !options.location.empty()) { + auto probed_type = + getTypeEnum(LocationParser(desc.location).type()); + auto caller_type = + getTypeEnum(LocationParser(options.location).type()); + if (caller_type == MTYPE_UNKNOWN) { + LOG(WARNING) + << "Ignoring unknown caller location '" << options.location + << "' for registered memory at " << addr_list[i] + << " (probed '" << desc.location + << "'); keeping probed location"; + } else if (caller_type != probed_type) { + LOG(WARNING) << "Ignoring caller location '" << options.location + << "' (type mismatch with probed '" + << desc.location << "') for registered memory at " + << addr_list[i] << "; keeping probed location"; + } else { + desc.location = options.location; + } + } if (options.internal) desc.internal = options.internal; desc_list.push_back(std::move(desc)); } auto status = local_segment_tracker_->addInBatch( desc_list, [&](std::vector& descs) -> Status { + const bool hp_tcp_required = + options.type == HP_TCP || + (options.type == UNSPEC && transports.size() == 1 && + transports.front() == HP_TCP); for (auto type : transports) { auto s = transport_list_[type]->addMemoryBuffer(descs, options); - if (!s.ok()) LOG(WARNING) << s.ToString(); + if (!s.ok()) { + if (type == HP_TCP && hp_tcp_required) return s; + LOG(WARNING) << s.ToString(); + } } // desc.transports lists the transports that actually registered // the buffer (each transport appends itself on success). @@ -876,11 +960,21 @@ Status TransferEngineImpl::unregisterLocalMemory(void* addr, size_t size) { auto status = local_segment_tracker_->remove( (uint64_t)addr, size, [&](BufferDesc& desc) -> Status { removed = true; - for (auto type : desc.transports) { - auto status = transport_list_[type]->removeMemoryBuffer(desc); + const auto registered_transports = desc.transports; + for (size_t type = 0; type < kSupportedTransportTypes; ++type) { + auto& transport = transport_list_[type]; + if (!transport) continue; + const auto transport_type = static_cast(type); + const bool advertised = + std::find(registered_transports.begin(), + registered_transports.end(), + transport_type) != registered_transports.end(); + if (!advertised && !transport->tracksLocalBuffer(desc)) + continue; + auto status = transport->removeMemoryBuffer(desc); if (!status.ok()) LOG(WARNING) << status.ToString(); } - for (auto type : desc.transports) { + for (auto type : registered_transports) { TentMetrics::instance().recordRegisteredBufferBytes( type, -static_cast(desc.length)); } @@ -904,11 +998,23 @@ Status TransferEngineImpl::unregisterLocalMemory( (uint64_t)addr_list[i], size_list.empty() ? 0 : size_list[i], [&](BufferDesc& desc) -> Status { removed = true; - for (auto type : desc.transports) { - auto s = transport_list_[type]->removeMemoryBuffer(desc); + const auto registered_transports = desc.transports; + for (size_t type = 0; type < kSupportedTransportTypes; ++type) { + auto& transport = transport_list_[type]; + if (!transport) continue; + const auto transport_type = + static_cast(type); + const bool advertised = + std::find(registered_transports.begin(), + registered_transports.end(), + transport_type) != + registered_transports.end(); + if (!advertised && !transport->tracksLocalBuffer(desc)) + continue; + auto s = transport->removeMemoryBuffer(desc); if (!s.ok()) LOG(WARNING) << s.ToString(); } - for (auto type : desc.transports) { + for (auto type : registered_transports) { TentMetrics::instance().recordRegisteredBufferBytes( type, -static_cast(desc.length)); } @@ -1656,7 +1762,8 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, // local HBM<->host executor), mirroring how the CUDA cases gate on NVLINK. // An empty stage location means "no staging needed on that side". if (transport_list_[TPU] && - (transport_list_[RDMA] || transport_list_[TCP])) { + (transport_list_[RDMA] || transport_list_[TCP] || + transport_list_[HP_TCP])) { if (local_mtype == MTYPE_TPU && remote_mtype == MTYPE_TPU) { policy.clear(); policy.push_back(server_addr); @@ -1711,7 +1818,7 @@ Status TransferEngineImpl::prepareSubmit( PreparedSubmit::Owner owner; owner.request = request; owner.route = resolveTransport(owner.request, 0); - if (owner.route.transport == TCP) { + if (owner.route.transport == TCP || owner.route.transport == HP_TCP) { findStagingPolicy(owner.request, owner.staging_params); owner.staging = !owner.staging_params.empty() && staging_proxy_; } @@ -2101,7 +2208,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { return finishQueuedOwner(owner_id, FAILED); } - if (task.type == TCP) { + if (task.type == TCP || task.type == HP_TCP) { std::vector staging_params; findStagingPolicy(task.request, staging_params); if (!staging_params.empty() && staging_proxy_) { @@ -2496,8 +2603,59 @@ Status TransferEngineImpl::pollTaskStatus(Batch* batch, size_t task_id, if (!transport || !sub_batch) { return Status::InvalidArgument("Transport not available" LOC_MARK); } - return transport->getTransferStatus(sub_batch, task.sub_task_id, - task_status); + // HP TCP classifies transport errors using the terminal output. Other + // transports retain their existing error-output contract. + if (task.type == HP_TCP) task_status = {PENDING, 0}; + Status result = + transport->getTransferStatus(sub_batch, task.sub_task_id, task_status); + if (result.ok() || task.type != HP_TCP || task_status.s != FAILED) { + return result; + } + + if (result.IsNeedsRefreshCache()) { + finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + if (task.metadata_refresh_retry_count >= + kMaxHpTcpMetadataRefreshRetries) { + task.suppress_failover = true; + return Status::OK(); + } + ++task.metadata_refresh_retry_count; + Status invalidated = metadata_->segmentManager().invalidateRemote( + task.request.target_id); + if (!invalidated.ok()) { + task.suppress_failover = true; + LOG(WARNING) << "HP TCP metadata cache invalidation failed: " + << invalidated.ToString(); + return Status::OK(); + } + + const auto retry_start = std::chrono::steady_clock::now(); + startTransportAttempt(task, HP_TCP, retry_start); + Status retried = transport->retryTransferTask( + sub_batch, task.sub_task_id, task.request); + if (!retried.ok()) { + finishTransportAttempt(task, FAILED, + std::chrono::steady_clock::now()); + task.suppress_failover = true; + LOG(WARNING) << "HP TCP metadata refresh retry failed: " + << retried.ToString(); + return Status::OK(); + } + task.status = PENDING; + task_status.s = PENDING; + task_status.transferred_bytes = 0; + return Status::OK(); + } + + // A valid remote permission/range/protocol rejection, or a WRITE whose + // remote outcome is unknown because its ACK was lost, is permanent for + // this logical request. ShuttingDown maps to TooManyRequests and remains + // transient, so the existing failover policy may act on it. + if (!result.IsTooManyRequests()) { + finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + task.suppress_failover = true; + } + return Status::OK(); } void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id, @@ -2513,8 +2671,8 @@ void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id, task_status.s == CANCELED) { if (task.failure_stage < 0) task.failure_stage = 1; } - if (!allow_failover || task.cancel_requested || task_status.s != FAILED || - task.type == UNSPEC) + if (!allow_failover || task.cancel_requested || task.suppress_failover || + task_status.s != FAILED || task.type == UNSPEC) return; // The current physical transport attempt has failed even if the logical @@ -2659,11 +2817,7 @@ Status TransferEngineImpl::getBatchStatus(BatchID batch_id, size_t total_tasks = 0; TransferStatusEnum worst_failure = PENDING; auto isWorse = [](TransferStatusEnum cur, TransferStatusEnum best) { - static const std::unordered_map severity = { - {INITIAL, 0}, {PENDING, 0}, {COMPLETED, 0}, {INVALID, 1}, - {CANCELED, 2}, {TIMEOUT, 3}, {FAILED, 4}, - }; - return severity.at(cur) > severity.at(best); + return transferStatusSeverity(cur) > transferStatusSeverity(best); }; for (size_t task_id = 0; task_id < batch->task_list.size(); ++task_id) { auto& task = batch->task_list[task_id]; diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp index a1444ce5c0..3708dc2112 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp @@ -15,11 +15,16 @@ #include "tent/runtime/transfer_engine_impl.h" #include "tent/transport/shm/shm_transport.h" #include "tent/transport/tcp/tcp_transport.h" +#include "tent/transport/hp_tcp/hp_tcp_transport.h" #ifdef USE_RDMA #include "tent/transport/rdma/rdma_transport.h" #endif +#ifdef USE_UB +#include "tent/transport/ub/ub_transport.h" +#endif + #ifdef USE_CUDA #include "tent/transport/nvlink/nvlink_transport.h" #include "tent/transport/mnnvl/mnnvl_transport.h" @@ -56,6 +61,10 @@ Status TransferEngineImpl::loadTransports() { if (conf_->get("transports/tcp/enable", true)) transport_list_[TCP] = std::make_shared(); + if (hp_tcp_transport_config_.enabled) + transport_list_[HP_TCP] = std::make_shared( + hp_tcp_transport_config_.params); + // SHM is opt-in: default false because the current path is not NUMA-aware // (see tent/config/transfer-engine.json for an example that enables it). if (conf_->get("transports/shm/enable", false)) @@ -68,6 +77,13 @@ Status TransferEngineImpl::loadTransports() { } #endif +#ifdef USE_UB + if (conf_->get("transports/ub/enable", false) && + topology_->getNicCount(Topology::NIC_UB)) { + transport_list_[UB] = std::make_shared(); + } +#endif + #ifdef USE_URING if (conf_->get("transports/io_uring/enable", true)) transport_list_[IOURING] = std::make_shared(); diff --git a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt index 4a3f3c9494..0ed51cfe18 100644 --- a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(tcp) +add_subdirectory(hp_tcp) add_subdirectory(rdma) add_subdirectory(shm) add_subdirectory(nvlink) @@ -23,6 +24,7 @@ foreach( tent_xport_rdma tent_xport_shm tent_xport_tcp + tent_xport_hp_tcp tent_xport_ub tent_xport_ascend_direct tent_xport_sunrise_link diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/hp_tcp/CMakeLists.txt new file mode 100644 index 0000000000..937a1e4839 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/CMakeLists.txt @@ -0,0 +1,9 @@ +add_library( + tent_xport_hp_tcp_core STATIC + hp_tcp_protocol.cpp hp_tcp_workers.cpp hp_tcp_task.cpp + hp_tcp_buffer_registry.cpp hp_tcp_client.cpp hp_tcp_server.cpp) +target_link_libraries(tent_xport_hp_tcp_core PUBLIC tent_common tent_interface) + +add_library(tent_xport_hp_tcp STATIC hp_tcp_transport.cpp) +target_link_libraries(tent_xport_hp_tcp PUBLIC tent_rpc tent_common + tent_xport_hp_tcp_core) diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_buffer_registry.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_buffer_registry.cpp new file mode 100644 index 0000000000..d4f39d5ca5 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_buffer_registry.cpp @@ -0,0 +1,265 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" + +#include +#include +#include +#include + +namespace mooncake::tent { +namespace { + +bool RangeEnd(uint64_t base, uint64_t length, uint64_t* end) { + if (end == nullptr || length == 0 || + base > std::numeric_limits::max() - length) { + return false; + } + *end = base + length; + return true; +} + +uint64_t MakeRegistrationNamespace() { + std::random_device device; + std::seed_seq seed{device(), device(), device(), device(), + device(), device(), device(), device()}; + std::mt19937_64 random(seed); + return random(); +} + +Status LeaseError(HighPerformanceTcpStatus failure, const char* message) { + switch (failure) { + case HighPerformanceTcpStatus::kStaleRegistration: + return Status::NeedsRefreshCache(std::string(message) + LOC_MARK); + case HighPerformanceTcpStatus::kShuttingDown: + return Status::TooManyRequests(std::string(message) + LOC_MARK); + case HighPerformanceTcpStatus::kPermissionDenied: + case HighPerformanceTcpStatus::kRangeRejected: + return Status::AddressNotRegistered(std::string(message) + + LOC_MARK); + default: + return Status::InternalError(std::string(message) + LOC_MARK); + } +} + +} // namespace + +HighPerformanceTcpBufferRegistry::HighPerformanceTcpBufferRegistry() + : registration_namespace_(MakeRegistrationNamespace()) {} + +HighPerformanceTcpBufferRegistry::Lease::Lease(std::shared_ptr entry) + : entry_(std::move(entry)) {} + +HighPerformanceTcpBufferRegistry::Lease::Lease(Lease&& other) noexcept + : entry_(std::move(other.entry_)) {} + +HighPerformanceTcpBufferRegistry::Lease& +HighPerformanceTcpBufferRegistry::Lease::operator=(Lease&& other) noexcept { + if (this != &other) { + reset(); + entry_ = std::move(other.entry_); + } + return *this; +} + +HighPerformanceTcpBufferRegistry::Lease::~Lease() { reset(); } + +void HighPerformanceTcpBufferRegistry::Lease::reset() { + if (!entry_) return; + { + std::lock_guard lock(entry_->mutex); + if (entry_->active_leases > 0) --entry_->active_leases; + } + entry_->drained.notify_all(); + entry_.reset(); +} + +void* HighPerformanceTcpBufferRegistry::Lease::data() const { + return entry_ ? reinterpret_cast(entry_->base) : nullptr; +} + +uint64_t HighPerformanceTcpBufferRegistry::Lease::base() const { + return entry_ ? entry_->base : 0; +} + +uint64_t HighPerformanceTcpBufferRegistry::Lease::length() const { + return entry_ ? entry_->length : 0; +} + +Status HighPerformanceTcpBufferRegistry::add(uint64_t base, uint64_t length, + Permission permission, + uint64_t* registration_id) { + uint64_t end = 0; + if (!RangeEnd(base, length, &end)) { + return Status::InvalidArgument("invalid HP TCP buffer range" LOC_MARK); + } + if (permission != kLocalReadWrite && permission != kGlobalReadOnly && + permission != kGlobalReadWrite) { + return Status::InvalidArgument( + "invalid HP TCP buffer permission" LOC_MARK); + } + + std::lock_guard lock(registry_mutex_); + if (closing_) { + return Status::TooManyRequests( + "HP TCP buffer registry is shutting down" LOC_MARK); + } + const auto next = entries_.upper_bound(base); + if (next != entries_.end() && end > next->second->base) { + return Status::InvalidArgument("overlapping HP TCP buffer" LOC_MARK); + } + if (next != entries_.begin()) { + const auto previous = std::prev(next); + uint64_t previous_end = 0; + if (!RangeEnd(previous->second->base, previous->second->length, + &previous_end) || + previous_end > base) { + return Status::InvalidArgument( + "overlapping HP TCP buffer" LOC_MARK); + } + } + + if (next_registration_sequence_ == 0 || + next_registration_sequence_ == std::numeric_limits::max()) { + return Status::InternalError( + "HP TCP registration id space exhausted" LOC_MARK); + } + uint64_t id = registration_namespace_ ^ next_registration_sequence_++; + if (id == 0) { + if (next_registration_sequence_ == + std::numeric_limits::max()) { + return Status::InternalError( + "HP TCP registration id space exhausted" LOC_MARK); + } + id = registration_namespace_ ^ next_registration_sequence_++; + } + + auto entry = std::make_shared(); + entry->base = base; + entry->length = length; + entry->permission = permission; + entry->registration_id = id; + entries_.emplace(base, entry); + if (registration_id != nullptr) *registration_id = id; + return Status::OK(); +} + +Status HighPerformanceTcpBufferRegistry::remove(uint64_t base, + uint64_t length) { + std::shared_ptr entry; + { + std::lock_guard registry_lock(registry_mutex_); + const auto it = entries_.find(base); + if (it == entries_.end() || it->second->length != length) { + return Status::AddressNotRegistered( + "HP TCP buffer not registered" LOC_MARK); + } + entry = it->second; + // Hide the range before waiting so no new lease can race in. + entries_.erase(it); + } + + std::unique_lock entry_lock(entry->mutex); + entry->drained.wait(entry_lock, [&] { return entry->active_leases == 0; }); + return Status::OK(); +} + +void HighPerformanceTcpBufferRegistry::close() { + std::lock_guard lock(registry_mutex_); + closing_ = true; +} + +Status HighPerformanceTcpBufferRegistry::reopen() { + std::lock_guard lock(registry_mutex_); + if (!entries_.empty()) { + return Status::InvalidArgument( + "cannot reopen HP TCP buffer registry while buffers remain " + "registered" LOC_MARK); + } + closing_ = false; + return Status::OK(); +} + +Status HighPerformanceTcpBufferRegistry::acquireLocalLease(uint64_t addr, + uint64_t length, + Lease* lease) { + return acquire(addr, length, 0, HighPerformanceTcpOpcode::kRead, false, + lease, nullptr); +} + +Status HighPerformanceTcpBufferRegistry::acquireRemoteLease( + uint64_t addr, uint64_t length, uint64_t registration_id, + HighPerformanceTcpOpcode opcode, Lease* lease, + HighPerformanceTcpStatus* failure) { + return acquire(addr, length, registration_id, opcode, true, lease, failure); +} + +Status HighPerformanceTcpBufferRegistry::acquire( + uint64_t addr, uint64_t length, uint64_t registration_id, + HighPerformanceTcpOpcode opcode, bool remote, Lease* lease, + HighPerformanceTcpStatus* failure) { + if (failure != nullptr) *failure = HighPerformanceTcpStatus::kOk; + if (lease == nullptr) { + return Status::InvalidArgument("HP TCP lease output is null" LOC_MARK); + } + lease->reset(); + + uint64_t end = 0; + if (!RangeEnd(addr, length, &end)) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kRangeRejected; + return Status::InvalidArgument("invalid HP TCP lease range" LOC_MARK); + } + + std::lock_guard registry_lock(registry_mutex_); + if (closing_) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kShuttingDown; + return LeaseError(HighPerformanceTcpStatus::kShuttingDown, + "HP TCP buffer registry is shutting down"); + } + const auto next = entries_.upper_bound(addr); + if (next == entries_.begin()) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kRangeRejected; + return LeaseError(HighPerformanceTcpStatus::kRangeRejected, + "HP TCP range not registered"); + } + const auto entry = std::prev(next)->second; + uint64_t entry_end = 0; + if (!RangeEnd(entry->base, entry->length, &entry_end) || + addr < entry->base || end > entry_end) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kRangeRejected; + return LeaseError(HighPerformanceTcpStatus::kRangeRejected, + "HP TCP range not registered"); + } + + std::lock_guard entry_lock(entry->mutex); + if (remote && entry->registration_id != registration_id) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kStaleRegistration; + return LeaseError(HighPerformanceTcpStatus::kStaleRegistration, + "stale HP TCP registration"); + } + if (remote && (entry->permission == kLocalReadWrite || + (opcode == HighPerformanceTcpOpcode::kWrite && + entry->permission != kGlobalReadWrite))) { + if (failure != nullptr) + *failure = HighPerformanceTcpStatus::kPermissionDenied; + return LeaseError(HighPerformanceTcpStatus::kPermissionDenied, + "HP TCP permission denied"); + } + + ++entry->active_leases; + *lease = Lease(entry); + return Status::OK(); +} + +bool HighPerformanceTcpBufferRegistry::tracks(uint64_t base, + uint64_t length) const { + std::lock_guard lock(registry_mutex_); + const auto it = entries_.find(base); + return it != entries_.end() && it->second->length == length; +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_client.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_client.cpp new file mode 100644 index 0000000000..5be2fcf925 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_client.cpp @@ -0,0 +1,603 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_client.h" + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace mooncake::tent { +class HighPerformanceTcpClient::Lane + : public std::enable_shared_from_this { + public: + Lane(HighPerformanceTcpClient* parent, asio::io_context& io, Config config, + LaneKey key) + : parent_(parent), + config_(std::move(config)), + key_(std::move(key)), + resolver_(io), + socket_(io), + timer_(io) {} + + void enqueue(Operation operation) { + try { + queue_.push_back(std::move(operation)); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP lane enqueue failed: " << error.what(); + completeStandalone(std::move(operation), FAILED, 0); + return; + } catch (...) { + LOG(ERROR) << "HP TCP lane enqueue failed"; + completeStandalone(std::move(operation), FAILED, 0); + return; + } + if (!current_) startNext(); + } + + void cancelAll(TransferStatusEnum terminal) { + while (!queue_.empty()) { + Operation operation = std::move(queue_.front()); + queue_.pop_front(); + completeStandalone(std::move(operation), terminal, 0); + } + if (!current_) { + closeDirty(); + return; + } + forced_terminal_ = terminal; + ++timer_generation_; + std::error_code ignored; + timer_.cancel(ignored); + resolver_.cancel(); + closeDirty(); + // The outstanding async callback owns the final completion. Releasing + // the operation (and its memory lease in the adapter) here would be a + // use-after-free risk because Asio still owns the user buffer. + } + + bool cancelRequest(uint64_t request_id) { + for (auto it = queue_.begin(); it != queue_.end(); ++it) { + if (it->request_id != request_id) continue; + Operation operation = std::move(*it); + queue_.erase(it); + completeStandalone(std::move(operation), CANCELED, 0); + return true; + } + if (!current_ || current_->request_id != request_id) return false; + forced_terminal_ = CANCELED; + cancelTimer(); + resolver_.cancel(); + closeDirty(); + // The live Asio callback owns final completion and buffer retirement. + return true; + } + + private: + template + void runHandler(uint64_t epoch, Function&& function) noexcept { + if (!matches(epoch)) return; + try { + function(); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP client handler failed: " << error.what(); + if (matches(epoch)) { + finishCurrent(FAILED, 0, false); + } + } catch (...) { + LOG(ERROR) << "HP TCP client handler failed"; + if (matches(epoch)) { + finishCurrent(FAILED, 0, false); + } + } + } + + void startNext() { + if (current_ || queue_.empty()) return; + current_.emplace(std::move(queue_.front())); + queue_.pop_front(); + ++operation_epoch_; + forced_terminal_.reset(); + body_offset_ = 0; + remote_write_outcome_unknown_ = false; + request_bytes_ = EncodeHighPerformanceTcpRequest( + {current_->opcode, current_->request_id, current_->registration_id, + current_->remote_addr, current_->length}); + + try { + if (socket_.is_open()) { + writeHeader(operation_epoch_); + } else { + resolve(operation_epoch_); + } + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP async initiation failed: " << error.what(); + finishCurrent(FAILED, 0, false); + } catch (...) { + LOG(ERROR) << "HP TCP async initiation failed"; + finishCurrent(FAILED, 0, false); + } + } + + void resolve(uint64_t epoch) { + armTimer(config_.connect_timeout_ms, epoch); + auto self = shared_from_this(); + resolver_.async_resolve( + key_.host, std::to_string(key_.port), + [self, epoch](const std::error_code& error, + asio::ip::tcp::resolver::results_type results) { + self->runHandler(epoch, [&] { + if (self->finishForcedIfAny()) return; + if (error) { + self->finishIoError(error); + return; + } + self->connect(epoch, std::move(results)); + }); + }); + } + + void connect(uint64_t epoch, + asio::ip::tcp::resolver::results_type results) { + std::error_code ignored; + socket_.close(ignored); + auto self = shared_from_this(); + asio::async_connect( + socket_, results, + [self, epoch](const std::error_code& error, + const asio::ip::tcp::endpoint&) { + self->runHandler(epoch, [&] { + if (self->finishForcedIfAny()) return; + if (error) { + self->finishIoError(error); + return; + } + self->cancelTimer(); + self->parent_->connections_created_.fetch_add( + 1, std::memory_order_relaxed); + self->writeHeader(epoch); + }); + }); + } + + void writeHeader(uint64_t epoch) { + armTimer(config_.progress_timeout_ms, epoch); + auto self = shared_from_this(); + asio::async_write( + socket_, asio::buffer(request_bytes_), + [self, epoch](const std::error_code& error, size_t bytes) { + self->runHandler(epoch, [&] { + if (self->current_->opcode == + HighPerformanceTcpOpcode::kWrite) { + // Any reported header byte makes replay conservative: + // the peer may subsequently observe a complete request. + self->remote_write_outcome_unknown_ = bytes > 0; + } + if (self->finishForcedIfAny()) return; + if (error || bytes != kHighPerformanceTcpRequestSize) { + self->finishIoError( + error ? error : asio::error::operation_aborted); + return; + } + self->armTimer(self->config_.progress_timeout_ms, epoch); + if (self->current_->opcode == + HighPerformanceTcpOpcode::kWrite) { + self->body_offset_ = 0; + self->writeBodyChunk(epoch); + } else { + self->readResponse(epoch); + } + }); + }); + } + + void writeBodyChunk(uint64_t epoch) { + if (body_offset_ == current_->length) { + readResponse(epoch); + return; + } + const size_t chunk = static_cast(std::min( + config_.chunk_size, current_->length - body_offset_)); + auto* data = static_cast(current_->local_addr) + body_offset_; + auto self = shared_from_this(); + asio::async_write( + socket_, asio::buffer(data, chunk), + [self, epoch, chunk](const std::error_code& error, size_t bytes) { + self->runHandler(epoch, [&] { + if (self->finishForcedIfAny()) return; + if (error || bytes != chunk) { + self->finishIoError( + error ? error : asio::error::operation_aborted); + return; + } + self->body_offset_ += bytes; + self->armTimer(self->config_.progress_timeout_ms, epoch); + self->writeBodyChunk(epoch); + }); + }); + } + + void readResponse(uint64_t epoch) { + auto self = shared_from_this(); + asio::async_read( + socket_, asio::buffer(response_bytes_), + [self, epoch](const std::error_code& error, size_t bytes) { + self->runHandler(epoch, [&] { + if (self->finishForcedIfAny()) return; + if (error || bytes != kHighPerformanceTcpResponseSize) { + self->finishIoError( + error ? error : asio::error::operation_aborted); + return; + } + self->armTimer(self->config_.progress_timeout_ms, epoch); + HighPerformanceTcpResponseFrame response; + const Status decoded = DecodeHighPerformanceTcpResponse( + self->response_bytes_.data(), + self->response_bytes_.size(), &response); + if (!decoded.ok()) { + self->finishProtocolError(); + return; + } + if (response.request_id != self->current_->request_id) { + self->finishProtocolError(); + return; + } + if (response.status != HighPerformanceTcpStatus::kOk) { + self->finishRemoteError(response.status); + return; + } + if (response.committed_bytes != self->current_->length) { + self->finishProtocolError(); + return; + } + if (self->current_->opcode == + HighPerformanceTcpOpcode::kWrite) { + self->finishClean(); + } else { + self->body_offset_ = 0; + self->readBodyChunk(epoch); + } + }); + }); + } + + void readBodyChunk(uint64_t epoch) { + if (body_offset_ == current_->length) { + finishClean(); + return; + } + const size_t chunk = static_cast(std::min( + config_.chunk_size, current_->length - body_offset_)); + auto* data = static_cast(current_->local_addr) + body_offset_; + auto self = shared_from_this(); + asio::async_read( + socket_, asio::buffer(data, chunk), + [self, epoch, chunk](const std::error_code& error, size_t bytes) { + self->runHandler(epoch, [&] { + if (self->finishForcedIfAny()) return; + if (error || bytes != chunk) { + self->finishIoError( + error ? error : asio::error::operation_aborted); + return; + } + self->body_offset_ += bytes; + self->armTimer(self->config_.progress_timeout_ms, epoch); + self->readBodyChunk(epoch); + }); + }); + } + + void armTimer(uint64_t timeout_ms, uint64_t epoch) { + const uint64_t generation = ++timer_generation_; + timer_.expires_after(std::chrono::milliseconds(timeout_ms)); + auto self = shared_from_this(); + timer_.async_wait( + [self, epoch, generation](const std::error_code& error) { + if (error == asio::error::operation_aborted) return; + if (error || !self->matches(epoch) || + generation != self->timer_generation_) { + return; + } + self->forced_terminal_ = TIMEOUT; + self->resolver_.cancel(); + self->closeDirty(); + // The active resolve/socket callback completes the operation. + // This makes timeout sticky without releasing the caller buffer + // before the canceled I/O handler has quiesced. + }); + } + + void cancelTimer() { + ++timer_generation_; + std::error_code ignored; + timer_.cancel(ignored); + } + + bool matches(uint64_t epoch) const { + return current_.has_value() && epoch == operation_epoch_; + } + + bool finishForcedIfAny() { + if (!forced_terminal_.has_value()) return false; + const TransferStatusEnum terminal = *forced_terminal_; + finishCurrent(terminal, 0, false); + return true; + } + + void closeDirty() { + std::error_code ignored; + if (socket_.is_open()) { + socket_.cancel(ignored); + socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored); + socket_.close(ignored); + } + } + + void finishProtocolError() { finishCurrent(FAILED, 0, false); } + + void finishRemoteError(HighPerformanceTcpStatus status) { + finishCurrent(FAILED, 0, false, status); + } + + void finishIoError(const std::error_code&) { + const TransferStatusEnum terminal = forced_terminal_.value_or(FAILED); + finishCurrent(terminal, 0, false); + } + + void finishClean() { + cancelTimer(); + remote_write_outcome_unknown_ = false; + finishCurrent(COMPLETED, current_->length, true); + } + + void finishCurrent( + TransferStatusEnum terminal, size_t bytes, bool keep_stream, + std::optional remote_status = std::nullopt) { + cancelTimer(); + if (!keep_stream) closeDirty(); + Operation operation = std::move(*current_); + if (terminal == FAILED && remote_write_outcome_unknown_ && + !remote_status.has_value()) { + // A WRITE may have reached the peer without a valid ACK. Reuse the + // existing terminal protocol classification so it cannot be + // replayed through another transport. + remote_status = HighPerformanceTcpStatus::kInternalError; + } + current_.reset(); + ++operation_epoch_; // invalidate late timer/cancel callbacks + forced_terminal_.reset(); + remote_write_outcome_unknown_ = false; + completeStandalone(std::move(operation), terminal, bytes, + remote_status); + startNext(); + } + + void completeStandalone( + Operation operation, TransferStatusEnum terminal, size_t bytes, + std::optional remote_status = std::nullopt) { + try { + if (operation.complete) + operation.complete(terminal, bytes, remote_status); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP completion callback threw: " << error.what(); + } catch (...) { + LOG(ERROR) << "HP TCP completion callback threw"; + } + parent_->operationFinished(); + } + + HighPerformanceTcpClient* parent_; + Config config_; + LaneKey key_; + asio::ip::tcp::resolver resolver_; + asio::ip::tcp::socket socket_; + asio::steady_timer timer_; + std::deque queue_; + std::optional current_; + std::array request_bytes_{}; + std::array response_bytes_{}; + uint64_t operation_epoch_{0}; + uint64_t timer_generation_{0}; + uint64_t body_offset_{0}; + bool remote_write_outcome_unknown_{false}; + std::optional forced_terminal_; +}; + +bool HighPerformanceTcpClient::LaneKey::operator==(const LaneKey& other) const { + return peer_id == other.peer_id && incarnation == other.incarnation && + host == other.host && port == other.port && lane_id == other.lane_id; +} + +size_t HighPerformanceTcpClient::LaneKeyHash::operator()( + const LaneKey& key) const { + size_t hash = std::hash{}(key.peer_id); + const auto mix = [&hash](size_t value) { + hash ^= value + static_cast(0x9e3779b97f4a7c15ULL) + + (hash << 6U) + (hash >> 2U); + }; + mix(std::hash{}(key.incarnation)); + mix(std::hash{}(key.host)); + mix(std::hash{}(key.port)); + mix(std::hash{}(key.lane_id)); + return hash; +} + +HighPerformanceTcpClient::HighPerformanceTcpClient( + Config config, HighPerformanceTcpWorkers* workers) + : config_(std::move(config)), workers_(workers) { + if (workers_ != nullptr) worker_states_.resize(workers_->workerCount()); +} + +HighPerformanceTcpClient::~HighPerformanceTcpClient() { + if (workers_ != nullptr && workers_->controlContextAvailable() && + !workers_->onWorkerThread()) { + (void)cancelAll(CANCELED); + } + DCHECK(workers_ == nullptr || + active_operations_.load(std::memory_order_acquire) == 0) + << "HP TCP client destroyed with active operations"; +} + +void HighPerformanceTcpClient::operationStarted() { + active_operations_.fetch_add(1, std::memory_order_acq_rel); +} + +void HighPerformanceTcpClient::operationFinished() { + const uint64_t previous = + active_operations_.fetch_sub(1, std::memory_order_acq_rel); + if (previous == 1) { + std::lock_guard lock(active_mutex_); + active_cv_.notify_all(); + } +} + +void HighPerformanceTcpClient::enqueueOnOwner(size_t owner_worker, + Operation operation) { + if (workers_ == nullptr || owner_worker >= worker_states_.size() || + !operation.complete) { + if (operation.complete) { + try { + operation.complete(FAILED, 0, std::nullopt); + } catch (...) { + LOG(ERROR) << "HP TCP rejected-operation callback threw"; + } + } + return; + } + + operationStarted(); + try { + if (stopping_.load(std::memory_order_acquire) || + operation.local_addr == nullptr || operation.length == 0 || + operation.length > config_.max_transfer_bytes || + operation.host.empty() || operation.port == 0 || + operation.lane_id >= config_.connections_per_peer) { + try { + operation.complete(CANCELED, 0, std::nullopt); + } catch (...) { + LOG(ERROR) << "HP TCP rejected-operation callback threw"; + } + operationFinished(); + return; + } + + auto& lanes = worker_states_[owner_worker].lanes; + for (auto it = lanes.begin(); it != lanes.end();) { + if (it->first.peer_id == operation.peer_id && + it->first.incarnation != operation.incarnation) { + it->second->cancelAll(CANCELED); + it = lanes.erase(it); + } else { + ++it; + } + } + + LaneKey key{operation.peer_id, operation.incarnation, operation.host, + operation.port, operation.lane_id}; + auto it = lanes.find(key); + if (it == lanes.end()) { + auto lane = std::make_shared( + this, workers_->ioContext(owner_worker), config_, key); + it = lanes.emplace(std::move(key), std::move(lane)).first; + } + // Lane::enqueue owns completion + active-operation retirement from + // this point forward, including allocation/initiation failures. + it->second->enqueue(std::move(operation)); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP lane setup failed: " << error.what(); + try { + operation.complete(FAILED, 0, std::nullopt); + } catch (...) { + LOG(ERROR) << "HP TCP failed-operation callback threw"; + } + operationFinished(); + } catch (...) { + LOG(ERROR) << "HP TCP lane setup failed"; + try { + operation.complete(FAILED, 0, std::nullopt); + } catch (...) { + LOG(ERROR) << "HP TCP failed-operation callback threw"; + } + operationFinished(); + } +} + +void HighPerformanceTcpClient::cancelWorker(size_t worker_id, + TransferStatusEnum terminal) { + if (worker_id >= worker_states_.size()) return; + auto& lanes = worker_states_[worker_id].lanes; + for (auto& [key, lane] : lanes) { + (void)key; + lane->cancelAll(terminal); + } + lanes.clear(); +} + +void HighPerformanceTcpClient::cancelRequestOnWorker(size_t worker_id, + uint64_t request_id) { + if (worker_id >= worker_states_.size()) return; + for (auto& [key, lane] : worker_states_[worker_id].lanes) { + (void)key; + if (lane->cancelRequest(request_id)) return; + } +} + +Status HighPerformanceTcpClient::cancelRequest(size_t owner_worker, + uint64_t request_id) { + if (workers_ == nullptr || owner_worker >= worker_states_.size() || + request_id == 0) { + return Status::InvalidArgument( + "invalid HP TCP cancellation request" LOC_MARK); + } + if (!workers_->controlContextAvailable()) { + return Status::InternalError( + "HP TCP worker contexts are unavailable" LOC_MARK); + } + try { + asio::post(workers_->ioContext(owner_worker), + [this, owner_worker, request_id] { + cancelRequestOnWorker(owner_worker, request_id); + }); + } catch (const std::exception& error) { + return Status::InternalError( + std::string("HP TCP cancellation post failed: ") + error.what() + + LOC_MARK); + } + return Status::OK(); +} + +Status HighPerformanceTcpClient::cancelAll(TransferStatusEnum terminal) { + if (workers_ == nullptr) return Status::OK(); + if (workers_->onWorkerThread()) { + return Status::InvalidArgument( + "HP TCP client cancelAll cannot block a worker" LOC_MARK); + } + stopping_.store(true, std::memory_order_release); + + if (workers_->controlContextAvailable()) { + try { + // Post cancellation to every owner before waiting for active I/O. + for (size_t i = 0; i < workers_->workerCount(); ++i) { + asio::post(workers_->ioContext(i), + [this, i, terminal] { cancelWorker(i, terminal); }); + } + } catch (const std::exception& error) { + return Status::InternalError( + std::string("HP TCP client cancellation post failed: ") + + error.what() + LOC_MARK); + } + CHECK_STATUS(workers_->barrier()); + } + + std::unique_lock lock(active_mutex_); + active_cv_.wait(lock, [&] { + return active_operations_.load(std::memory_order_acquire) == 0; + }); + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_protocol.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_protocol.cpp new file mode 100644 index 0000000000..52907e7164 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_protocol.cpp @@ -0,0 +1,294 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" + +#include +#include +#include +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake::tent { +namespace { + +using json = nlohmann::json; + +void Put16(uint8_t* out, uint16_t value) { + out[0] = static_cast(value >> 8U); + out[1] = static_cast(value); +} + +void Put32(uint8_t* out, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + out[i] = static_cast(value >> ((3U - i) * 8U)); + } +} + +void Put64(uint8_t* out, uint64_t value) { + for (size_t i = 0; i < 8; ++i) { + out[i] = static_cast(value >> ((7U - i) * 8U)); + } +} + +uint16_t Get16(const uint8_t* in) { + return static_cast((static_cast(in[0]) << 8U) | + static_cast(in[1])); +} + +uint32_t Get32(const uint8_t* in) { + uint32_t value = 0; + for (size_t i = 0; i < 4; ++i) { + value = (value << 8U) | in[i]; + } + return value; +} + +uint64_t Get64(const uint8_t* in) { + uint64_t value = 0; + for (size_t i = 0; i < 8; ++i) { + value = (value << 8U) | in[i]; + } + return value; +} + +Status InvalidFrame(const std::string& detail, HighPerformanceTcpStatus error, + HighPerformanceTcpStatus* wire_error) { + if (wire_error != nullptr) *wire_error = error; + return Status::InvalidArgument( + "Invalid high-performance TCP frame: " + detail + LOC_MARK); +} + +bool IsHex128(const std::string& value) { + return value.size() == 32 && + std::all_of(value.begin(), value.end(), + [](unsigned char ch) { return std::isxdigit(ch) != 0; }); +} + +Status InvalidAttr(const std::string& detail) { + return Status::InvalidArgument("Invalid HP TCP attribute: " + detail + + LOC_MARK); +} + +bool ReadPositiveUint64(const json& object, const char* key, uint64_t* out) { + auto it = object.find(key); + if (it == object.end()) return false; + if (it->is_number_unsigned()) { + const auto value = it->get(); + if (value == 0) return false; + *out = value; + return true; + } + if (!it->is_number_integer()) return false; + const auto value = it->get(); + if (value <= 0) return false; + *out = static_cast(value); + return true; +} + +} // namespace + +std::array +EncodeHighPerformanceTcpRequest(const HighPerformanceTcpRequestFrame& frame) { + std::array bytes{}; + Put32(bytes.data(), kHighPerformanceTcpMagic); + Put16(bytes.data() + 4, kHighPerformanceTcpVersion); + bytes[6] = static_cast(frame.opcode); + bytes[7] = 0; // flags, reserved in v1 + Put64(bytes.data() + 8, frame.request_id); + Put64(bytes.data() + 16, frame.registration_id); + Put64(bytes.data() + 24, frame.remote_addr); + Put64(bytes.data() + 32, frame.length); + Put64(bytes.data() + 40, 0); // reserved + return bytes; +} + +Status DecodeHighPerformanceTcpRequest(const uint8_t* bytes, size_t size, + HighPerformanceTcpRequestFrame* frame, + HighPerformanceTcpStatus* wire_error) { + if (wire_error != nullptr) { + *wire_error = HighPerformanceTcpStatus::kInternalError; + } + if (bytes == nullptr || frame == nullptr || + size != kHighPerformanceTcpRequestSize) { + return InvalidFrame("request size", + HighPerformanceTcpStatus::kInternalError, + wire_error); + } + if (Get32(bytes) != kHighPerformanceTcpMagic) { + return InvalidFrame("magic", HighPerformanceTcpStatus::kInternalError, + wire_error); + } + if (Get16(bytes + 4) != kHighPerformanceTcpVersion) { + return InvalidFrame("version", HighPerformanceTcpStatus::kBadVersion, + wire_error); + } + if (bytes[7] != 0 || Get64(bytes + 40) != 0) { + return InvalidFrame("flags/reserved", + HighPerformanceTcpStatus::kInternalError, + wire_error); + } + if (bytes[6] != static_cast(HighPerformanceTcpOpcode::kRead) && + bytes[6] != static_cast(HighPerformanceTcpOpcode::kWrite)) { + return InvalidFrame("opcode", HighPerformanceTcpStatus::kBadOpcode, + wire_error); + } + + const uint64_t remote_addr = Get64(bytes + 24); + const uint64_t length = Get64(bytes + 32); + if (length == 0 || + remote_addr > std::numeric_limits::max() - length) { + return InvalidFrame("length/range", + HighPerformanceTcpStatus::kBadLength, wire_error); + } + + frame->opcode = static_cast(bytes[6]); + frame->request_id = Get64(bytes + 8); + frame->registration_id = Get64(bytes + 16); + frame->remote_addr = remote_addr; + frame->length = length; + return Status::OK(); +} + +std::array +EncodeHighPerformanceTcpResponse(const HighPerformanceTcpResponseFrame& frame) { + std::array bytes{}; + Put32(bytes.data(), kHighPerformanceTcpMagic); + Put16(bytes.data() + 4, kHighPerformanceTcpVersion); + Put16(bytes.data() + 6, static_cast(frame.status)); + Put64(bytes.data() + 8, frame.request_id); + Put64(bytes.data() + 16, frame.committed_bytes); + Put64(bytes.data() + 24, 0); + return bytes; +} + +Status DecodeHighPerformanceTcpResponse( + const uint8_t* bytes, size_t size, HighPerformanceTcpResponseFrame* frame) { + if (bytes == nullptr || frame == nullptr || + size != kHighPerformanceTcpResponseSize) { + return Status::InvalidArgument( + "Invalid high-performance TCP response size" LOC_MARK); + } + if (Get32(bytes) != kHighPerformanceTcpMagic || + Get16(bytes + 4) != kHighPerformanceTcpVersion || + Get64(bytes + 24) != 0) { + return Status::InvalidArgument( + "Invalid high-performance TCP response header" LOC_MARK); + } + const uint16_t status = Get16(bytes + 6); + if (status > + static_cast(HighPerformanceTcpStatus::kInternalError)) { + return Status::InvalidArgument( + "Invalid high-performance TCP response status" LOC_MARK); + } + frame->status = static_cast(status); + frame->request_id = Get64(bytes + 8); + frame->committed_bytes = Get64(bytes + 16); + return Status::OK(); +} + +const char* HighPerformanceTcpPermissionName(Permission permission) { + switch (permission) { + case kLocalReadWrite: + return "local_read_write"; + case kGlobalReadOnly: + return "global_read_only"; + case kGlobalReadWrite: + return "global_read_write"; + } + return "unknown"; +} + +Status EncodeHighPerformanceTcpEndpointAttr( + const HighPerformanceTcpEndpointAttr& attr, std::string* encoded) { + if (encoded == nullptr || !IsHex128(attr.incarnation) || + attr.host.empty() || attr.port == 0 || attr.max_transfer_bytes == 0) { + return InvalidAttr("endpoint"); + } + + json object = { + {"protocol", "tent_hp_tcp"}, + {"version", kHighPerformanceTcpVersion}, + {"incarnation", attr.incarnation}, + {"endpoints", + json::array({{{"host", attr.host}, {"port", attr.port}}})}, + {"max_transfer_bytes", attr.max_transfer_bytes}, + }; + *encoded = object.dump(); + return Status::OK(); +} + +Status DecodeHighPerformanceTcpEndpointAttr( + const std::string& encoded, HighPerformanceTcpEndpointAttr* attr) { + if (attr == nullptr) return InvalidAttr("null endpoint output"); + try { + const json object = json::parse(encoded); + if (!object.is_object() || + object.value("protocol", "") != "tent_hp_tcp" || + object.value("version", 0) != kHighPerformanceTcpVersion) { + return InvalidAttr("endpoint protocol/version"); + } + const std::string incarnation = object.value("incarnation", ""); + const auto endpoints = object.find("endpoints"); + if (!IsHex128(incarnation) || endpoints == object.end() || + !endpoints->is_array() || endpoints->size() != 1 || + !(*endpoints)[0].is_object()) + return InvalidAttr("endpoint identity/list"); + const json& endpoint = (*endpoints)[0]; + const std::string host = endpoint.value("host", ""); + uint64_t port = 0; + uint64_t max_transfer_bytes = 0; + if (host.empty() || !ReadPositiveUint64(endpoint, "port", &port) || + port > 65535 || + !ReadPositiveUint64(object, "max_transfer_bytes", + &max_transfer_bytes)) + return InvalidAttr("endpoint address/limit"); + *attr = {incarnation, host, static_cast(port), + max_transfer_bytes}; + return Status::OK(); + } catch (const std::exception& error) { + return Status::MalformedJson( + std::string("Invalid HP TCP endpoint attribute: ") + error.what() + + LOC_MARK); + } +} + +Status EncodeHighPerformanceTcpBufferAttr( + const HighPerformanceTcpBufferAttr& attr, std::string* encoded) { + if (encoded == nullptr || attr.registration_id == 0 || + (attr.permission != "global_read_only" && + attr.permission != "global_read_write")) { + return InvalidAttr("buffer"); + } + *encoded = json{{"protocol", "tent_hp_tcp"}, + {"version", kHighPerformanceTcpVersion}, + {"registration_id", attr.registration_id}, + {"permission", attr.permission}} + .dump(); + return Status::OK(); +} + +Status DecodeHighPerformanceTcpBufferAttr(const std::string& encoded, + HighPerformanceTcpBufferAttr* attr) { + if (attr == nullptr) return InvalidAttr("null buffer output"); + try { + const json object = json::parse(encoded); + if (!object.is_object() || + object.value("protocol", "") != "tent_hp_tcp" || + object.value("version", 0) != kHighPerformanceTcpVersion) + return InvalidAttr("buffer protocol/version"); + uint64_t registration = 0; + const std::string permission = object.value("permission", ""); + if (!ReadPositiveUint64(object, "registration_id", ®istration) || + (permission != "global_read_only" && + permission != "global_read_write")) + return InvalidAttr("buffer registration/permission"); + *attr = {registration, permission}; + return Status::OK(); + } catch (const std::exception& error) { + return Status::MalformedJson( + std::string("Invalid HP TCP buffer attribute: ") + error.what() + + LOC_MARK); + } +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_server.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_server.cpp new file mode 100644 index 0000000000..70ee0cfd6a --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_server.cpp @@ -0,0 +1,627 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_server.h" + +#include +#include +#include +#include +#include +#include + +#include + +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" + +namespace mooncake::tent { + +class HighPerformanceTcpServer::Session + : public std::enable_shared_from_this { + public: + Session(HighPerformanceTcpServer* parent, size_t worker_id, + std::shared_ptr socket, Config config, + HighPerformanceTcpBufferRegistry* registry) + : parent_(parent), + worker_id_(worker_id), + socket_(std::move(socket)), + config_(std::move(config)), + registry_(registry), + timer_(socket_->get_executor()) {} + + void start() { readHeader(); } + + void cancel() { + forced_close_ = true; + cancelTimer(); + std::error_code ignored; + socket_->cancel(ignored); + socket_->shutdown(asio::ip::tcp::socket::shutdown_both, ignored); + socket_->close(ignored); + // Every live session always has exactly one async socket operation + // outstanding. Its callback performs the final lease release and + // registry removal, so cancel() never publishes closure early. + } + + private: + template + void runHandler(Function&& function) noexcept { + try { + function(); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP server handler failed: " << error.what(); + finishClosed(); + } catch (...) { + LOG(ERROR) << "HP TCP server handler failed"; + finishClosed(); + } + } + + void readHeader() { + if (forced_close_) { + finishClosed(); + return; + } + lease_.reset(); + body_offset_ = 0; + header_offset_ = 0; + ++request_epoch_; + cancelTimer(); + // A newly accepted peer must send its first byte within the progress + // deadline so an empty connection cannot hold a bounded session slot. + // Once a request has completed, waiting for the next request is + // connection-pool idle time rather than stalled I/O. + if (request_epoch_ == 1) armProgressTimer(request_epoch_); + readHeaderChunk(); + } + + void readHeaderChunk() { + auto self = shared_from_this(); + socket_->async_read_some( + asio::buffer(request_bytes_.data() + header_offset_, + request_bytes_.size() - header_offset_), + [self](const std::error_code& error, size_t bytes) { + self->runHandler([&] { + if (self->forced_close_) { + self->finishClosed(); + return; + } + if (error || bytes == 0) { + self->finishClosed(); + return; + } + self->header_offset_ += bytes; + // A reused stream may remain idle indefinitely. Once any + // byte of its next frame arrives, each partial read must + // make progress before the deadline. + if (self->header_offset_ == self->request_bytes_.size()) { + self->handleHeader(); + } else { + self->armProgressTimer(self->request_epoch_); + self->readHeaderChunk(); + } + }); + }); + } + + void handleHeader() { + HighPerformanceTcpStatus wire_error = + HighPerformanceTcpStatus::kInternalError; + const Status decoded = DecodeHighPerformanceTcpRequest( + request_bytes_.data(), request_bytes_.size(), &request_, + &wire_error); + if (!decoded.ok()) { + request_.request_id = 0; + sendResponse(wire_error, 0, true); + return; + } + if (request_.length > config_.max_transfer_bytes) { + sendResponse(HighPerformanceTcpStatus::kBadLength, 0, true); + return; + } + + HighPerformanceTcpStatus wire_status = + HighPerformanceTcpStatus::kInternalError; + const Status lease_status = registry_->acquireRemoteLease( + request_.remote_addr, request_.length, request_.registration_id, + request_.opcode, &lease_, &wire_status); + if (!lease_status.ok()) { + if (request_.opcode == HighPerformanceTcpOpcode::kWrite) { + // The client sends the complete WRITE body before reading the + // response. Consume exactly this rejected frame under the + // normal progress deadline, then reply and close the stream. + body_offset_ = 0; + discardRejectedWriteBody(wire_status); + } else { + sendResponse(wire_status, 0, true); + } + return; + } + + armProgressTimer(request_epoch_); + if (request_.opcode == HighPerformanceTcpOpcode::kRead) { + // The response says the entire requested range is valid. The + // client still completes only after all payload bytes arrive. + sendResponse(HighPerformanceTcpStatus::kOk, request_.length, false, + [self = shared_from_this()] { + self->body_offset_ = 0; + self->writeReadBodyChunk(); + }); + } else { + body_offset_ = 0; + readWriteBodyChunk(); + } + } + + void discardRejectedWriteBody(HighPerformanceTcpStatus status) { + if (body_offset_ == request_.length) { + sendResponse(status, 0, true); + return; + } + const size_t chunk = static_cast(std::min( + discard_bytes_.size(), request_.length - body_offset_)); + armProgressTimer(request_epoch_); + auto self = shared_from_this(); + asio::async_read( + *socket_, asio::buffer(discard_bytes_.data(), chunk), + [self, status, chunk](const std::error_code& error, size_t bytes) { + self->runHandler([&] { + if (self->forced_close_) { + self->finishClosed(); + return; + } + if (error || bytes != chunk) { + self->finishClosed(); + return; + } + self->body_offset_ += bytes; + self->discardRejectedWriteBody(status); + }); + }); + } + + uint8_t* remoteDataAt(uint64_t offset) { + return static_cast(lease_.data()) + + (request_.remote_addr - lease_.base()) + offset; + } + + void writeReadBodyChunk() { + if (body_offset_ == request_.length) { + lease_.reset(); + readHeader(); + return; + } + const size_t chunk = static_cast(std::min( + config_.chunk_size, request_.length - body_offset_)); + auto self = shared_from_this(); + asio::async_write( + *socket_, asio::buffer(remoteDataAt(body_offset_), chunk), + [self, chunk](const std::error_code& error, size_t bytes) { + self->runHandler([&] { + if (self->forced_close_) { + self->finishClosed(); + return; + } + if (error || bytes != chunk) { + self->finishClosed(); + return; + } + self->body_offset_ += bytes; + self->armProgressTimer(self->request_epoch_); + self->writeReadBodyChunk(); + }); + }); + } + + void readWriteBodyChunk() { + if (body_offset_ == request_.length) { + // The final async_read callback has returned: every byte is now in + // destination DRAM and no socket callback can touch the range. + lease_.reset(); + armProgressTimer(request_epoch_); + sendResponse(HighPerformanceTcpStatus::kOk, request_.length, false, + [self = shared_from_this()] { self->readHeader(); }); + return; + } + const size_t chunk = static_cast(std::min( + config_.chunk_size, request_.length - body_offset_)); + auto self = shared_from_this(); + asio::async_read( + *socket_, asio::buffer(remoteDataAt(body_offset_), chunk), + [self, chunk](const std::error_code& error, size_t bytes) { + self->runHandler([&] { + if (self->forced_close_) { + self->finishClosed(); + return; + } + if (error || bytes != chunk) { + self->finishClosed(); + return; + } + self->body_offset_ += bytes; + self->armProgressTimer(self->request_epoch_); + self->readWriteBodyChunk(); + }); + }); + } + + void sendResponse(HighPerformanceTcpStatus status, uint64_t committed, + bool close_after, + std::function continuation = {}) { + response_bytes_ = EncodeHighPerformanceTcpResponse( + {status, request_.request_id, committed}); + // Error responses are I/O too: arm a deadline even when request + // validation failed before a lease was acquired. + armProgressTimer(request_epoch_); + auto self = shared_from_this(); + asio::async_write( + *socket_, asio::buffer(response_bytes_), + [self, close_after, continuation = std::move(continuation)]( + const std::error_code& error, size_t bytes) mutable { + self->runHandler([&] { + if (self->forced_close_) { + self->finishClosed(); + return; + } + if (error || bytes != kHighPerformanceTcpResponseSize) { + self->finishClosed(); + return; + } + if (close_after) { + self->finishClosed(); + return; + } + self->armProgressTimer(self->request_epoch_); + if (continuation) { + continuation(); + } else { + self->readHeader(); + } + }); + }); + } + + void armProgressTimer(uint64_t epoch) { + const uint64_t generation = ++timer_generation_; + timer_.expires_after( + std::chrono::milliseconds(config_.progress_timeout_ms)); + auto self = shared_from_this(); + timer_.async_wait( + [self, epoch, generation](const std::error_code& error) { + if (error == asio::error::operation_aborted) return; + if (error || epoch != self->request_epoch_ || + generation != self->timer_generation_) { + return; + } + self->forced_close_ = true; + std::error_code ignored; + self->socket_->cancel(ignored); + self->socket_->shutdown(asio::ip::tcp::socket::shutdown_both, + ignored); + self->socket_->close(ignored); + // Finalization waits for the canceled body/response callback. + }); + } + + void cancelTimer() { + ++timer_generation_; + std::error_code ignored; + timer_.cancel(ignored); + } + + void finishClosed() { + if (closed_) return; + closed_ = true; + cancelTimer(); + std::error_code ignored; + socket_->cancel(ignored); + socket_->shutdown(asio::ip::tcp::socket::shutdown_both, ignored); + socket_->close(ignored); + lease_.reset(); + parent_->onSessionClosed(worker_id_, shared_from_this()); + } + + HighPerformanceTcpServer* parent_; + size_t worker_id_; + std::shared_ptr socket_; + Config config_; + HighPerformanceTcpBufferRegistry* registry_; + asio::steady_timer timer_; + + std::array request_bytes_{}; + std::array response_bytes_{}; + std::array discard_bytes_{}; + HighPerformanceTcpRequestFrame request_; + HighPerformanceTcpBufferRegistry::Lease lease_; + size_t header_offset_{0}; + uint64_t body_offset_{0}; + uint64_t request_epoch_{0}; + uint64_t timer_generation_{0}; + bool forced_close_{false}; + bool closed_{false}; +}; + +HighPerformanceTcpServer::HighPerformanceTcpServer( + Config config, HighPerformanceTcpBufferRegistry* registry, + HighPerformanceTcpWorkers* workers) + : config_(std::move(config)), registry_(registry), workers_(workers) { + if (workers_ != nullptr) sessions_.resize(workers_->workerCount()); +} + +HighPerformanceTcpServer::~HighPerformanceTcpServer() { + (void)stop(); + DCHECK(workers_ == nullptr || + active_sessions_.load(std::memory_order_acquire) == 0) + << "HP TCP server destroyed with active sessions"; +} + +bool HighPerformanceTcpServer::reserveConnection() { + size_t current = active_sessions_.load(std::memory_order_acquire); + while (current < config_.max_connections) { + if (active_sessions_.compare_exchange_weak(current, current + 1, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return true; + } + } + return false; +} + +Status HighPerformanceTcpServer::start(uint16_t* bound_port) { + if (bound_port == nullptr || registry_ == nullptr || workers_ == nullptr || + workers_->workerCount() == 0 || config_.chunk_size == 0 || + config_.max_transfer_bytes == 0 || config_.progress_timeout_ms == 0 || + config_.max_connections == 0) { + return Status::InvalidArgument( + "invalid HP TCP server configuration" LOC_MARK); + } + if (started_.exchange(true, std::memory_order_acq_rel)) { + return Status::InvalidArgument( + "HP TCP server already started" LOC_MARK); + } + + try { + asio::ip::tcp::endpoint endpoint; + if (config_.bind_address.empty()) { + endpoint = + asio::ip::tcp::endpoint(asio::ip::tcp::v4(), config_.port); + } else { + endpoint = asio::ip::tcp::endpoint( + asio::ip::make_address(config_.bind_address), config_.port); + } + acceptor_ = std::make_unique(accept_io_); + acceptor_->open(endpoint.protocol()); + acceptor_->set_option(asio::socket_base::reuse_address(true)); + acceptor_->bind(endpoint); + acceptor_->listen(asio::socket_base::max_listen_connections); + *bound_port = acceptor_->local_endpoint().port(); + accept_guard_.emplace(asio::make_work_guard(accept_io_)); + stopping_.store(false, std::memory_order_release); + Status accept_status = startAccept(); + if (!accept_status.ok()) { + std::error_code ignored; + acceptor_->close(ignored); + accept_guard_->reset(); + started_.store(false, std::memory_order_release); + acceptor_.reset(); + accept_guard_.reset(); + return accept_status; + } + accept_thread_ = std::thread([this] { + try { + accept_io_.run(); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP accept loop failed: " << error.what(); + stopping_.store(true, std::memory_order_release); + } catch (...) { + LOG(ERROR) << "HP TCP accept loop failed"; + stopping_.store(true, std::memory_order_release); + } + }); + return Status::OK(); + } catch (const std::exception& error) { + started_.store(false, std::memory_order_release); + acceptor_.reset(); + accept_guard_.reset(); + return Status::InternalError( + std::string("HP TCP listener start failed: ") + error.what() + + LOC_MARK); + } +} + +Status HighPerformanceTcpServer::startAccept() { + if (stopping_.load(std::memory_order_acquire) || !acceptor_ || + !acceptor_->is_open()) { + return Status::OK(); + } + try { + const size_t worker_id = + next_worker_.fetch_add(1, std::memory_order_relaxed) % + workers_->workerCount(); + auto socket = std::make_shared( + workers_->ioContext(worker_id)); + acceptor_->async_accept(*socket, [this, worker_id, socket]( + const std::error_code& error) { + if (!error && !stopping_.load(std::memory_order_acquire)) { + if (reserveConnection()) { + try { + asio::post(workers_->ioContext(worker_id), + [this, worker_id, socket] { + installAcceptedSocket(worker_id, socket); + }); + } catch (const std::exception& post_error) { + LOG(ERROR) << "HP TCP accepted-socket dispatch failed: " + << post_error.what(); + if (active_sessions_.fetch_sub( + 1, std::memory_order_acq_rel) == 1) { + sessions_wait_cv_.notify_all(); + } + std::error_code ignored; + socket->close(ignored); + } catch (...) { + LOG(ERROR) << "HP TCP accepted-socket dispatch failed"; + if (active_sessions_.fetch_sub( + 1, std::memory_order_acq_rel) == 1) { + sessions_wait_cv_.notify_all(); + } + std::error_code ignored; + socket->close(ignored); + } + } else { + std::error_code ignored; + socket->close(ignored); + } + } + if (!stopping_.load(std::memory_order_acquire)) { + Status next = startAccept(); + if (!next.ok()) { + LOG(ERROR) + << "HP TCP accept re-arm failed: " << next.ToString(); + stopping_.store(true, std::memory_order_release); + std::error_code ignored; + if (acceptor_) acceptor_->close(ignored); + if (accept_guard_) accept_guard_->reset(); + } + } + }); + return Status::OK(); + } catch (const std::exception& error) { + return Status::InternalError( + std::string("HP TCP accept initiation failed: ") + error.what() + + LOC_MARK); + } catch (...) { + return Status::InternalError( + "HP TCP accept initiation failed" LOC_MARK); + } +} + +void HighPerformanceTcpServer::installAcceptedSocket( + size_t worker_id, std::shared_ptr socket) { + if (stopping_.load(std::memory_order_acquire)) { + std::error_code ignored; + socket->close(ignored); + if (active_sessions_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + sessions_wait_cv_.notify_all(); + } + return; + } + std::shared_ptr session; + try { + session = std::make_shared(this, worker_id, socket, config_, + registry_); + sessions_[worker_id].insert(session); + session->start(); + } catch (const std::exception& start_error) { + LOG(ERROR) << "HP TCP server session start failed: " + << start_error.what(); + if (session) sessions_[worker_id].erase(session); + std::error_code ignored; + socket->cancel(ignored); + socket->close(ignored); + if (active_sessions_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + sessions_wait_cv_.notify_all(); + } + } catch (...) { + LOG(ERROR) << "HP TCP server session start failed"; + if (session) sessions_[worker_id].erase(session); + std::error_code ignored; + socket->cancel(ignored); + socket->close(ignored); + if (active_sessions_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + sessions_wait_cv_.notify_all(); + } + } +} + +void HighPerformanceTcpServer::onSessionClosed( + size_t worker_id, const std::shared_ptr& session) { + if (worker_id < sessions_.size()) sessions_[worker_id].erase(session); + if (active_sessions_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + std::lock_guard lock(sessions_wait_mutex_); + sessions_wait_cv_.notify_all(); + } +} + +Status HighPerformanceTcpServer::stopAccepting() { + if (!started_.load(std::memory_order_acquire)) return Status::OK(); + if (stopping_.exchange(true, std::memory_order_acq_rel)) { + if (accept_thread_.joinable()) accept_thread_.join(); + return Status::OK(); + } + + if (accept_thread_.joinable()) { + auto done = std::make_shared>(); + auto future = done->get_future(); + try { + asio::post(accept_io_, [this, done] { + std::error_code ignored; + if (acceptor_) { + acceptor_->cancel(ignored); + acceptor_->close(ignored); + } + if (accept_guard_) accept_guard_->reset(); + done->set_value(); + }); + future.wait(); + } catch (...) { + std::error_code ignored; + if (acceptor_) acceptor_->close(ignored); + if (accept_guard_) accept_guard_->reset(); + } + accept_thread_.join(); + } + return Status::OK(); +} + +void HighPerformanceTcpServer::cancelWorkerSessions(size_t worker_id) { + if (worker_id >= sessions_.size()) return; + // Copy because Session::finishClosed erases from this set later. + std::vector> copy(sessions_[worker_id].begin(), + sessions_[worker_id].end()); + for (const auto& session : copy) session->cancel(); +} + +Status HighPerformanceTcpServer::cancelAll() { + if (workers_ == nullptr || active_sessions_.load() == 0) { + return Status::OK(); + } + if (workers_->onWorkerThread()) { + return Status::InvalidArgument( + "HP TCP server cancelAll cannot block a worker" LOC_MARK); + } + if (!workers_->controlContextAvailable()) { + return active_sessions_.load() == 0 + ? Status::OK() + : Status::InternalError( + "HP TCP worker contexts unavailable with live " + "sessions" LOC_MARK); + } + + try { + for (size_t i = 0; i < workers_->workerCount(); ++i) { + asio::post(workers_->ioContext(i), + [this, i] { cancelWorkerSessions(i); }); + } + } catch (const std::exception& error) { + return Status::InternalError( + std::string("HP TCP server cancellation post failed: ") + + error.what() + LOC_MARK); + } + CHECK_STATUS(workers_->barrier()); + + std::unique_lock lock(sessions_wait_mutex_); + sessions_wait_cv_.wait(lock, [&] { + return active_sessions_.load(std::memory_order_acquire) == 0; + }); + return Status::OK(); +} + +Status HighPerformanceTcpServer::stop() { + Status first = stopAccepting(); + Status canceled = cancelAll(); + acceptor_.reset(); + accept_guard_.reset(); + started_.store(false, std::memory_order_release); + if (!first.ok()) return first; + return canceled; +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_task.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_task.cpp new file mode 100644 index 0000000000..c40bccb71e --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_task.cpp @@ -0,0 +1,64 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_task.h" + +#include + +namespace mooncake::tent { + +bool HighPerformanceTcpTaskState::completeOnce( + TransferStatusEnum terminal, size_t bytes, + std::optional remote_status) noexcept { + if (terminal == INITIAL || terminal == PENDING || terminal == INVALID) { + terminal = FAILED; + bytes = 0; + } + + bool expected = false; + if (!completion_claimed_.compare_exchange_strong( + expected, true, std::memory_order_acq_rel, + std::memory_order_acquire)) { + return false; + } + + // This function is called only after the final socket callback that could + // touch the local user buffer has retired (or before the operation reached + // a socket). Retire memory ownership and budget before publishing terminal + // status so teardown observing terminal cannot race unregister/free. + bytes_.store(bytes, std::memory_order_relaxed); + if (remote_status.has_value() && + *remote_status != HighPerformanceTcpStatus::kOk) { + remote_status_.store(*remote_status, std::memory_order_relaxed); + } + local_lease_.reset(); + + if (reservation_active_.exchange(false, std::memory_order_acq_rel)) { + if (admission_ != nullptr) admission_->release(1, reserved_bytes_); + } + + status_.store(terminal, std::memory_order_release); + + try { + if (notify_progress_) notify_progress_(progress_batch_id_); + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP progress notification threw: " << error.what(); + } catch (...) { + LOG(ERROR) << "HP TCP progress notification threw"; + } + return true; +} + +TransferStatus HighPerformanceTcpTaskState::snapshot() const noexcept { + TransferStatus result; + result.s = status_.load(std::memory_order_acquire); + result.transferred_bytes = bytes_.load(std::memory_order_acquire); + return result; +} + +std::optional +HighPerformanceTcpTaskState::remoteStatus() const noexcept { + const auto status = remote_status_.load(std::memory_order_acquire); + if (status == HighPerformanceTcpStatus::kOk) return std::nullopt; + return status; +} + +} // namespace mooncake::tent 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 new file mode 100644 index 0000000000..ebfb2c3a6a --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp @@ -0,0 +1,802 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_transport.h" + +#include +#include +#include +#include +#include +#include + +#include + +#include "tent/common/config.h" +#include "tent/runtime/platform.h" +#include "tent/runtime/slab.h" + +namespace mooncake::tent { +namespace { + +constexpr uint64_t kIoProgressStepBytes = 1ULL << 20; + +std::string HostFromRpc(const std::string& address) { + if (address.empty()) return {}; + if (address.front() == '[') { + const auto close = address.find(']'); + return close == std::string::npos ? std::string{} + : address.substr(1, close - 1); + } + const auto colon = address.rfind(':'); + if (colon == std::string::npos) return address; + // An unbracketed address with multiple colons is IPv6 and cannot be + // safely split as host:port. RPC addresses produced by TENT are bracketed, + // so reject rather than guessing. + if (address.find(':') != colon) return {}; + return address.substr(0, colon); +} + +bool HasTransport(const BufferDesc& buffer, TransportType type) { + return std::find(buffer.transports.begin(), buffer.transports.end(), + type) != buffer.transports.end(); +} + +bool RemotePermissionAllows(const HighPerformanceTcpBufferAttr& attr, + Request::OpCode opcode) { + if (opcode == Request::READ) { + return attr.permission == "global_read_only" || + attr.permission == "global_read_write"; + } + return attr.permission == "global_read_write"; +} + +Status FirstError(Status first, const Status& next) { + if (first.ok() && !next.ok()) return next; + return first; +} + +Status NeedsRefresh(const std::string& message) { + return Status::NeedsRefreshCache(message + LOC_MARK); +} + +Status CheckRuntimeHealth( + const HighPerformanceTcpWorkers* workers, + const HighPerformanceTcpAdmissionController* admission) { + if (workers != nullptr && workers->hasFailedWorker()) { + return Status::InternalError( + "HP TCP worker runtime has failed" LOC_MARK); + } + if (admission != nullptr && admission->failed()) { + return Status::InternalError( + "HP TCP admission accounting has failed" LOC_MARK); + } + return Status::OK(); +} + +Status RemoteWireStatus(HighPerformanceTcpStatus status) { + switch (status) { + case HighPerformanceTcpStatus::kStaleRegistration: + return Status::NeedsRefreshCache( + "remote HP TCP registration is stale" LOC_MARK); + case HighPerformanceTcpStatus::kPermissionDenied: + return Status::AddressNotRegistered( + "remote HP TCP permission denied" LOC_MARK); + case HighPerformanceTcpStatus::kRangeRejected: + return Status::AddressNotRegistered( + "remote HP TCP range rejected" LOC_MARK); + case HighPerformanceTcpStatus::kShuttingDown: + return Status::TooManyRequests( + "remote HP TCP transport is shutting down" LOC_MARK); + case HighPerformanceTcpStatus::kOk: + return Status::OK(); + default: + return Status::InvalidEntry( + "remote HP TCP protocol status " + + std::to_string(static_cast(status)) + LOC_MARK); + } +} + +} // namespace + +struct HighPerformanceTcpTransport::TaskPlan { + std::shared_ptr task; + HighPerformanceTcpWorkers::Command command; +}; + +HighPerformanceTcpTransport::HighPerformanceTcpTransport() + : HighPerformanceTcpTransport(HighPerformanceTcpParams{}) {} + +HighPerformanceTcpTransport::HighPerformanceTcpTransport( + HighPerformanceTcpParams params) + : params_(std::move(params)) { + caps.dram_to_dram = true; +} + +HighPerformanceTcpTransport::~HighPerformanceTcpTransport() { + const Status status = uninstall(); + if (!status.ok()) { + LOG(ERROR) << "HP TCP destructor uninstall failed: " + << status.ToString(); + } +} + +Status HighPerformanceTcpTransport::validateParams() const { + if (params_.worker_count == 0 || params_.connections_per_peer == 0 || + params_.max_outstanding_tasks == 0 || + params_.max_outstanding_bytes == 0 || params_.max_transfer_bytes == 0 || + params_.connect_timeout_ms == 0 || params_.progress_timeout_ms == 0) { + return Status::InvalidArgument( + "invalid high-performance TCP limits" LOC_MARK); + } + return Status::OK(); +} + +std::string HighPerformanceTcpTransport::makeIncarnation() const { + std::random_device device; + std::mt19937_64 random(device()); + std::ostringstream out; + for (int i = 0; i < 2; ++i) { + out << std::hex << std::setw(16) << std::setfill('0') << random(); + } + return out.str(); +} + +Status HighPerformanceTcpTransport::rollbackPublishedEndpoint( + const std::optional& previous_attr) { + if (!metadata_) return Status::OK(); + Status first = metadata_->segmentManager().updateLocal( + [&](SegmentDesc& desc) -> Status { + if (desc.type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "local segment is not memory while rolling back HP " + "TCP" LOC_MARK); + } + auto& attrs = + std::get(desc.detail).transport_attrs; + if (previous_attr.has_value()) { + attrs[static_cast(TransportType::HP_TCP)] = *previous_attr; + } else { + attrs.erase(static_cast(TransportType::HP_TCP)); + } + return Status::OK(); + }); + if (first.ok()) { + first = metadata_->segmentManager().synchronizeLocal(); + } + return first; +} + +Status HighPerformanceTcpTransport::install( + std::string&, std::shared_ptr metadata, + std::shared_ptr, std::shared_ptr) { + std::lock_guard lifecycle_lock(lifecycle_mutex_); + if (installed_.load(std::memory_order_acquire) || workers_ || server_ || + client_ || admission_) { + return Status::InvalidArgument( + "HP TCP transport is already installed or was not fully torn " + "down" LOC_MARK); + } + CHECK_STATUS(validateParams()); + CHECK_STATUS(registry_.reopen()); + if (!metadata) { + return Status::InvalidArgument("HP TCP metadata is null" LOC_MARK); + } + + metadata_ = std::move(metadata); + stopping_.store(false, std::memory_order_release); + + std::optional previous_attr; + { + const SegmentDescRef local = metadata_->segmentManager().getLocal(); + if (!local || local->type != SegmentType::Memory) { + metadata_.reset(); + return Status::InvalidMetadataType( + "HP TCP requires a local memory segment" LOC_MARK); + } + const auto& attrs = local->getMemory().transport_attrs; + const auto it = attrs.find(static_cast(TransportType::HP_TCP)); + if (it != attrs.end()) previous_attr = it->second; + } + + admission_ = std::make_unique( + params_.max_outstanding_tasks, params_.max_outstanding_bytes); + workers_ = std::make_unique( + HighPerformanceTcpWorkers::Config{params_.worker_count}); + + Status status = workers_->start(); + if (!status.ok()) { + admission_.reset(); + workers_.reset(); + metadata_.reset(); + return status; + } + + client_ = std::make_unique( + HighPerformanceTcpClient::Config{ + params_.max_transfer_bytes, + static_cast(std::min(kIoProgressStepBytes, + params_.max_transfer_bytes)), + params_.connect_timeout_ms, params_.progress_timeout_ms, + params_.connections_per_peer}, + workers_.get()); + + const uint64_t max_connections_u64 = std::max( + params_.connections_per_peer, params_.max_outstanding_tasks); + const size_t max_connections = static_cast(std::min( + max_connections_u64, std::numeric_limits::max())); + server_ = std::make_unique( + HighPerformanceTcpServer::Config{ + params_.bind_address, params_.port, params_.max_transfer_bytes, + static_cast(std::min(kIoProgressStepBytes, + params_.max_transfer_bytes)), + params_.progress_timeout_ms, max_connections}, + ®istry_, workers_.get()); + + uint16_t bound_port = 0; + status = server_->start(&bound_port); + if (!status.ok()) { + (void)server_->stop(); + (void)client_->cancelAll(CANCELED); + (void)workers_->stop(); + server_.reset(); + client_.reset(); + workers_.reset(); + admission_.reset(); + metadata_.reset(); + return status; + } + + const SegmentDescRef local = metadata_->segmentManager().getLocal(); + const std::string host = params_.advertise_address.empty() + ? HostFromRpc(local->rpc_server_addr) + : params_.advertise_address; + if (host.empty()) { + status = Status::InvalidArgument( + "unable to derive HP TCP advertise address" LOC_MARK); + } else { + const std::string incarnation = makeIncarnation(); + std::string encoded; + status = EncodeHighPerformanceTcpEndpointAttr( + {incarnation, host, bound_port, params_.max_transfer_bytes}, + &encoded); + if (status.ok()) { + status = metadata_->segmentManager().updateLocal( + [&](SegmentDesc& desc) -> Status { + if (desc.type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "HP TCP local segment changed type" LOC_MARK); + } + std::get(desc.detail) + .transport_attrs[static_cast( + TransportType::HP_TCP)] = encoded; + return Status::OK(); + }); + } + if (status.ok()) { + status = metadata_->segmentManager().synchronizeLocal(); + } + } + + if (!status.ok()) { + const Status rollback = rollbackPublishedEndpoint(previous_attr); + if (!rollback.ok()) { + LOG(ERROR) << "HP TCP install metadata rollback failed: " + << rollback.ToString(); + } + (void)server_->stopAccepting(); + (void)client_->cancelAll(CANCELED); + (void)server_->cancelAll(); + (void)server_->stop(); + (void)workers_->stop(); + server_.reset(); + client_.reset(); + workers_.reset(); + admission_.reset(); + metadata_.reset(); + return status; + } + + metadata_->setNotifyCallback([this](const Notification& notification) { + RWSpinlock::WriteGuard guard(notify_lock_); + notifications_.push_back(notification); + return 0; + }); + installed_.store(true, std::memory_order_release); + return Status::OK(); +} + +Status HighPerformanceTcpTransport::stopRuntime() { + // Lifecycle invariant: this is the sole normal teardown order. Client and + // server async callbacks hold raw parent pointers, so they must quiesce + // before their owners are destroyed and before worker contexts disappear. + Status first = Status::OK(); + if (admission_) admission_->close(); + if (server_) first = FirstError(std::move(first), server_->stopAccepting()); + registry_.close(); + + if (workers_) first = FirstError(std::move(first), workers_->barrier()); + if (client_) + first = FirstError(std::move(first), client_->cancelAll(CANCELED)); + if (server_) first = FirstError(std::move(first), server_->cancelAll()); + + if (admission_) { + if (workers_ && workers_->hasFailedWorker() && + (admission_->outstandingTasks() != 0 || + admission_->outstandingBytes() != 0)) { + first = FirstError( + std::move(first), + Status::InternalError("HP TCP admission cannot drain after " + "worker failure" LOC_MARK)); + } else { + first = FirstError(std::move(first), admission_->waitForZero()); + } + } + if (server_) first = FirstError(std::move(first), server_->stop()); + if (first.ok()) { + DCHECK(!admission_ || (admission_->outstandingTasks() == 0 && + admission_->outstandingBytes() == 0)); + DCHECK(!client_ || client_->activeOperations() == 0); + DCHECK(!server_ || server_->activeSessionsForTest() == 0); + } + if (workers_) first = FirstError(std::move(first), workers_->stop()); + return first; +} + +Status HighPerformanceTcpTransport::quiesce() { + std::lock_guard lifecycle_lock(lifecycle_mutex_); + if (!workers_ && !server_ && !client_ && !admission_) return Status::OK(); + if (stopping_.exchange(true, std::memory_order_acq_rel)) { + // A previous quiesce under the same lifecycle lock completed all + // cancellation/join work before returning. + return Status::OK(); + } + return stopRuntime(); +} + +Status HighPerformanceTcpTransport::uninstall() { + std::lock_guard lifecycle_lock(lifecycle_mutex_); + if (!installed_.load(std::memory_order_acquire) && !workers_ && !server_ && + !client_ && !admission_) { + return Status::OK(); + } + + Status first = Status::OK(); + if (!stopping_.exchange(true, std::memory_order_acq_rel)) { + first = stopRuntime(); + } + + if (metadata_) { + metadata_->setNotifyCallback(nullptr); + Status removed = metadata_->segmentManager().updateLocal( + [](SegmentDesc& desc) -> Status { + if (desc.type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "HP TCP local segment is not memory during " + "uninstall" LOC_MARK); + } + std::get(desc.detail) + .transport_attrs.erase( + static_cast(TransportType::HP_TCP)); + return Status::OK(); + }); + first = FirstError(std::move(first), removed); + if (removed.ok()) { + first = FirstError(std::move(first), + metadata_->segmentManager().synchronizeLocal()); + } + } + + server_.reset(); + client_.reset(); + workers_.reset(); + admission_.reset(); + metadata_.reset(); + installed_.store(false, std::memory_order_release); + return first; +} + +Status HighPerformanceTcpTransport::allocateSubBatch(SubBatchRef& batch, + size_t max_size) { + if (batch != nullptr) { + return Status::InvalidArgument( + "HP TCP SubBatch output must be null" LOC_MARK); + } + auto* result = Slab::Get().allocate(); + if (result == nullptr) { + return Status::InternalError( + "unable to allocate HP TCP SubBatch" LOC_MARK); + } + result->max_size = max_size; + result->tasks.clear(); + try { + result->tasks.reserve(max_size); + } catch (...) { + Slab::Get().deallocate(result); + return Status::InternalError( + "unable to reserve HP TCP SubBatch storage" LOC_MARK); + } + batch = result; + return Status::OK(); +} + +Status HighPerformanceTcpTransport::freeSubBatch(SubBatchRef& batch) { + auto* hp_batch = dynamic_cast(batch); + if (hp_batch == nullptr) { + return Status::InvalidArgument("invalid HP TCP SubBatch" LOC_MARK); + } + for (const auto& task : hp_batch->tasks) { + const TransferStatusEnum state = task->snapshot().s; + if (state == INITIAL || state == PENDING) { + return Status::InvalidArgument( + "cannot free an HP TCP SubBatch with pending tasks" LOC_MARK); + } + } + hp_batch->tasks.clear(); + hp_batch->max_size = 0; + Slab::Get().deallocate(hp_batch); + batch = nullptr; + return Status::OK(); +} + +Status HighPerformanceTcpTransport::planTask(const Request& request, + HighPerformanceTcpSubBatch* batch, + TaskPlan* plan) { + if (batch == nullptr || plan == nullptr) { + return Status::InvalidArgument( + "invalid HP TCP task plan output" LOC_MARK); + } + if (request.source == nullptr || request.length == 0 || + request.length > params_.max_transfer_bytes || + (request.opcode != Request::READ && request.opcode != Request::WRITE)) { + return Status::InvalidArgument("invalid HP TCP request" LOC_MARK); + } + + HighPerformanceTcpBufferRegistry::Lease local_lease; + const uint64_t local_addr = + static_cast(reinterpret_cast(request.source)); + CHECK_STATUS( + registry_.acquireLocalLease(local_addr, request.length, &local_lease)); + + HighPerformanceTcpEndpointAttr endpoint_attr; + HighPerformanceTcpBufferAttr buffer_attr; + SegmentDescRef pin; + Status resolved = metadata_->segmentManager().withCachedSegment( + request.target_id, pin, [&](SegmentDesc* segment) -> Status { + if (segment == nullptr || segment->type != SegmentType::Memory) { + return NeedsRefresh("HP TCP target is not a memory segment"); + } + BufferDesc* buffer = + segment->findBuffer(request.target_offset, request.length); + if (buffer == nullptr || + !HasTransport(*buffer, TransportType::HP_TCP)) { + return NeedsRefresh("HP TCP target buffer is not advertised"); + } + const auto endpoint_it = segment->getMemory().transport_attrs.find( + static_cast(TransportType::HP_TCP)); + if (endpoint_it == segment->getMemory().transport_attrs.end()) { + return NeedsRefresh("HP TCP endpoint metadata is missing"); + } + Status decoded = DecodeHighPerformanceTcpEndpointAttr( + endpoint_it->second, &endpoint_attr); + if (!decoded.ok()) { + return NeedsRefresh("HP TCP endpoint metadata is incompatible"); + } + const auto registration_it = + buffer->transport_attrs.find(TransportType::HP_TCP); + if (registration_it == buffer->transport_attrs.end()) { + return NeedsRefresh("HP TCP buffer registration is missing"); + } + decoded = DecodeHighPerformanceTcpBufferAttr( + registration_it->second, &buffer_attr); + if (!decoded.ok()) { + return NeedsRefresh("HP TCP buffer metadata is incompatible"); + } + return Status::OK(); + }); + if (!resolved.ok()) return resolved; + + if (request.length > endpoint_attr.max_transfer_bytes) { + return Status::InvalidArgument( + "HP TCP request exceeds remote endpoint capability" LOC_MARK); + } + if (!RemotePermissionAllows(buffer_attr, request.opcode)) { + return Status::AddressNotRegistered( + "HP TCP remote permission does not allow requested " + "operation" LOC_MARK); + } + + uint64_t request_id = + next_request_id_.fetch_add(1, std::memory_order_relaxed); + if (request_id == 0) { + return Status::InternalError( + "HP TCP request id space exhausted" LOC_MARK); + } + const uint32_t lane_id = static_cast( + request_id % static_cast(params_.connections_per_peer)); + const size_t owner_worker = + workers_->affinityOwner(request.target_id, lane_id); + + auto task = std::make_shared( + request.length, batch->progress_batch_id, batch->notify_progress, + std::move(local_lease)); + task->setDispatchIdentity(owner_worker, request_id); + + HighPerformanceTcpClient::Operation operation; + operation.peer_id = request.target_id; + operation.incarnation = endpoint_attr.incarnation; + operation.host = endpoint_attr.host; + operation.port = endpoint_attr.port; + operation.lane_id = lane_id; + operation.registration_id = buffer_attr.registration_id; + operation.remote_addr = request.target_offset; + operation.local_addr = request.source; + operation.length = request.length; + operation.opcode = request.opcode == Request::READ + ? HighPerformanceTcpOpcode::kRead + : HighPerformanceTcpOpcode::kWrite; + operation.request_id = request_id; + operation.complete = + [task](TransferStatusEnum terminal, size_t bytes, + std::optional remote_status) { + (void)task->completeOnce(terminal, bytes, remote_status); + }; + + HighPerformanceTcpWorkers::Command command; + command.worker_id = owner_worker; + command.run = [this, task, + operation = std::move(operation)](size_t worker_id) mutable { + if (task->cancelRequested() || + stopping_.load(std::memory_order_acquire)) { + (void)task->completeOnce(CANCELED, 0); + return; + } + client_->enqueueOnOwner(worker_id, std::move(operation)); + }; + command.cancel = [task] { (void)task->completeOnce(CANCELED, 0); }; + + plan->task = std::move(task); + plan->command = std::move(command); + return Status::OK(); +} + +Status HighPerformanceTcpTransport::submitTransferTasks( + SubBatchRef batch, const std::vector& requests) { + auto* hp_batch = dynamic_cast(batch); + if (hp_batch == nullptr) { + return Status::InvalidArgument("invalid HP TCP SubBatch" LOC_MARK); + } + if (!installed_.load(std::memory_order_acquire) || + stopping_.load(std::memory_order_acquire) || workers_ == nullptr || + client_ == nullptr || admission_ == nullptr) { + return Status::InvalidArgument( + "HP TCP transport is unavailable" LOC_MARK); + } + CHECK_STATUS(CheckRuntimeHealth(workers_.get(), admission_.get())); + if (requests.empty()) return Status::OK(); + if (requests.size() > hp_batch->max_size - hp_batch->tasks.size()) { + return Status::TooManyRequests( + "HP TCP SubBatch capacity exceeded" LOC_MARK); + } + + uint64_t total_bytes = 0; + for (const Request& request : requests) { + if (request.length > + std::numeric_limits::max() - total_bytes || + request.length == 0) { + return Status::InvalidArgument( + "HP TCP batch byte count overflow" LOC_MARK); + } + total_bytes += request.length; + } + + std::vector plans; + std::vector commands; + try { + plans.resize(requests.size()); + commands.reserve(requests.size()); + } catch (...) { + return Status::InternalError( + "unable to allocate HP TCP task planning storage" LOC_MARK); + } + + for (size_t i = 0; i < requests.size(); ++i) { + CHECK_STATUS(planTask(requests[i], hp_batch, &plans[i])); + } + for (auto& plan : plans) { + commands.push_back(std::move(plan.command)); + } + + // SubBatch capacity was reserved at allocateSubBatch(). The callback below + // therefore performs only noexcept shared_ptr moves and reservation flag + // stores while dispatch ownership is being committed. + const size_t old_size = hp_batch->tasks.size(); + Status committed = workers_->tryCommitBatch( + commands, admission_.get(), requests.size(), total_bytes, [&] { + for (auto& plan : plans) { + plan.task->activateReservation(admission_.get()); + hp_batch->tasks.push_back(std::move(plan.task)); + } + }); + if (!committed.ok()) { + // The worker transaction promises not to invoke on_commit on failure. + // Keep a defensive assertion in debug logs without mutating the batch. + if (hp_batch->tasks.size() != old_size) { + LOG(FATAL) << "HP TCP atomic admission violated SubBatch rollback"; + } + return committed; + } + return Status::OK(); +} + +Status HighPerformanceTcpTransport::getTransferStatus(SubBatchRef batch, + int task_id, + TransferStatus& status) { + auto* hp_batch = dynamic_cast(batch); + if (hp_batch == nullptr || task_id < 0 || + static_cast(task_id) >= hp_batch->tasks.size()) { + return Status::InvalidArgument("invalid HP TCP task id" LOC_MARK); + } + const auto& task = hp_batch->tasks[static_cast(task_id)]; + status = task->snapshot(); + if (status.s == PENDING) { + return CheckRuntimeHealth(workers_.get(), admission_.get()); + } + const auto remote_status = task->remoteStatus(); + if (remote_status.has_value()) return RemoteWireStatus(*remote_status); + return Status::OK(); +} + +Status HighPerformanceTcpTransport::retryTransferTask(SubBatchRef batch, + int task_id, + const Request& request) { + auto* hp_batch = dynamic_cast(batch); + if (hp_batch == nullptr || task_id < 0 || + static_cast(task_id) >= hp_batch->tasks.size()) { + return Status::InvalidArgument("invalid HP TCP retry task id" LOC_MARK); + } + CHECK_STATUS(CheckRuntimeHealth(workers_.get(), admission_.get())); + if (hp_batch->tasks[static_cast(task_id)]->snapshot().s != FAILED) { + return Status::InvalidArgument( + "HP TCP retry requires a failed attempt" LOC_MARK); + } + + TaskPlan plan; + CHECK_STATUS(planTask(request, hp_batch, &plan)); + std::vector commands; + try { + commands.push_back(std::move(plan.command)); + } catch (...) { + return Status::InternalError( + "unable to allocate HP TCP retry command" LOC_MARK); + } + + Status committed = workers_->tryCommitBatch( + commands, admission_.get(), 1, request.length, [&] { + plan.task->activateReservation(admission_.get()); + hp_batch->tasks[static_cast(task_id)] = + std::move(plan.task); + }); + return committed; +} + +Status HighPerformanceTcpTransport::cancelTransferTask(SubBatchRef batch, + int task_id) { + auto* hp_batch = dynamic_cast(batch); + if (hp_batch == nullptr || task_id < 0 || + static_cast(task_id) >= hp_batch->tasks.size()) { + return Status::InvalidArgument("invalid HP TCP task id" LOC_MARK); + } + const auto& task = hp_batch->tasks[static_cast(task_id)]; + const TransferStatusEnum state = task->snapshot().s; + if (state != INITIAL && state != PENDING) return Status::OK(); + + task->requestCancel(); + if (client_ == nullptr || workers_ == nullptr) return Status::OK(); + const Status canceled = + client_->cancelRequest(task->ownerWorker(), task->requestId()); + // A request still in the worker dispatch queue has no client lane yet; + // its command observes cancelRequested() and settles it. Treat inability + // to find/post a lane cancellation during shutdown as best effort. + if (canceled.IsInternalError() && + stopping_.load(std::memory_order_acquire)) { + return Status::OK(); + } + return canceled; +} + +Status HighPerformanceTcpTransport::addMemoryBuffer( + BufferDesc& desc, const MemoryOptions& options) { + if (stopping_.load(std::memory_order_acquire)) { + return Status::TooManyRequests( + "HP TCP transport is shutting down" LOC_MARK); + } + const LocationParser location(desc.location); + if ((location.type() != "cpu" && location.type() != kWildcardLocation) || + Platform::getLoader().getMemoryType( + reinterpret_cast(desc.addr)) != MTYPE_CPU) { + return Status::InvalidArgument( + "HP TCP v1 supports CPU DRAM only" LOC_MARK); + } + + uint64_t registration_id = 0; + CHECK_STATUS( + registry_.add(desc.addr, desc.length, options.perm, ®istration_id)); + if (options.perm == kLocalReadWrite) return Status::OK(); + + std::string encoded; + Status status = EncodeHighPerformanceTcpBufferAttr( + {registration_id, HighPerformanceTcpPermissionName(options.perm)}, + &encoded); + if (!status.ok()) { + (void)registry_.remove(desc.addr, desc.length); + return status; + } + desc.transport_attrs[TransportType::HP_TCP] = std::move(encoded); + if (!HasTransport(desc, TransportType::HP_TCP)) { + desc.transports.push_back(TransportType::HP_TCP); + } + return Status::OK(); +} + +Status HighPerformanceTcpTransport::addMemoryBuffer( + std::vector& desc_list, const MemoryOptions& options) { + std::vector created; + try { + created.reserve(desc_list.size()); + } catch (...) { + return Status::InternalError( + "unable to allocate HP TCP registration rollback state" LOC_MARK); + } + + for (size_t i = 0; i < desc_list.size(); ++i) { + const bool tracked_before = + registry_.tracks(desc_list[i].addr, desc_list[i].length); + Status status = addMemoryBuffer(desc_list[i], options); + if (status.ok()) { + if (!tracked_before) created.push_back(i); + continue; + } + for (auto it = created.rbegin(); it != created.rend(); ++it) { + const Status rollback = removeMemoryBuffer(desc_list[*it]); + if (!rollback.ok()) { + LOG(ERROR) << "HP TCP registration rollback failed: " + << rollback.ToString(); + } + } + return status; + } + return Status::OK(); +} + +Status HighPerformanceTcpTransport::removeMemoryBuffer(BufferDesc& desc) { + if (!registry_.tracks(desc.addr, desc.length)) return Status::OK(); + Status status = registry_.remove(desc.addr, desc.length); + if (!status.ok()) return status; + desc.transport_attrs.erase(TransportType::HP_TCP); + desc.transports.erase( + std::remove(desc.transports.begin(), desc.transports.end(), + TransportType::HP_TCP), + desc.transports.end()); + return Status::OK(); +} + +Status HighPerformanceTcpTransport::sendNotification( + SegmentID target_id, const Notification& notification) { + if (!metadata_) { + return Status::InternalError("HP TCP metadata is unavailable" LOC_MARK); + } + return metadata_->segmentManager().withCachedSegment( + target_id, [&](SegmentDesc* segment) { + return ControlClient::notify(segment->rpc_server_addr, + notification); + }); +} + +Status HighPerformanceTcpTransport::receiveNotification( + std::vector& notifications) { + RWSpinlock::WriteGuard guard(notify_lock_); + notifications.clear(); + notifications.swap(notifications_); + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_workers.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_workers.cpp new file mode 100644 index 0000000000..bf2ddeed33 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_workers.cpp @@ -0,0 +1,308 @@ +// Copyright 2026 KVCache.AI +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +#include +#include +#include + +#include + +namespace mooncake::tent { + +Status HighPerformanceTcpAdmissionController::tryReserve(uint64_t tasks, + uint64_t bytes) { + if (tasks == 0 || bytes == 0) { + return Status::InvalidArgument( + "HP TCP admission reservation must be non-zero" LOC_MARK); + } + std::lock_guard lock(mutex_); + if (failed_.load(std::memory_order_acquire)) { + return Status::InternalError( + "HP TCP admission accounting is failed" LOC_MARK); + } + if (!accepting_) { + return Status::TooManyRequests("HP TCP admission is closed" LOC_MARK); + } + if (tasks > max_tasks_ || bytes > max_bytes_ || + tasks_ > max_tasks_ - tasks || bytes_ > max_bytes_ - bytes) { + return Status::TooManyRequests( + "HP TCP admission limit exceeded" LOC_MARK); + } + tasks_ += tasks; + bytes_ += bytes; + return Status::OK(); +} + +void HighPerformanceTcpAdmissionController::release(uint64_t tasks, + uint64_t bytes) { + bool notify_waiters = false; + { + std::lock_guard lock(mutex_); + if (tasks > tasks_ || bytes > bytes_) { + LOG(ERROR) << "HP TCP admission release underflow"; + // Do not manufacture a drained state. A double release is a + // lifecycle invariant violation; preserving the counters keeps + // shutdown from treating live work as retired. + failed_.store(true, std::memory_order_release); + notify_waiters = true; + } else { + tasks_ -= tasks; + bytes_ -= bytes; + notify_waiters = tasks_ == 0 && bytes_ == 0; + } + } + if (notify_waiters) zero_cv_.notify_all(); +} + +void HighPerformanceTcpAdmissionController::close() { + std::lock_guard lock(mutex_); + accepting_ = false; +} + +Status HighPerformanceTcpAdmissionController::waitForZero() { + std::unique_lock lock(mutex_); + zero_cv_.wait(lock, [&] { + return failed_.load(std::memory_order_acquire) || + (tasks_ == 0 && bytes_ == 0); + }); + if (failed_.load(std::memory_order_acquire)) { + return Status::InternalError( + "HP TCP admission accounting is failed" LOC_MARK); + } + return Status::OK(); +} + +bool HighPerformanceTcpAdmissionController::failed() const { + return failed_.load(std::memory_order_acquire); +} + +uint64_t HighPerformanceTcpAdmissionController::outstandingTasks() const { + std::lock_guard lock(mutex_); + return tasks_; +} + +uint64_t HighPerformanceTcpAdmissionController::outstandingBytes() const { + std::lock_guard lock(mutex_); + return bytes_; +} + +HighPerformanceTcpWorkers::HighPerformanceTcpWorkers() + : HighPerformanceTcpWorkers(Config{}) {} + +HighPerformanceTcpWorkers::HighPerformanceTcpWorkers(Config config) + : config_(config) {} + +HighPerformanceTcpWorkers::~HighPerformanceTcpWorkers() { (void)stop(); } + +Status HighPerformanceTcpWorkers::start() { + std::lock_guard lock(lifecycle_mutex_); + if (running()) return Status::OK(); + if (started_ || !workers_.empty() || config_.worker_count == 0) { + return Status::InvalidArgument("invalid HP TCP worker state" LOC_MARK); + } + + try { + started_ = true; + failed_.store(false, std::memory_order_release); + workers_.reserve(config_.worker_count); + for (size_t i = 0; i < config_.worker_count; ++i) { + workers_.push_back(std::make_unique()); + } + running_.store(true, std::memory_order_release); + for (auto& worker : workers_) { + worker->thread = std::thread([this, context = worker.get()] { + for (;;) { + try { + context->io.run(); + break; + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP io_context handler failed: " + << error.what(); + markWorkerFailed(); + } catch (...) { + LOG(ERROR) << "HP TCP io_context handler failed"; + markWorkerFailed(); + } + } + }); + } + } catch (const std::exception& error) { + running_.store(false, std::memory_order_release); + for (auto& worker : workers_) { + worker->guard.reset(); + worker->io.stop(); + if (worker->thread.joinable()) worker->thread.join(); + } + workers_.clear(); + return Status::InternalError( + std::string("failed to start HP TCP workers: ") + error.what() + + LOC_MARK); + } + return Status::OK(); +} + +void HighPerformanceTcpWorkers::markWorkerFailed() noexcept { + failed_.store(true, std::memory_order_release); + // New batches observe failed_ and are rejected. The owner loop deliberately + // keeps running so already-committed work and teardown cancellation retain + // their original affinity and can retire every task, session and lease. +} + +bool HighPerformanceTcpWorkers::onWorkerThread() const { + const auto current = std::this_thread::get_id(); + return std::any_of(workers_.begin(), workers_.end(), [&](const auto& w) { + return w->thread.joinable() && w->thread.get_id() == current; + }); +} + +Status HighPerformanceTcpWorkers::stop() { + std::lock_guard lock(lifecycle_mutex_); + if (workers_.empty()) return Status::OK(); + if (onWorkerThread()) { + return Status::InvalidArgument( + "HP TCP worker cannot synchronously stop itself" LOC_MARK); + } + const bool failed = hasFailedWorker(); + running_.store(false, std::memory_order_release); + for (auto& worker : workers_) worker->guard.reset(); + for (auto& worker : workers_) worker->io.stop(); + for (auto& worker : workers_) { + if (worker->thread.joinable()) worker->thread.join(); + } + // Keep stopped contexts alive until their client/server owners are + // destroyed. Those owners contain socket objects and cancellation handlers + // that refer to these contexts, even after no worker can execute them. + if (failed) { + return Status::InternalError( + "HP TCP workers stopped after an owner thread failure" LOC_MARK); + } + return Status::OK(); +} + +size_t HighPerformanceTcpWorkers::affinityOwner(uint64_t peer, + uint32_t lane) const { + size_t hash = std::hash{}(peer); + const auto mix = [&hash](size_t value) { + hash ^= value + static_cast(0x9e3779b97f4a7c15ULL) + + (hash << 6U) + (hash >> 2U); + }; + // HP TCP v1 has one endpoint. Keep its zero index in the hash so this + // interface simplification does not change the existing worker mapping. + mix(std::hash{}(0)); + mix(std::hash{}(lane)); + return config_.worker_count == 0 ? 0 : hash % config_.worker_count; +} + +Status HighPerformanceTcpWorkers::tryCommitBatch( + std::vector& commands, + HighPerformanceTcpAdmissionController* admission, uint64_t reserve_tasks, + uint64_t reserve_bytes, const std::function& on_commit) { + if (commands.empty()) return Status::OK(); + + std::lock_guard submit_lock(submit_mutex_); + if (hasFailedWorker()) { + return Status::InternalError( + "HP TCP worker owner thread has failed" LOC_MARK); + } + if (!running()) { + return Status::InternalError("HP TCP workers are not running" LOC_MARK); + } + + for (const auto& command : commands) { + if (!command.run || command.worker_id >= workers_.size()) { + return Status::InvalidArgument( + "invalid HP TCP worker command" LOC_MARK); + } + } + + bool reserved = false; + if (admission != nullptr) { + CHECK_STATUS(admission->tryReserve(reserve_tasks, reserve_bytes)); + reserved = true; + } + try { + if (on_commit) on_commit(); + } catch (...) { + if (reserved) admission->release(reserve_tasks, reserve_bytes); + return Status::InternalError("HP TCP ownership commit failed" LOC_MARK); + } + + for (size_t i = 0; i < commands.size(); ++i) { + Command command = std::move(commands[i]); + const size_t owner = command.worker_id; + const auto cancel = command.cancel; + try { + asio::post(workers_[owner]->io, + [this, command = std::move(command)]() mutable { + runCommand(std::move(command)); + }); + } catch (...) { + if (cancel) cancel(); + for (++i; i < commands.size(); ++i) { + auto& pending = commands[i]; + if (pending.cancel) pending.cancel(); + } + return Status::OK(); + } + } + return Status::OK(); +} + +void HighPerformanceTcpWorkers::runCommand(Command command) { + const size_t owner = command.worker_id; + try { + if (running()) { + command.run(owner); + } else if (command.cancel) { + command.cancel(); + } + } catch (const std::exception& error) { + LOG(ERROR) << "HP TCP worker command failed: " << error.what(); + if (command.cancel) command.cancel(); + } catch (...) { + LOG(ERROR) << "HP TCP worker command failed"; + if (command.cancel) command.cancel(); + } +} + +Status HighPerformanceTcpWorkers::barrier() { + if (onWorkerThread()) { + return Status::InvalidArgument( + "HP TCP barrier cannot run on a worker thread" LOC_MARK); + } + if (workers_.empty()) return Status::OK(); + if (!running()) { + return Status::InternalError( + "HP TCP barrier cannot run after workers stop" LOC_MARK); + } + + struct Latch { + std::mutex mutex; + std::condition_variable cv; + size_t remaining{0}; + }; + auto latch = std::make_shared(); + latch->remaining = workers_.size(); + try { + for (auto& worker : workers_) { + asio::post(worker->io, [latch] { + std::lock_guard lock(latch->mutex); + if (--latch->remaining == 0) latch->cv.notify_all(); + }); + } + } catch (...) { + return Status::InternalError("HP TCP barrier post failed" LOC_MARK); + } + std::unique_lock lock(latch->mutex); + latch->cv.wait(lock, [&] { return latch->remaining == 0; }); + return Status::OK(); +} + +asio::io_context& HighPerformanceTcpWorkers::ioContext(size_t worker_id) { + if (worker_id >= workers_.size()) { + throw std::out_of_range("HP TCP worker id out of range"); + } + return workers_[worker_id]->io; +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp index 9a861cc9a7..0afda3af8f 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp @@ -851,6 +851,7 @@ int RdmaContext::openDevice(const std::string& device_name, uint8_t port) { native_context_ = context.release(); lid_ = port_attr.lid; recordPortSpeed(port_attr); + queryEffectiveSpeed(); return 0; } @@ -877,12 +878,64 @@ int RdmaContext::refreshPortAttributes() { return -1; } recordPortSpeed(port_attr); + queryEffectiveSpeed(); return 0; } +int RdmaContext::queryPortState(ibv_port_state* state) const { + if (!state || !native_context_ || !params_) return -1; + ibv_port_attr port_attr; + if (verbs_.ibv_query_port_default(native_context_, params_->device.port, + &port_attr)) { + PLOG(WARNING) << "Failed to query state of port " + << static_cast(params_->device.port) << " on " + << device_name_; + return -1; + } + // No recordPortSpeed() here on purpose -- see the header. + *state = port_attr.state; + return 0; +} + +void RdmaContext::queryEffectiveSpeed() { + if (!native_context_ || !verbs_.ibv_query_port_speed) { + effective_speed_mbps_.store(0, std::memory_order_relaxed); + return; + } + uint64_t speed = 0; + int rc = verbs_.ibv_query_port_speed(native_context_, params_->device.port, + &speed); + if (rc != 0) { + // Keep the last known value: on a degraded LAG, dropping to the + // encoded rate would overstate the port until the next successful + // query. Log once per failure episode, not per query. + effective_speed_query_failures_.fetch_add(1, std::memory_order_relaxed); + if (!effective_speed_query_failing_.exchange( + true, std::memory_order_relaxed)) { + LOG(WARNING) << "ibv_query_port_speed failed on " << device_name_ + << " (rc " << rc << "), keeping " + << effective_speed_mbps_.load( + std::memory_order_relaxed) + << " Mb/s (" + << effective_speed_query_failures_.load( + std::memory_order_relaxed) + << " failures so far)"; + } + return; + } + if (effective_speed_query_failing_.exchange(false, + std::memory_order_relaxed)) { + LOG(INFO) << "ibv_query_port_speed recovered on " << device_name_; + } + // The verb reports 100 Mb/s units; store plain Mb/s. + effective_speed_mbps_.store(speed * 100, std::memory_order_relaxed); +} + double RdmaContext::linkSpeedGbps() const { - return ibLinkSpeedGbps(active_speed_.load(std::memory_order_relaxed), - active_width_.load(std::memory_order_relaxed)); + return ibPortSpeedGbps( + effective_speed_mbps_.load(std::memory_order_relaxed), + active_speed_.load(std::memory_order_relaxed), + active_width_.load(std::memory_order_relaxed)); } } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp index f9df5330b3..ac4501e617 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include #include #include @@ -177,9 +179,18 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, return -1; } - // Pre-post recv buffers for notification - notify_recv_buffers_.resize(kNotifyMaxPendingSends); - notify_recv_mrs_.resize(kNotifyMaxPendingSends); + // Allocate one contiguous recv buffer, logically split into + // kNotifyMaxPendingSends slots. One MR covers every slot so high + // endpoint counts do not explode ibv_reg_mr. + notify_recv_buffer_.resize(kNotifyBufferSize * kNotifyMaxPendingSends); + notify_recv_mr_ = context_->verbs_.ibv_reg_mr_default( + context_->nativePD(), notify_recv_buffer_.data(), + notify_recv_buffer_.size(), IBV_ACCESS_LOCAL_WRITE); + if (!notify_recv_mr_) { + PLOG(ERROR) << "Failed to register notification recv buffer"; + deconstruct(); + return -1; + } // Allocate one contiguous send buffer, logically split into // kNotifyMaxPendingSends slots to avoid DMA-vs-overwrite races. @@ -193,20 +204,6 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, return -1; } - for (size_t i = 0; i < kNotifyMaxPendingSends; ++i) { - notify_recv_buffers_[i].resize(kNotifyBufferSize); - - // Register memory for recv buffer - notify_recv_mrs_[i] = context_->verbs_.ibv_reg_mr_default( - context_->nativePD(), notify_recv_buffers_[i].data(), - kNotifyBufferSize, IBV_ACCESS_LOCAL_WRITE); - if (!notify_recv_mrs_[i]) { - PLOG(ERROR) << "Failed to register notification recv buffer"; - deconstruct(); - return -1; - } - } - return 0; } @@ -254,20 +251,15 @@ int RdmaEndPoint::deconstructUnlocked() { // A live QP may still reference the notification MRs. Keep both MRs // and backing buffers intact until QP destruction succeeds. if (!notify_qp_) { - for (auto& mr : notify_recv_mrs_) { - if (!mr) continue; - if (context_->verbs_.ibv_dereg_mr(mr)) { + if (notify_recv_mr_) { + if (context_->verbs_.ibv_dereg_mr(notify_recv_mr_)) { PLOG(ERROR) << "Failed to deregister notification recv MR"; result = -1; } else { - mr = nullptr; + notify_recv_mr_ = nullptr; } } - if (std::all_of(notify_recv_mrs_.begin(), notify_recv_mrs_.end(), - [](ibv_mr* mr) { return mr == nullptr; })) { - notify_recv_mrs_.clear(); - notify_recv_buffers_.clear(); - } + if (!notify_recv_mr_) notify_recv_buffer_.clear(); if (notify_send_mr_) { if (context_->verbs_.ibv_dereg_mr(notify_send_mr_)) { @@ -306,8 +298,8 @@ int RdmaEndPoint::deconstructUnlocked() { wr_depth_list_ = nullptr; } - if (result == 0 && !notify_qp_ && notify_recv_mrs_.empty() && - !notify_send_mr_ && qp_list_.empty()) { + if (result == 0 && !notify_qp_ && !notify_recv_mr_ && !notify_send_mr_ && + qp_list_.empty()) { peer_server_name_.clear(); peer_nic_name_.clear(); status_.store(EP_DESTROYED, std::memory_order_release); @@ -1036,14 +1028,57 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid, return 0; } +char* RdmaEndPoint::notifySlotPtr(char* base, size_t idx) { + return base + idx * kNotifyBufferSize; +} + +bool RdmaEndPoint::encodeNotifyPayload(char* slot, const std::string& name, + const std::string& msg, + uint32_t* out_len) { + if (!slot || !out_len) return false; + if (name.size() > UINT32_MAX || msg.size() > UINT32_MAX) return false; + uint32_t name_len = static_cast(name.size()); + uint32_t msg_len = static_cast(msg.size()); + const size_t total_size = + sizeof(name_len) + name_len + sizeof(msg_len) + msg_len; + if (total_size > kNotifyBufferSize) return false; + + auto* ptr = slot; + std::memcpy(ptr, &name_len, sizeof(name_len)); + ptr += sizeof(name_len); + std::memcpy(ptr, name.data(), name_len); + ptr += name_len; + std::memcpy(ptr, &msg_len, sizeof(msg_len)); + ptr += sizeof(msg_len); + std::memcpy(ptr, msg.data(), msg_len); + *out_len = static_cast(total_size); + return true; +} + +bool RdmaEndPoint::decodeNotifyPayload(const char* data, size_t byte_len, + std::string* name, std::string* msg) { + if (!data || !name || !msg || byte_len < 8) return false; + uint32_t name_len = 0; + std::memcpy(&name_len, data, sizeof(name_len)); + if (name_len > byte_len - 8) return false; + uint32_t msg_len = 0; + std::memcpy(&msg_len, data + 4 + name_len, sizeof(msg_len)); + if (msg_len > byte_len - 8 - name_len) return false; + name->assign(data + 4, name_len); + msg->assign(data + 4 + name_len + 4, msg_len); + return true; +} + void RdmaEndPoint::postNotifyRecv(size_t idx) { - if (idx >= notify_recv_buffers_.size() || idx >= notify_recv_mrs_.size()) + if (!notify_qp_ || !notify_recv_mr_ || idx >= kNotifyMaxPendingSends || + notify_recv_buffer_.size() < (idx + 1) * kNotifyBufferSize) return; ibv_sge sge = {}; - sge.addr = reinterpret_cast(notify_recv_buffers_[idx].data()); - sge.length = notify_recv_buffers_[idx].size(); - sge.lkey = notify_recv_mrs_[idx]->lkey; + sge.addr = reinterpret_cast( + notifySlotPtr(notify_recv_buffer_.data(), idx)); + sge.length = kNotifyBufferSize; + sge.lkey = notify_recv_mr_->lkey; ibv_recv_wr wr = {}; wr.wr_id = idx; @@ -1169,30 +1204,13 @@ bool RdmaEndPoint::sendNotification(const std::string& name, // Pick the next send slot — flow control guarantees this slot's previous // DMA has completed (at most kNotifyMaxPendingSends-1 in-flight). size_t slot = notify_send_wr_id_ % kNotifyMaxPendingSends; - char* slot_ptr = notify_send_buffer_.data() + slot * kNotifyBufferSize; - - // Serialize: [name_len(4)][name][msg_len(4)][msg] - if (name.size() > UINT32_MAX || msg.size() > UINT32_MAX) { - LOG(ERROR) << "Notification field exceeds uint32 limit"; - return false; - } - uint32_t name_len = static_cast(name.size()); - uint32_t msg_len = static_cast(msg.size()); - size_t total_size = sizeof(name_len) + name_len + sizeof(msg_len) + msg_len; - if (total_size > kNotifyBufferSize) { - LOG(ERROR) << "Notification message too large: " << total_size; + char* slot_ptr = notifySlotPtr(notify_send_buffer_.data(), slot); + uint32_t total_size = 0; + if (!encodeNotifyPayload(slot_ptr, name, msg, &total_size)) { + LOG(ERROR) << "Failed to encode notification payload"; return false; } - auto* ptr = slot_ptr; - std::memcpy(ptr, &name_len, sizeof(name_len)); - ptr += sizeof(name_len); - std::memcpy(ptr, name.data(), name_len); - ptr += name_len; - std::memcpy(ptr, &msg_len, sizeof(msg_len)); - ptr += sizeof(msg_len); - std::memcpy(ptr, msg.data(), msg_len); - // Post send ibv_sge sge = {}; sge.addr = reinterpret_cast(slot_ptr); @@ -1223,7 +1241,8 @@ bool RdmaEndPoint::sendNotification(const std::string& name, bool RdmaEndPoint::handleNotifyRecv(size_t buffer_idx, size_t byte_len) { std::lock_guard resource_guard(notify_resource_mutex_); - if (buffer_idx >= notify_recv_buffers_.size()) { + if (!notify_recv_mr_ || buffer_idx >= kNotifyMaxPendingSends || + notify_recv_buffer_.size() < (buffer_idx + 1) * kNotifyBufferSize) { LOG(ERROR) << "Invalid recv buffer index: " << buffer_idx; return false; } @@ -1234,33 +1253,16 @@ bool RdmaEndPoint::handleNotifyRecv(size_t buffer_idx, size_t byte_len) { return false; } - char* data = notify_recv_buffers_[buffer_idx].data(); - size_t len = byte_len; - - // Deserialize: [name_len(4)][name][msg_len(4)][msg] - if (len < 8) { - LOG(ERROR) << "Invalid notification message size: " << len; - postNotifyRecv(buffer_idx); - return false; - } - - uint32_t name_len = *reinterpret_cast(data); - if (name_len > len - 8) { - LOG(ERROR) << "Invalid notification message format (name too long)"; - postNotifyRecv(buffer_idx); - return false; - } - - std::string name(data + 4, name_len); - uint32_t msg_len = *reinterpret_cast(data + 4 + name_len); - if (msg_len > len - 8 - name_len) { - LOG(ERROR) << "Invalid notification message format (msg too long)"; + char* data = notifySlotPtr(notify_recv_buffer_.data(), buffer_idx); + std::string name; + std::string msg; + if (!decodeNotifyPayload(data, byte_len, &name, &msg)) { + LOG(ERROR) << "Invalid notification message size or format: " + << byte_len; postNotifyRecv(buffer_idx); return false; } - std::string msg(data + 4 + name_len + 4, msg_len); - // Add directly to transport queue (skip endpoint queue for lower latency) context_->transport_.addNotificationToQueue(name, msg); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/ibv_loader.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/ibv_loader.cpp index 9f4bee3c98..5cea729053 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/ibv_loader.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/ibv_loader.cpp @@ -32,6 +32,19 @@ bool LoadSymbol(void* handle, const char* name, Fn& out) { return true; } +// A symbol newer libibverbs add; its absence must not disable RDMA. +template +void LoadOptionalSymbol(void* handle, const char* name, Fn& out) { + void* sym = dlsym(handle, name); + if (!sym) { + LOG(INFO) << "libibverbs lacks optional symbol " << name + << "; the fallback path will be used"; + out = nullptr; + return; + } + out = reinterpret_cast(sym); +} + IbvLoader& IbvLoader::Instance() { static IbvLoader instance; return instance; @@ -55,6 +68,8 @@ IbvLoader::IbvLoader() { ok &= LoadSymbol(handle_, "ibv_query_gid", symbols_.ibv_query_gid); ok &= LoadSymbol(handle_, "ibv_query_port", symbols_.ibv_query_port_default); + LoadOptionalSymbol(handle_, "ibv_query_port_speed", + symbols_.ibv_query_port_speed); ok &= LoadSymbol(handle_, "ibv_get_device_name", symbols_.ibv_get_device_name); @@ -116,4 +131,4 @@ IbvLoader::~IbvLoader() { } } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp index 9a81373bb1..e94f2df8f2 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp @@ -129,6 +129,7 @@ Status DeviceSelector::allocate(uint64_t total_length, uint32_t num_slices, for (int dev_id : entry->device_list[rank]) { if (!usable(dev_id)) continue; if ((device_mask & (1ULL << dev_id)) == 0) continue; + if (!isNumaEligible(entry, dev_id)) continue; tl_eligible.push_back(dev_id); } if (tl_eligible.empty()) continue; @@ -147,7 +148,7 @@ Status DeviceSelector::allocate(uint64_t total_length, uint32_t num_slices, } return Status::OK(); } - return Status::DeviceNotFound("no eligible devices"); + return Status::DeviceNotFound(noEligibleDeviceReason()); } std::vector tl_candidates; @@ -168,6 +169,39 @@ Status DeviceSelector::allocate(uint64_t total_length, uint32_t num_slices, return Status::OK(); } +void DeviceSelector::auditStrictLocalNuma() const { + if (!sched_params_.strict_local_numa || !local_topology_) return; + + size_t excludable = 0; + for (const auto& mem : local_topology_->mem_list_) { + bool has_nic = false, has_local_nic = false; + for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { + for (int dev_id : mem.device_list[rank]) { + if (devices_.find(dev_id) == devices_.end()) continue; + has_nic = true; + if (local_topology_->isCrossNuma(mem, dev_id)) + excludable++; + else + has_local_nic = true; + } + } + if (has_nic && !has_local_nic) { + LOG(WARNING) << "strict_local_numa: location " << mem.name + << " (NUMA " << mem.numa_node + << ") has no same-NUMA RDMA NIC; transfers from it " + "will fail with DeviceNotFound"; + } + } + + if (excludable == 0) { + LOG(WARNING) << "strict_local_numa is enabled but no NIC can be " + "classified as cross-NUMA on this host (custom " + "priority matrix, VM, or sysfs without NUMA info), so " + "the flag has no effect and cross-NUMA NICs keep the " + "numa_penalties soft penalty"; + } +} + int DeviceSelector::getDeviceRank(const std::string& location, int dev_id) const { auto entry = local_topology_->getMemEntry(location); @@ -211,6 +245,7 @@ Status DeviceSelector::buildCandidates(const Topology::MemEntry* entry, for (int dev_id : entry->device_list[rank]) { if (!usable(dev_id)) continue; if ((device_mask & (1ULL << dev_id)) == 0) continue; + if (!isNumaEligible(entry, dev_id)) continue; // QoS: Get device's current priority slot (local, per-process) // Device accepts request if dev_priority >= request_priority int dev_priority = PRIO_LOW; // Default: accept all @@ -222,20 +257,20 @@ Status DeviceSelector::buildCandidates(const Topology::MemEntry* entry, } } - // If no devices after priority filtering, fall back to every usable - // device. Availability is not a QoS filter and is never relaxed here. + // Retry without QoS filtering; availability and NUMA exclusion stay. if (candidates.empty()) { for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { for (int dev_id : entry->device_list[rank]) { if (!usable(dev_id)) continue; if ((device_mask & (1ULL << dev_id)) == 0) continue; + if (!isNumaEligible(entry, dev_id)) continue; add_candidate(dev_id, rank); } } } if (candidates.empty()) { - return Status::DeviceNotFound("no eligible devices"); + return Status::DeviceNotFound(noEligibleDeviceReason()); } std::sort( diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 13e973af86..d7be5e4cb9 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -111,6 +112,9 @@ Workers::Workers(RdmaTransport* transport) } } + params.strict_local_numa = + conf->get("transports/rdma/strict_local_numa", false); + // ============================================================ // Bandwidth Estimation (EWMA) // ============================================================ @@ -187,6 +191,7 @@ Workers::Workers(RdmaTransport* transport) params.max_bandwidth_gbps); device_selector_->setSchedulingParams(params); + device_selector_->auditStrictLocalNuma(); // Seed each device from its context. context_set_ is indexed by NicID, // the same id the selector uses. Three cases: @@ -861,14 +866,35 @@ void Workers::workerThread(int thread_id) { int Workers::handleContextEvents(int dev_id, std::shared_ptr& context) { - ibv_async_event event; - if (ibv_get_async_event(context->nativeContext(), &event) < 0) return -1; - LOG(WARNING) << "Received context async event " - << ibv_event_type_str(event.event_type) << " for context " - << context->name(); - applyContextEvent(dev_id, *context, event); - ibv_ack_async_event(&event); - return 0; + // The async fd is non-blocking and edge-triggered + // (joinNonblockingPollList), and ibv_get_async_event() dequeues one record + // per call, so every queued event has to be consumed here: epoll only + // reports readiness again once a *new* event arrives. Bursts are routine + // (IBV_EVENT_COMM_EST fires once per connection), and a PORT_ACTIVE + // stranded behind one keeps the context paused until some unrelated event + // happens to release it -- which may be never. + while (true) { + ibv_async_event event; + errno = 0; + if (ibv_get_async_event(context->nativeContext(), &event) < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) return 0; // drained + if (errno == EINTR) continue; + PLOG(ERROR) << "ibv_get_async_event for context " + << context->name(); + return -1; + } + if (event.event_type == IBV_EVENT_COMM_EST) { + VLOG(1) << "Received context async event " + << ibv_event_type_str(event.event_type) << " for context " + << context->name(); + } else { + LOG(WARNING) << "Received context async event " + << ibv_event_type_str(event.event_type) + << " for context " << context->name(); + } + applyContextEvent(dev_id, *context, event); + ibv_ack_async_event(&event); + } } void Workers::applyContextEvent(int dev_id, RdmaContext& context, @@ -909,12 +935,7 @@ void Workers::applyContextEvent(int dev_id, RdmaContext& context, device_selector_->setDeviceAvailable(dev_id, false); LOG(WARNING) << "Action: " << context.name() << " down"; } else { - context.resume(); - // The link may have renegotiated while down: re-seed before - // the device becomes selectable so no worker scores it on - // the old rate. - refreshLinkSpeed(dev_id, context); - device_selector_->setDeviceAvailable(dev_id, true); + activateContext(dev_id, context); LOG(WARNING) << "Action: " << context.name() << " up"; } break; @@ -933,6 +954,14 @@ void Workers::applyContextEvent(int dev_id, RdmaContext& context, } } +void Workers::activateContext(int dev_id, RdmaContext& context) { + context.resume(); + // The link may have renegotiated while down: re-seed before the device + // becomes selectable so no worker scores it on the old rate. + refreshLinkSpeed(dev_id, context); + if (device_selector_) device_selector_->setDeviceAvailable(dev_id, true); +} + void Workers::refreshLinkSpeed(int dev_id, RdmaContext& context) { if (!device_selector_) return; const double before = context.linkSpeedGbps(); @@ -944,7 +973,10 @@ void Workers::refreshLinkSpeed(int dev_id, RdmaContext& context) { // configured default and warns, the same as at startup. if (after == before) return; LOG(WARNING) << context.name() << " link speed " << before << " -> " - << after << " Gbps, re-seeding its bandwidth estimate"; + << after << " Gbps (" + << (context.effectiveSpeedKnown() ? "effective speed" + : "encoded rate") + << "), re-seeding its bandwidth estimate"; device_selector_->setDeviceBandwidth(dev_id, after); } @@ -956,6 +988,27 @@ void Workers::reclaimEndpoints() { } } +void Workers::resumePausedContexts() { + for (size_t dev_id = 0; dev_id < transport_->context_set_.size(); + ++dev_id) { + auto& context = transport_->context_set_[dev_id]; + // Only a paused context is waiting for a recovery event; this also + // filters out inert slots, which never leave DEVICE_UNINIT. + if (!context || context->status() != RdmaContext::DEVICE_PAUSED) + continue; + ibv_port_state state; + if (context->queryPortState(&state) != 0) continue; // already logged + // Only a fully active port carries traffic. Intermediate states + // (INIT/ARMED/ACTIVE_DEFER) mean the link is still settling, so leave + // the context paused and re-check on the next tick. + if (state != IBV_PORT_ACTIVE) continue; + LOG(WARNING) << "Action: " << context->name() + << " up (port reports ACTIVE without an " + "IBV_EVENT_PORT_ACTIVE event)"; + activateContext(static_cast(dev_id), *context); + } +} + void Workers::monitorThread() { // Track time for periodic endpoint reclaim (1 Hz heartbeat) auto last_reclaim_time = std::chrono::steady_clock::now(); @@ -971,6 +1024,8 @@ void Workers::monitorThread() { if (time_since_last_reclaim >= 1000) { // 1 second = 1000 ms reclaimEndpoints(); + // Safety net for a recovery event that never reached us. + resumePausedContexts(); last_reclaim_time = current_time; } @@ -1086,7 +1141,18 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { const auto& list = target.topo_entry->device_list[rank]; if (list.empty()) continue; - slice->target_dev_id = list[SimpleRandom::Get().next(list.size())]; + size_t start = SimpleRandom::Get().next(list.size()); + slice->target_dev_id = list[start]; + // Prefer a same-NUMA peer NIC; do not fail if none exist. + if (strictLocalNuma()) { + for (size_t i = 0; i < list.size(); ++i) { + int tdev = list[(start + i) % list.size()]; + if (!target.topo->isCrossNuma(*target.topo_entry, tdev)) { + slice->target_dev_id = tdev; + break; + } + } + } break; } } @@ -1137,6 +1203,11 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, return Status::OK(); } +bool Workers::strictLocalNuma() const { + return device_selector_ && + device_selector_->getSchedulingParams().strict_local_numa; +} + int Workers::getDeviceByFlatIndex(const RouteHint& hint, size_t flat_idx) { for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { auto& list = hint.topo_entry->device_list[rank]; @@ -1219,6 +1290,9 @@ Status Workers::selectFallbackDevice(RouteHint& source, RouteHint& target, int tdev = getDeviceByFlatIndex(target, dst_idx); if (sdev < 0 || sdev >= 64 || (device_mask & (1ULL << sdev)) == 0) continue; + if (strictLocalNuma() && + source.topo->isCrossNuma(*source.topo_entry, sdev)) + continue; bool reachable = same_machine ? (sdev == tdev) // loopback is safe : rail_mon->available(sdev, tdev); diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/tcp/CMakeLists.txt index 3b49319c32..0942f9e522 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/transport/tcp/CMakeLists.txt @@ -1,3 +1,2 @@ -file(GLOB XPORT_SOURCES "*.cpp") -add_library(tent_xport_tcp STATIC ${XPORT_SOURCES}) +add_library(tent_xport_tcp STATIC tcp_transport.cpp) target_link_libraries(tent_xport_tcp PUBLIC tent_rpc tent_common) diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp new file mode 100644 index 0000000000..a9556dcaae --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp @@ -0,0 +1,472 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/endpoint.h" + +#include +#include +#include +#include +#include +#include + +namespace mooncake::tent::ub { +namespace { + +uint64_t generationSeed() { + const auto wall = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + const auto steady = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + const uint64_t seed = + (wall ^ (steady << 13) ^ (steady >> 7)) & 0x7fffffffffffffffULL; + return seed == 0 ? 1 : seed; +} + +std::atomic g_next_endpoint_generation{generationSeed()}; + +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; +} + +} // namespace + +size_t UbEndpointKeyHash::operator()(const UbEndpointKey& key) const noexcept { + size_t seed = std::hash{}(key.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(key.remote_segment_id)); + combine(std::hash{}(key.remote_topology_id)); + combine(std::hash{}(key.peer_nic_path)); + return seed; +} + +uint64_t UbEndpoint::allocateGeneration() noexcept { + const uint64_t generation = + g_next_endpoint_generation.fetch_add(1, std::memory_order_relaxed); + // Generation zero is reserved for "no endpoint". Exhausting the 64-bit + // process-wide sequence is not recoverable without reusing generations. + if (generation == 0 || generation == std::numeric_limits::max()) { + std::terminate(); + } + return generation; +} + +UbEndpoint::UbEndpoint(UbEndpointKey key, UbContextPtr context, + std::shared_ptr adapter, + uint32_t jetty_count, JettyOptions jetty_options) + : key_(std::move(key)), + context_(std::move(context)), + adapter_(std::move(adapter)), + jetty_count_(jetty_count), + jetty_options_(jetty_options), + generation_(allocateGeneration()) {} + +UbEndpoint::~UbEndpoint() { + auto status = retire(); + if (!status.ok() || outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + // A caller that drops the final endpoint reference after a failed + // native fence must not let Jetty destructors force RESET/delete. A + // process-lifetime quarantine is safer than DMA-after-free. + static auto* leaked = new std::vector(); + static auto* leaked_mutex = new std::mutex(); + std::scoped_lock lock(lifecycle_mutex_, *leaked_mutex); + leaked->insert(leaked->end(), std::make_move_iterator(jetties_.begin()), + std::make_move_iterator(jetties_.end())); + jetties_.clear(); + } +} + +void UbEndpoint::rememberFirstError(const Status& candidate, Status& first) { + if (first.ok() && !candidate.ok()) first = candidate; +} + +Status UbEndpoint::resetAndUnbindLocked() { + Status first_error = Status::OK(); + if (!adapter_) return first_error; + + for (auto it = jetties_.rbegin(); it != jetties_.rend(); ++it) { + if (!*it) continue; + auto reset_status = adapter_->resetJetty(*it); + rememberFirstError(reset_status, first_error); + // Unimporting a peer before RESET succeeds can sever resources still + // referenced by hardware. Keep this Jetty intact for a retry. + if (!reset_status.ok()) continue; + rememberFirstError(adapter_->unbindJetty(*it), first_error); + } + return first_error; +} + +Status UbEndpoint::deleteJettysLocked() { + Status first_error = Status::OK(); + if (adapter_) { + for (auto it = jetties_.rbegin(); it != jetties_.rend(); ++it) { + if (!*it) continue; + rememberFirstError(adapter_->deleteJetty(*it), first_error); + } + } + jetties_.clear(); + jfc_indices_.clear(); + return first_error; +} + +Status UbEndpoint::failLocked(Status status) { + if (status.ok()) { + status = Status::InternalError( + "UB endpoint entered failed state without an error" LOC_MARK); + } + lifecycle_status_ = status; + state_.store(State::kFailed, std::memory_order_release); + + rememberFirstError(resetAndUnbindLocked(), retire_status_); + rememberFirstError(deleteJettysLocked(), retire_status_); + return lifecycle_status_; +} + +Status UbEndpoint::prepare() { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kPrepared || current == State::kReady) { + return Status::OK(); + } + if (current == State::kFailed) return lifecycle_status_; + if (current != State::kUninitialized) { + return Status::InvalidArgument( + "UB endpoint cannot be prepared in its current state" LOC_MARK); + } + + state_.store(State::kHandshaking, std::memory_order_release); + if (!key_.valid()) { + return failLocked( + Status::InvalidArgument("Invalid UB endpoint key" LOC_MARK)); + } + if (!adapter_ || !context_ || !context_->active() || !context_->handle() || + !context_->handle()->valid()) { + return failLocked(Status::InvalidArgument( + "UB endpoint requires an active context and adapter" LOC_MARK)); + } + if (context_->topologyId() != key_.local_topology_id) { + return failLocked(Status::InvalidArgument( + "UB endpoint key does not match its local context" LOC_MARK)); + } + if (jetty_count_ == 0 || context_->jfcs().empty()) { + return failLocked(Status::InvalidArgument( + "UB endpoint requires at least one Jetty and JFC" LOC_MARK)); + } + const uint32_t max_jetty = context_->deviceInfo().capabilities.max_jetty; + if (max_jetty != 0 && jetty_count_ > max_jetty) { + return failLocked( + Status::InvalidArgument("Requested endpoint Jetty count exceeds " + "device capability" LOC_MARK)); + } + + jetties_.reserve(jetty_count_); + jfc_indices_.reserve(jetty_count_); + for (uint32_t index = 0; index < jetty_count_; ++index) { + const size_t jfc_index = index % context_->jfcs().size(); + auto jfc = context_->jfc(jfc_index); + if (!jfc || !jfc->valid()) { + return failLocked(Status::InvalidArgument( + "UB endpoint selected an inactive JFC" LOC_MARK)); + } + + JettyPtr jetty; + auto status = adapter_->createJetty(context_->handle(), jfc->handle(), + jetty_options_, jetty); + if (!status.ok()) return failLocked(std::move(status)); + if (!jetty || !jetty->valid() || jetty->id() == 0) { + if (jetty) jetties_.push_back(std::move(jetty)); + return failLocked(Status::InternalError( + "URMA adapter returned an invalid Jetty" LOC_MARK)); + } + jetties_.push_back(std::move(jetty)); + jfc_indices_.push_back(jfc_index); + } + + lifecycle_status_ = Status::OK(); + state_.store(State::kPrepared, std::memory_order_release); + return Status::OK(); +} + +Status UbEndpoint::bind(const UbBootstrapDesc& peer) { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kReady) { + if (samePeer(peer_, peer)) return Status::OK(); + return Status::InvalidArgument( + "UB endpoint is already bound to another peer generation" LOC_MARK); + } + if (current == State::kFailed) return lifecycle_status_; + if (current != State::kPrepared) { + return Status::InvalidArgument( + "UB endpoint must be prepared before bind" LOC_MARK); + } + + state_.store(State::kBinding, std::memory_order_release); + if (!peer.reply_msg.empty()) { + return failLocked(Status::RpcServiceError( + std::string("Peer rejected UB bootstrap: ") + peer.reply_msg)); + } + if (peer.protocol_version != 1) { + return failLocked(Status::InvalidArgument( + "Unsupported UB bootstrap protocol version" LOC_MARK)); + } + if (peer.endpoint_generation == 0 || peer.local_eid.empty()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap is missing EID or " + "endpoint generation" LOC_MARK)); + } + if (peer.jetty_ids.size() != jetties_.size()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap Jetty count does not " + "match local endpoint" LOC_MARK)); + } + if (!peer.jetty_uasids.empty() && + peer.jetty_uasids.size() != jetties_.size()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap UASID count does not " + "match local endpoint" LOC_MARK)); + } + + for (size_t index = 0; index < jetties_.size(); ++index) { + if (peer.jetty_ids[index] == 0) { + return failLocked(Status::InvalidArgument( + "Peer UB bootstrap contains a zero Jetty ID" LOC_MARK)); + } + RemoteJettyInfo remote; + remote.eid = peer.local_eid; + remote.id = peer.jetty_ids[index]; + if (!peer.jetty_uasids.empty()) { + remote.uasid = peer.jetty_uasids[index]; + } + auto status = adapter_->bindJetty(jetties_[index], remote); + if (!status.ok()) return failLocked(std::move(status)); + } + + peer_ = peer; + peer_generation_.store(peer.endpoint_generation, std::memory_order_release); + lifecycle_status_ = Status::OK(); + state_.store(State::kReady, std::memory_order_release); + return Status::OK(); +} + +Status UbEndpoint::makeBootstrapDesc(const std::string& segment_name, + const std::string& local_nic_path, + const std::string& peer_nic_path, + uint64_t segment_generation, + UbBootstrapDesc& output) const { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current != State::kPrepared && current != State::kReady) { + return Status::InvalidArgument( + "UB endpoint must be prepared before bootstrap" LOC_MARK); + } + + output = UbBootstrapDesc{}; + output.protocol_version = 1; + output.segment_name = segment_name; + output.local_nic_path = local_nic_path.empty() + ? context_->deviceInfo().native_device_path + : local_nic_path; + output.peer_nic_path = + peer_nic_path.empty() ? key_.peer_nic_path : peer_nic_path; + output.local_device_name = context_->deviceInfo().native_device_name; + output.local_device_id = context_->topologyId(); + output.local_eid_index = static_cast(context_->deviceInfo().eid_index); + output.local_eid = context_->deviceInfo().eid; + output.jetty_ids.reserve(jetties_.size()); + output.jetty_uasids.reserve(jetties_.size()); + for (const auto& jetty : jetties_) { + if (!jetty || !jetty->valid() || jetty->id() == 0) { + return Status::InternalError( + "UB endpoint contains an invalid Jetty" LOC_MARK); + } + output.jetty_ids.push_back(jetty->id()); + output.jetty_uasids.push_back(jetty->uasid()); + } + output.endpoint_generation = generation_; + output.segment_generation = segment_generation; + output.capabilities = {"read", "write", "endpoint_generation"}; + return Status::OK(); +} + +bool UbEndpoint::tryAcquireOutstanding(uint64_t bytes) noexcept { + std::lock_guard lock(lifecycle_mutex_); + if (state_.load(std::memory_order_relaxed) != State::kReady || !context_ || + !context_->active()) { + return false; + } + + outstanding_wrs_.fetch_add(1, std::memory_order_relaxed); + outstanding_bytes_.fetch_add(bytes, std::memory_order_relaxed); + context_->addInflight(bytes); + return true; +} + +void UbEndpoint::releaseOutstanding(uint64_t bytes) noexcept { + std::lock_guard lock(lifecycle_mutex_); + const uint64_t current_wrs = + outstanding_wrs_.load(std::memory_order_relaxed); + if (current_wrs == 0) return; + + outstanding_wrs_.store(current_wrs - 1, std::memory_order_relaxed); + const uint64_t current_bytes = + outstanding_bytes_.load(std::memory_order_relaxed); + outstanding_bytes_.store(current_bytes >= bytes ? current_bytes - bytes : 0, + std::memory_order_relaxed); + context_->removeInflight(bytes); + + if (current_wrs == 1 && + state_.load(std::memory_order_relaxed) == State::kDestroying) { + (void)finishRetireLocked(); + } +} + +Status UbEndpoint::finishRetireLocked() { + if (state_.load(std::memory_order_relaxed) == State::kDestroyed) { + return retire_status_; + } + if (outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + return retire_status_; + } + + // With no outstanding WR, RESET itself is a sufficient fence. When + // quiesce() ran earlier this is an idempotent cleanup pass. + auto reset_status = resetAndUnbindLocked(); + if (!reset_status.ok()) { + retire_status_ = reset_status; + return reset_status; + } + native_quiesced_ = true; + auto delete_status = deleteJettysLocked(); + if (!delete_status.ok()) { + retire_status_ = delete_status; + return delete_status; + } + retire_status_ = Status::OK(); + peer_ = UbBootstrapDesc{}; + peer_generation_.store(0, std::memory_order_release); + state_.store(State::kDestroyed, std::memory_order_release); + return retire_status_; +} + +Status UbEndpoint::retire() { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kDestroyed) return retire_status_; + + if (current != State::kDestroying) { + state_.store(State::kDestroying, std::memory_order_release); + } + if (outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + return retire_status_; + } + return finishRetireLocked(); +} + +Status UbEndpoint::quiesce(uint32_t timeout_ms, + std::vector& completions) { + completions.clear(); + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kDestroyed) return retire_status_; + if (!adapter_ || timeout_ms == 0) { + return Status::InvalidArgument( + "UB endpoint quiesce requires an adapter and timeout" LOC_MARK); + } + + state_.store(State::kDestroying, std::memory_order_release); + Status fence_error = Status::OK(); + if (!native_quiesced_) { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::milliseconds(timeout_ms); + for (const auto& jetty : jetties_) { + if (!jetty) continue; + const auto now = std::chrono::steady_clock::now(); + if (now >= deadline) { + if (fence_error.ok()) { + fence_error = Status::RdmaError( + "UB endpoint Jetty drain budget exhausted"); + } + break; + } + const auto remaining = + std::chrono::duration_cast(deadline - + now); + const uint32_t remaining_ms = + static_cast(std::max(1, remaining.count())); + std::vector drained; + auto status = adapter_->quiesceJetty(jetty, remaining_ms, drained); + completions.insert(completions.end(), + std::make_move_iterator(drained.begin()), + std::make_move_iterator(drained.end())); + if (!status.ok() && fence_error.ok()) fence_error = status; + } + if (!fence_error.ok()) { + // Do not RESET, unbind, or delete anything without a fence for + // every Jetty. A later shutdown attempt can safely retry. + return fence_error; + } + native_quiesced_ = true; + } + + auto reset_status = resetAndUnbindLocked(); + if (!reset_status.ok()) return reset_status; + if (outstanding_wrs_.load(std::memory_order_relaxed) == 0) { + return finishRetireLocked(); + } + return Status::OK(); +} + +bool UbEndpoint::reusable() const noexcept { + switch (state()) { + case State::kUninitialized: + case State::kHandshaking: + case State::kPrepared: + case State::kBinding: + case State::kReady: + return true; + case State::kFailed: + case State::kDestroying: + case State::kDestroyed: + return false; + } + return false; +} + +size_t UbEndpoint::jettyCount() const { + std::lock_guard lock(lifecycle_mutex_); + return jetties_.size(); +} + +JettyPtr UbEndpoint::jetty(size_t index) const { + std::lock_guard lock(lifecycle_mutex_); + if (jetties_.empty()) return nullptr; + return jetties_[index % jetties_.size()]; +} + +size_t UbEndpoint::jfcIndex(size_t jetty_index) const { + std::lock_guard lock(lifecycle_mutex_); + if (jfc_indices_.empty()) return 0; + return jfc_indices_[jetty_index % jfc_indices_.size()]; +} + +std::vector UbEndpoint::jetties() const { + std::lock_guard lock(lifecycle_mutex_); + return jetties_; +} + +Status UbEndpoint::lifecycleStatus() const { + std::lock_guard lock(lifecycle_mutex_); + return lifecycle_status_; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp new file mode 100644 index 0000000000..186bd3f254 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp @@ -0,0 +1,207 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/endpoint_store.h" + +#include +#include +#include +#include + +namespace mooncake::tent::ub { + +EndpointStore::EndpointStore(std::shared_ptr adapter, + size_t max_size, uint32_t jetty_count, + JettyOptions jetty_options) + : adapter_(std::move(adapter)), + max_size_(max_size), + jetty_count_(jetty_count), + jetty_options_(jetty_options) {} + +EndpointStore::~EndpointStore() { + auto status = clear(); + if (!status.ok()) { + // Standalone owners may ignore clear(); preserve unsafe-to-destroy + // endpoints for process lifetime rather than letting native handle + // destructors bypass a failed drain fence. + static auto* leaked = new std::vector>(); + static auto* leaked_mutex = new std::mutex(); + std::scoped_lock lock(mutex_, *leaked_mutex); + leaked->insert(leaked->end(), + std::make_move_iterator(quarantined_.begin()), + std::make_move_iterator(quarantined_.end())); + quarantined_.clear(); + } +} + +std::shared_ptr EndpointStore::get(const UbEndpointKey& key) { + std::shared_ptr retired; + std::shared_ptr result; + { + std::lock_guard lock(mutex_); + auto it = endpoints_.find(key); + if (it == endpoints_.end()) return nullptr; + if (!it->second.endpoint || !it->second.endpoint->reusable()) { + retired = std::move(it->second.endpoint); + endpoints_.erase(it); + } else { + result = it->second.endpoint; + } + } + if (retired) { + auto status = retired->retire(); + if (!status.ok() || retired->state() != UbEndpoint::State::kDestroyed) { + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(retired)); + } + } + return result; +} + +Status EndpointStore::getOrCreate(const UbEndpointKey& key, + const UbContextPtr& context, + std::shared_ptr& endpoint) { + endpoint.reset(); + if (!key.valid() || !context || + context->topologyId() != key.local_topology_id || max_size_ == 0 || + jetty_count_ == 0) { + return Status::InvalidArgument( + "Invalid UB endpoint store request" LOC_MARK); + } + + for (;;) { + std::shared_ptr evicted; + { + std::lock_guard lock(mutex_); + auto existing = endpoints_.find(key); + if (existing != endpoints_.end()) { + if (existing->second.endpoint && + existing->second.endpoint->reusable()) { + endpoint = existing->second.endpoint; + } else { + evicted = std::move(existing->second.endpoint); + endpoints_.erase(existing); + } + } + + if (!endpoint && !evicted && + endpoints_.size() + quarantined_.size() >= max_size_) { + auto victim = endpoints_.end(); + uint64_t oldest = std::numeric_limits::max(); + for (auto it = endpoints_.begin(); it != endpoints_.end(); + ++it) { + if (it->second.endpoint && + it->second.endpoint->outstandingWrs() == 0 && + it->second.insertion_order < oldest) { + victim = it; + oldest = it->second.insertion_order; + } + } + if (victim == endpoints_.end()) { + return Status::TooManyRequests( + "All UB endpoint cache entries are in flight" LOC_MARK); + } + evicted = std::move(victim->second.endpoint); + endpoints_.erase(victim); + } + + if (!endpoint && !evicted) { + endpoint = std::make_shared( + key, context, adapter_, jetty_count_, jetty_options_); + endpoints_.emplace(key, + Entry{endpoint, next_insertion_order_++}); + } + } + + if (evicted) { + auto status = evicted->retire(); + if (!status.ok() || + evicted->state() != UbEndpoint::State::kDestroyed) { + if (status.ok()) { + status = Status::TooManyRequests( + "UB endpoint cleanup is still in progress" LOC_MARK); + } + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(evicted)); + return status; + } + continue; + } + + auto status = endpoint->prepare(); + if (!status.ok()) { + (void)retire(key, endpoint->generation()); + endpoint.reset(); + return status; + } + return Status::OK(); + } +} + +bool EndpointStore::retire(const UbEndpointKey& key, uint64_t generation) { + std::shared_ptr endpoint; + { + std::lock_guard lock(mutex_); + auto it = endpoints_.find(key); + if (it == endpoints_.end() || !it->second.endpoint || + it->second.endpoint->generation() != generation) { + return false; + } + endpoint = std::move(it->second.endpoint); + endpoints_.erase(it); + } + auto status = endpoint->retire(); + if (!status.ok() || endpoint->state() != UbEndpoint::State::kDestroyed) { + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(endpoint)); + } + return true; +} + +bool EndpointStore::retire(const std::shared_ptr& endpoint) { + return endpoint && retire(endpoint->key(), endpoint->generation()); +} + +Status EndpointStore::clear() { + std::vector> endpoints; + { + std::lock_guard lock(mutex_); + endpoints.reserve(endpoints_.size()); + for (auto& [_, entry] : endpoints_) { + if (entry.endpoint) endpoints.push_back(std::move(entry.endpoint)); + } + endpoints_.clear(); + for (auto& endpoint : quarantined_) { + if (endpoint) endpoints.push_back(std::move(endpoint)); + } + quarantined_.clear(); + } + Status first_error = Status::OK(); + std::vector> failed; + for (auto& endpoint : endpoints) { + auto status = endpoint->retire(); + if (!status.ok() || + endpoint->state() != UbEndpoint::State::kDestroyed) { + if (status.ok()) { + status = Status::TooManyRequests( + "UB endpoint cleanup is still in progress" LOC_MARK); + } + if (first_error.ok()) first_error = status; + failed.push_back(std::move(endpoint)); + } + } + if (!failed.empty()) { + std::lock_guard lock(mutex_); + quarantined_.insert(quarantined_.end(), + std::make_move_iterator(failed.begin()), + std::make_move_iterator(failed.end())); + } + return first_error; +} + +size_t EndpointStore::size() const { + std::lock_guard lock(mutex_); + return endpoints_.size(); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp b/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp new file mode 100644 index 0000000000..dab6d0da1b --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp @@ -0,0 +1,371 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/ub/quota.h" + +#include + +namespace mooncake::tent::ub { + +QuotaManager::QuotaManager(QuotaLimits default_device_limits, + QuotaLimits default_path_limits) + : default_device_limits_(default_device_limits), + default_path_limits_(default_path_limits) {} + +void QuotaManager::setDefaultDeviceLimits(const QuotaLimits& limits) { + std::lock_guard lock(mutex_); + default_device_limits_ = limits; +} + +void QuotaManager::setDefaultPathLimits(const QuotaLimits& limits) { + std::lock_guard lock(mutex_); + default_path_limits_ = limits; +} + +QuotaLimits QuotaManager::defaultDeviceLimits() const { + std::lock_guard lock(mutex_); + return default_device_limits_; +} + +QuotaLimits QuotaManager::defaultPathLimits() const { + std::lock_guard lock(mutex_); + return default_path_limits_; +} + +bool QuotaManager::setDeviceLimits(Topology::NicID local_topology_id, + const QuotaLimits& limits) { + if (local_topology_id < 0) return false; + std::lock_guard lock(mutex_); + devices_[local_topology_id].override_limits = limits; + return true; +} + +bool QuotaManager::clearDeviceLimits(Topology::NicID local_topology_id) { + if (local_topology_id < 0) return false; + std::lock_guard lock(mutex_); + auto it = devices_.find(local_topology_id); + if (it == devices_.end() || !it->second.override_limits) return false; + it->second.override_limits.reset(); + return true; +} + +bool QuotaManager::setPathLimits(const UbPostPath& path, + const QuotaLimits& limits) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + auto& record = paths_[UbRailKey::fromPath(path)]; + record.latest_path = path; + record.override_limits = limits; + return true; +} + +bool QuotaManager::clearPathLimits(const UbPostPath& path) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + auto it = paths_.find(UbRailKey::fromPath(path)); + if (it == paths_.end() || !it->second.override_limits) return false; + it->second.override_limits.reset(); + return true; +} + +std::optional QuotaManager::tryAcquire(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) { + std::lock_guard lock(mutex_); + return tryAcquireLocked(path, bytes, wrs, true); +} + +std::optional QuotaManager::tryAcquireFirst( + const std::vector& paths, uint64_t bytes, uint64_t wrs) { + std::lock_guard lock(mutex_); + for (const auto& path : paths) { + auto reservation = tryAcquireLocked(path, bytes, wrs, false); + if (reservation) return reservation; + } + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + return std::nullopt; +} + +QuotaAvailability QuotaManager::availability(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) const { + std::lock_guard lock(mutex_); + return availabilityLocked(path, bytes, wrs); +} + +std::optional QuotaManager::tryAcquireLocked( + const UbPostPath& path, uint64_t bytes, uint64_t wrs, + bool count_aggregate_reject) { + if (!path.valid() || wrs == 0) { + if (count_aggregate_reject) { + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + } + return std::nullopt; + } + + auto& device = devices_[path.local_topology_id]; + auto& rail = paths_[UbRailKey::fromPath(path)]; + if (!rail.latest_path.valid() || + path.endpoint_generation > rail.latest_path.endpoint_generation) { + rail.latest_path = path; + } + const auto device_limits = effectiveLimits(device, default_device_limits_); + const auto path_limits = effectiveLimits(rail, default_path_limits_); + const bool device_fits = fits(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes) && + fits(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs); + const bool path_fits = + fits(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes) && + fits(rail.usage.outstanding_wrs, wrs, path_limits.max_outstanding_wrs); + if (!device_fits || !path_fits) { + if (!device_fits) { + device.rejected_acquisitions = + saturatingAdd(device.rejected_acquisitions, 1); + } + if (!path_fits) { + rail.rejected_acquisitions = + saturatingAdd(rail.rejected_acquisitions, 1); + } + if (count_aggregate_reject) { + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + } + return std::nullopt; + } + + const uint64_t id = nextReservationIdLocked(); + addUsage(device.usage, bytes, wrs); + addUsage(rail.usage, bytes, wrs); + updatePeak(device.usage, device.peak_usage); + updatePeak(rail.usage, rail.peak_usage); + device.total_acquisitions = saturatingAdd(device.total_acquisitions, 1); + rail.total_acquisitions = saturatingAdd(rail.total_acquisitions, 1); + + addUsage(aggregate_stats_.usage, bytes, wrs); + updatePeak(aggregate_stats_.usage, aggregate_stats_.peak_usage); + aggregate_stats_.total_acquisitions = + saturatingAdd(aggregate_stats_.total_acquisitions, 1); + active_reservations_.emplace(id, ActiveReservation{path, bytes, wrs}); + aggregate_stats_.active_reservations = active_reservations_.size(); + return QuotaReservation{id, path, bytes, wrs}; +} + +QuotaAvailability QuotaManager::availabilityLocked(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) const { + if (!path.valid() || wrs == 0) return {}; + + const auto device_it = devices_.find(path.local_topology_id); + const auto path_it = paths_.find(UbRailKey::fromPath(path)); + const QuotaRecord empty; + const auto& device = + device_it == devices_.end() ? empty : device_it->second; + const auto& rail = path_it == paths_.end() ? empty : path_it->second; + const auto device_limits = effectiveLimits(device, default_device_limits_); + const auto path_limits = effectiveLimits(rail, default_path_limits_); + + QuotaAvailability result; + result.can_acquire = + fits(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes) && + fits(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs) && + fits(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes) && + fits(rail.usage.outstanding_wrs, wrs, path_limits.max_outstanding_wrs); + result.normalized_inflight = + std::max(normalizedUsage(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes), + normalizedUsage(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes)); + result.normalized_outstanding_wrs = + std::max(normalizedUsage(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs), + normalizedUsage(rail.usage.outstanding_wrs, wrs, + path_limits.max_outstanding_wrs)); + return result; +} + +bool QuotaManager::release(const QuotaReservation& reservation) { + std::lock_guard lock(mutex_); + const auto active = active_reservations_.find(reservation.id); + if (reservation.id == 0 || active == active_reservations_.end()) { + aggregate_stats_.duplicate_release_attempts = + saturatingAdd(aggregate_stats_.duplicate_release_attempts, 1); + return false; + } + + // Always release the manager-owned record. The caller's copy may have + // been changed or may refer to an earlier copy of this reservation. + const ActiveReservation charge = active->second; + active_reservations_.erase(active); + + auto device = devices_.find(charge.path.local_topology_id); + if (device != devices_.end()) { + releaseUsage(device->second.usage, charge.bytes, charge.wrs); + device->second.total_releases = + saturatingAdd(device->second.total_releases, 1); + } + auto path = paths_.find(UbRailKey::fromPath(charge.path)); + if (path != paths_.end()) { + releaseUsage(path->second.usage, charge.bytes, charge.wrs); + path->second.total_releases = + saturatingAdd(path->second.total_releases, 1); + } + + releaseUsage(aggregate_stats_.usage, charge.bytes, charge.wrs); + aggregate_stats_.total_releases = + saturatingAdd(aggregate_stats_.total_releases, 1); + aggregate_stats_.active_reservations = active_reservations_.size(); + return true; +} + +DeviceQuotaStats QuotaManager::deviceStats( + Topology::NicID local_topology_id) const { + std::lock_guard lock(mutex_); + DeviceQuotaStats result; + result.local_topology_id = local_topology_id; + auto it = devices_.find(local_topology_id); + if (it == devices_.end()) { + static_cast(result).limits = default_device_limits_; + } else { + static_cast(result) = + makeStats(it->second, default_device_limits_); + } + return result; +} + +PathQuotaStats QuotaManager::pathStats(const UbPostPath& path) const { + std::lock_guard lock(mutex_); + PathQuotaStats result; + result.path = path; + auto it = paths_.find(UbRailKey::fromPath(path)); + if (it == paths_.end()) { + static_cast(result).limits = default_path_limits_; + } else { + static_cast(result) = + makeStats(it->second, default_path_limits_); + } + return result; +} + +std::vector QuotaManager::allDeviceStats() const { + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(devices_.size()); + for (const auto& [id, record] : devices_) { + DeviceQuotaStats stats; + static_cast(stats) = + makeStats(record, default_device_limits_); + stats.local_topology_id = id; + result.push_back(stats); + } + return result; +} + +std::vector QuotaManager::allPathStats() const { + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(paths_.size()); + for (const auto& [_, record] : paths_) { + PathQuotaStats stats; + static_cast(stats) = + makeStats(record, default_path_limits_); + stats.path = record.latest_path; + result.push_back(stats); + } + return result; +} + +AggregateQuotaStats QuotaManager::aggregateStats() const { + std::lock_guard lock(mutex_); + return aggregate_stats_; +} + +size_t QuotaManager::activeReservationCount() const { + std::lock_guard lock(mutex_); + return active_reservations_.size(); +} + +bool QuotaManager::fits(uint64_t current, uint64_t charge, uint64_t limit) { + return current <= limit && charge <= limit - current; +} + +double QuotaManager::normalizedUsage(uint64_t current, uint64_t charge, + uint64_t limit) { + if (limit == std::numeric_limits::max()) return 0.0; + if (limit == 0) return current == 0 && charge == 0 ? 0.0 : 1.0; + const long double projected = std::min( + static_cast(limit), + static_cast(current) + static_cast(charge)); + return static_cast(projected / static_cast(limit)); +} + +uint64_t QuotaManager::saturatingAdd(uint64_t lhs, uint64_t rhs) { + if (std::numeric_limits::max() - lhs < rhs) { + return std::numeric_limits::max(); + } + return lhs + rhs; +} + +void QuotaManager::addUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs) { + usage.inflight_bytes = saturatingAdd(usage.inflight_bytes, bytes); + usage.outstanding_wrs = saturatingAdd(usage.outstanding_wrs, wrs); +} + +void QuotaManager::releaseUsage(QuotaUsage& usage, uint64_t bytes, + uint64_t wrs) { + usage.inflight_bytes = + bytes >= usage.inflight_bytes ? 0 : usage.inflight_bytes - bytes; + usage.outstanding_wrs = + wrs >= usage.outstanding_wrs ? 0 : usage.outstanding_wrs - wrs; +} + +void QuotaManager::updatePeak(const QuotaUsage& usage, QuotaUsage& peak) { + peak.inflight_bytes = std::max(peak.inflight_bytes, usage.inflight_bytes); + peak.outstanding_wrs = + std::max(peak.outstanding_wrs, usage.outstanding_wrs); +} + +QuotaLimits QuotaManager::effectiveLimits(const QuotaRecord& record, + const QuotaLimits& defaults) { + return record.override_limits.value_or(defaults); +} + +QuotaStats QuotaManager::makeStats(const QuotaRecord& record, + const QuotaLimits& defaults) { + return QuotaStats{effectiveLimits(record, defaults), + record.usage, + record.peak_usage, + record.total_acquisitions, + record.total_releases, + record.rejected_acquisitions}; +} + +uint64_t QuotaManager::nextReservationIdLocked() { + // IDs are never reused while active, including across uint64_t wrap. + while (next_reservation_id_ == 0 || + active_reservations_.contains(next_reservation_id_)) { + ++next_reservation_id_; + } + const uint64_t result = next_reservation_id_++; + if (next_reservation_id_ == 0) ++next_reservation_id_; + return result; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp b/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp new file mode 100644 index 0000000000..02ae248945 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp @@ -0,0 +1,279 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/ub/rail_monitor.h" + +#include +#include +#include + +namespace mooncake::tent::ub { + +namespace { + +double updateEwma(double current, double sample, double alpha) { + if (current < 0.0) return sample; + return alpha * sample + (1.0 - alpha) * current; +} + +} // namespace + +RailMonitor::RailMonitor(RailMonitorConfig config) { + if (config.valid()) config_ = config; +} + +bool RailMonitor::configure(const RailMonitorConfig& config) { + if (!config.valid()) return false; + std::lock_guard lock(mutex_); + config_ = config; + return true; +} + +RailMonitorConfig RailMonitor::config() const { + std::lock_guard lock(mutex_); + return config_; +} + +bool RailMonitor::registerPath(const UbPostPath& path) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + getOrCreateLocked(path); + return true; +} + +bool RailMonitor::available(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return false; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + return !state.stats.paused; +} + +void RailMonitor::recordSuccess(const UbPostPath& path, uint64_t bytes, + uint64_t latency_ns, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + + ++state.stats.successful_completions; + if (std::numeric_limits::max() - state.stats.completed_bytes < + bytes) { + state.stats.completed_bytes = std::numeric_limits::max(); + } else { + state.stats.completed_bytes += bytes; + } + state.stats.last_success_ns = std::max(state.stats.last_success_ns, now_ns); + + // Zero-byte or zero-latency completions are valid completions but are not + // usable bandwidth observations. + if (bytes == 0 || latency_ns == 0) return; + const double bandwidth = static_cast(bytes) * 1'000'000'000.0 / + static_cast(latency_ns); + state.stats.ewma_bandwidth_bytes_per_second = + updateEwma(state.stats.ewma_bandwidth_bytes_per_second, bandwidth, + config_.ewma_alpha); + state.stats.ewma_latency_ns = + updateEwma(state.stats.ewma_latency_ns, static_cast(latency_ns), + config_.ewma_alpha); +} + +void RailMonitor::recordError(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + recordFailureLocked(path, now_ns, false); +} + +void RailMonitor::recordTimeout(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + recordFailureLocked(path, now_ns, true); +} + +bool RailMonitor::recordEndpointRebuild(const UbPostPath& path, + uint64_t now_ns) { + if (!path.valid()) return false; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + if (path.endpoint_generation <= state.recorded_rebuild_generation) { + return false; + } + state.recorded_rebuild_generation = path.endpoint_generation; + ++state.stats.endpoint_rebuilds; + return true; +} + +RailStats RailMonitor::stats(const UbPostPath& path, uint64_t now_ns) { + RailStats result; + result.key = UbRailKey::fromPath(path); + if (!path.valid()) return result; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + pruneErrorsLocked(state, now_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + return state.stats; +} + +std::vector RailMonitor::allStats(uint64_t now_ns) { + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(rails_.size()); + for (auto& [key, state] : rails_) { + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + pruneErrorsLocked(state, observed_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + result.push_back(state.stats); + } + return result; +} + +double RailMonitor::aggregateBandwidth(uint64_t now_ns) { + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + + bool has_sample = false; + std::unordered_map best_per_device; + for (auto& [key, state] : rails_) { + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + const double bandwidth = state.stats.ewma_bandwidth_bytes_per_second; + if (bandwidth < 0.0) continue; + has_sample = true; + if (state.stats.paused) continue; + auto [it, inserted] = + best_per_device.emplace(key.local_topology_id, bandwidth); + if (!inserted) it->second = std::max(it->second, bandwidth); + } + + if (!has_sample) return -1.0; + double aggregate = 0.0; + for (const auto& entry : best_per_device) { + aggregate += entry.second; + } + return aggregate; +} + +size_t RailMonitor::pathCount() const { + std::lock_guard lock(mutex_); + return rails_.size(); +} + +uint64_t RailMonitor::normalizedNow(uint64_t now_ns) { + return now_ns == 0 ? steadyNowNs() : now_ns; +} + +uint64_t RailMonitor::deadlineAfter(uint64_t now_ns, uint64_t duration_ns) { + if (std::numeric_limits::max() - now_ns < duration_ns) { + return std::numeric_limits::max(); + } + return now_ns + duration_ns; +} + +uint64_t RailMonitor::observeTimeLocked(RailState& state, uint64_t event_ns) { + state.observed_through_ns = std::max(state.observed_through_ns, event_ns); + return state.observed_through_ns; +} + +RailMonitor::RailState& RailMonitor::getOrCreateLocked(const UbPostPath& path) { + const auto key = UbRailKey::fromPath(path); + auto [it, inserted] = rails_.try_emplace(key); + if (inserted) it->second.stats.key = key; + it->second.stats.latest_endpoint_generation = std::max( + it->second.stats.latest_endpoint_generation, path.endpoint_generation); + return it->second; +} + +void RailMonitor::insertErrorLocked(RailState& state, uint64_t event_ns) { + const auto position = std::upper_bound(state.recent_errors.begin(), + state.recent_errors.end(), event_ns); + state.recent_errors.insert(position, event_ns); +} + +void RailMonitor::pruneErrorsLocked(RailState& state, uint64_t now_ns) { + while (!state.recent_errors.empty()) { + const uint64_t error_ns = state.recent_errors.front(); + if (error_ns > now_ns || now_ns - error_ns < config_.error_window_ns) { + break; + } + state.recent_errors.pop_front(); + } +} + +void RailMonitor::refreshCooldownLocked(RailState& state, uint64_t now_ns) { + pruneErrorsLocked(state, now_ns); + if (!state.stats.paused || now_ns < state.stats.cooldown_until_ns) return; + state.ignore_errors_through_ns = + std::max(state.ignore_errors_through_ns, state.stats.cooldown_until_ns); + state.stats.paused = false; + state.stats.cooldown_until_ns = 0; + state.recent_errors.clear(); + state.stats.errors_in_window = 0; + ++state.stats.recoveries; +} + +void RailMonitor::recordFailureLocked(const UbPostPath& path, uint64_t now_ns, + bool timeout) { + auto& state = getOrCreateLocked(path); + const uint64_t event_ns = now_ns; + const uint64_t observed_ns = observeTimeLocked(state, event_ns); + refreshCooldownLocked(state, observed_ns); + ++state.stats.completion_errors; + if (timeout) ++state.stats.timeouts; + state.stats.last_error_ns = std::max(state.stats.last_error_ns, event_ns); + + if (event_ns <= state.ignore_errors_through_ns) return; + insertErrorLocked(state, event_ns); + pruneErrorsLocked(state, observed_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + + if (state.recent_errors.size() < config_.error_threshold) return; + const uint64_t cooldown_until = + deadlineAfter(state.recent_errors.back(), config_.cooldown_ns); + if (cooldown_until <= observed_ns) { + // This entire failure burst and its cooldown are already in the past + // relative to the watermark. Treat it as a completed health epoch + // instead of resurrecting an expired pause because an event was late. + state.ignore_errors_through_ns = + std::max(state.ignore_errors_through_ns, cooldown_until); + state.recent_errors.clear(); + state.stats.errors_in_window = 0; + return; + } + if (!state.stats.paused) { + state.stats.paused = true; + state.stats.pause_started_ns = state.recent_errors.back(); + ++state.stats.pauses; + } + state.stats.cooldown_until_ns = + std::max(state.stats.cooldown_until_ns, cooldown_until); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp new file mode 100644 index 0000000000..130e799ca3 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -0,0 +1,783 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/ub_transport.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "tent/common/utils/string_builder.h" +#include "tent/runtime/segment.h" +#include "tent/thirdparty/nlohmann/json.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/device_selection.h" +#include "tent/transport/ub/endpoint.h" +#include "tent/transport/ub/endpoint_store.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" +#include "tent/transport/ub/urma_adapter.h" +#include "tent/transport/ub/workers.h" + +namespace mooncake::tent { +namespace { + +bool filterAllows(const std::vector& filter, + const ub::DeviceInfo& device) { + if (filter.empty()) return true; + return std::find(filter.begin(), filter.end(), device.topology_name) != + filter.end() || + std::find(filter.begin(), filter.end(), device.native_device_name) != + filter.end(); +} + +uint64_t generationSeed() { + const auto wall = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + const auto steady = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + const uint64_t seed = + (wall ^ (steady << 19) ^ (steady >> 5)) & 0x7fffffffffffffffULL; + return seed == 0 ? 1 : seed; +} + +std::string encodeDeviceMetadata(Topology::NicID topology_id, + const ub::DeviceInfo& device) { + const auto& caps = device.capabilities; + return nlohmann::json{{"schema_version", 1}, + {"topology_id", topology_id}, + {"native_device_name", device.native_device_name}, + {"native_device_path", device.native_device_path}, + {"eid_index", device.eid_index}, + {"eid", device.eid}, + {"active", device.active}, + {"capabilities", + {{"max_jfc", caps.max_jfc}, + {"max_jfc_depth", caps.max_jfc_depth}, + {"max_jfr_depth", caps.max_jfr_depth}, + {"max_jetty", caps.max_jetty}, + {"max_jetty_depth", caps.max_jetty_depth}, + {"max_send_sge", caps.max_send_sge}, + {"max_remote_sge", caps.max_remote_sge}, + {"max_message_size", caps.max_message_size}, + {"max_read_size", caps.max_read_size}, + {"max_write_size", caps.max_write_size}, + {"feature_flags", caps.feature_flags}, + {"transport_modes", caps.transport_modes}}}} + .dump(); +} + +} // namespace + +struct UbTransport::Impl { + explicit Impl(std::shared_ptr injected) + : adapter(injected ? std::move(injected) + : ub::createDefaultUrmaAdapter()) {} + + Status install(const std::string& segment_name, + std::shared_ptr control, + std::shared_ptr topology, + std::shared_ptr config) { + std::lock_guard lock(lifecycle_mutex); + if (installed.load(std::memory_order_acquire)) { + return Status::InvalidArgument( + "UB transport has already been installed" LOC_MARK); + } + if (shutting_down.load(std::memory_order_acquire) || workers || + adapter_initialized) { + return Status::InvalidArgument( + "A previous UB uninstall has not drained yet" LOC_MARK); + } + if (segment_name.empty() || !control || !topology) { + return Status::InvalidArgument( + "UB install requires segment, control service and " + "topology" LOC_MARK); + } + if (!adapter || !adapter->available()) { + return Status::DeviceNotFound( + "A real or injected URMA adapter is unavailable" LOC_MARK); + } + if (!config) config = std::make_shared(); + ub::UbParams parsed; + CHECK_STATUS(ub::UbParams::FromConfig(*config, parsed)); + if (!parsed.enable) { + return Status::InvalidArgument( + "UB transport is disabled by configuration" LOC_MARK); + } + if (parsed.enable_notifications) { + return Status::NotImplemented( + "UB notifications are not supported by protocol version " + "1" LOC_MARK); + } + + shutting_down.store(false, std::memory_order_release); + local_segment_name = segment_name; + metadata = std::move(control); + local_topology = std::move(topology); + conf = std::move(config); + params = parsed; + + auto status = adapter->initialize(); + if (!status.ok()) return failInstall(status); + adapter_initialized = true; + + std::vector discovered; + status = adapter->discoverDevices(discovered); + if (!status.ok()) return failInstall(status); + + const bool explicit_filter = !params.device_filter.empty(); + const auto preferred_devices = + ub::preferBondingDevicesIfPresent(discovered, explicit_filter); + const bool prefer_bonding = + !explicit_filter && + std::any_of(discovered.begin(), discovered.end(), + [](const ub::DeviceInfo& device) { + return ub::isBondingDevice(device); + }); + std::unordered_set preferred_topology_names; + preferred_topology_names.reserve(preferred_devices.size()); + for (const auto& device : preferred_devices) { + preferred_topology_names.insert(device.topology_name); + } + + std::unordered_map by_topology_name; + std::unordered_map jfc_depth_by_device; + for (auto& device : discovered) { + by_topology_name.emplace(device.topology_name, std::move(device)); + } + + std::vector preferred_names; + std::vector skipped_names; + + for (size_t id = 0; id < local_topology->getNicCount(); ++id) { + const auto* nic = local_topology->getNicEntry(static_cast(id)); + if (!nic || nic->type != Topology::NIC_UB) continue; + auto found = by_topology_name.find(nic->name); + if (found == by_topology_name.end() || !found->second.active) { + continue; + } + if (!filterAllows(params.device_filter, found->second)) continue; + if (preferred_topology_names.find(found->second.topology_name) == + preferred_topology_names.end()) { + skipped_names.push_back(found->second.native_device_name); + continue; + } + + ub::JfcOptions jfc_options; + const auto& device_caps = found->second.capabilities; + if (device_caps.max_jfc_depth != 0) { + jfc_options.depth = + std::min(jfc_options.depth, device_caps.max_jfc_depth); + jfc_options.receiver_depth = std::min( + jfc_options.receiver_depth, device_caps.max_jfc_depth); + } + if (device_caps.max_jfr_depth != 0) { + jfc_options.receiver_depth = std::min( + jfc_options.receiver_depth, device_caps.max_jfr_depth); + } + auto context = std::make_shared( + static_cast(id), found->second, adapter); + status = context->initialize(params.jfc_per_context, jfc_options); + if (!status.ok()) { + LOG(WARNING) << "Disable UB device " << nic->name << ": " + << status.ToString(); + continue; + } + preferred_names.push_back(found->second.native_device_name); + context_by_topology_id.emplace(static_cast(id), context); + context_by_topology_name.emplace(nic->name, context); + jfc_depth_by_device.emplace(static_cast(id), + jfc_options.depth); + contexts.push_back(std::move(context)); + } + + { + std::ostringstream selection_log; + if (explicit_filter) { + selection_log << "UB device selection: mode=explicit-filter"; + } else if (prefer_bonding) { + selection_log + << "UB device selection: mode=auto-prefer-bonding"; + } else { + selection_log << "UB device selection: mode=all-devices"; + } + if (!preferred_names.empty()) { + selection_log << " preferred=["; + for (size_t i = 0; i < preferred_names.size(); ++i) { + if (i != 0) selection_log << ", "; + selection_log << preferred_names[i]; + } + selection_log << "]"; + } + if (!skipped_names.empty()) { + selection_log << " skipped=["; + for (size_t i = 0; i < skipped_names.size(); ++i) { + if (i != 0) selection_log << ", "; + selection_log << skipped_names[i]; + } + selection_log << "]"; + } + LOG(INFO) << selection_log.str(); + } + + if (contexts.empty()) { + return failInstall(Status::DeviceNotFound( + "No UB context initialized successfully" LOC_MARK)); + } + + size_t safe_slice_size = params.slice_size; + ub::JettyOptions jetty_options; + for (const auto& context : contexts) { + const auto& device_caps = context->deviceInfo().capabilities; + auto clamp_slice = [&safe_slice_size](uint64_t limit) { + if (limit != 0) { + safe_slice_size = static_cast( + std::min(safe_slice_size, limit)); + } + }; + clamp_slice(device_caps.max_message_size); + clamp_slice(device_caps.max_read_size); + clamp_slice(device_caps.max_write_size); + if (device_caps.max_jetty_depth != 0) { + jetty_options.depth = + std::min(jetty_options.depth, device_caps.max_jetty_depth); + } + if (device_caps.max_send_sge != 0) { + jetty_options.max_sge = static_cast(std::min( + jetty_options.max_sge, device_caps.max_send_sge)); + } + if (device_caps.max_remote_sge != 0) { + jetty_options.max_sge = static_cast(std::min( + jetty_options.max_sge, device_caps.max_remote_sge)); + } + } + if (safe_slice_size == 0 || jetty_options.depth == 0 || + jetty_options.max_sge == 0) { + return failInstall(Status::InvalidArgument( + "UB device capabilities cannot support the configured data " + "path" LOC_MARK)); + } + if (safe_slice_size != params.slice_size) { + LOG(WARNING) << "Clamp UB slice_size from " << params.slice_size + << " to device limit " << safe_slice_size; + params.slice_size = safe_slice_size; + } + + buffers = std::make_unique(adapter, contexts); + endpoints = std::make_unique( + adapter, params.max_endpoints, params.jetty_per_endpoint, + jetty_options); + ub::RailMonitorConfig rail_config; + rail_config.cooldown_ns = + static_cast(params.endpoint_cooldown_ms) * 1'000'000ULL; + rails = std::make_unique(rail_config); + auto saturatedProduct = [](uint64_t lhs, uint64_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return std::numeric_limits::max(); + } + return lhs * rhs; + }; + const uint64_t path_wrs = + saturatedProduct(jetty_options.depth, params.jetty_per_endpoint); + ub::QuotaLimits path_limits{ + saturatedProduct(path_wrs, params.slice_size), path_wrs}; + quota = std::make_unique(path_limits, path_limits); + for (const auto& context : contexts) { + const auto depth = jfc_depth_by_device.at(context->topologyId()); + const uint64_t device_wrs = + saturatedProduct(depth, context->jfcs().size()); + (void)quota->setDeviceLimits( + context->topologyId(), + ub::QuotaLimits{saturatedProduct(device_wrs, params.slice_size), + device_wrs}); + } + + status = publishLocalDevices(); + if (!status.ok()) return failInstall(status); + + metadata->setBootstrapUbCallback( + [this](const UbBootstrapDesc& request, UbBootstrapDesc& response) { + return onBootstrap(request, response); + }); + callback_installed = true; + + workers = std::make_unique( + adapter, contexts, local_topology, &metadata->segmentManager(), + buffers.get(), rails.get(), quota.get(), params, + [this](const ub::EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + return resolveEndpoint(request, endpoint); + }, + [this](const std::shared_ptr& endpoint) { + if (endpoints) (void)endpoints->retire(endpoint); + }); + status = workers->start(); + if (!status.ok()) return failInstall(status); + + installed.store(true, std::memory_order_release); + return Status::OK(); + } + + Status uninstall() { + std::lock_guard lock(lifecycle_mutex); + return shutdownUnlocked(); + } + + Status failInstall(Status failure) { + (void)shutdownUnlocked(); + return failure; + } + + Status shutdownUnlocked() { + installed.store(false, std::memory_order_release); + shutting_down.store(true, std::memory_order_release); + Status first_error = Status::OK(); + auto remember = [&first_error](const Status& status) { + if (first_error.ok() && !status.ok()) first_error = status; + }; + + // Callback replacement waits for a currently executing UB bootstrap + // handler, fencing all control-plane access before resources retire. + if (callback_installed && metadata) { + metadata->setBootstrapUbCallback({}); + callback_installed = false; + } + if (workers) { + auto status = workers->stop(); + if (!status.ok()) { + // A failed native drain fence is not permission to destroy + // memory registrations, JFCs, Contexts, or the adapter. Keep + // the complete ownership graph alive so uninstall can be + // retried after a late completion/provider recovery. + return status; + } + workers.reset(); + } + if (endpoints) { + auto status = endpoints->clear(); + if (!status.ok()) return status; + endpoints.reset(); + } + if (buffers) { + auto status = buffers->clear(); + if (!status.ok()) return status; + buffers.reset(); + } + for (auto it = contexts.rbegin(); it != contexts.rend(); ++it) { + if (*it) remember((*it)->shutdown()); + } + contexts.clear(); + context_by_topology_id.clear(); + context_by_topology_name.clear(); + rails.reset(); + quota.reset(); + if (adapter_initialized && adapter) { + remember(adapter->shutdown()); + adapter_initialized = false; + } + metadata.reset(); + local_topology.reset(); + conf.reset(); + local_segment_name.clear(); + shutting_down.store(false, std::memory_order_release); + return first_error; + } + + Status publishLocalDevices() { + auto& manager = metadata->segmentManager(); + CHECK_STATUS( + manager.updateLocal([this](SegmentDesc& segment) -> Status { + if (segment.type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "Local segment is not memory-backed" LOC_MARK); + } + auto& detail = std::get(segment.detail); + detail.transport_attrs[TransportType::UB] = nlohmann::json{ + {"schema_version", 1}, + {"protocol", "urma"}, + {"notifications", false}}.dump(); + std::unordered_set existing; + for (const auto& device : detail.devices) { + existing.insert(device.name); + } + for (const auto& context : contexts) { + if (!context || !context->active()) continue; + if (existing.insert(context->deviceInfo().topology_name) + .second) { + DeviceDesc device; + device.name = context->deviceInfo().topology_name; + device.lid = 0; + device.gid.clear(); + device.transport_attrs[TransportType::UB] = + encodeDeviceMetadata(context->topologyId(), + context->deviceInfo()); + detail.devices.push_back(std::move(device)); + } + } + return Status::OK(); + })); + return manager.synchronizeLocal(); + } + + Status resolveEndpoint(const ub::EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + endpoint.reset(); + if (shutting_down.load(std::memory_order_acquire) || !endpoints || + !request.local_context || !request.remote_segment) { + return Status::InvalidArgument( + "UB transport is shutting down or route is invalid" LOC_MARK); + } + const auto* remote_nic = + request.remote_segment->getMemory().topology.getNicEntry( + request.remote_topology_id); + if (!remote_nic || remote_nic->type != Topology::NIC_UB) { + return Status::DeviceNotFound( + "Remote segment does not advertise a UB topology " + "device" LOC_MARK); + } + const std::string peer_path = + MakeNicPath(request.remote_segment->name, remote_nic->name); + ub::UbEndpointKey key{request.local_context->topologyId(), + request.remote_segment_id, + request.remote_topology_id, peer_path}; + CHECK_STATUS( + endpoints->getOrCreate(key, request.local_context, endpoint)); + if (endpoint->ready()) return Status::OK(); + + const std::string local_path = + MakeNicPath(local_segment_name, + request.local_context->deviceInfo().topology_name); + UbBootstrapDesc bootstrap; + auto status = endpoint->makeBootstrapDesc( + local_segment_name, local_path, peer_path, + request.segment_generation, bootstrap); + if (!status.ok()) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return status; + } + UbBootstrapDesc response; + status = ControlClient::bootstrapUb( + request.remote_segment->rpc_server_addr, bootstrap, response); + if (status.ok()) status = endpoint->bind(response); + if (!status.ok()) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return status; + } + return Status::OK(); + } + + int onBootstrap(const UbBootstrapDesc& request, UbBootstrapDesc& response) { + if (shutting_down.load(std::memory_order_acquire) || !endpoints) { + response.reply_msg = "UB transport is shutting down"; + return -1; + } + const std::string local_name = + getNicNameFromNicPath(request.peer_nic_path); + auto context_it = context_by_topology_name.find(local_name); + if (local_name.empty() || + context_it == context_by_topology_name.end()) { + response.reply_msg = "UB bootstrap selected an unknown local NIC"; + return -1; + } + if (request.local_device_id < 0 || request.local_eid.empty() || + request.jetty_ids.empty() || request.endpoint_generation == 0) { + response.reply_msg = "UB bootstrap request is incomplete"; + return -1; + } + + ub::UbEndpointKey key{context_it->second->topologyId(), + LOCAL_SEGMENT_ID, request.local_device_id, + request.local_nic_path}; + std::shared_ptr endpoint; + auto status = endpoints->getOrCreate(key, context_it->second, endpoint); + if (status.ok()) status = endpoint->bind(request); + if (status.ok()) { + status = endpoint->makeBootstrapDesc( + local_segment_name, request.peer_nic_path, + request.local_nic_path, currentSegmentGeneration(), response); + } + if (!status.ok()) { + if (endpoint) (void)endpoints->retire(endpoint); + response = UbBootstrapDesc{}; + response.reply_msg = status.ToString(); + return -1; + } + return 0; + } + + uint64_t currentSegmentGeneration() const { + // Buffer generations are carried in buffer metadata. The bootstrap + // field is a peer restart/refresh hint and must still be nonzero even + // before the first user buffer is registered. + return segment_generation.load(std::memory_order_relaxed); + } + + mutable std::mutex lifecycle_mutex; + std::shared_ptr adapter; + bool adapter_initialized{false}; + bool callback_installed{false}; + std::atomic installed{false}; + std::atomic shutting_down{false}; + std::atomic segment_generation{generationSeed()}; + ub::UbParams params; + std::string local_segment_name; + std::shared_ptr metadata; + std::shared_ptr local_topology; + std::shared_ptr conf; + std::vector contexts; + std::unordered_map + context_by_topology_id; + std::unordered_map context_by_topology_name; + std::unique_ptr buffers; + std::unique_ptr endpoints; + std::unique_ptr rails; + std::unique_ptr quota; + std::unique_ptr workers; +}; + +UbTransport::UbTransport(std::shared_ptr adapter) + : impl_(std::make_unique(std::move(adapter))) {} + +UbTransport::~UbTransport() { + auto status = uninstall(); + if (!status.ok()) { + // There is no caller left to retry an explicit uninstall. Leaking the + // still-live ownership graph is the only safe failure mode: destroying + // joinable pollers or registered memory after a failed device fence + // would terminate the process or permit DMA-after-free. The OS/provider + // reclaims these resources at process exit. + LOG(ERROR) << "Preserve undrained UB resources during destruction: " + << status.ToString(); + (void)impl_.release(); + } +} + +Status UbTransport::install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf) { + auto status = impl_->install(local_segment_name, std::move(metadata), + std::move(local_topology), std::move(conf)); + if (status.ok()) { + caps = Capabilities{}; + caps.dram_to_dram = true; + } + return status; +} + +Status UbTransport::uninstall() { + auto status = impl_->uninstall(); + caps = Capabilities{}; + return status; +} + +Status UbTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) { + batch = nullptr; + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->installed.load(std::memory_order_acquire)) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto* ub_batch = new (std::nothrow) UbSubBatch(); + if (!ub_batch) { + return Status::InternalError( + "Unable to allocate UB sub-batch" LOC_MARK); + } + ub_batch->max_size = max_size; + ub_batch->task_list.reserve(max_size); + batch = ub_batch; + return Status::OK(); +} + +Status UbTransport::freeSubBatch(SubBatchRef& batch) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument("Invalid UB sub-batch" LOC_MARK); + } + ub_batch->task_list.clear(); + delete ub_batch; + batch = nullptr; + return Status::OK(); +} + +Status UbTransport::submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument("Invalid UB sub-batch" LOC_MARK); + } + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->installed.load(std::memory_order_acquire) || !impl_->workers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + if (request_list.size() > ub_batch->max_size - ub_batch->task_list.size()) { + return Status::TooManyRequests("Exceed UB batch capacity" LOC_MARK); + } + for (const auto& request : request_list) { + const auto local_address = reinterpret_cast(request.source); + if (!request.source || request.length == 0 || + request.length > + std::numeric_limits::max() - request.target_offset || + request.length > + std::numeric_limits::max() - local_address) { + return Status::InvalidArgument( + "UB request range is empty or overflows" LOC_MARK); + } + } + + const auto notify_progress = ub_batch->notify_progress; + const auto progress_batch_id = ub_batch->progress_batch_id; + std::vector new_tasks; + new_tasks.reserve(request_list.size()); + for (const auto& request : request_list) { + const size_t slice_count = + request.length / impl_->params.slice_size + + (request.length % impl_->params.slice_size != 0 ? 1 : 0); + if (slice_count > impl_->params.max_slices_per_task) { + return Status::InvalidArgument( + "UB request exceeds " + "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); + }); + size_t offset = 0; + while (offset < request.length) { + const size_t length = + std::min(impl_->params.slice_size, request.length - offset); + ub::UbSliceSpec spec; + spec.local_address = static_cast(request.source) + offset; + spec.remote_address = request.target_offset + offset; + spec.length = length; + spec.request_offset = offset; + spec.max_retries = impl_->params.max_retries; + if (!task->addSlice(spec)) { + return Status::InternalError( + "Unable to construct UB slice" LOC_MARK); + } + offset += length; + } + (void)task->seal(); + 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; + } + } + return Status::OK(); +} + +Status UbTransport::getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch || task_id < 0 || + static_cast(task_id) >= ub_batch->task_list.size()) { + return Status::InvalidArgument("Invalid UB task ID" LOC_MARK); + } + status = ub_batch->task_list[task_id]->transferStatus(); + return Status::OK(); +} + +Status UbTransport::cancelTransferTask(SubBatchRef batch, int task_id) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch || task_id < 0 || + static_cast(task_id) >= ub_batch->task_list.size()) { + return Status::InvalidArgument("Invalid UB task ID" LOC_MARK); + } + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->workers) { + return Status::InvalidArgument("UB workers are not running" LOC_MARK); + } + return impl_->workers->cancel(ub_batch->task_list[task_id]); +} + +Status UbTransport::addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto status = impl_->buffers->addBuffer(desc, options); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +Status UbTransport::addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto status = impl_->buffers->addBuffers(desc_list, options); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +Status UbTransport::removeMemoryBuffer(BufferDesc& desc) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) return Status::OK(); + auto status = impl_->buffers->removeBuffer(desc); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +bool UbTransport::warmupMemory(void* addr, size_t length) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!addr || length == 0 || !impl_->adapter || impl_->contexts.empty()) { + return false; + } + ub::LocalSegmentPtr segment; + ub::SegmentOptions options; + auto status = impl_->adapter->registerLocalSegment( + impl_->contexts.front()->handle(), reinterpret_cast(addr), + length, options, segment); + if (!status.ok()) return false; + return impl_->adapter->unregisterLocalSegment(segment).ok(); +} + +double UbTransport::getEstimatedBandwidth() const { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->params.enable_bandwidth_estimation || !impl_->rails) + return -1.0; + return impl_->rails->aggregateBandwidth(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp index 0c3f4b62b8..eaaac17088 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp @@ -30,8 +30,13 @@ #include #include +#include + +#include "tent/transport/ub/device_selection.h" + #if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA #include +#include #endif namespace mooncake { @@ -352,10 +357,11 @@ class RuntimeLease { class RealContext final : public Context { public: RealContext(std::shared_ptr runtime, DeviceInfo info, - urma_context_t* native) + urma_context_t* native, uint32_t ctp_priority) : runtime_(std::move(runtime)), info_(std::move(info)), - native_(native) {} + native_(native), + ctp_priority_(ctp_priority) {} ~RealContext() override { (void)close(); } @@ -367,6 +373,12 @@ class RealContext final : public Context { urma_context_t* native() const noexcept { return native_; } + bool isBondingDevice() const noexcept { + return ::mooncake::tent::ub::isBondingDevice(info_); + } + + uint32_t ctpPriority() const noexcept { return ctp_priority_; } + Status close() { if (native_ == nullptr) return Status::OK(); const int rc = urma_delete_context(native_); @@ -379,6 +391,7 @@ class RealContext final : public Context { std::shared_ptr runtime_; DeviceInfo info_; urma_context_t* native_ = nullptr; + uint32_t ctp_priority_ = 15; }; class RealJfc final : public Jfc { @@ -554,12 +567,14 @@ class RealLocalSegment final : public LocalSegment { public: RealLocalSegment(std::shared_ptr context, urma_target_seg_t* native, uint64_t address, - uint64_t length, SegmentDescriptor descriptor) + uint64_t length, SegmentDescriptor descriptor, + urma_token_id_t* token_id) : context_(std::move(context)), native_(native), address_(address), length_(length), - descriptor_(std::move(descriptor)) {} + descriptor_(std::move(descriptor)), + token_id_(token_id) {} ~RealLocalSegment() override { (void)close(); } @@ -576,11 +591,22 @@ class RealLocalSegment final : public LocalSegment { } Status close() { - if (native_ == nullptr) return Status::OK(); - const int rc = urma_unregister_seg(native_); - if (rc != URMA_SUCCESS) return nativeError("urma_unregister_seg", rc); - native_ = nullptr; - return Status::OK(); + Status status = Status::OK(); + if (native_ != nullptr) { + const int rc = urma_unregister_seg(native_); + if (rc != URMA_SUCCESS) { + status = nativeError("urma_unregister_seg", rc); + } else { + native_ = nullptr; + } + } + + if (token_id_ != nullptr) { + (void)urma_free_token_id(token_id_); + token_id_ = nullptr; + } + + return status; } private: @@ -589,6 +615,7 @@ class RealLocalSegment final : public LocalSegment { uint64_t address_ = 0; uint64_t length_ = 0; SegmentDescriptor descriptor_; + urma_token_id_t* token_id_ = nullptr; }; class RealRemoteSegment final : public RemoteSegment { @@ -645,7 +672,12 @@ class RealJetty final : public Jetty { config.flag.bs.share_jfr = 1; config.jfs_cfg.depth = options.depth; config.jfs_cfg.trans_mode = URMA_TM_RC; - config.jfs_cfg.priority = options.priority; + config.jfs_cfg.priority = context_->isBondingDevice() + ? context_->ctpPriority() + : options.priority; + if (context_->isBondingDevice()) { + config.jfs_cfg.flag.bs.multi_path = 1; + } config.jfs_cfg.max_sge = options.max_sge; config.jfs_cfg.max_rsge = options.max_sge; config.jfs_cfg.rnr_retry = options.rnr_retry; @@ -1109,6 +1141,47 @@ class RealUrmaAdapter final : public UrmaAdapter { return nativePointerError("urma_create_context"); } + if (isBondingDevice(requested)) { + bondp_set_bonding_mode_in_t bonding_mode{}; + bonding_mode.bonding_mode = BONDP_BONDING_MODE_STANDALONE; + bonding_mode.bonding_level = BONDP_BONDING_LEVEL_IODIE; + + urma_user_ctl_in_t ctl_in{}; + ctl_in.addr = reinterpret_cast(&bonding_mode); + ctl_in.len = sizeof(bonding_mode); + ctl_in.opcode = BONDP_USER_CTL_SET_BONDING_MODE; + + urma_user_ctl_out_t ctl_out{}; + const int ctl_rc = + urma_user_ctl(native_context, &ctl_in, &ctl_out); + if (ctl_rc != URMA_SUCCESS) { + (void)urma_delete_context(native_context); + return nativeError("urma_user_ctl(SET_BONDING_MODE)", + ctl_rc); + } + } + + uint32_t ctp_priority = 15; + bool ctp_priority_found = false; + for (uint32_t priority_index = 0; + priority_index < URMA_MAX_PRIORITY_CNT; ++priority_index) { + const auto& priority_info = + attributes.dev_cap.priority_info[priority_index]; + if (priority_info.tp_type.bs.ctp != 0) { + ctp_priority = priority_index; + LOG(INFO) << "UB_PRIO_SELECTED priority=" << ctp_priority + << " SL=" << static_cast(priority_info.SL); + ctp_priority_found = true; + break; + } + } + + if (isBondingDevice(requested) && !ctp_priority_found) { + (void)urma_delete_context(native_context); + return Status::InvalidArgument( + "bonding URMA device has no CTP priority"); + } + DeviceInfo current = requested; current.native_device_path = boundedString(native_device->path, URMA_MAX_PATH); @@ -1120,7 +1193,8 @@ class RealUrmaAdapter final : public UrmaAdapter { std::to_string(current.eid_index); } output = std::make_shared( - std::move(runtime), std::move(current), native_context); + std::move(runtime), std::move(current), native_context, + ctp_priority); return Status::OK(); } return Status::DeviceNotFound("URMA device not found: " + @@ -1191,20 +1265,28 @@ class RealUrmaAdapter final : public UrmaAdapter { return Status::InvalidArgument("invalid local segment range"); } + urma_token_id_t* token_id = urma_alloc_token_id(real_context->native()); + if (token_id == nullptr) { + return nativePointerError("urma_alloc_token_id"); + } + urma_reg_seg_flag_t flags{}; flags.bs.token_policy = URMA_TOKEN_NONE; flags.bs.cacheable = options.cacheable ? URMA_CACHEABLE : URMA_NON_CACHEABLE; flags.bs.access = nativeAccess(options.access); + flags.bs.token_id_valid = 1; urma_seg_cfg_t config{}; config.va = address; config.len = length; + config.token_id = token_id; config.token_value.token = options.token; config.flag = flags; urma_target_seg_t* native_segment = urma_register_seg(real_context->native(), &config); if (native_segment == nullptr) { + (void)urma_free_token_id(token_id); return nativePointerError("urma_register_seg"); } @@ -1214,13 +1296,17 @@ class RealUrmaAdapter final : public UrmaAdapter { wire_descriptor.attr = native_segment->seg.attr; wire_descriptor.token_id = native_segment->seg.token_id; + LOG(INFO) << "UB_SEG_TOKEN_ID token_id_valid=1" + << " va=0x" << std::hex << address << " len=0x" << length + << std::dec; + SegmentDescriptor descriptor; descriptor.urma_api_version = URMA_API_VERSION; descriptor.urma_abi_size = sizeof(urma_seg_t); descriptor.hex = encodeHex(&wire_descriptor, sizeof(wire_descriptor)); output = std::make_shared( std::move(real_context), native_segment, address, length, - std::move(descriptor)); + std::move(descriptor), token_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 new file mode 100644 index 0000000000..47808fdace --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -0,0 +1,932 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/workers.h" + +#include +#include +#include +#include +#include + +#include + +#include "tent/runtime/platform.h" +#include "tent/transport/ub/endpoint.h" + +namespace mooncake::tent::ub { +namespace { + +uint64_t deadlineAfter(uint64_t now_ns, uint64_t timeout_ns) { + return timeout_ns > std::numeric_limits::max() - now_ns + ? std::numeric_limits::max() + : now_ns + timeout_ns; +} + +bool retryableStatus(const Status& status) { + return status.IsDeviceNotFound() || status.IsNeedsRefreshCache() || + status.IsRpcServiceError() || status.IsInternalError() || + status.IsRdmaError(); +} + +} // namespace + +struct UbWorkers::Route { + SegmentDescRef pin; + BufferDesc* remote_buffer{nullptr}; + UbBufferMetadata metadata; + std::vector remote_devices; +}; + +struct UbWorkers::Inflight { + uint64_t completion_token{0}; + PendingSlice pending; + UbAttemptToken attempt; + UbPostPath path; + std::shared_ptr endpoint; + LocalSegmentPtr local_segment; + RemoteSegmentPtr remote_segment; + QuotaReservation quota; + uint64_t posted_ns{0}; + uint64_t deadline_ns{0}; + std::atomic timed_out{false}; + std::atomic timeout_recorded{false}; + std::atomic resources_released{false}; +}; + +UbWorkers::UbWorkers(std::shared_ptr adapter, + std::vector contexts, + std::shared_ptr local_topology, + SegmentManager* segment_manager, UbBufferManager* buffers, + RailMonitor* rail_monitor, QuotaManager* quota, + UbParams params, EndpointResolver endpoint_resolver, + EndpointRetirer endpoint_retirer) + : adapter_(std::move(adapter)), + contexts_(std::move(contexts)), + local_topology_(std::move(local_topology)), + segment_manager_(segment_manager), + buffers_(buffers), + rail_monitor_(rail_monitor), + quota_(quota), + params_(std::move(params)), + endpoint_resolver_(std::move(endpoint_resolver)), + endpoint_retirer_(std::move(endpoint_retirer)) { + for (const auto& context : contexts_) { + if (!context) continue; + context_by_topology_id_[context->topologyId()] = context; + for (const auto& jfc : context->jfcs()) { + if (jfc) { + all_jfcs_.push_back(jfc); + context_by_jfc_[jfc.get()] = context; + } + } + } +} + +UbWorkers::~UbWorkers() { (void)stop(); } + +Status UbWorkers::start() { + bool expected = false; + if (!accepting_.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + return Status::InvalidArgument("UB workers already started" LOC_MARK); + } + if (!adapter_ || contexts_.empty() || all_jfcs_.empty() || + !segment_manager_ || !buffers_ || !rail_monitor_ || !quota_ || + !endpoint_resolver_ || params_.worker_count == 0 || + params_.poller_count == 0) { + accepting_.store(false, std::memory_order_release); + return Status::InvalidArgument( + "UB workers have incomplete dependencies" LOC_MARK); + } + + posting_.store(true, std::memory_order_release); + polling_.store(true, std::memory_order_release); + timeout_scans_enabled_.store(true, std::memory_order_release); + try { + posting_threads_.reserve(params_.worker_count); + for (uint32_t i = 0; i < params_.worker_count; ++i) { + posting_threads_.emplace_back(&UbWorkers::postingLoop, this, i); + } + const size_t poller_count = + std::min(params_.poller_count, all_jfcs_.size()); + polling_threads_.reserve(poller_count); + for (size_t i = 0; i < poller_count; ++i) { + polling_threads_.emplace_back(&UbWorkers::pollingLoop, this, i); + } + } catch (const std::exception& error) { + accepting_.store(false, std::memory_order_release); + posting_.store(false, std::memory_order_release); + polling_.store(false, std::memory_order_release); + timeout_scans_enabled_.store(false, std::memory_order_release); + queue_cv_.notify_all(); + for (auto& thread : posting_threads_) { + if (thread.joinable()) thread.join(); + } + for (auto& thread : polling_threads_) { + if (thread.joinable()) thread.join(); + } + posting_threads_.clear(); + polling_threads_.clear(); + return Status::InternalError(std::string("Cannot start UB workers: ") + + error.what() + LOC_MARK); + } + return Status::OK(); +} + +Status UbWorkers::stop() { + accepting_.store(false, std::memory_order_release); + // Fence retry enqueue before draining the queues. A completion racing + // with shutdown will turn its retry-pending slice into CANCELED instead of + // leaving a new queue entry behind the drain pass. + posting_.store(false, std::memory_order_release); + timeout_scans_enabled_.store(false, std::memory_order_release); + queue_cv_.notify_all(); + + std::vector abandoned; + { + std::lock_guard lock(queue_mutex_); + for (auto& queue : queues_) { + while (!queue.empty()) { + abandoned.push_back(std::move(queue.front())); + queue.pop_front(); + } + } + } + for (auto& pending : abandoned) { + if (pending.slice) pending.slice->requestCancellation(); + } + + for (auto& thread : posting_threads_) { + if (thread.joinable()) thread.join(); + } + posting_threads_.clear(); + + // Wait for a timeout scan that began before the flag change. scanTimeouts + // rechecks under this mutex, so no new endpoint can enter the drain set + // after this barrier and escape the shutdown snapshot. + { + std::lock_guard lock(endpoint_drain_mutex_); + } + + // Snapshot all native work after posting threads have joined; no new WR + // can cross the adapter boundary beyond this point. Quiesce each endpoint + // before stopping pollers. A failed fence leaves pollers, tokens, native + // segments, quota, and endpoints intact so uninstall can be retried + // safely; fixed sleeps are never treated as proof that DMA stopped. + std::vector> still_inflight; + std::vector> endpoints_to_fence; + { + std::lock_guard lock(inflight_mutex_); + still_inflight.reserve(inflight_.size()); + for (const auto& [_, inflight] : inflight_) { + still_inflight.push_back(inflight); + if (inflight && inflight->endpoint) { + endpoints_to_fence.push_back(inflight->endpoint); + } + } + for (const auto& [_, endpoint] : draining_endpoints_) { + if (endpoint) endpoints_to_fence.push_back(endpoint); + } + } + for (const auto& inflight : still_inflight) { + if (inflight && inflight->pending.slice) { + inflight->pending.slice->requestCancellation(); + } + } + + Status fence_error = Status::OK(); + std::unordered_set fenced; + for (const auto& endpoint : endpoints_to_fence) { + if (!endpoint || !fenced.insert(endpoint.get()).second) { + continue; + } + std::vector drained; + auto status = endpoint->quiesce(params_.slice_timeout_ms, drained); + for (const auto& completion : drained) { + if (completion.token != 0) handleCompletion(completion); + } + if (!status.ok()) { + rememberEndpointDrain(endpoint); + if (fence_error.ok()) fence_error = status; + continue; + } + forgetEndpointDrain(endpoint); + if (endpoint_retirer_) endpoint_retirer_(endpoint); + } + if (!fence_error.ok()) { + return fence_error; + } + { + std::lock_guard lock(inflight_mutex_); + if (!draining_endpoints_.empty()) { + return Status::InternalError( + "UB endpoint drain set changed during shutdown" LOC_MARK); + } + } + + // Successful endpoint fences prove that any token still missing from the + // JFC can no longer touch memory. Let pollers consume already-queued CRs, + // then stop them and safely resolve any provider-lost token. + if (!still_inflight.empty()) { + std::unique_lock lock(inflight_mutex_); + (void)inflight_cv_.wait_for(lock, std::chrono::milliseconds(100), + [this] { return inflight_.empty(); }); + } + polling_.store(false, std::memory_order_release); + for (auto& thread : polling_threads_) { + if (thread.joinable()) thread.join(); + } + polling_threads_.clear(); + + still_inflight.clear(); + { + std::lock_guard lock(inflight_mutex_); + for (auto& [_, inflight] : inflight_) { + still_inflight.push_back(std::move(inflight)); + } + inflight_.clear(); + } + for (const auto& inflight : still_inflight) { + releaseInflight(inflight); + (void)inflight->pending.slice->resolveAttempt(inflight->attempt, FAILED, + 0, false); + } + inflight_cv_.notify_all(); + return Status::OK(); +} + +Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { + if (!task) { + return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK); + } + 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()) { + pending.push_back(PendingSlice{task, slice, device_mask, priority, + task->request().target_id, + task->request().opcode}); + } + } + bool stopped_during_submit = false; + { + 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)); + } + } + } + if (stopped_during_submit) { + for (auto& item : pending) item.slice->requestCancellation(); + return Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + if (!pending.empty()) queue_cv_.notify_all(); + return Status::OK(); +} + +Status UbWorkers::cancel(const UbTask::Ptr& task) { + if (!task) { + return Status::InvalidArgument("Cannot cancel a null UB task" LOC_MARK); + } + (void)task->requestCancellation(); + queue_cv_.notify_all(); + return Status::OK(); +} + +size_t UbWorkers::queuedCount() const { + std::lock_guard lock(queue_mutex_); + size_t count = 0; + for (const auto& queue : queues_) count += queue.size(); + return count; +} + +size_t UbWorkers::inflightCount() const { + std::lock_guard lock(inflight_mutex_); + return inflight_.size(); +} + +bool UbWorkers::popPending(PendingSlice& pending) { + std::unique_lock lock(queue_mutex_); + queue_cv_.wait(lock, [this] { + if (!posting_.load(std::memory_order_acquire)) return true; + for (const auto& queue : queues_) { + if (!queue.empty()) return true; + } + return false; + }); + if (!posting_.load(std::memory_order_acquire)) return false; + for (auto& queue : queues_) { + if (!queue.empty()) { + pending = std::move(queue.front()); + queue.pop_front(); + return true; + } + } + return false; +} + +void UbWorkers::postingLoop(size_t worker_index) { + while (posting_.load(std::memory_order_acquire)) { + PendingSlice pending; + if (!popPending(pending)) continue; + if (pending.slice) processPending(pending, worker_index); + } +} + +void UbWorkers::pollingLoop(size_t poller_index) { + const size_t poller_count = std::max( + 1, std::min(params_.poller_count, all_jfcs_.size())); + uint64_t last_timeout_scan = 0; + while (polling_.load(std::memory_order_acquire)) { + bool progressed = false; + for (size_t index = poller_index; index < all_jfcs_.size(); + index += poller_count) { + std::vector completions; + auto status = all_jfcs_[index]->poll(64, completions); + if (!status.ok()) { + auto context = context_by_jfc_.find(all_jfcs_[index].get()); + if (context != context_by_jfc_.end() && context->second) { + (void)context->second->markUnavailable(); + } + LOG_EVERY_N(WARNING, 1000) + << "UB JFC poll failed: " << status.ToString(); + continue; + } + progressed = progressed || !completions.empty(); + for (const auto& completion : completions) { + if (completion.token != 0) handleCompletion(completion); + } + } + const uint64_t now = steadyNowNs(); + if (timeout_scans_enabled_.load(std::memory_order_acquire) && + now - last_timeout_scan >= 1'000'000ULL) { + scanTimeouts(); + last_timeout_scan = now; + } + if (!progressed) + std::this_thread::sleep_for(std::chrono::microseconds(20)); + } +} + +void UbWorkers::enqueueRetry(const PendingSlice& pending) { + if (!pending.slice) return; + if (!posting_.load(std::memory_order_acquire) || + !pending.slice->markQueued()) { + pending.slice->requestCancellation(); + return; + } + bool stopped = false; + { + std::lock_guard lock(queue_mutex_); + if (!posting_.load(std::memory_order_relaxed)) { + stopped = true; + } else { + queues_[std::clamp(pending.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))] + .push_back(pending); + } + } + if (stopped) { + pending.slice->requestCancellation(); + return; + } + queue_cv_.notify_one(); +} + +void UbWorkers::deferPending(const PendingSlice& pending) { + if (!pending.slice || pending.slice->cancellationRequested() || + !posting_.load(std::memory_order_acquire)) { + if (pending.slice) pending.slice->requestCancellation(); + return; + } + // Capacity pressure is not a transfer failure. A short bounded backoff + // prevents a posting lane from exhausting CPU while completions release + // quota; releaseInflight() also wakes the queue condition variable. + { + std::unique_lock lock(queue_mutex_); + (void)queue_cv_.wait_for(lock, std::chrono::microseconds(50), [&] { + return !posting_.load(std::memory_order_relaxed) || + pending.slice->cancellationRequested(); + }); + if (posting_.load(std::memory_order_relaxed) && + !pending.slice->cancellationRequested()) { + queues_[std::clamp(pending.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))] + .push_back(pending); + lock.unlock(); + queue_cv_.notify_one(); + return; + } + } + pending.slice->requestCancellation(); +} + +Status UbWorkers::buildRoute(const PendingSlice& pending, Route& route) { + return segment_manager_->withCachedSegment( + pending.target_id, route.pin, [&](SegmentDesc* segment) -> Status { + if (!segment || segment->type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "UB target is not a memory segment" LOC_MARK); + } + auto* buffer = + segment->findBuffer(pending.slice->spec().remote_address, + pending.slice->spec().length); + if (!buffer) { + return Status::NeedsRefreshCache( + "UB target range is not registered" LOC_MARK); + } + auto attr = buffer->transport_attrs.find(TransportType::UB); + if (attr == buffer->transport_attrs.end()) { + return Status::NeedsRefreshCache( + "UB target buffer has no transport metadata" LOC_MARK); + } + UbBufferMetadata metadata; + auto status = decodeBufferMetadata(attr->second, metadata); + if (!status.ok()) return status; + route.remote_buffer = buffer; + route.metadata = std::move(metadata); + route.remote_devices = orderedRemoteDevices(*segment, *buffer); + if (route.remote_devices.empty()) { + return Status::DeviceNotFound( + "UB target has no active advertised device" LOC_MARK); + } + return Status::OK(); + }); +} + +std::vector UbWorkers::orderedLocalDevices( + const PendingSlice& pending) const { + std::vector ordered; + std::unordered_set seen; + std::unordered_map device_rank; + std::string location = kWildcardLocation; + auto locations = Platform::getLoader().getLocation( + pending.slice->spec().local_address, 1, true); + if (!locations.empty()) location = locations.front().location; + + auto append = [&](Topology::NicID id) { + auto found = context_by_topology_id_.find(id); + if (found == context_by_topology_id_.end() || !found->second || + !found->second->active()) { + return; + } + const bool allowed = pending.device_mask == ~0ULL || + (id >= 0 && id < 64 && + (pending.device_mask & (uint64_t{1} << id)) != 0); + if (allowed && seen.insert(id).second) ordered.push_back(id); + }; + if (local_topology_) { + if (const auto* memory = local_topology_->getMemEntry(location)) { + for (size_t rank = 0; rank < Topology::DevicePriorityRanks; + ++rank) { + for (auto id : memory->device_list[rank]) { + auto [it, inserted] = device_rank.emplace(id, rank); + if (!inserted) it->second = std::min(it->second, rank); + append(id); + } + } + } + } + for (const auto& context : contexts_) { + if (context) { + device_rank.try_emplace(context->topologyId(), + Topology::DevicePriorityRanks); + append(context->topologyId()); + } + } + std::stable_sort(ordered.begin(), ordered.end(), [&](auto lhs, auto rhs) { + const auto lhs_rank = device_rank.at(lhs); + const auto rhs_rank = device_rank.at(rhs); + if (lhs_rank != rhs_rank) return lhs_rank < rhs_rank; + return context_by_topology_id_.at(lhs)->inflightBytes() < + context_by_topology_id_.at(rhs)->inflightBytes(); + }); + return ordered; +} + +std::vector UbWorkers::orderedRemoteDevices( + const SegmentDesc& segment, const BufferDesc& buffer) { + std::vector ordered; + auto attr = buffer.transport_attrs.find(TransportType::UB); + if (attr == buffer.transport_attrs.end()) return ordered; + UbBufferMetadata metadata; + if (!decodeBufferMetadata(attr->second, metadata).ok()) return ordered; + std::unordered_set advertised; + for (const auto& item : metadata.segments) { + advertised.insert(item.topology_id); + } + std::unordered_set seen; + const auto& topology = segment.getMemory().topology; + auto append = [&](Topology::NicID id) { + const auto* nic = topology.getNicEntry(id); + if (nic && nic->type == Topology::NIC_UB && advertised.count(id) && + seen.insert(id).second) { + ordered.push_back(id); + } + }; + if (const auto* memory = topology.getMemEntry(buffer.location)) { + for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { + for (auto id : memory->device_list[rank]) append(id); + } + } + for (const auto& item : metadata.segments) append(item.topology_id); + return ordered; +} + +Status UbWorkers::chooseAndResolveEndpoint( + const PendingSlice& pending, Route& route, + std::shared_ptr& endpoint, UbPostPath& path) { + const auto local_devices = orderedLocalDevices(pending); + if (local_devices.empty() || route.remote_devices.empty()) { + return Status::DeviceNotFound("No usable UB posting path" LOC_MARK); + } + const auto snapshot = pending.slice->snapshot(); + const size_t combinations = + local_devices.size() * route.remote_devices.size(); + const size_t start = + combinations == 0 ? 0 : snapshot.retry_count % combinations; + Status first_error = Status::DeviceNotFound( + "No ready UB endpoint for any posting path" LOC_MARK); + for (size_t candidate = 0; candidate < combinations; ++candidate) { + const size_t flat = (start + candidate) % combinations; + const auto local_id = local_devices[flat / route.remote_devices.size()]; + const auto remote_id = + route.remote_devices[flat % route.remote_devices.size()]; + auto context = context_by_topology_id_.at(local_id); + EndpointResolveRequest request{context, pending.target_id, + route.pin.get(), remote_id, + route.metadata.generation}; + std::shared_ptr candidate_endpoint; + auto status = endpoint_resolver_(request, candidate_endpoint); + if (!status.ok() || !candidate_endpoint || + !candidate_endpoint->ready()) { + if (first_error.IsDeviceNotFound() && !status.ok()) + first_error = status; + continue; + } + UbPostPath candidate_path{local_id, pending.target_id, remote_id, + candidate_endpoint->generation()}; + if (!rail_monitor_->available(candidate_path)) continue; + endpoint = std::move(candidate_endpoint); + path = candidate_path; + return Status::OK(); + } + return first_error; +} + +void UbWorkers::processPending(const PendingSlice& pending, + size_t worker_index) { + if (pending.slice->cancellationRequested()) { + pending.slice->requestCancellation(); + return; + } + Route route; + auto status = buildRoute(pending, route); + if (!status.ok()) { + const auto resolution = pending.slice->resolveBeforePost( + FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + std::shared_ptr endpoint; + UbPostPath path; + status = chooseAndResolveEndpoint(pending, route, endpoint, path); + if (!status.ok()) { + const auto resolution = pending.slice->resolveBeforePost( + FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + LocalSegmentRef local; + status = buffers_->findLocal( + reinterpret_cast(pending.slice->spec().local_address), + pending.slice->spec().length, path.local_topology_id, local); + if (!status.ok()) { + (void)pending.slice->resolveBeforePost(FAILED, 0, false); + return; + } + ImportedSegmentRef remote; + status = buffers_->importRemote(pending.target_id, path.local_topology_id, + path.remote_device_id, *route.remote_buffer, + pending.opcode, + pending.slice->spec().remote_address, + pending.slice->spec().length, remote); + if (!status.ok()) { + const bool retryable = status.IsNeedsRefreshCache(); + if (retryable && pending.target_id != LOCAL_SEGMENT_ID) { + (void)segment_manager_->invalidateRemote(pending.target_id); + } + const auto resolution = + pending.slice->resolveBeforePost(FAILED, 0, retryable); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + auto reservation = quota_->tryAcquire(path, pending.slice->spec().length); + if (!reservation) { + deferPending(pending); + return; + } + if (!endpoint->tryAcquireOutstanding(pending.slice->spec().length)) { + (void)quota_->release(*reservation); + deferPending(pending); + return; + } + + auto attempt = pending.slice->beginAttempt(path); + if (!attempt) { + endpoint->releaseOutstanding(pending.slice->spec().length); + (void)quota_->release(*reservation); + return; + } + + const uint64_t completion_token = nextCompletionToken(); + auto inflight = std::make_shared(); + inflight->completion_token = completion_token; + inflight->pending = pending; + inflight->attempt = *attempt; + inflight->path = path; + inflight->endpoint = endpoint; + inflight->local_segment = local.segment; + inflight->remote_segment = remote.segment; + inflight->quota = *reservation; + inflight->posted_ns = steadyNowNs(); + inflight->deadline_ns = deadlineAfter( + inflight->posted_ns, + static_cast(params_.slice_timeout_ms) * 1'000'000ULL); + + if (!pending.slice->tryCommitPost(*attempt, inflight->posted_ns)) { + endpoint->releaseOutstanding(pending.slice->spec().length); + (void)quota_->release(*reservation); + return; + } + { + std::lock_guard lock(inflight_mutex_); + inflight_.emplace(completion_token, inflight); + } + + WorkRequest work; + work.operation = + pending.opcode == Request::READ ? Operation::READ : Operation::WRITE; + work.local_address = + reinterpret_cast(pending.slice->spec().local_address); + work.remote_address = pending.slice->spec().remote_address; + work.length = pending.slice->spec().length; + work.token = completion_token; + work.local_segment = local.segment; + work.remote_segment = remote.segment; + auto jetty = endpoint->jetty(worker_index); + size_t posted_count = 0; + status = jetty ? adapter_->post(jetty, {work}, posted_count) + : Status::InternalError("UB endpoint has no Jetty" LOC_MARK); + if (posted_count == 1) { + rail_monitor_->registerPath(path); + if (!status.ok()) { + LOG_EVERY_N(WARNING, 1000) + << "UB post returned an error after accepting the WR: " + << status.ToString(); + } + return; + } + + if (status.ok()) { + status = Status::InternalError( + "URMA adapter accepted an unexpected WR count" LOC_MARK); + } + + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(completion_token); + if (it != inflight_.end() && it->second == inflight) + inflight_.erase(it); + } + inflight_cv_.notify_all(); + releaseInflight(inflight); + rail_monitor_->recordError(path); + if (endpoint_retirer_) endpoint_retirer_(endpoint); + const auto resolution = pending.slice->resolveAttempt( + *attempt, FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } +} + +void UbWorkers::handleCompletion(const Completion& completion) { + std::shared_ptr inflight; + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(completion.token); + if (it == inflight_.end()) return; + inflight = std::move(it->second); + inflight_.erase(it); + } + inflight_cv_.notify_all(); + releaseInflight(inflight); + if (inflight->timed_out.load(std::memory_order_acquire)) { + // A natural or flush completion proves that this particular WR can no + // longer touch memory. Whichever side wins the timeout/drain race may + // advance the logical attempt; attempt matching makes the other call + // an idempotent no-op. + const uint64_t now = steadyNowNs(); + recordTimeoutOnce(inflight, now); + resolveInflight(inflight, TIMEOUT, 0, true); + return; + } + + const uint64_t now = steadyNowNs(); + const uint64_t latency = + now >= inflight->posted_ns ? now - inflight->posted_ns : 0; + switch (completion.category) { + case CompletionCategory::SUCCESS: + rail_monitor_->recordSuccess(inflight->path, + inflight->pending.slice->spec().length, + latency, now); + resolveInflight(inflight, COMPLETED, + inflight->pending.slice->spec().length, false); + break; + case CompletionCategory::TIMEOUT: + rail_monitor_->recordTimeout(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, TIMEOUT, 0, true); + break; + case CompletionCategory::LOCAL_DEVICE_ERROR: + if (inflight->endpoint && inflight->endpoint->context()) { + (void)inflight->endpoint->context()->markUnavailable(); + } + rail_monitor_->recordError(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, FAILED, 0, true); + break; + case CompletionCategory::REMOTE_PATH_ERROR: + case CompletionCategory::ENDPOINT_ERROR: + rail_monitor_->recordError(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, FAILED, 0, true); + break; + case CompletionCategory::MEMORY_ERROR: + case CompletionCategory::UNKNOWN_ERROR: + rail_monitor_->recordError(inflight->path, now); + resolveInflight(inflight, FAILED, 0, false); + break; + } +} + +void UbWorkers::scanTimeouts() { + std::lock_guard drain_lock(endpoint_drain_mutex_); + if (!timeout_scans_enabled_.load(std::memory_order_acquire)) return; + const uint64_t now = steadyNowNs(); + std::vector> expired; + { + std::lock_guard lock(inflight_mutex_); + for (const auto& [_, inflight] : inflight_) { + bool expected = false; + if (inflight->deadline_ns <= now && + inflight->timed_out.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + expired.push_back(inflight); + } + } + } + for (const auto& inflight : expired) { + std::vector drained; + auto status = + inflight->endpoint + ? inflight->endpoint->quiesce(params_.slice_timeout_ms, drained) + : Status::InvalidArgument( + "Timed-out UB WR has no endpoint" LOC_MARK); + // Even a partial/failed fence may have returned valid WR completions; + // they must never be lost. + for (const auto& completion : drained) { + if (completion.token != 0) handleCompletion(completion); + } + if (!status.ok()) { + rememberEndpointDrain(inflight->endpoint); + LOG_EVERY_N(ERROR, 100) + << "Cannot establish UB timeout drain fence: " + << status.ToString(); + // Leave the token and all resource references alive. A natural + // completion can still resolve it safely; otherwise a later scan + // retries the native fence. + bool still_inflight = false; + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(inflight->completion_token); + still_inflight = + it != inflight_.end() && it->second == inflight; + } + if (still_inflight) { + inflight->timed_out.store(false, std::memory_order_release); + } + continue; + } + forgetEndpointDrain(inflight->endpoint); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + recordTimeoutOnce(inflight, now); + // quiesce() is the proof that retry cannot overlap old DMA. Reclaim a + // provider-lost token here: the successful fence proves the old WR can + // no longer touch memory, so releasing outstanding/quota accounting is + // safe. releaseInflight() is idempotent against a racing completion. + bool erased = false; + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(inflight->completion_token); + if (it != inflight_.end() && it->second == inflight) { + inflight_.erase(it); + erased = true; + } + } + if (erased) { + inflight_cv_.notify_all(); + releaseInflight(inflight); + } + resolveInflight(inflight, TIMEOUT, 0, true); + } +} + +void UbWorkers::releaseInflight(const std::shared_ptr& inflight) { + bool expected = false; + if (!inflight || !inflight->resources_released.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + return; + } + if (inflight->endpoint) { + inflight->endpoint->releaseOutstanding( + inflight->pending.slice->spec().length); + } + if (inflight->quota.valid()) (void)quota_->release(inflight->quota); + queue_cv_.notify_all(); +} + +void UbWorkers::recordTimeoutOnce(const std::shared_ptr& inflight, + uint64_t now_ns) { + bool expected = false; + if (inflight && inflight->timeout_recorded.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + rail_monitor_->recordTimeout(inflight->path, now_ns); + } +} + +void UbWorkers::rememberEndpointDrain( + const std::shared_ptr& endpoint) { + if (!endpoint) return; + std::lock_guard lock(inflight_mutex_); + draining_endpoints_[endpoint->generation()] = endpoint; +} + +void UbWorkers::forgetEndpointDrain( + const std::shared_ptr& endpoint) { + if (!endpoint) return; + std::lock_guard lock(inflight_mutex_); + auto it = draining_endpoints_.find(endpoint->generation()); + if (it != draining_endpoints_.end() && it->second == endpoint) { + draining_endpoints_.erase(it); + } +} + +void UbWorkers::resolveInflight(const std::shared_ptr& inflight, + TransferStatusEnum outcome, size_t bytes, + bool retryable) { + const auto resolution = inflight->pending.slice->resolveAttempt( + inflight->attempt, outcome, bytes, retryable); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(inflight->pending); + } +} + +void UbWorkers::failUnposted(const PendingSlice& pending, + TransferStatusEnum outcome) { + if (pending.slice) (void)pending.slice->tryResolveBeforePost(outcome); +} + +uint64_t UbWorkers::nextCompletionToken() { + uint64_t token = next_token_.fetch_add(1, std::memory_order_relaxed); + if (token == 0) token = next_token_.fetch_add(1, std::memory_order_relaxed); + return token; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 875a111470..754020340e 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -1,6 +1,6 @@ -# TENT Tests and Examples -# tent_metrics_example and deadline_promotion_bench are not add_test targets. -# tent/benchmark/hip_bandwidth_bench.cpp is a standalone hipcc program. +# TENT Tests and Examples tent_metrics_example and deadline_promotion_bench are +# not add_test targets. tent/benchmark/hip_bandwidth_bench.cpp is a standalone +# hipcc program. # TENT Metrics Example add_executable(tent_metrics_example tent_metrics_example.cpp) @@ -125,9 +125,9 @@ add_test(NAME tent_ip_utils_test COMMAND tent_ip_utils_test) # values with embedded NULs. Only built when hiredis is available. if(TARGET metastore_redis) add_executable(tent_redis_metastore_test redis_metastore_test.cpp) - target_link_libraries(tent_redis_metastore_test - PRIVATE metastore_redis tent_common gtest gtest_main - glog) + target_link_libraries( + tent_redis_metastore_test PRIVATE metastore_redis tent_common gtest + gtest_main glog) target_include_directories(tent_redis_metastore_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_redis_metastore_test COMMAND tent_redis_metastore_test) @@ -184,6 +184,62 @@ target_include_directories(tent_tcp_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_tcp_transport_test COMMAND tent_tcp_transport_test) +add_executable(tent_hp_tcp_transport_config_test + hp_tcp_transport_config_test.cpp) +target_link_libraries(tent_hp_tcp_transport_config_test + PRIVATE gtest gtest_main tent_link_group) +target_include_directories(tent_hp_tcp_transport_config_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_transport_config_test + COMMAND tent_hp_tcp_transport_config_test) + +add_executable(tent_hp_tcp_workers_test hp_tcp_workers_test.cpp) +target_link_libraries(tent_hp_tcp_workers_test PRIVATE gtest gtest_main + tent_xport_hp_tcp_core) +target_include_directories(tent_hp_tcp_workers_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_workers_test COMMAND tent_hp_tcp_workers_test) + +add_executable(tent_hp_tcp_protocol_test hp_tcp_protocol_test.cpp) +target_link_libraries(tent_hp_tcp_protocol_test PRIVATE gtest gtest_main + tent_xport_hp_tcp_core) +target_include_directories(tent_hp_tcp_protocol_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_protocol_test COMMAND tent_hp_tcp_protocol_test) + +add_executable(tent_hp_tcp_buffer_registry_test hp_tcp_buffer_registry_test.cpp) +target_link_libraries(tent_hp_tcp_buffer_registry_test + PRIVATE gtest gtest_main tent_xport_hp_tcp_core) +target_include_directories(tent_hp_tcp_buffer_registry_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_buffer_registry_test + COMMAND tent_hp_tcp_buffer_registry_test) + +add_executable(tent_hp_tcp_socket_test hp_tcp_socket_test.cpp) +target_link_libraries(tent_hp_tcp_socket_test PRIVATE gtest gtest_main + tent_xport_hp_tcp_core) +target_include_directories(tent_hp_tcp_socket_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_socket_test COMMAND tent_hp_tcp_socket_test) + +add_executable(tent_hp_tcp_transport_test hp_tcp_transport_test.cpp) +target_link_libraries(tent_hp_tcp_transport_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_hp_tcp_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_hp_tcp_transport_test COMMAND tent_hp_tcp_transport_test) + +add_executable(tent_hp_tcp_e2e_test hp_tcp_e2e_test.cpp) +target_link_libraries(tent_hp_tcp_e2e_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_hp_tcp_e2e_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test( + NAME tent_hp_tcp_e2e_test_concurrency_16 + COMMAND tent_hp_tcp_e2e_test + --gtest_filter=HighPerformanceTcpE2eTest.WriteThenReadConcurrency16) +set_tests_properties(tent_hp_tcp_e2e_test_concurrency_16 PROPERTIES TIMEOUT 120) + add_executable(tent_tcp_datapath_roundtrip_test tcp_datapath_roundtrip_test.cpp) target_link_libraries(tent_tcp_datapath_roundtrip_test PRIVATE gtest gtest_main tent_link_group) @@ -283,6 +339,25 @@ target_include_directories(tent_rdma_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rdma_transport_test COMMAND tent_rdma_transport_test) +# Regression test for the edge-triggered async event fd: one epoll wakeup must +# drain the whole queue. Wraps ibv_get_async_event to script the event source, +# so it needs no RDMA device. -Wl,--wrap is a GNU ld / lld feature that Apple's +# linker lacks, and the test is meaningless without it, so skip the target there +# rather than build one that calls the real symbols. +if(UNIX AND NOT APPLE) + add_executable(tent_rdma_async_event_drain_test + rdma_async_event_drain_test.cpp) + target_link_libraries(tent_rdma_async_event_drain_test + PRIVATE gtest gtest_main tent_link_group ibverbs) + target_include_directories(tent_rdma_async_event_drain_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_link_options( + tent_rdma_async_event_drain_test PRIVATE "-Wl,--wrap=ibv_get_async_event" + "-Wl,--wrap=ibv_ack_async_event") + add_test(NAME tent_rdma_async_event_drain_test + COMMAND tent_rdma_async_event_drain_test) +endif() + if(USE_HIP) find_package(HIP REQUIRED) add_executable(tent_rocm_platform_test rocm_platform_test.cpp) @@ -382,7 +457,22 @@ target_include_directories(tent_topology_priority_matrix_test add_test(NAME tent_topology_priority_matrix_test COMMAND tent_topology_priority_matrix_test) +add_executable(tent_strict_local_numa_test strict_local_numa_test.cpp) +target_link_libraries(tent_strict_local_numa_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_strict_local_numa_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_strict_local_numa_test COMMAND tent_strict_local_numa_test) + if(USE_UB) + add_executable( + tent_ub_core_test ub_core_test.cpp ../src/transport/ub/quota.cpp + ../src/transport/ub/rail_monitor.cpp) + target_link_libraries(tent_ub_core_test PRIVATE tent_common gtest gtest_main) + target_include_directories(tent_ub_core_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_test(NAME tent_ub_core_test COMMAND tent_ub_core_test) + add_executable( tent_ub_teardown_test ub_teardown_test.cpp ../src/transport/ub/buffers.cpp @@ -392,6 +482,14 @@ if(USE_UB) target_include_directories(tent_ub_teardown_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_ub_teardown_test COMMAND tent_ub_teardown_test) + + add_executable(tent_ub_native_data_path_test ub_native_data_path_test.cpp) + target_link_libraries(tent_ub_native_data_path_test PRIVATE gtest gtest_main + tent_link_group) + target_include_directories(tent_ub_native_data_path_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_test(NAME tent_ub_native_data_path_test + COMMAND tent_ub_native_data_path_test) endif() # End-to-end failover test: drives real TransferEngineImpl with @@ -508,6 +606,18 @@ target_include_directories(tent_runtime_queue_dispatch_test add_test(NAME tent_runtime_queue_dispatch_test COMMAND tent_runtime_queue_dispatch_test) +add_executable(tent_local_memory_lifecycle_test + local_memory_lifecycle_test.cpp) +target_link_libraries(tent_local_memory_lifecycle_test + PRIVATE gtest gtest_main tent_link_group) +if(TARGET asio_shared) + target_link_libraries(tent_local_memory_lifecycle_test PRIVATE asio_shared) +endif() +target_include_directories(tent_local_memory_lifecycle_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_local_memory_lifecycle_test + COMMAND tent_local_memory_lifecycle_test) + # Causal chain stage decomposition: validates that dispatch_time and post_time # timestamps are populated on both queue and direct-commit paths. add_executable(causal_chain_test causal_chain_test.cpp) @@ -557,9 +667,9 @@ if(USE_TPU) add_test(NAME tent_tpu_transport_test COMMAND tent_tpu_transport_test) endif() -# NOT registered via add_test: the test needs a real CUDA device (it exports -# a real CUDA IPC handle), so it must be run manually on a GPU host. Only -# built when CUDA is enabled: the test includes CUDA headers directly. +# NOT registered via add_test: the test needs a real CUDA device (it exports a +# real CUDA IPC handle), so it must be run manually on a GPU host. Only built +# when CUDA is enabled: the test includes CUDA headers directly. if(USE_CUDA) add_executable(tent_nvlink_ipc_handle_sharing_test nvlink_ipc_handle_sharing_test.cpp) diff --git a/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp index 32ae7c72b3..34c70dd7b6 100644 --- a/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp +++ b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp @@ -14,7 +14,10 @@ #include +#include #include +#include +#include #include "tent/common/utils/string_builder.h" #include "tent/transport/rdma/endpoint.h" @@ -45,6 +48,28 @@ class EndpointTestAccess { static bool notifyConnected(const RdmaEndPoint& endpoint) { return endpoint.notify_connected_.load(std::memory_order_relaxed); } + + static constexpr size_t notifyBufferSize() { + return RdmaEndPoint::kNotifyBufferSize; + } + + static constexpr size_t notifyMaxPendingSends() { + return RdmaEndPoint::kNotifyMaxPendingSends; + } + + static char* notifySlotPtr(char* base, size_t idx) { + return RdmaEndPoint::notifySlotPtr(base, idx); + } + + static bool encodeNotifyPayload(char* slot, const std::string& name, + const std::string& msg, uint32_t* out_len) { + return RdmaEndPoint::encodeNotifyPayload(slot, name, msg, out_len); + } + + static bool decodeNotifyPayload(const char* data, size_t byte_len, + std::string* name, std::string* msg) { + return RdmaEndPoint::decodeNotifyPayload(data, byte_len, name, msg); + } }; namespace { @@ -113,6 +138,77 @@ TEST(EndpointLifecycleTest, NotificationFailsWhenEndpointIsNotConnected) { EXPECT_FALSE(endpoint.sendNotification("name", "message")); } +TEST(EndpointLifecycleTest, NotifySlotsAreAdjacentAndNonOverlapping) { + std::vector buf(EndpointTestAccess::notifyMaxPendingSends() * + EndpointTestAccess::notifyBufferSize()); + char* slot0 = EndpointTestAccess::notifySlotPtr(buf.data(), 0); + char* slot1 = EndpointTestAccess::notifySlotPtr(buf.data(), 1); + char* last = EndpointTestAccess::notifySlotPtr( + buf.data(), EndpointTestAccess::notifyMaxPendingSends() - 1); + EXPECT_EQ(slot0, buf.data()); + EXPECT_EQ(static_cast(slot1 - slot0), + EndpointTestAccess::notifyBufferSize()); + EXPECT_EQ(static_cast(last - slot0), + (EndpointTestAccess::notifyMaxPendingSends() - 1) * + EndpointTestAccess::notifyBufferSize()); + EXPECT_EQ(last + EndpointTestAccess::notifyBufferSize(), + buf.data() + buf.size()); +} + +TEST(EndpointLifecycleTest, NotifyAdjacentSlotsDoNotCrosstalk) { + // Recv used to be 256 independent vectors. After coalescing, a wrong + // offset would mix payload from a neighbor slot. + const size_t slot_size = EndpointTestAccess::notifyBufferSize(); + const size_t nslots = EndpointTestAccess::notifyMaxPendingSends(); + std::vector buf(nslots * slot_size, '\xee'); + const size_t indices[] = {0, 1, 2, 127, 254, 255}; + uint32_t encoded[6] = {}; + for (size_t i = 0; i < 6; ++i) { + const size_t idx = indices[i]; + std::string name = "n" + std::to_string(idx); + std::string msg = + "payload-" + std::to_string(idx) + std::string(2048, 'x'); + ASSERT_TRUE(EndpointTestAccess::encodeNotifyPayload( + EndpointTestAccess::notifySlotPtr(buf.data(), idx), name, msg, + &encoded[i])); + ASSERT_GT(encoded[i], 8u); + ASSERT_LE(encoded[i], slot_size); + } + for (size_t i = 0; i < 6; ++i) { + const size_t idx = indices[i]; + std::string name; + std::string msg; + ASSERT_TRUE(EndpointTestAccess::decodeNotifyPayload( + EndpointTestAccess::notifySlotPtr(buf.data(), idx), encoded[i], + &name, &msg)) + << "slot " << idx; + EXPECT_EQ(name, "n" + std::to_string(idx)); + EXPECT_EQ(msg, + "payload-" + std::to_string(idx) + std::string(2048, 'x')); + // Trailing bytes in the 64KiB slot stay the fill pattern, so a + // too-long read would have picked up 0xee as name_len. + ASSERT_TRUE(EndpointTestAccess::decodeNotifyPayload( + EndpointTestAccess::notifySlotPtr(buf.data(), idx), slot_size, + &name, &msg)); + EXPECT_EQ(name, "n" + std::to_string(idx)); + } + // Slot 3 was never written; decoding the fill pattern must fail rather + // than returning a neighbor's payload. + std::string name; + std::string msg; + EXPECT_FALSE(EndpointTestAccess::decodeNotifyPayload( + EndpointTestAccess::notifySlotPtr(buf.data(), 3), slot_size, &name, + &msg)); +} + +TEST(EndpointLifecycleTest, NotifyPayloadRejectedWhenLargerThanSlot) { + std::vector slot(EndpointTestAccess::notifyBufferSize()); + uint32_t encoded = 0; + std::string too_big(EndpointTestAccess::notifyBufferSize(), 'z'); + EXPECT_FALSE(EndpointTestAccess::encodeNotifyPayload(slot.data(), "n", + too_big, &encoded)); +} + TEST(EndpointLifecycleTest, NotifyLocalFaultKeepsEndpointServingData) { // A fault confined to the notify QP must not retire the endpoint: doing so // moves every data QP to ERR and flushes in-flight transfers. diff --git a/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp b/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp index e5149dcdbc..9fe8946661 100644 --- a/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp +++ b/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp @@ -202,6 +202,100 @@ class FakeTransport : public Transport { bool force_submit_fail_; }; +class HpTcpRecoveryTransport : public FakeTransport { + public: + explicit HpTcpRecoveryTransport(bool permanent_failure = false) + : FakeTransport(HP_TCP), permanent_failure_(permanent_failure) {} + + std::atomic retry_calls{0}; + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& /*options*/) override { + if (std::find(desc.transports.begin(), desc.transports.end(), HP_TCP) == + desc.transports.end()) { + // HP TCP is first so an unhinted request exercises the HP TCP + // failure classification before the fallback transport. + desc.transports.insert(desc.transports.begin(), HP_TCP); + } + return Status::OK(); + } + + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override { + for (auto& desc : desc_list) { + CHECK_STATUS(addMemoryBuffer(desc, options)); + } + return Status::OK(); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + ++status_calls; + auto* hp_batch = static_cast(batch); + if (task_id < 0 || + task_id >= static_cast(hp_batch->statuses.size())) { + return Status::InvalidArgument("bad HP TCP task_id" LOC_MARK); + } + + if (permanent_failure_) { + status = {FAILED, 0}; + return Status::InvalidArgument( + "HP TCP WRITE outcome is unknown" LOC_MARK); + } + if (retry_calls.load(std::memory_order_acquire) != 0) { + status = {COMPLETED, hp_batch->requests[task_id].length}; + return Status::OK(); + } + + status = {FAILED, 0}; + return Status::NeedsRefreshCache( + "remote HP TCP metadata is stale" LOC_MARK); + } + + Status retryTransferTask(SubBatchRef batch, int task_id, + const Request& request) override { + auto* hp_batch = static_cast(batch); + if (task_id < 0 || + task_id >= static_cast(hp_batch->statuses.size())) { + return Status::InvalidArgument("bad HP TCP retry task_id" LOC_MARK); + } + if (request.source != hp_batch->requests[task_id].source || + request.target_id != hp_batch->requests[task_id].target_id || + request.target_offset != + hp_batch->requests[task_id].target_offset || + request.length != hp_batch->requests[task_id].length) { + return Status::InvalidArgument( + "HP TCP retry changed the logical request" LOC_MARK); + } + ++retry_calls; + return Status::OK(); + } + + const char* getName() const override { return ""; } + + private: + bool permanent_failure_; +}; + +class HpTcpStatusErrorOnceTransport : public FakeTransport { + public: + HpTcpStatusErrorOnceTransport() : FakeTransport(HP_TCP) {} + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + if (poll_count_.fetch_add(1, std::memory_order_relaxed) == 0) { + // Deliberately leave status untouched: Transport does not promise + // a valid output when it returns an error. + return Status::InternalError( + "injected HP TCP status poll failure" LOC_MARK); + } + return FakeTransport::getTransferStatus(batch, task_id, status); + } + + private: + std::atomic poll_count_{0}; +}; + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -255,6 +349,50 @@ TransferStatus pollUntilDone( return ts; } +struct HpTcpRecoveryBatch { + std::shared_ptr hp_tcp; + std::shared_ptr fallback_tcp; + std::vector buffer; + BatchID batch_id{0}; +}; + +void submitHpTcpRecoveryBatch(TransferEngineImpl& engine, + HpTcpRecoveryBatch& batch, + bool permanent_failure = false) { + batch.hp_tcp = std::make_shared(permanent_failure); + batch.fallback_tcp = std::make_shared(TCP); + + std::string segment_name = engine.getSegmentName(); + ASSERT_TRUE(batch.hp_tcp->install(segment_name, nullptr, nullptr).ok()); + ASSERT_TRUE( + batch.fallback_tcp->install(segment_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(HP_TCP, batch.hp_tcp); + engine.swapTransportForTest(TCP, batch.fallback_tcp); + + batch.buffer.assign(4096, 0xA5); + ASSERT_TRUE( + engine.registerLocalMemory(batch.buffer.data(), batch.buffer.size()) + .ok()); + batch.batch_id = engine.allocateBatch(1); + ASSERT_NE(batch.batch_id, static_cast(0)); + + Request request; + request.opcode = Request::WRITE; + request.source = batch.buffer.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(batch.buffer.data()); + request.length = batch.buffer.size(); + ASSERT_TRUE(engine.submitTransfer(batch.batch_id, {request}).ok()); +} + +void releaseHpTcpRecoveryBatch(TransferEngineImpl& engine, + HpTcpRecoveryBatch& batch) { + EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(batch.buffer.data(), batch.buffer.size()) + .ok()); +} + struct CorruptedRdmaBatch { std::shared_ptr fake_rdma; std::shared_ptr fake_tcp; @@ -301,6 +439,81 @@ void submitCorruptedRdmaBatch(TransferEngineImpl& engine, // mid-transfer). Engine must failover. // --------------------------------------------------------------------------- +TEST(EngineFailoverE2E, HpTcpStaleMetadataRetriesSameTransportOnce) { + auto config = makeMinimalP2PConfig(); + TransferEngineImpl engine(config); + ASSERT_TRUE(engine.available()); + + HpTcpRecoveryBatch batch; + submitHpTcpRecoveryBatch(engine, batch); + + const TransferStatus final_status = + pollUntilDone(engine, batch.batch_id, 0); + EXPECT_EQ(final_status.s, COMPLETED); + EXPECT_EQ(batch.hp_tcp->submit_calls.load(), 1); + EXPECT_EQ(batch.hp_tcp->retry_calls.load(), 1); + EXPECT_EQ(batch.fallback_tcp->submit_calls.load(), 0); + + releaseHpTcpRecoveryBatch(engine, batch); +} + +TEST(EngineFailoverE2E, HpTcpPermanentFailureDoesNotFailOver) { + auto config = makeMinimalP2PConfig(); + TransferEngineImpl engine(config); + ASSERT_TRUE(engine.available()); + + HpTcpRecoveryBatch batch; + submitHpTcpRecoveryBatch(engine, batch, /*permanent_failure=*/true); + + const TransferStatus final_status = + pollUntilDone(engine, batch.batch_id, 0); + EXPECT_EQ(final_status.s, FAILED); + EXPECT_EQ(batch.hp_tcp->submit_calls.load(), 1); + EXPECT_EQ(batch.hp_tcp->retry_calls.load(), 0); + EXPECT_EQ(batch.fallback_tcp->submit_calls.load(), 0); + + releaseHpTcpRecoveryBatch(engine, batch); +} + +TEST(EngineFailoverE2E, HpTcpPollErrorDoesNotInspectStaleStatusOutput) { + auto config = makeMinimalP2PConfig(); + TransferEngineImpl engine(config); + ASSERT_TRUE(engine.available()); + + auto hp_tcp = std::make_shared(); + std::string segment_name = engine.getSegmentName(); + ASSERT_TRUE(hp_tcp->install(segment_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(HP_TCP, hp_tcp); + + std::vector buffer(4096, 0x5A); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + const BatchID batch_id = engine.allocateBatch(1); + ASSERT_NE(batch_id, static_cast(0)); + + Request request{}; + request.opcode = Request::WRITE; + request.source = buffer.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(buffer.data()); + request.length = buffer.size(); + request.transport_hint = HP_TCP; + ASSERT_TRUE(engine.submitTransfer(batch_id, {request}).ok()); + + TransferStatus status{FAILED, 123}; + const Status first = engine.getTransferStatus(batch_id, 0, status); + EXPECT_TRUE(first.IsInternalError()) << first.ToString(); + EXPECT_EQ(status.s, PENDING); + EXPECT_EQ(status.transferred_bytes, 0U); + + ASSERT_TRUE(engine.getTransferStatus(batch_id, 0, status).ok()); + EXPECT_EQ(status.s, COMPLETED); + EXPECT_EQ(status.transferred_bytes, request.length); + + EXPECT_TRUE(engine.freeBatch(batch_id).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + TEST(EngineFailoverE2E, StatusCorruptionTriggersFailoverToSecondary) { auto cfg = makeMinimalP2PConfig(); TransferEngineImpl engine(cfg); @@ -622,6 +835,22 @@ TEST(EngineFailoverE2E, OverallStatusUsesWorstFailureNotGenericFailed) { engine.unregisterLocalMemory(completed_buf.data(), kBufLen).ok()); } +TEST(TransferStatusSeverityTest, KnownRanksMatchFormerMap) { + EXPECT_EQ(transferStatusSeverity(INITIAL), 0); + EXPECT_EQ(transferStatusSeverity(PENDING), 0); + EXPECT_EQ(transferStatusSeverity(COMPLETED), 0); + EXPECT_EQ(transferStatusSeverity(INVALID), 1); + EXPECT_EQ(transferStatusSeverity(CANCELED), 2); + EXPECT_EQ(transferStatusSeverity(TIMEOUT), 3); + EXPECT_EQ(transferStatusSeverity(FAILED), 4); +} + +TEST(TransferStatusSeverityTest, UnknownValueRanksWithFailedAndDoesNotThrow) { + auto unknown = static_cast(99); + EXPECT_EQ(transferStatusSeverity(unknown), transferStatusSeverity(FAILED)); + EXPECT_GT(transferStatusSeverity(unknown), transferStatusSeverity(TIMEOUT)); +} + TEST(EngineFailoverE2E, WaitTransferCompletionUsesProgressBatchWhenPollDisabled) { auto cfg = makeMinimalP2PConfig(); diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_buffer_registry_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_buffer_registry_test.cpp new file mode 100644 index 0000000000..5a337bc6b6 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_buffer_registry_test.cpp @@ -0,0 +1,111 @@ +// Copyright 2026 KVCache.AI +#include + +#include +#include +#include + +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" + +namespace mooncake::tent { +namespace { + +TEST(HighPerformanceTcpBufferRegistryTest, EnforcesPermissionAndRegistration) { + HighPerformanceTcpBufferRegistry registry; + std::array data{}; + const uint64_t base = reinterpret_cast(data.data()); + uint64_t id = 0; + ASSERT_TRUE(registry.add(base, data.size(), kGlobalReadOnly, &id).ok()); + + HighPerformanceTcpBufferRegistry::Lease lease; + HighPerformanceTcpStatus failure; + ASSERT_TRUE(registry + .acquireRemoteLease(base, data.size(), id, + HighPerformanceTcpOpcode::kRead, &lease, + &failure) + .ok()); + lease.reset(); + EXPECT_FALSE(registry + .acquireRemoteLease(base, data.size(), id, + HighPerformanceTcpOpcode::kWrite, + &lease, &failure) + .ok()); + EXPECT_EQ(failure, HighPerformanceTcpStatus::kPermissionDenied); + EXPECT_FALSE(registry + .acquireRemoteLease(base, data.size(), id + 1, + HighPerformanceTcpOpcode::kRead, + &lease, &failure) + .ok()); + EXPECT_EQ(failure, HighPerformanceTcpStatus::kStaleRegistration); +} + +TEST(HighPerformanceTcpBufferRegistryTest, + RejectsRegistrationFromPreviousRegistryIncarnation) { + std::array data{}; + const uint64_t base = reinterpret_cast(data.data()); + + uint64_t stale_id = 0; + { + HighPerformanceTcpBufferRegistry previous; + ASSERT_TRUE( + previous.add(base, data.size(), kGlobalReadWrite, &stale_id).ok()); + } + + HighPerformanceTcpBufferRegistry current; + uint64_t current_id = 0; + ASSERT_TRUE( + current.add(base, data.size(), kGlobalReadWrite, ¤t_id).ok()); + ASSERT_NE(stale_id, current_id); + + HighPerformanceTcpBufferRegistry::Lease lease; + HighPerformanceTcpStatus failure; + EXPECT_FALSE(current + .acquireRemoteLease(base, data.size(), stale_id, + HighPerformanceTcpOpcode::kRead, + &lease, &failure) + .ok()); + EXPECT_EQ(failure, HighPerformanceTcpStatus::kStaleRegistration); + EXPECT_TRUE(current + .acquireRemoteLease(base, data.size(), current_id, + HighPerformanceTcpOpcode::kRead, &lease, + &failure) + .ok()); +} + +TEST(HighPerformanceTcpBufferRegistryTest, UnregisterWaitsForActiveLease) { + HighPerformanceTcpBufferRegistry registry; + std::array data{}; + const uint64_t base = reinterpret_cast(data.data()); + uint64_t id = 0; + ASSERT_TRUE(registry.add(base, data.size(), kGlobalReadWrite, &id).ok()); + HighPerformanceTcpBufferRegistry::Lease lease; + ASSERT_TRUE(registry.acquireLocalLease(base, data.size(), &lease).ok()); + + std::atomic done{false}; + std::thread remover([&] { + EXPECT_TRUE(registry.remove(base, data.size()).ok()); + done = true; + }); + while (registry.tracks(base, data.size())) std::this_thread::yield(); + EXPECT_FALSE(done.load()); + lease.reset(); + remover.join(); + EXPECT_TRUE(done.load()); +} + +TEST(HighPerformanceTcpBufferRegistryTest, CloseRejectsNewWork) { + HighPerformanceTcpBufferRegistry registry; + std::array data{}; + uint64_t id = 0; + ASSERT_TRUE(registry + .add(reinterpret_cast(data.data()), data.size(), + kGlobalReadWrite, &id) + .ok()); + registry.close(); + EXPECT_TRUE( + registry.add(0x1000, 8, kGlobalReadWrite, nullptr).IsTooManyRequests()); + EXPECT_FALSE(registry.reopen().ok()); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_e2e_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_e2e_test.cpp new file mode 100644 index 0000000000..862d94508c --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_e2e_test.cpp @@ -0,0 +1,263 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/transfer_engine.h" + +namespace mooncake::tent { +namespace { + +constexpr size_t kDataLength = 256 * 1024; + +bool ReadExactly(int fd, void* buffer, size_t length) { + auto* cursor = static_cast(buffer); + size_t received = 0; + while (received < length) { + ssize_t read_result = 0; + do { + read_result = read(fd, cursor + received, length - received); + } while (read_result < 0 && errno == EINTR); + if (read_result <= 0) return false; + received += static_cast(read_result); + } + return true; +} + +class ChildProcessGuard { + public: + ChildProcessGuard(pid_t pid, int stop_fd) : pid_(pid), stop_fd_(stop_fd) {} + + ~ChildProcessGuard() { + if (pid_ <= 0) return; + close(stop_fd_); + (void)waitpid(pid_, nullptr, 0); + } + + int finish() { + close(stop_fd_); + int status = 0; + (void)waitpid(pid_, &status, 0); + pid_ = -1; + stop_fd_ = -1; + return status; + } + + private: + pid_t pid_; + int stop_fd_; +}; + +std::shared_ptr MakeHpConfig() { + auto config = std::make_shared(); + config->set("metadata_type", "p2p"); + config->set("metadata_servers", "P2PHANDSHAKE"); + config->set("transports/tcp/enable", false); + config->set("transports/hp_tcp/enable", true); + config->set("transports/hp_tcp/bind_address", "127.0.0.1"); + config->set("transports/hp_tcp/advertise_address", "127.0.0.1"); + config->set("transports/hp_tcp/port", 0); + config->set("transports/hp_tcp/worker_count", 4); + config->set("transports/hp_tcp/connections_per_peer", 4); + config->set("transports/hp_tcp/max_outstanding_tasks", 128); + config->set("transports/hp_tcp/max_outstanding_bytes", 64ULL << 20); + config->set("transports/hp_tcp/max_transfer_bytes", 8ULL << 20); + config->set("transports/hp_tcp/connect_timeout_ms", 2000); + config->set("transports/hp_tcp/progress_timeout_ms", 5000); + config->set("transports/rdma/enable", false); + config->set("transports/shm/enable", false); + config->set("rpc_server_threads", 1); + return config; +} + +bool WaitBatchDone(TransferEngine& engine, BatchID batch) { + TransferStatus status; + for (int i = 0; i < 10000; ++i) { + const Status result = engine.getTransferStatus(batch, status); + if (!result.ok() || status.s == FAILED || status.s == CANCELED || + status.s == TIMEOUT || status.s == INVALID) { + return false; + } + if (status.s == COMPLETED) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return false; +} + +void RunWriteThenReadAcrossProcesses(size_t task_count) { + const size_t remote_buffer_length = task_count * kDataLength; + const size_t local_buffer_length = 2 * remote_buffer_length; + + int ready_pipe[2]; + int stop_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0); + ASSERT_EQ(pipe(stop_pipe), 0); + + const pid_t child = fork(); + ASSERT_GE(child, 0); + if (child == 0) { + close(ready_pipe[0]); + close(stop_pipe[1]); + + int exit_code = 0; + { + TransferEngine server(MakeHpConfig()); + if (!server.available()) { + exit_code = 2; + } else { + std::vector remote(remote_buffer_length, 0); + const Status registered = server.registerLocalMemory( + remote.data(), remote.size(), kGlobalReadWrite); + if (!registered.ok()) { + exit_code = 3; + } else { + const std::string segment = server.getSegmentName(); + const uint32_t length = + static_cast(segment.size()); + if (write(ready_pipe[1], &length, sizeof(length)) != + static_cast(sizeof(length))) { + exit_code = 4; + } else if (write(ready_pipe[1], segment.data(), length) != + static_cast(length)) { + exit_code = 5; + } else { + char stop = 0; + (void)ReadExactly(stop_pipe[0], &stop, 1); + } + if (!server + .unregisterLocalMemory(remote.data(), + remote.size()) + .ok() && + exit_code == 0) { + exit_code = 6; + } + } + } + } + close(ready_pipe[1]); + close(stop_pipe[0]); + _exit(exit_code); + } + + close(ready_pipe[1]); + close(stop_pipe[0]); + ChildProcessGuard child_guard(child, stop_pipe[1]); + + uint32_t segment_length = 0; + if (!ReadExactly(ready_pipe[0], &segment_length, sizeof(segment_length))) { + close(ready_pipe[0]); + const int status = child_guard.finish(); + FAIL() << "HP TCP server initialization failed, child status " + << status; + } + std::string server_segment(segment_length, '\0'); + if (!ReadExactly(ready_pipe[0], server_segment.data(), segment_length)) { + close(ready_pipe[0]); + const int status = child_guard.finish(); + FAIL() << "Failed to read the HP TCP server segment, child status " + << status; + } + close(ready_pipe[0]); + + TransferEngine client(MakeHpConfig()); + ASSERT_TRUE(client.available()); + std::vector local(local_buffer_length, 0); + for (size_t task = 0; task < task_count; ++task) { + uint8_t* source = local.data() + task * kDataLength; + for (size_t i = 0; i < kDataLength; ++i) { + source[i] = static_cast((task * 131 + i * 7) & 0xff); + } + } + ASSERT_TRUE( + client.registerLocalMemory(local.data(), local.size(), kGlobalReadWrite) + .ok()); + + SegmentID segment = 0; + Status result; + for (int i = 0; i < 100; ++i) { + result = client.openSegment(segment, server_segment); + if (result.ok()) break; + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_TRUE(result.ok()) << result.ToString(); + + SegmentInfo info; + ASSERT_TRUE(client.getSegmentInfo(segment, info).ok()); + ASSERT_FALSE(info.buffers.empty()); + + std::vector writes; + std::vector reads; + writes.reserve(task_count); + reads.reserve(task_count); + for (size_t task = 0; task < task_count; ++task) { + Request write_request{}; + write_request.opcode = Request::WRITE; + write_request.source = local.data() + task * kDataLength; + write_request.target_id = segment; + write_request.target_offset = info.buffers[0].base + task * kDataLength; + write_request.length = kDataLength; + write_request.transport_hint = HP_TCP; + writes.push_back(write_request); + + Request read_request{}; + read_request.opcode = Request::READ; + read_request.source = + local.data() + remote_buffer_length + task * kDataLength; + read_request.target_id = segment; + read_request.target_offset = info.buffers[0].base + task * kDataLength; + read_request.length = kDataLength; + read_request.transport_hint = HP_TCP; + reads.push_back(read_request); + } + + BatchID batch = client.allocateBatch(task_count); + ASSERT_TRUE(client.submitTransfer(batch, writes).ok()); + ASSERT_TRUE(WaitBatchDone(client, batch)); + ASSERT_TRUE(client.freeBatch(batch).ok()); + + batch = client.allocateBatch(task_count); + ASSERT_TRUE(client.submitTransfer(batch, reads).ok()); + ASSERT_TRUE(WaitBatchDone(client, batch)); + ASSERT_TRUE(client.freeBatch(batch).ok()); + + for (size_t task = 0; task < task_count; ++task) { + const uint8_t* written = local.data() + task * kDataLength; + const uint8_t* read_back = + local.data() + remote_buffer_length + task * kDataLength; + EXPECT_EQ(std::memcmp(written, read_back, kDataLength), 0) + << "task " << task << " mismatch"; + } + + EXPECT_TRUE(client.closeSegment(segment).ok()); + EXPECT_TRUE(client.unregisterLocalMemory(local.data(), local.size()).ok()); + + const int status = child_guard.finish(); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 0); +} + +TEST(HighPerformanceTcpE2eTest, WriteThenReadConcurrency16) { + RunWriteThenReadAcrossProcesses(16); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_protocol_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_protocol_test.cpp new file mode 100644 index 0000000000..1c4bd0ad20 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_protocol_test.cpp @@ -0,0 +1,93 @@ +// Copyright 2026 KVCache.AI +#include + +#include + +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" + +namespace mooncake::tent { +namespace { + +TEST(HighPerformanceTcpProtocolTest, RequestAndResponseRoundTrip) { + const HighPerformanceTcpRequestFrame request{ + HighPerformanceTcpOpcode::kWrite, 11, 22, 33, 44}; + const auto request_bytes = EncodeHighPerformanceTcpRequest(request); + HighPerformanceTcpRequestFrame decoded_request; + ASSERT_TRUE(DecodeHighPerformanceTcpRequest(request_bytes.data(), + request_bytes.size(), + &decoded_request) + .ok()); + EXPECT_EQ(decoded_request.opcode, request.opcode); + EXPECT_EQ(decoded_request.request_id, request.request_id); + EXPECT_EQ(decoded_request.registration_id, request.registration_id); + EXPECT_EQ(decoded_request.remote_addr, request.remote_addr); + EXPECT_EQ(decoded_request.length, request.length); + + const HighPerformanceTcpResponseFrame response{ + HighPerformanceTcpStatus::kOk, request.request_id, request.length}; + const auto response_bytes = EncodeHighPerformanceTcpResponse(response); + HighPerformanceTcpResponseFrame decoded_response; + ASSERT_TRUE(DecodeHighPerformanceTcpResponse(response_bytes.data(), + response_bytes.size(), + &decoded_response) + .ok()); + EXPECT_EQ(decoded_response.status, response.status); + EXPECT_EQ(decoded_response.request_id, response.request_id); + EXPECT_EQ(decoded_response.committed_bytes, response.committed_bytes); +} + +TEST(HighPerformanceTcpProtocolTest, RejectsMalformedWireFrames) { + auto request = EncodeHighPerformanceTcpRequest( + {HighPerformanceTcpOpcode::kRead, 1, 2, 3, 4}); + HighPerformanceTcpRequestFrame decoded; + HighPerformanceTcpStatus wire_error = HighPerformanceTcpStatus::kOk; + request[5] = 2; + EXPECT_FALSE(DecodeHighPerformanceTcpRequest(request.data(), request.size(), + &decoded, &wire_error) + .ok()); + EXPECT_EQ(wire_error, HighPerformanceTcpStatus::kBadVersion); + request[5] = 1; + request[6] = 0xff; + EXPECT_FALSE(DecodeHighPerformanceTcpRequest(request.data(), request.size(), + &decoded, &wire_error) + .ok()); + EXPECT_EQ(wire_error, HighPerformanceTcpStatus::kBadOpcode); +} + +TEST(HighPerformanceTcpProtocolTest, MetadataAttributesRoundTrip) { + HighPerformanceTcpEndpointAttr endpoint{"00112233445566778899aabbccddeeff", + "127.0.0.1", 1234, 4096}; + std::string encoded; + ASSERT_TRUE(EncodeHighPerformanceTcpEndpointAttr(endpoint, &encoded).ok()); + HighPerformanceTcpEndpointAttr decoded_endpoint; + ASSERT_TRUE( + DecodeHighPerformanceTcpEndpointAttr(encoded, &decoded_endpoint).ok()); + EXPECT_EQ(decoded_endpoint.incarnation, endpoint.incarnation); + EXPECT_EQ(decoded_endpoint.host, endpoint.host); + EXPECT_EQ(decoded_endpoint.port, endpoint.port); + + ASSERT_TRUE( + EncodeHighPerformanceTcpBufferAttr({42, "global_read_write"}, &encoded) + .ok()); + HighPerformanceTcpBufferAttr decoded_buffer; + ASSERT_TRUE( + DecodeHighPerformanceTcpBufferAttr(encoded, &decoded_buffer).ok()); + EXPECT_EQ(decoded_buffer.registration_id, 42u); + EXPECT_EQ(decoded_buffer.permission, "global_read_write"); + + const uint64_t max_registration_id = std::numeric_limits::max(); + ASSERT_TRUE(EncodeHighPerformanceTcpBufferAttr( + {max_registration_id, "global_read_only"}, &encoded) + .ok()); + ASSERT_TRUE( + DecodeHighPerformanceTcpBufferAttr(encoded, &decoded_buffer).ok()); + EXPECT_EQ(decoded_buffer.registration_id, max_registration_id); + EXPECT_EQ(decoded_buffer.permission, "global_read_only"); + + EXPECT_FALSE( + DecodeHighPerformanceTcpEndpointAttr("not-json", &decoded_endpoint) + .ok()); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_socket_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_socket_test.cpp new file mode 100644 index 0000000000..24f0dc996b --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_socket_test.cpp @@ -0,0 +1,372 @@ +// Copyright 2026 KVCache.AI +#include + +#include + +#include +#include +#include +#include +#include + +#include "tent/transport/hp_tcp/hp_tcp_client.h" +#include "tent/transport/hp_tcp/hp_tcp_server.h" + +namespace mooncake::tent { +namespace { +using namespace std::chrono_literals; +constexpr char kIncarnation[] = "00112233445566778899aabbccddeeff"; + +struct Completion { + std::atomic done{false}; + TransferStatusEnum status{PENDING}; + size_t bytes{0}; + std::optional protocol_status; + + auto callback() { + return [this](TransferStatusEnum value, size_t count, + std::optional result) { + status = value; + bytes = count; + protocol_status = result; + done.store(true, std::memory_order_release); + }; + } + bool wait() { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (!done.load(std::memory_order_acquire) && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(1ms); + } + return done.load(std::memory_order_acquire); + } +}; + +template +bool WaitUntil(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (!predicate() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(1ms); + } + return predicate(); +} + +// Kernel TCP orders the remote write before its response, but that ordering is +// not visible to ThreadSanitizer as a C++ happens-before edge. +template +__attribute__((no_sanitize("thread"))) bool SocketOrderedEqual( + const std::array& left, const std::array& right) { + return left == right; +} + +class Runtime { + public: + explicit Runtime(size_t max_connections = 16, uint64_t timeout_ms = 500) + : workers({.worker_count = 2}), + client({.max_transfer_bytes = 1 << 20, + .chunk_size = 128, + .connect_timeout_ms = 500, + .progress_timeout_ms = timeout_ms, + .connections_per_peer = 2}, + &workers), + server({.bind_address = "127.0.0.1", + .port = 0, + .max_transfer_bytes = 1 << 20, + .chunk_size = 128, + .progress_timeout_ms = timeout_ms, + .max_connections = max_connections}, + ®istry, &workers) {} + + ~Runtime() { + (void)server.stopAccepting(); + (void)client.cancelAll(); + (void)server.stop(); + (void)workers.stop(); + } + void start() { + ASSERT_TRUE(workers.start().ok()); + ASSERT_TRUE(server.start(&port).ok()); + } + Status submit(HighPerformanceTcpClient::Operation operation) { + const size_t owner = + workers.affinityOwner(operation.peer_id, operation.lane_id); + std::vector commands; + commands.push_back({.worker_id = owner, + .run = + [this, owner, operation = std::move(operation)]( + size_t) mutable { + client.enqueueOnOwner(owner, + std::move(operation)); + }, + .cancel = {}}); + return workers.tryCommitBatch(commands, nullptr, 0, 0, [] {}); + } + + HighPerformanceTcpWorkers workers; + HighPerformanceTcpBufferRegistry registry; + HighPerformanceTcpClient client; + HighPerformanceTcpServer server; + uint16_t port{0}; +}; + +HighPerformanceTcpClient::Operation Operation( + void* local, size_t length, uint64_t remote, uint64_t registration, + uint64_t request_id, HighPerformanceTcpOpcode opcode, + Completion* completion, uint16_t port) { + HighPerformanceTcpClient::Operation operation; + operation.peer_id = 7; + operation.incarnation = kIncarnation; + operation.host = "127.0.0.1"; + operation.port = port; + operation.registration_id = registration; + operation.remote_addr = remote; + operation.local_addr = local; + operation.length = length; + operation.opcode = opcode; + operation.request_id = request_id; + operation.complete = completion->callback(); + return operation; +} + +TEST(HighPerformanceTcpSocketTest, WriteReadAndReuseConnection) { + Runtime runtime(/*max_connections=*/16, /*timeout_ms=*/100); + runtime.start(); + std::array remote{}; + std::array source{}; + std::array destination{}; + source.fill(0x5a); + uint64_t registration = 0; + ASSERT_TRUE(runtime.registry + .add(reinterpret_cast(remote.data()), + remote.size(), kGlobalReadWrite, ®istration) + .ok()); + + Completion write; + ASSERT_TRUE( + runtime + .submit(Operation(source.data(), source.size(), + reinterpret_cast(remote.data()), + registration, 1, HighPerformanceTcpOpcode::kWrite, + &write, runtime.port)) + .ok()); + ASSERT_TRUE(write.wait()); + EXPECT_EQ(write.status, COMPLETED); + EXPECT_TRUE(SocketOrderedEqual(remote, source)); + + std::this_thread::sleep_for(300ms); + ASSERT_EQ(runtime.server.activeSessionsForTest(), 1u); + + Completion read; + ASSERT_TRUE( + runtime + .submit(Operation(destination.data(), destination.size(), + reinterpret_cast(remote.data()), + registration, 2, HighPerformanceTcpOpcode::kRead, + &read, runtime.port)) + .ok()); + ASSERT_TRUE(read.wait()); + EXPECT_EQ(read.status, COMPLETED); + EXPECT_TRUE(SocketOrderedEqual(destination, source)); + EXPECT_EQ(runtime.client.connectionsCreatedForTest(), 1u); +} + +TEST(HighPerformanceTcpSocketTest, RejectedWriteBodyCannotBecomeNextFrame) { + Runtime runtime; + runtime.start(); + std::array remote{}; + uint64_t registration = 0; + ASSERT_TRUE(runtime.registry + .add(reinterpret_cast(remote.data()), + remote.size(), kGlobalReadWrite, ®istration) + .ok()); + + asio::io_context io; + asio::ip::tcp::socket socket(io); + socket.connect({asio::ip::make_address("127.0.0.1"), runtime.port}); + const auto rejected = EncodeHighPerformanceTcpRequest( + {HighPerformanceTcpOpcode::kWrite, 20, registration + 1, + reinterpret_cast(remote.data()), + kHighPerformanceTcpRequestSize + 1}); + const auto hidden = EncodeHighPerformanceTcpRequest( + {HighPerformanceTcpOpcode::kWrite, 21, registration, + reinterpret_cast(remote.data()), 1}); + const uint8_t hidden_body = 0x5a; + asio::write(socket, asio::buffer(rejected)); + asio::write(socket, asio::buffer(hidden)); + asio::write(socket, asio::buffer(&hidden_body, 1)); + + std::array response_bytes{}; + ASSERT_EQ(asio::read(socket, asio::buffer(response_bytes)), + response_bytes.size()); + HighPerformanceTcpResponseFrame response; + ASSERT_TRUE(DecodeHighPerformanceTcpResponse( + response_bytes.data(), response_bytes.size(), &response) + .ok()); + EXPECT_EQ(response.status, HighPerformanceTcpStatus::kStaleRegistration); + EXPECT_EQ(response.request_id, 20u); + EXPECT_EQ(response.committed_bytes, 0u); + EXPECT_EQ(remote[0], 0u); + EXPECT_TRUE( + WaitUntil([&] { return runtime.server.activeSessionsForTest() == 0; })); +} + +TEST(HighPerformanceTcpSocketTest, + RejectedWritePartialPayloadTimesOutAndReleasesSlot) { + Runtime runtime(/*max_connections=*/1, /*timeout_ms=*/100); + runtime.start(); + std::array remote{}; + uint64_t registration = 0; + ASSERT_TRUE(runtime.registry + .add(reinterpret_cast(remote.data()), + remote.size(), kGlobalReadWrite, ®istration) + .ok()); + + asio::io_context io; + asio::ip::tcp::socket socket(io); + socket.connect({asio::ip::make_address("127.0.0.1"), runtime.port}); + const auto rejected = EncodeHighPerformanceTcpRequest( + {HighPerformanceTcpOpcode::kWrite, 22, registration + 1, + reinterpret_cast(remote.data()), 1024}); + const uint8_t one_body_byte = 0; + asio::write(socket, asio::buffer(rejected)); + asio::write(socket, asio::buffer(&one_body_byte, 1)); + ASSERT_TRUE( + WaitUntil([&] { return runtime.server.activeSessionsForTest() == 1; })); + EXPECT_TRUE( + WaitUntil([&] { return runtime.server.activeSessionsForTest() == 0; })); + EXPECT_EQ(remote[0], 0u); +} + +TEST(HighPerformanceTcpSocketTest, ClientProgressTimeoutCompletesTask) { + asio::io_context peer_io; + asio::ip::tcp::acceptor acceptor(peer_io, {asio::ip::tcp::v4(), 0}); + std::thread peer([&] { + asio::ip::tcp::socket socket(peer_io); + acceptor.accept(socket); + std::array request{}; + asio::read(socket, asio::buffer(request)); + std::this_thread::sleep_for(200ms); + }); + + HighPerformanceTcpWorkers workers({.worker_count = 1}); + ASSERT_TRUE(workers.start().ok()); + HighPerformanceTcpClient client({4096, 128, 100, 30, 1}, &workers); + std::array local{}; + Completion completion; + auto operation = Operation(local.data(), local.size(), 0x1000, 1, 3, + HighPerformanceTcpOpcode::kRead, &completion, + acceptor.local_endpoint().port()); + std::vector commands; + commands.push_back({.worker_id = 0, + .run = + [&](size_t) mutable { + client.enqueueOnOwner(0, std::move(operation)); + }, + .cancel = {}}); + ASSERT_TRUE(workers.tryCommitBatch(commands, nullptr, 0, 0, [] {}).ok()); + ASSERT_TRUE(completion.wait()); + EXPECT_EQ(completion.status, TIMEOUT); + EXPECT_FALSE(completion.protocol_status.has_value()); + EXPECT_TRUE(client.cancelAll().ok()); + EXPECT_TRUE(workers.stop().ok()); + peer.join(); +} + +TEST(HighPerformanceTcpSocketTest, + WriteWithoutCompletionAckMarksRemoteOutcomeUnknown) { + asio::io_context peer_io; + asio::ip::tcp::acceptor acceptor(peer_io, {asio::ip::tcp::v4(), 0}); + std::thread peer([&] { + asio::ip::tcp::socket socket(peer_io); + acceptor.accept(socket); + std::array request{}; + std::array body{}; + asio::read(socket, asio::buffer(request)); + asio::read(socket, asio::buffer(body)); + socket.close(); // The payload arrived, but no completion ACK did. + }); + + HighPerformanceTcpWorkers workers({.worker_count = 1}); + ASSERT_TRUE(workers.start().ok()); + HighPerformanceTcpClient client({4096, 128, 100, 100, 1}, &workers); + std::array local{}; + Completion completion; + auto operation = Operation(local.data(), local.size(), 0x1000, 1, 31, + HighPerformanceTcpOpcode::kWrite, &completion, + acceptor.local_endpoint().port()); + std::vector commands; + commands.push_back({.worker_id = 0, + .run = + [&](size_t) mutable { + client.enqueueOnOwner(0, std::move(operation)); + }, + .cancel = {}}); + ASSERT_TRUE(workers.tryCommitBatch(commands, nullptr, 0, 0, [] {}).ok()); + ASSERT_TRUE(completion.wait()); + EXPECT_EQ(completion.status, FAILED); + EXPECT_EQ(completion.protocol_status, + HighPerformanceTcpStatus::kInternalError); + EXPECT_TRUE(client.cancelAll().ok()); + EXPECT_TRUE(workers.stop().ok()); + peer.join(); +} + +TEST(HighPerformanceTcpSocketTest, EmptyAndPartialHeadersReleaseSlot) { + Runtime runtime(/*max_connections=*/1, /*timeout_ms=*/100); + runtime.start(); + asio::io_context io; + { + asio::ip::tcp::socket socket(io); + socket.connect({asio::ip::make_address("127.0.0.1"), runtime.port}); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 1; })); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 0; })); + } + { + asio::ip::tcp::socket socket(io); + socket.connect({asio::ip::make_address("127.0.0.1"), runtime.port}); + const uint8_t byte = 0; + asio::write(socket, asio::buffer(&byte, 1)); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 1; })); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 0; })); + } + + std::array remote{}; + std::array local{}; + uint64_t registration = 0; + ASSERT_TRUE(runtime.registry + .add(reinterpret_cast(remote.data()), + remote.size(), kGlobalReadOnly, ®istration) + .ok()); + Completion completion; + ASSERT_TRUE( + runtime + .submit(Operation(local.data(), local.size(), + reinterpret_cast(remote.data()), + registration, 4, HighPerformanceTcpOpcode::kRead, + &completion, runtime.port)) + .ok()); + ASSERT_TRUE(completion.wait()); + EXPECT_EQ(completion.status, COMPLETED); +} + +TEST(HighPerformanceTcpSocketTest, ClosedSessionsAreReaped) { + Runtime runtime(/*max_connections=*/2); + runtime.start(); + for (int i = 0; i < 8; ++i) { + asio::io_context io; + asio::ip::tcp::socket socket(io); + socket.connect({asio::ip::make_address("127.0.0.1"), runtime.port}); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 1; })); + socket.close(); + ASSERT_TRUE(WaitUntil( + [&] { return runtime.server.activeSessionsForTest() == 0; })); + } +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_config_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_config_test.cpp new file mode 100644 index 0000000000..39fe7be713 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_config_test.cpp @@ -0,0 +1,88 @@ +// Copyright 2026 KVCache.AI +#include + +#include + +#include "tent/common/config.h" +#include "tent/runtime/hp_tcp_transport_config.h" + +namespace mooncake::tent { +namespace { + +TEST(HpTcpTransportConfigTest, DefaultsToDisabled) { + Config config; + HpTcpTransportConfig parsed; + ASSERT_TRUE(ParseHpTcpTransportConfig(config, &parsed).ok()); + EXPECT_FALSE(parsed.enabled); +} + +TEST(HpTcpTransportConfigTest, RejectsWrongLeafTypes) { + Config config; + ASSERT_TRUE( + config.load(R"({"transports":{"hp_tcp":{"enable":"yes"}}})").ok()); + HpTcpTransportConfig parsed; + EXPECT_TRUE(ParseHpTcpTransportConfig(config, &parsed).IsInvalidArgument()); + + ASSERT_TRUE( + config.load(R"({"transports":{"hp_tcp":{"worker_count":"16"}}})").ok()); + EXPECT_TRUE(ParseHpTcpTransportConfig(config, &parsed).IsInvalidArgument()); +} + +TEST(HpTcpTransportConfigTest, ParsesAndValidatesLimits) { + Config config; + ASSERT_TRUE( + config + .load( + R"({"transports":{"tcp":{"enable":false},"hp_tcp":{"enable":true,"port":0,"worker_count":4,"connections_per_peer":2,"max_outstanding_tasks":16,"max_outstanding_bytes":4096,"max_transfer_bytes":1024,"connect_timeout_ms":1,"progress_timeout_ms":2}}})") + .ok()); + HpTcpTransportConfig parsed; + ASSERT_TRUE(ParseHpTcpTransportConfig(config, &parsed).ok()); + EXPECT_TRUE(parsed.enabled); + EXPECT_EQ(parsed.params.worker_count, 4U); + + ASSERT_TRUE( + config + .load( + R"({"transports":{"tcp":{"enable":false},"hp_tcp":{"enable":true,"max_transfer_bytes":0}}})") + .ok()); + EXPECT_TRUE(ParseHpTcpTransportConfig(config, &parsed).IsInvalidArgument()); +} + +TEST(HpTcpTransportConfigTest, AcceptsFullWidthUnsignedLimit) { + Config config; + ASSERT_TRUE( + config + .load( + R"({"transports":{"tcp":{"enable":false},"hp_tcp":{"enable":true,"max_outstanding_bytes":18446744073709551615}}})") + .ok()); + HpTcpTransportConfig parsed; + ASSERT_TRUE(ParseHpTcpTransportConfig(config, &parsed).ok()); + EXPECT_EQ(parsed.params.max_outstanding_bytes, + std::numeric_limits::max()); +} + +TEST(HpTcpTransportConfigTest, DisabledHpTcpIgnoresInactiveLimits) { + Config config; + ASSERT_TRUE( + config + .load( + R"({"transports":{"hp_tcp":{"enable":false,"worker_count":0}}})") + .ok()); + HpTcpTransportConfig parsed; + ASSERT_TRUE(ParseHpTcpTransportConfig(config, &parsed).ok()); + EXPECT_FALSE(parsed.enabled); +} + +TEST(HpTcpTransportConfigTest, RejectsTcpAndHpTcpTogether) { + Config config; + ASSERT_TRUE( + config + .load( + R"({"transports":{"tcp":{"enable":true},"hp_tcp":{"enable":true}}})") + .ok()); + HpTcpTransportConfig parsed; + EXPECT_TRUE(ParseHpTcpTransportConfig(config, &parsed).IsInvalidArgument()); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp new file mode 100644 index 0000000000..56db4c1aea --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp @@ -0,0 +1,361 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/control_plane.h" +#include "tent/transport/hp_tcp/hp_tcp_protocol.h" +#include "tent/transport/hp_tcp/hp_tcp_transport.h" + +namespace mooncake::tent { + +class HighPerformanceTcpTransportTestPeer { + public: + static void failWorker(HighPerformanceTcpTransport& transport) { + asio::post(transport.workers_->ioContext(0), [] { + throw std::runtime_error("injected HP TCP worker failure"); + }); + } + + static bool hasFailedWorker(const HighPerformanceTcpTransport& transport) { + return transport.workers_->hasFailedWorker(); + } +}; + +namespace { + +template +bool WaitUntil(Predicate predicate) { + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!predicate() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + return predicate(); +} + +std::shared_ptr MakeLocalMetadata() { + auto metadata = std::make_shared("p2p", "", nullptr); + EXPECT_TRUE(metadata->segmentManager() + .updateLocal([](SegmentDesc& segment) -> Status { + segment.name = "hp_transport_test"; + segment.machine_id = "hp_transport_test_machine"; + segment.rpc_server_addr = "127.0.0.1:40000"; + segment.type = SegmentType::Memory; + std::get(segment.detail) = + MemorySegmentDesc{}; + return Status::OK(); + }) + .ok()); + return metadata; +} + +HighPerformanceTcpParams MakeParams() { + HighPerformanceTcpParams params; + params.bind_address = "127.0.0.1"; + params.advertise_address = "127.0.0.1"; + params.port = 0; + params.worker_count = 2; + params.connections_per_peer = 2; + params.max_outstanding_tasks = 16; + params.max_outstanding_bytes = 1 << 20; + params.max_transfer_bytes = 1 << 20; + params.connect_timeout_ms = 1000; + params.progress_timeout_ms = 1000; + return params; +} + +bool ContainsTransport(const BufferDesc& desc, TransportType type) { + return std::find(desc.transports.begin(), desc.transports.end(), type) != + desc.transports.end(); +} + +TEST(HighPerformanceTcpTransportTest, + WorkerFailureDoesNotHideCommittedTerminalStatus) { + auto metadata = MakeLocalMetadata(); + HighPerformanceTcpTransport transport(MakeParams()); + std::string segment_name = "hp_transport_test"; + ASSERT_TRUE( + transport.install(segment_name, metadata, nullptr, nullptr).ok()); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 1).ok()); + auto* hp_batch = dynamic_cast(batch); + ASSERT_NE(hp_batch, nullptr); + + Request request{}; + request.length = 37; + auto task = std::make_shared( + request.length, 0, [](BatchID) {}, + HighPerformanceTcpBufferRegistry::Lease{}); + ASSERT_TRUE(task->completeOnce(COMPLETED, request.length)); + hp_batch->tasks.push_back(std::move(task)); + + HighPerformanceTcpTransportTestPeer::failWorker(transport); + ASSERT_TRUE(WaitUntil([&] { + return HighPerformanceTcpTransportTestPeer::hasFailedWorker(transport); + })); + + TransferStatus status{FAILED, 0}; + const Status result = transport.getTransferStatus(batch, 0, status); + EXPECT_TRUE(result.ok()) << result.ToString(); + EXPECT_EQ(status.s, COMPLETED); + EXPECT_EQ(status.transferred_bytes, request.length); + + EXPECT_TRUE(transport.freeSubBatch(batch).ok()); + EXPECT_TRUE(transport.quiesce().IsInternalError()); + EXPECT_TRUE(transport.uninstall().ok()); +} + +TEST(HighPerformanceTcpTransportTest, UnknownWriteOutcomeIsPermanent) { + HighPerformanceTcpTransport transport(MakeParams()); + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 1).ok()); + auto* hp_batch = dynamic_cast(batch); + ASSERT_NE(hp_batch, nullptr); + + auto uncertain_write = std::make_shared( + 0, 0, [](BatchID) {}, HighPerformanceTcpBufferRegistry::Lease{}); + ASSERT_TRUE(uncertain_write->completeOnce( + FAILED, 0, HighPerformanceTcpStatus::kInternalError)); + hp_batch->tasks.push_back(std::move(uncertain_write)); + + TransferStatus status{}; + const Status unsafe_replay = transport.getTransferStatus(batch, 0, status); + EXPECT_TRUE(unsafe_replay.IsInvalidEntry()) << unsafe_replay.ToString(); + EXPECT_EQ(status.s, FAILED); + EXPECT_TRUE(transport.freeSubBatch(batch).ok()); +} + +TEST(HighPerformanceTcpTransportTest, + PublishesEndpointAndSeparatesLocalOnlyCapabilities) { + auto metadata = MakeLocalMetadata(); + HighPerformanceTcpTransport transport(MakeParams()); + std::string segment_name = "hp_transport_test"; + ASSERT_TRUE( + transport.install(segment_name, metadata, nullptr, nullptr).ok()); + + const SegmentDescRef local = metadata->segmentManager().getLocal(); + const auto attr_it = local->getMemory().transport_attrs.find( + static_cast(TransportType::HP_TCP)); + ASSERT_NE(attr_it, local->getMemory().transport_attrs.end()); + HighPerformanceTcpEndpointAttr endpoint; + ASSERT_TRUE( + DecodeHighPerformanceTcpEndpointAttr(attr_it->second, &endpoint).ok()); + EXPECT_EQ(endpoint.host, "127.0.0.1"); + EXPECT_NE(endpoint.port, 0); + + std::array local_only_storage{}; + BufferDesc local_only; + local_only.addr = reinterpret_cast(local_only_storage.data()); + local_only.length = local_only_storage.size(); + local_only.location = "cpu:0"; + MemoryOptions local_options; + local_options.perm = kLocalReadWrite; + ASSERT_TRUE(transport.addMemoryBuffer(local_only, local_options).ok()); + EXPECT_TRUE(transport.tracksLocalBuffer(local_only)); + EXPECT_FALSE(ContainsTransport(local_only, TransportType::HP_TCP)); + EXPECT_EQ(local_only.transport_attrs.count(TransportType::HP_TCP), 0U); + + std::array global_storage{}; + BufferDesc global; + global.addr = reinterpret_cast(global_storage.data()); + global.length = global_storage.size(); + global.location = "cpu:0"; + MemoryOptions global_options; + global_options.perm = kGlobalReadWrite; + ASSERT_TRUE(transport.addMemoryBuffer(global, global_options).ok()); + EXPECT_TRUE(ContainsTransport(global, TransportType::HP_TCP)); + const auto buffer_attr = global.transport_attrs.find(TransportType::HP_TCP); + ASSERT_NE(buffer_attr, global.transport_attrs.end()); + HighPerformanceTcpBufferAttr decoded_buffer; + ASSERT_TRUE( + DecodeHighPerformanceTcpBufferAttr(buffer_attr->second, &decoded_buffer) + .ok()); + EXPECT_NE(decoded_buffer.registration_id, 0U); + EXPECT_EQ(decoded_buffer.permission, "global_read_write"); + + ASSERT_TRUE(transport.quiesce().ok()); + ASSERT_TRUE(transport.removeMemoryBuffer(local_only).ok()); + ASSERT_TRUE(transport.removeMemoryBuffer(global).ok()); + ASSERT_TRUE(transport.uninstall().ok()); +} + +Status PublishBuffers(const std::shared_ptr& metadata, + const std::vector& buffers) { + return metadata->segmentManager().updateLocal( + [&](SegmentDesc& segment) -> Status { + std::get(segment.detail).buffers = buffers; + return Status::OK(); + }); +} + +Status WaitForTransportResult(HighPerformanceTcpTransport& transport, + Transport::SubBatchRef batch, + TransferStatus& transfer_status) { + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + Status result = Status::OK(); + while (std::chrono::steady_clock::now() < deadline) { + result = transport.getTransferStatus(batch, 0, transfer_status); + if (transfer_status.s != PENDING) return result; + std::this_thread::yield(); + } + return Status::InternalError("HP TCP test transfer did not finish"); +} + +TEST(HighPerformanceTcpTransportTest, + StaleRegistrationRefreshRetriesSameTransportWithFreshMetadata) { + auto server_metadata = MakeLocalMetadata(); + uint16_t rpc_port = 0; + ASSERT_TRUE(server_metadata->start(rpc_port).ok()); + ASSERT_NE(rpc_port, 0); + const std::string server_name = "127.0.0.1:" + std::to_string(rpc_port); + ASSERT_TRUE(server_metadata->segmentManager() + .updateLocal([&](SegmentDesc& segment) -> Status { + segment.name = server_name; + segment.rpc_server_addr = server_name; + return Status::OK(); + }) + .ok()); + + HighPerformanceTcpTransport server(MakeParams()); + std::string installed_server_name = server_name; + ASSERT_TRUE( + server.install(installed_server_name, server_metadata, nullptr, nullptr) + .ok()); + std::array remote_storage{}; + BufferDesc registration_a; + registration_a.addr = reinterpret_cast(remote_storage.data()); + registration_a.length = remote_storage.size(); + registration_a.location = "cpu:0"; + MemoryOptions remote_options; + remote_options.type = HP_TCP; + remote_options.perm = kGlobalReadWrite; + ASSERT_TRUE(server.addMemoryBuffer(registration_a, remote_options).ok()); + HighPerformanceTcpBufferAttr attr_a; + ASSERT_TRUE( + DecodeHighPerformanceTcpBufferAttr( + registration_a.transport_attrs.at(TransportType::HP_TCP), &attr_a) + .ok()); + ASSERT_TRUE(PublishBuffers(server_metadata, {registration_a}).ok()); + + auto client_metadata = MakeLocalMetadata(); + ASSERT_TRUE(client_metadata->segmentManager() + .updateLocal([](SegmentDesc& segment) -> Status { + // Deliberately omit a callback address. The remote + // cache stays on A until this test invalidates it. + segment.rpc_server_addr.clear(); + return Status::OK(); + }) + .ok()); + HighPerformanceTcpTransport client(MakeParams()); + std::string client_name = "hp_transport_client"; + ASSERT_TRUE( + client.install(client_name, client_metadata, nullptr, nullptr).ok()); + SegmentID target = 0; + ASSERT_TRUE( + client_metadata->segmentManager().openRemote(target, server_name).ok()); + SegmentDescRef cached_a; + ASSERT_TRUE(client_metadata->segmentManager() + .getRemoteCached(cached_a, target) + .ok()); + const BufferDesc* cached_buffer_a = + cached_a->findBuffer(registration_a.addr, registration_a.length); + ASSERT_NE(cached_buffer_a, nullptr); + + ASSERT_TRUE(server.removeMemoryBuffer(registration_a).ok()); + BufferDesc registration_b; + registration_b.addr = reinterpret_cast(remote_storage.data()); + registration_b.length = remote_storage.size(); + registration_b.location = "cpu:0"; + ASSERT_TRUE(server.addMemoryBuffer(registration_b, remote_options).ok()); + HighPerformanceTcpBufferAttr attr_b; + ASSERT_TRUE( + DecodeHighPerformanceTcpBufferAttr( + registration_b.transport_attrs.at(TransportType::HP_TCP), &attr_b) + .ok()); + ASSERT_NE(attr_a.registration_id, attr_b.registration_id); + ASSERT_TRUE(PublishBuffers(server_metadata, {registration_b}).ok()); + + std::array local_storage{}; + BufferDesc local; + local.addr = reinterpret_cast(local_storage.data()); + local.length = local_storage.size(); + local.location = "cpu:0"; + MemoryOptions local_options; + local_options.type = HP_TCP; + local_options.perm = kLocalReadWrite; + ASSERT_TRUE(client.addMemoryBuffer(local, local_options).ok()); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(client.allocateSubBatch(batch, 1).ok()); + Request request{}; + request.opcode = Request::READ; + request.source = local_storage.data(); + request.target_id = target; + request.target_offset = registration_b.addr; + request.length = registration_b.length; + request.transport_hint = HP_TCP; + ASSERT_TRUE(client.submitTransferTasks(batch, {request}).ok()); + + TransferStatus transfer_status; + Status first_result = + WaitForTransportResult(client, batch, transfer_status); + EXPECT_EQ(transfer_status.s, FAILED); + EXPECT_TRUE(first_result.IsNeedsRefreshCache()) << first_result.ToString(); + + int metadata_refresh_retry_count = 0; + ASSERT_LT(metadata_refresh_retry_count, 1); + ++metadata_refresh_retry_count; + ASSERT_TRUE( + client_metadata->segmentManager().invalidateRemote(target).ok()); + ASSERT_TRUE(client.retryTransferTask(batch, 0, request).ok()); + + Status retry_result = + WaitForTransportResult(client, batch, transfer_status); + EXPECT_TRUE(retry_result.ok()) << retry_result.ToString(); + EXPECT_EQ(transfer_status.s, COMPLETED); + EXPECT_EQ(transfer_status.transferred_bytes, request.length); + EXPECT_EQ(metadata_refresh_retry_count, 1); + + SegmentDescRef refreshed; + ASSERT_TRUE(client_metadata->segmentManager() + .getRemoteCached(refreshed, target) + .ok()); + const BufferDesc* refreshed_buffer = + refreshed->findBuffer(registration_b.addr, registration_b.length); + ASSERT_NE(refreshed_buffer, nullptr); + HighPerformanceTcpBufferAttr refreshed_attr; + ASSERT_TRUE(DecodeHighPerformanceTcpBufferAttr( + refreshed_buffer->transport_attrs.at(TransportType::HP_TCP), + &refreshed_attr) + .ok()); + EXPECT_EQ(refreshed_attr.registration_id, attr_b.registration_id); + + ASSERT_TRUE(client.freeSubBatch(batch).ok()); + ASSERT_TRUE(client.removeMemoryBuffer(local).ok()); + ASSERT_TRUE(server.removeMemoryBuffer(registration_b).ok()); + ASSERT_TRUE(client.quiesce().ok()); + ASSERT_TRUE(server.quiesce().ok()); + ASSERT_TRUE(client.uninstall().ok()); + ASSERT_TRUE(server.uninstall().ok()); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_workers_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_workers_test.cpp new file mode 100644 index 0000000000..109f53d1e8 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_workers_test.cpp @@ -0,0 +1,150 @@ +// Copyright 2026 KVCache.AI +#include + +#include +#include +#include +#include +#include +#include + +#include "tent/transport/hp_tcp/hp_tcp_buffer_registry.h" +#include "tent/transport/hp_tcp/hp_tcp_task.h" +#include "tent/transport/hp_tcp/hp_tcp_workers.h" + +namespace mooncake::tent { +namespace { +using namespace std::chrono_literals; + +template +bool WaitUntil(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (!predicate() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(1ms); + } + return predicate(); +} + +Status SubmitToWorker(HighPerformanceTcpWorkers& workers, size_t owner, + HighPerformanceTcpWorkers::Task task) { + std::vector commands{ + {.worker_id = owner, .run = std::move(task), .cancel = {}}}; + return workers.tryCommitBatch(commands, nullptr, 0, 0, [] {}); +} + +TEST(HighPerformanceTcpWorkersTest, PreservesAffinity) { + HighPerformanceTcpWorkers workers({.worker_count = 2}); + ASSERT_TRUE(workers.start().ok()); + std::atomic completed{0}; + for (size_t owner : {0u, 1u, 0u, 1u}) { + ASSERT_TRUE(SubmitToWorker(workers, owner, [&, owner](size_t actual) { + EXPECT_EQ(actual, owner); + ++completed; + }).ok()); + } + EXPECT_TRUE(WaitUntil([&] { return completed.load() == 4; })); + EXPECT_TRUE(workers.stop().ok()); +} + +TEST(HighPerformanceTcpWorkersTest, BatchAdmissionHasNoPartialCommit) { + HighPerformanceTcpWorkers workers({.worker_count = 1}); + ASSERT_TRUE(workers.start().ok()); + HighPerformanceTcpAdmissionController admission(1, 512); + ASSERT_TRUE(admission.tryReserve(1, 512).ok()); + bool committed = false; + std::vector commands{ + {.worker_id = 0, .run = [](size_t) {}, .cancel = [] {}}}; + EXPECT_TRUE(workers + .tryCommitBatch(commands, &admission, 1, 512, + [&] { committed = true; }) + .IsTooManyRequests()); + EXPECT_FALSE(committed); + EXPECT_EQ(admission.outstandingTasks(), 1u); + admission.release(1, 512); + EXPECT_TRUE(workers.stop().ok()); +} + +TEST(HighPerformanceTcpAdmissionControllerTest, + UnderflowFailsClosedWithoutForgingDrain) { + HighPerformanceTcpAdmissionController admission(2, 1024); + ASSERT_TRUE(admission.tryReserve(1, 512).ok()); + + std::promise result; + auto result_future = result.get_future(); + std::thread waiter( + [&] { result.set_value(admission.waitForZero().IsInternalError()); }); + + admission.release(2, 512); + + EXPECT_TRUE(admission.failed()); + EXPECT_EQ(admission.outstandingTasks(), 1u); + EXPECT_EQ(admission.outstandingBytes(), 512u); + const auto wait_status = result_future.wait_for(std::chrono::seconds(1)); + if (wait_status != std::future_status::ready) { + // Ensure a broken implementation cannot strand the test thread. + admission.release(1, 512); + } + EXPECT_EQ(wait_status, std::future_status::ready); + EXPECT_TRUE(result_future.get()); + waiter.join(); + EXPECT_TRUE(admission.tryReserve(1, 1).IsInternalError()); +} + +TEST(HighPerformanceTcpWorkersTest, + FailedWorkerRejectsAdmissionButKeepsTeardownLive) { + HighPerformanceTcpWorkers workers({.worker_count = 1}); + ASSERT_TRUE(workers.start().ok()); + + std::atomic owner_loop_continued{false}; + + asio::post(workers.ioContext(0), + [] { throw std::runtime_error("test worker failure"); }); + asio::post(workers.ioContext(0), [&] { owner_loop_continued.store(true); }); + + ASSERT_TRUE(WaitUntil([&] { return workers.hasFailedWorker(); })); + EXPECT_TRUE(workers.barrier().ok()); + EXPECT_TRUE(owner_loop_continued.load()); + EXPECT_TRUE(SubmitToWorker(workers, 0, [](size_t) {}).IsInternalError()); + EXPECT_TRUE(workers.stop().IsInternalError()); +} + +TEST(HighPerformanceTcpTaskTest, CompletionReleasesLeaseAndBudgetOnce) { + HighPerformanceTcpBufferRegistry registry; + std::array memory{}; + uint64_t registration = 0; + ASSERT_TRUE(registry + .add(reinterpret_cast(memory.data()), + memory.size(), kGlobalReadWrite, ®istration) + .ok()); + HighPerformanceTcpBufferRegistry::Lease lease; + ASSERT_TRUE( + registry + .acquireLocalLease(reinterpret_cast(memory.data()), + memory.size(), &lease) + .ok()); + HighPerformanceTcpAdmissionController admission(1, memory.size()); + ASSERT_TRUE(admission.tryReserve(1, memory.size()).ok()); + + auto task = std::make_shared( + memory.size(), 1, [](BatchID) {}, std::move(lease)); + task->activateReservation(&admission); + EXPECT_TRUE(task->completeOnce(COMPLETED, memory.size())); + EXPECT_FALSE(task->completeOnce(FAILED, 0)); + EXPECT_EQ(admission.outstandingTasks(), 0u); + EXPECT_TRUE( + registry + .remove(reinterpret_cast(memory.data()), memory.size()) + .ok()); +} + +TEST(HighPerformanceTcpWorkersTest, StopIsIdempotentAndNotRestartable) { + HighPerformanceTcpWorkers workers({.worker_count = 1}); + ASSERT_TRUE(workers.start().ok()); + EXPECT_TRUE(workers.stop().ok()); + EXPECT_TRUE(workers.stop().ok()); + EXPECT_TRUE(workers.start().IsInvalidArgument()); + EXPECT_TRUE(SubmitToWorker(workers, 0, [](size_t) {}).IsInternalError()); +} + +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/local_memory_lifecycle_test.cpp b/mooncake-transfer-engine/tent/tests/local_memory_lifecycle_test.cpp new file mode 100644 index 0000000000..4a3859f64d --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/local_memory_lifecycle_test.cpp @@ -0,0 +1,95 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/runtime/transfer_engine_impl.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { +namespace { + +class FailOnceFreeTransport : public Transport { + public: + Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions&) override { + *addr = std::malloc(size); + return *addr ? Status::OK() + : Status::InternalError("malloc failed" LOC_MARK); + } + + Status freeLocalMemory(void* addr, size_t) override { + ++free_calls; + if (free_calls == 1) { + return Status::InternalError("injected free failure" LOC_MARK); + } + std::free(addr); + return Status::OK(); + } + + const char* getName() const override { return "fail-once-free"; } + + int free_calls{0}; +}; + +std::shared_ptr makeConfig() { + auto config = std::make_shared(); + config->set("metadata_type", "p2p"); + config->set("metadata_servers", ""); + config->set("rpc_server_hostname", "127.0.0.1"); + config->set("rpc_server_port", "0"); + config->set("log_level", "warning"); + + for (const char* transport : {"tcp", "shm", "rdma", "io_uring", "nvlink", + "mnnvl", "gds", "ascend_direct"}) { + config->set(std::string("transports/") + transport + "/enable", false); + } + return config; +} + +TEST(LocalMemoryLifecycle, RetainsOwnershipWhenTransportFreeFails) { + TransferEngineImpl engine(makeConfig()); + ASSERT_TRUE(engine.available()); + + auto transport = std::make_shared(); + engine.swapTransportForTest(RDMA, transport); + + MemoryOptions options; + options.type = RDMA; + options.location = "cpu:0"; + void* addr = nullptr; + ASSERT_TRUE(engine.allocateLocalMemory(&addr, 4096, options).ok()); + + auto first = engine.freeLocalMemory(addr); + EXPECT_FALSE(first.ok()); + EXPECT_EQ(transport->free_calls, 1); + + auto second = engine.freeLocalMemory(addr); + EXPECT_TRUE(second.ok()) << second.ToString(); + EXPECT_EQ(transport->free_calls, 2); + + auto third = engine.freeLocalMemory(addr); + EXPECT_TRUE(third.IsInvalidArgument()); + EXPECT_EQ(transport->free_calls, 2); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp index 1e715134b2..051a1b533c 100644 --- a/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp +++ b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp @@ -260,7 +260,12 @@ class FakeTransport : public Transport { return Status::OK(); } - Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); } + Status removeMemoryBuffer(BufferDesc& desc) override { + desc.transports.erase(std::remove(desc.transports.begin(), + desc.transports.end(), self_type_), + desc.transports.end()); + return Status::OK(); + } Status allocateLocalMemory(void** addr, size_t size, MemoryOptions&) override { diff --git a/mooncake-transfer-engine/tent/tests/rdma_async_event_drain_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_async_event_drain_test.cpp new file mode 100644 index 0000000000..06aa280cde --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/rdma_async_event_drain_test.cpp @@ -0,0 +1,219 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Workers::handleContextEvents() must empty the async event fd on every call: +// the fd is edge-triggered, so anything left queued is stranded until an +// unrelated later event releases it -- and a stranded IBV_EVENT_PORT_ACTIVE +// leaves its context paused, silently failing every transfer on that NIC. +// These tests replace the event source with a scripted queue via linker +// wrapping, so no RDMA device is needed. + +#include +#include + +#include +#include +#include +#include + +#include "tent/runtime/topology.h" +#include "tent/transport/rdma/context.h" +#include "tent/transport/rdma/params.h" +#include "tent/transport/rdma/rdma_transport.h" +#include "tent/transport/rdma/workers.h" + +namespace { + +// Stands in for the kernel's async event queue behind async_fd. +struct AsyncEventScript { + std::deque pending; + // Reported once `pending` runs dry. EAGAIN mimics a drained non-blocking + // fd; anything else mimics a real read failure. + int drained_errno = EAGAIN; + // Fail this many reads with EINTR before serving the queue. + int pending_eintr = 0; + int get_calls = 0; + int ack_calls = 0; +}; + +AsyncEventScript g_script; + +} // namespace + +extern "C" int __wrap_ibv_get_async_event(struct ibv_context* context, + struct ibv_async_event* event) { + (void)context; // Never dereferenced by the code under test. + g_script.get_calls++; + if (g_script.pending_eintr > 0) { + g_script.pending_eintr--; + errno = EINTR; + return -1; + } + if (g_script.pending.empty()) { + errno = g_script.drained_errno; + return -1; + } + memset(event, 0, sizeof(*event)); + event->event_type = g_script.pending.front(); + g_script.pending.pop_front(); + return 0; +} + +extern "C" void __wrap_ibv_ack_async_event(struct ibv_async_event* event) { + (void)event; + g_script.ack_calls++; +} + +namespace mooncake { +namespace tent { + +// Friend accessor for driving the event loop without a full install(). +class RdmaTransportTestPeer { + public: + static void bindTopology(RdmaTransport& transport, + std::shared_ptr topology) { + transport.local_topology_ = topology; + transport.local_buffer_manager_.setTopology(topology); + transport.params_ = std::make_shared(); + transport.conf_ = std::make_shared(); + } + + static size_t initializeContexts(RdmaTransport& transport) { + return transport.initializeContexts(); + } + + static std::unique_ptr makeWorkers(RdmaTransport& transport) { + return std::make_unique(&transport); + } + + static int handleContextEvents(Workers& workers, int dev_id, + std::shared_ptr& context) { + return workers.handleContextEvents(dev_id, context); + } + + static RdmaContextSet& contextSet(RdmaTransport& transport) { + return transport.context_set_; + } +}; + +namespace { + +// Event types whose handlers touch nothing a device-less context lacks. +// IBV_EVENT_COMM_EST falls through applyContextEvent()'s default label; +// IBV_EVENT_PORT_ACTIVE takes the recovery path, where resume() is a no-op on +// an inert context and the link-speed refresh declines without a device. +// Together they cover a handled and an unhandled event. +constexpr ibv_event_type kUnhandledEvent = IBV_EVENT_COMM_EST; +constexpr ibv_event_type kHandledEvent = IBV_EVENT_PORT_ACTIVE; + +class AsyncEventDrainTest : public ::testing::Test { + protected: + void SetUp() override { + topology_ = std::make_shared(); + // No device is named "mc-absent-rnic-0" on any host, so the context + // stays inert: construct() still builds the endpoint store, but no + // real async fd is ever opened and only the scripted queue answers. + ASSERT_TRUE( + topology_ + ->parse( + R"({"nics":[{"name":"mc-absent-rnic-0","type":0,"numa_node":0}]})") + .ok()); + RdmaTransportTestPeer::bindTopology(transport_, topology_); + ASSERT_EQ(RdmaTransportTestPeer::initializeContexts(transport_), 0u); + workers_ = RdmaTransportTestPeer::makeWorkers(transport_); + + g_script = AsyncEventScript{}; + } + + int drain() { + auto& context = RdmaTransportTestPeer::contextSet(transport_)[kDev]; + return RdmaTransportTestPeer::handleContextEvents(*workers_, kDev, + context); + } + + static constexpr int kDev = 0; + std::shared_ptr topology_; + RdmaTransport transport_; + std::unique_ptr workers_; +}; + +// The regression: a burst arriving between two epoll wakeups must be consumed +// in full, not one event at a time. +TEST_F(AsyncEventDrainTest, DrainsEveryQueuedEvent) { + constexpr int kBurst = 5; + for (int i = 0; i < kBurst; ++i) + g_script.pending.push_back(kUnhandledEvent); + + EXPECT_EQ(drain(), 0); + + EXPECT_TRUE(g_script.pending.empty()) + << "one epoll wakeup must drain the whole async event queue; " + << g_script.pending.size() << " event(s) were left stranded"; + // kBurst reads plus the trailing EAGAIN that ends the drain. + EXPECT_EQ(g_script.get_calls, kBurst + 1); + EXPECT_EQ(g_script.ack_calls, kBurst); +} + +// A PORT_ACTIVE queued behind a burst is the case that wedged a NIC for good, +// so it must be reached and acked like any other event. +TEST_F(AsyncEventDrainTest, ReachesHandledEventQueuedBehindOthers) { + g_script.pending.push_back(kUnhandledEvent); + g_script.pending.push_back(kUnhandledEvent); + g_script.pending.push_back(kHandledEvent); + + EXPECT_EQ(drain(), 0); + + EXPECT_TRUE(g_script.pending.empty()); + EXPECT_EQ(g_script.get_calls, 4); + EXPECT_EQ(g_script.ack_calls, 3); +} + +// A wakeup with nothing pending is a drained fd, not a failure. +TEST_F(AsyncEventDrainTest, EmptyQueueIsNotAnError) { + EXPECT_EQ(drain(), 0); + + EXPECT_EQ(g_script.get_calls, 1); + EXPECT_EQ(g_script.ack_calls, 0); +} + +// A signal must not end the drain early, or the events behind it are stranded +// exactly as they were before the fix. +TEST_F(AsyncEventDrainTest, InterruptedReadIsRetried) { + g_script.pending.push_back(kUnhandledEvent); + g_script.pending.push_back(kUnhandledEvent); + g_script.pending_eintr = 1; + + EXPECT_EQ(drain(), 0); + + EXPECT_TRUE(g_script.pending.empty()); + // One EINTR, two reads, one trailing EAGAIN. + EXPECT_EQ(g_script.get_calls, 4); + EXPECT_EQ(g_script.ack_calls, 2); +} + +// A genuine read failure still aborts and propagates, so the loop cannot spin +// forever on a broken fd. +TEST_F(AsyncEventDrainTest, RealReadErrorStopsTheDrain) { + g_script.pending.push_back(kUnhandledEvent); + g_script.drained_errno = EIO; + + EXPECT_EQ(drain(), -1); + + EXPECT_EQ(g_script.get_calls, 2); + EXPECT_EQ(g_script.ack_calls, 1); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp index d4b70a89eb..b2b96cc16f 100644 --- a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include #include @@ -34,6 +36,7 @@ #include "tent/transport/rdma/params.h" #include "tent/transport/rdma/quota.h" #include "tent/transport/rdma/rdma_transport.h" +#include "tent/transport/rdma/ibv_loader.h" #include "tent/transport/rdma/workers.h" namespace mooncake { @@ -75,6 +78,12 @@ class RdmaTransportTestPeer { workers.applyContextEvent(dev_id, context, event); } + // Runs the monitorThread() 1 Hz safety net for contexts whose + // IBV_EVENT_PORT_ACTIVE never arrived, without starting any threads. + static void resumePausedContexts(Workers& workers) { + workers.resumePausedContexts(); + } + static const RdmaContextSet& contextSet(const RdmaTransport& transport) { return transport.context_set_; } @@ -89,6 +98,27 @@ class RdmaTransportTestPeer { } }; +// Friend accessor for RdmaContext: TENT reaches libibverbs through a table of +// function pointers the context copies from IbvLoader, so a test can replace +// individual entries and hand the context a placeholder device instead of +// needing an RNIC. +class RdmaContextTestPeer { + public: + static IbvSymbols& verbs(RdmaContext& context) { return context.verbs_; } + + // Make the context look opened on `native` (never dereferenced by the + // port-attribute paths, only passed back to the verbs) with `params`. + static void bindDevice(RdmaContext& context, ibv_context* native, + std::shared_ptr params) { + context.native_context_ = native; + context.params_ = std::move(params); + } + + static void unbindDevice(RdmaContext& context) { + context.native_context_ = nullptr; + } +}; + namespace { bool hasRdmaDevice() { @@ -479,6 +509,202 @@ TEST_F(RdmaContextEventTest, CqErrLeavesAvailabilityAlone) { EXPECT_TRUE(selector_->isDeviceAvailable(kDev)); } +// ibv_query_port_speed() exists only in rdma-core >= 62. It must be resolved +// as an optional symbol: present -> non-null, absent -> null, and either way +// the mandatory verbs are still there (an older libibverbs must not lose +// RDMA over it). Compared against a direct dlsym so the expectation is +// whatever this host's library actually has. +TEST(RdmaContextPortSpeedTest, EffectiveSpeedVerbIsOptional) { + void* lib = dlopen("libibverbs.so.1", RTLD_NOW | RTLD_LOCAL); + if (!lib) GTEST_SKIP() << "libibverbs.so.1 not loadable"; + const bool host_has_verb = dlsym(lib, "ibv_query_port_speed") != nullptr; + dlclose(lib); + + const auto& sym = IbvLoader::Instance().sym(); + EXPECT_EQ(sym.ibv_query_port_speed != nullptr, host_has_verb); + // Mandatory symbols resolve regardless of the optional one. + EXPECT_NE(sym.ibv_query_port_default, nullptr); + EXPECT_NE(sym.ibv_open_device, nullptr); +} + +// Verbs stand-ins wired through RdmaContextTestPeer::verbs(). Plain function +// pointers, so state lives in one static block. +struct FakePortVerbs { + ibv_context native{}; // placeholder handle, never dereferenced + uint8_t active_speed = 0; // what ibv_query_port reports + uint8_t active_width = 0; + int query_port_rc = 0; + uint64_t speed_100mbps = 0; // what ibv_query_port_speed reports + int query_speed_rc = 0; + int query_speed_calls = 0; +}; +FakePortVerbs fake_port; + +int fakeQueryPort(ibv_context* context, uint8_t, ibv_port_attr* attr) { + if (context != &fake_port.native) return EINVAL; + if (fake_port.query_port_rc) return fake_port.query_port_rc; + *attr = {}; + attr->state = IBV_PORT_ACTIVE; + attr->active_speed = fake_port.active_speed; + attr->active_width = fake_port.active_width; + return 0; +} + +int fakeQueryPortSpeed(ibv_context* context, uint32_t, uint64_t* speed) { + ++fake_port.query_speed_calls; + if (context != &fake_port.native) return EINVAL; + if (fake_port.query_speed_rc) return fake_port.query_speed_rc; + *speed = fake_port.speed_100mbps; + return 0; +} + +// A context whose port-attribute verbs are the fakes above, "opened" on the +// placeholder device. Exercises refreshPortAttributes()/linkSpeedGbps() +// exactly as the monitor thread does, without an RNIC. +class RdmaContextFakeVerbsTest : public ::testing::Test { + protected: + void SetUp() override { + fake_port = FakePortVerbs{}; + fake_port.active_speed = 128; // NDR + fake_port.active_width = 2; // 4x -> 400 Gbps encoded + context_ = std::make_unique(transport_); + auto& verbs = RdmaContextTestPeer::verbs(*context_); + verbs.ibv_query_port_default = fakeQueryPort; + verbs.ibv_query_port_speed = fakeQueryPortSpeed; + RdmaContextTestPeer::bindDevice(*context_, &fake_port.native, + std::make_shared()); + } + + void TearDown() override { + // The context never owned the placeholder; keep its destructor away + // from it. + RdmaContextTestPeer::unbindDevice(*context_); + } + + RdmaTransport transport_; + std::unique_ptr context_; +}; + +TEST_F(RdmaContextFakeVerbsTest, EffectiveSpeedPreferredWhenVerbReportsIt) { + fake_port.speed_100mbps = 2000; // LAG down to one 200G PF + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 200.0); + EXPECT_EQ(fake_port.query_speed_calls, 1); +} + +// A transient verb failure must not revert a degraded LAG to the higher +// encoded rate: the last known effective speed is held until a query +// succeeds again (a real recovery re-fires PORT_ACTIVE / SPEED_CHANGE). +TEST_F(RdmaContextFakeVerbsTest, EffectiveSpeedHeldWhenVerbFails) { + fake_port.speed_100mbps = 2000; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + ASSERT_DOUBLE_EQ(context_->linkSpeedGbps(), 200.0); + fake_port.query_speed_rc = EIO; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 200.0); + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 200.0); + EXPECT_EQ(context_->effectiveSpeedQueryFailures(), 2u); + // Recovery at a new speed is picked up again. + fake_port.query_speed_rc = 0; + fake_port.speed_100mbps = 4000; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 400.0); + EXPECT_EQ(context_->effectiveSpeedQueryFailures(), 2u); +} + +// A verb that succeeds but reports 0 means "nothing to say", not a failure: +// the encodings decide, as when the verb is absent. +TEST_F(RdmaContextFakeVerbsTest, EncodedRateWhenVerbReportsZero) { + fake_port.speed_100mbps = 2000; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + ASSERT_DOUBLE_EQ(context_->linkSpeedGbps(), 200.0); + fake_port.speed_100mbps = 0; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 400.0); + EXPECT_EQ(context_->effectiveSpeedQueryFailures(), 0u); +} + +TEST_F(RdmaContextFakeVerbsTest, EncodedRateWhenVerbAbsent) { + fake_port.speed_100mbps = 2000; + RdmaContextTestPeer::verbs(*context_).ibv_query_port_speed = nullptr; + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 400.0); + EXPECT_EQ(fake_port.query_speed_calls, 0); +} + +TEST_F(RdmaContextFakeVerbsTest, RefreshSeesRenegotiatedLink) { + ASSERT_EQ(context_->refreshPortAttributes(), 0); + ASSERT_DOUBLE_EQ(context_->linkSpeedGbps(), 400.0); + fake_port.active_speed = 32; // came back as EDR 4x + ASSERT_EQ(context_->refreshPortAttributes(), 0); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 100.0); +} + +TEST_F(RdmaContextFakeVerbsTest, RefreshFailsCleanlyWhenQueryPortFails) { + ASSERT_EQ(context_->refreshPortAttributes(), 0); + fake_port.query_port_rc = EIO; + fake_port.active_speed = 32; + EXPECT_EQ(context_->refreshPortAttributes(), -1); + EXPECT_DOUBLE_EQ(context_->linkSpeedGbps(), 400.0); // cached values kept +} + +// The whole runtime chain: a port event reaches Workers::applyContextEvent, +// the context re-reads its (fake) port, and the selector is re-seeded only +// when the speed actually changed -- with the device marked available again +// on the new rate, not the old one. +TEST(RdmaContextEventChainTest, PortActiveReseedsOnlyWhenTheSpeedChanged) { + auto topology = std::make_shared(); + ASSERT_TRUE(topology + ->parse(R"({"nics":[ + {"name":"mc-tcp-0","type":1,"numa_node":0}, + {"name":"mc-absent-rnic-1","type":0,"numa_node":0}]})") + .ok()); + RdmaTransport transport; + RdmaTransportTestPeer::bindTopology(transport, topology); + ASSERT_EQ(RdmaTransportTestPeer::initializeContexts(transport), 0u); + auto workers = RdmaTransportTestPeer::makeWorkers(transport); + auto* selector = workers->getDeviceSelector(); + constexpr int kDev = 1; + auto& context = *RdmaTransportTestPeer::contextSet(transport)[kDev]; + + fake_port = FakePortVerbs{}; + fake_port.active_speed = 128; + fake_port.active_width = 2; // 400G + auto& verbs = RdmaContextTestPeer::verbs(context); + verbs.ibv_query_port_default = fakeQueryPort; + verbs.ibv_query_port_speed = fakeQueryPortSpeed; + RdmaContextTestPeer::bindDevice(context, &fake_port.native, + std::make_shared()); + + // Pretend init seeded it at 400G and it learned ~45 GB/s since. + ASSERT_EQ(context.refreshPortAttributes(), 0); + ASSERT_TRUE( + selector->setDeviceBandwidth(kDev, context.linkSpeedGbps()).ok()); + ASSERT_TRUE(selector->setDeviceAvailable(kDev, true).ok()); + for (int i = 0; i < 64; ++i) + ASSERT_TRUE(selector->release(kDev, 1 << 20, (1 << 20) / 45e9).ok()); + ASSERT_NEAR(selector->getAggregateEwmaBandwidth(), 45e9, 45e9 * 0.02); + + ibv_async_event event{}; + event.event_type = IBV_EVENT_PORT_ACTIVE; + event.element.port_num = context.portNum(); + + // Same speed after the flap: keep what was learned. + RdmaTransportTestPeer::applyContextEvent(*workers, kDev, context, event); + EXPECT_NEAR(selector->getAggregateEwmaBandwidth(), 45e9, 45e9 * 0.02); + EXPECT_TRUE(selector->isDeviceAvailable(kDev)); + + // LAG lost a PF: the effective speed halves, the seed and clamp follow. + fake_port.speed_100mbps = 2000; + RdmaTransportTestPeer::applyContextEvent(*workers, kDev, context, event); + EXPECT_DOUBLE_EQ(context.linkSpeedGbps(), 200.0); + EXPECT_DOUBLE_EQ(selector->getAggregateEwmaBandwidth(), 25e9); + EXPECT_TRUE(selector->isDeviceAvailable(kDev)); + + RdmaContextTestPeer::unbindDevice(context); +} + TEST(RdmaContextPortSpeedTest, RefreshOnInertContextIsRejected) { RdmaTransport transport; RdmaContext context(transport); @@ -487,6 +713,75 @@ TEST(RdmaContextPortSpeedTest, RefreshOnInertContextIsRejected) { EXPECT_DOUBLE_EQ(context.linkSpeedGbps(), 0.0); } +// The 1 Hz recovery poll only reactivates contexts that are actually paused. +// An inert slot never opened a device, so there is no port to consult and +// nothing to hand back to the selector. +TEST_F(RdmaContextEventTest, RecoveryPollLeavesInertContextsAlone) { + fire(IBV_EVENT_PORT_ERR, ourPort()); + ASSERT_FALSE(selector_->isDeviceAvailable(kDev)); + ASSERT_EQ(context().status(), RdmaContext::DEVICE_UNINIT); + + RdmaTransportTestPeer::resumePausedContexts(*workers_); + + EXPECT_FALSE(selector_->isDeviceAvailable(kDev)); +} + +TEST(RdmaContextPortStateTest, QueryOnInertContextIsRejected) { + RdmaTransport transport; + RdmaContext context(transport); + ASSERT_EQ(context.status(), RdmaContext::DEVICE_UNINIT); + + ibv_port_state state = IBV_PORT_DOWN; + EXPECT_EQ(context.queryPortState(&state), -1); + EXPECT_EQ(state, IBV_PORT_DOWN) << "a failed query must not invent a state"; + EXPECT_EQ(context.queryPortState(nullptr), -1); +} + +// A context paused by IBV_EVENT_PORT_ERR whose IBV_EVENT_PORT_ACTIVE never +// arrives (the async fd is edge-triggered, so a queued event can be stranded) +// used to stay DEVICE_PAUSED for the rest of the process' life, failing every +// transfer routed to that NIC. The monitor thread's poll must notice that the +// hardware reports the port as up and reactivate the context. +TEST(RdmaPausedContextRecoveryTest, PollActivatesPausedContextWithLivePort) { + if (!hasRdmaDevice()) GTEST_SKIP() << "no RDMA device detected"; + + int count = 0; + ibv_device** devices = ibv_get_device_list(&count); + ASSERT_NE(devices, nullptr); + ASSERT_GT(count, 0); + const std::string device_name = ibv_get_device_name(devices[0]); + ibv_free_device_list(devices); + + auto topology = std::make_shared(); + const std::string spec = R"({"nics":[{"name":")" + device_name + + R"(","type":0,"numa_node":0}]})"; + ASSERT_TRUE(topology->parse(spec).ok()); + + RdmaTransport transport; + RdmaTransportTestPeer::bindTopology(transport, topology); + RdmaTransportTestPeer::initializeContexts(transport); + const auto& contexts = RdmaTransportTestPeer::contextSet(transport); + ASSERT_EQ(contexts.size(), 1u); + RdmaContext& context = *contexts[0]; + if (context.status() != RdmaContext::DEVICE_ENABLED) + GTEST_SKIP() << device_name << " has no usable active port"; + + auto workers = RdmaTransportTestPeer::makeWorkers(transport); + auto* selector = workers->getDeviceSelector(); + ASSERT_NE(selector, nullptr); + + // Exactly the state IBV_EVENT_PORT_ERR leaves behind, minus the event + // that would have undone it. + context.pause(); + ASSERT_EQ(context.status(), RdmaContext::DEVICE_PAUSED); + ASSERT_TRUE(selector->setDeviceAvailable(0, false).ok()); + + RdmaTransportTestPeer::resumePausedContexts(*workers); + + EXPECT_EQ(context.status(), RdmaContext::DEVICE_ENABLED); + EXPECT_TRUE(selector->isDeviceAvailable(0)); +} + TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { if (!hasRdmaDevice()) GTEST_SKIP() << "no RDMA device detected"; diff --git a/mooncake-transfer-engine/tent/tests/strict_local_numa_test.cpp b/mooncake-transfer-engine/tent/tests/strict_local_numa_test.cpp new file mode 100644 index 0000000000..e381c9c48a --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/strict_local_numa_test.cpp @@ -0,0 +1,294 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/runtime/topology.h" +#include "tent/transport/rdma/quota.h" + +namespace mooncake { +namespace tent { +namespace { + +// cuda:0 (NUMA 0): local dev 0 + remote dev 1. cuda:1 (NUMA 2): only remote +// dev 1. +constexpr const char* kTwoTierTopology = R"json( +{ + "nics": [ + {"name": "mlx5_local", "type": 0, "numa_node": 0}, + {"name": "mlx5_remote", "type": 0, "numa_node": 1} + ], + "mems": [ + {"name": "cuda:0", "type": 1, "numa_node": 0, + "device_list": {"rank0": [0], "rank2": [1]}}, + {"name": "cuda:1", "type": 1, "numa_node": 2, + "device_list": {"rank2": [1]}} + ] +} +)json"; + +// Unknown NIC or memory NUMA (-1) must stay selectable. +constexpr const char* kUnknownNumaTopology = R"json( +{ + "nics": [ + {"name": "mlx5_known", "type": 0, "numa_node": 1}, + {"name": "mlx5_bond_0", "type": 0, "numa_node": -1} + ], + "mems": [ + {"name": "cuda:0", "type": 1, "numa_node": 0, + "device_list": {"rank2": [1]}}, + {"name": "*", "type": 4, "numa_node": -1, + "device_list": {"rank2": [0, 1]}} + ] +} +)json"; + +// Remote NIC in rank 0 (ROCm / MemoryProber closest-PCIe shape). +constexpr const char* kRemoteNicInFirstRankTopology = R"json( +{ + "nics": [ + {"name": "mlx5_close_remote", "type": 0, "numa_node": 1}, + {"name": "mlx5_far_local", "type": 0, "numa_node": 0} + ], + "mems": [ + {"name": "hip:0", "type": 2, "numa_node": 0, + "device_list": {"rank0": [0], "rank1": [1]}} + ] +} +)json"; + +std::shared_ptr makeTopology(const char* json = kTwoTierTopology) { + auto topo = std::make_shared(); + EXPECT_TRUE(topo->parse(json).ok()); + return topo; +} + +std::unique_ptr makeSelector(std::shared_ptr topo, + bool strict, bool smart) { + auto selector = std::make_unique(); + EXPECT_TRUE(selector->loadTopology(topo).ok()); + DeviceSelector::SchedulingParams params; + params.strict_local_numa = strict; + // Isolate NUMA behavior from QoS priority rotation. + params.enable_priority_filtering = false; + selector->setSchedulingParams(params); + selector->setSmartSelection(smart); + return selector; +} + +TEST(StrictLocalNumaTest, SmartModeNeverPicksCrossNumaWhenLocalExists) { + auto selector = + makeSelector(makeTopology(), /*strict=*/true, /*smart=*/true); + for (int i = 0; i < 64; ++i) { + std::vector devs; + ASSERT_TRUE(selector->allocate(4096, 4, 1024, "cuda:0", devs).ok()); + ASSERT_FALSE(devs.empty()); + for (int dev : devs) EXPECT_EQ(dev, 0) << "iteration " << i; + } +} + +TEST(StrictLocalNumaTest, SmartModeReturnsDeviceNotFoundWithoutLocalNic) { + auto selector = + makeSelector(makeTopology(), /*strict=*/true, /*smart=*/true); + std::vector devs; + auto status = selector->allocate(4096, 4, 1024, "cuda:1", devs); + EXPECT_TRUE(status.IsDeviceNotFound()) << status.ToString(); + EXPECT_TRUE(devs.empty()); +} + +TEST(StrictLocalNumaTest, NonStrictModeFallsBackToCrossNuma) { + auto selector = + makeSelector(makeTopology(), /*strict=*/false, /*smart=*/true); + std::vector devs; + ASSERT_TRUE(selector->allocate(4096, 4, 1024, "cuda:1", devs).ok()); + ASSERT_FALSE(devs.empty()); + for (int dev : devs) EXPECT_EQ(dev, 1); +} + +TEST(StrictLocalNumaTest, BaselineModeExcludesCrossNuma) { + auto selector = + makeSelector(makeTopology(), /*strict=*/true, /*smart=*/false); + for (int i = 0; i < 16; ++i) { + std::vector devs; + ASSERT_TRUE(selector->allocate(4096, 4, 1024, "cuda:0", devs).ok()); + ASSERT_FALSE(devs.empty()); + for (int dev : devs) EXPECT_EQ(dev, 0); + } + std::vector devs; + EXPECT_TRUE( + selector->allocate(4096, 4, 1024, "cuda:1", devs).IsDeviceNotFound()); +} + +TEST(StrictLocalNumaTest, BaselineNonStrictFallsBackToCrossNuma) { + auto selector = + makeSelector(makeTopology(), /*strict=*/false, /*smart=*/false); + std::vector devs; + ASSERT_TRUE(selector->allocate(4096, 4, 1024, "cuda:1", devs).ok()); + ASSERT_FALSE(devs.empty()); + for (int dev : devs) EXPECT_EQ(dev, 1); +} + +// Masking the local NIC must not fall back to the remote one under strict mode. +TEST(StrictLocalNumaTest, FallbackDoesNotReintroduceCrossNuma) { + const uint64_t mask = ~(1ULL << 0); // exclude the local NIC (dev 0) + + auto strict = makeSelector(makeTopology(), /*strict=*/true, /*smart=*/true); + std::vector devs; + auto status = + strict->allocate(4096, 4, 1024, "cuda:0", devs, PRIO_HIGH, mask); + EXPECT_TRUE(status.IsDeviceNotFound()) << status.ToString(); + EXPECT_TRUE(devs.empty()); + + // Same mask without strict still reaches the remote NIC. + auto relaxed = + makeSelector(makeTopology(), /*strict=*/false, /*smart=*/true); + std::vector relaxed_devs; + ASSERT_TRUE( + relaxed + ->allocate(4096, 4, 1024, "cuda:0", relaxed_devs, PRIO_HIGH, mask) + .ok()); + ASSERT_FALSE(relaxed_devs.empty()); + for (int dev : relaxed_devs) EXPECT_EQ(dev, 1); +} + +TEST(StrictLocalNumaTest, UnknownNicNumaIsNotExcluded) { + for (bool smart : {true, false}) { + auto selector = makeSelector(makeTopology(kUnknownNumaTopology), + /*strict=*/true, smart); + std::vector devs; + auto status = selector->allocate(4096, 4, 1024, "cuda:0", devs); + ASSERT_TRUE(status.ok()) << status.ToString() << ", smart=" << smart; + ASSERT_FALSE(devs.empty()); + for (int dev : devs) EXPECT_EQ(dev, 1); + } +} + +TEST(StrictLocalNumaTest, UnknownMemoryNumaIsNotExcluded) { + for (bool smart : {true, false}) { + auto selector = makeSelector(makeTopology(kUnknownNumaTopology), + /*strict=*/true, smart); + std::vector devs; + auto status = + selector->allocate(4096, 4, 1024, kWildcardLocation, devs); + ASSERT_TRUE(status.ok()) << status.ToString() << ", smart=" << smart; + EXPECT_EQ(devs.size(), 4u); + } +} + +TEST(StrictLocalNumaTest, RemoteNicPromotedToFirstRankIsStillExcluded) { + for (bool smart : {true, false}) { + auto selector = + makeSelector(makeTopology(kRemoteNicInFirstRankTopology), + /*strict=*/true, smart); + std::vector devs; + auto status = selector->allocate(4096, 4, 1024, "hip:0", devs); + ASSERT_TRUE(status.ok()) << status.ToString() << ", smart=" << smart; + ASSERT_FALSE(devs.empty()); + // dev 0 is closest by PCIe but on another NUMA node. + for (int dev : devs) EXPECT_EQ(dev, 1) << "smart=" << smart; + } +} + +// Priority matrix has no NUMA ids, so every NIC stays selectable. +TEST(StrictLocalNumaTest, PriorityMatrixTopologyKeepsSoftPenalty) { + auto topo = std::make_shared(); + ASSERT_TRUE(topo->parsePriorityMatrix( + R"json({"cpu:0": [["mlx5_0"], ["mlx5_1"]]})json") + .ok()); + + for (bool smart : {true, false}) { + auto selector = makeSelector(topo, /*strict=*/true, smart); + std::vector devs; + auto status = selector->allocate(4096, 4, 1024, "cpu:0", devs); + ASSERT_TRUE(status.ok()) << status.ToString() << ", smart=" << smart; + EXPECT_EQ(devs.size(), 4u); + } +} + +TEST(StrictLocalNumaTest, IsCrossNumaRequiresBothSidesKnown) { + auto topo = makeTopology(kUnknownNumaTopology); + const auto* cuda0 = topo->getMemEntry("cuda:0"); + const auto* wildcard = topo->getMemEntry(kWildcardLocation); + ASSERT_NE(cuda0, nullptr); + ASSERT_NE(wildcard, nullptr); + + EXPECT_TRUE(topo->isCrossNuma(*cuda0, 0)); // both known, differ + EXPECT_FALSE(topo->isCrossNuma(*cuda0, 1)); // NIC unknown + EXPECT_FALSE(topo->isCrossNuma(*wildcard, 0)); + EXPECT_FALSE(topo->isCrossNuma(*wildcard, 1)); + EXPECT_FALSE(topo->isCrossNuma(*cuda0, 42)); +} + +class StrictLocalNumaEnvTest : public ::testing::Test { + protected: + // Isolate from an ambient MC_TENT_CONF. + void SetUp() override { unsetenv("MC_TENT_CONF"); } + + void TearDown() override { + unsetenv("MC_STRICT_LOCAL_NUMA"); + unsetenv("MC_IB_PORT"); + } + + static bool loadStrictFlag(const char* env_value, bool default_value) { + setenv("MC_STRICT_LOCAL_NUMA", env_value, 1); + Config config; + ConfigHelper helper; + EXPECT_TRUE(helper.loadFromEnv(config).ok()); + return config.get("transports/rdma/strict_local_numa", default_value); + } +}; + +TEST_F(StrictLocalNumaEnvTest, AcceptsNumericAndTextualBooleans) { + for (const char* on : {"1", "true", "TRUE", "yes", "on"}) { + EXPECT_TRUE(loadStrictFlag(on, false)) << "value: " << on; + } + for (const char* off : {"0", "false", "no", "off"}) { + EXPECT_FALSE(loadStrictFlag(off, true)) << "value: " << off; + } +} + +// Unparsable env must not overwrite an existing JSON value. +TEST_F(StrictLocalNumaEnvTest, UnrecognizedValueKeepsConfiguredValue) { + setenv("MC_STRICT_LOCAL_NUMA", "maybe", 1); + ConfigHelper helper; + + Config preset; + preset.set("transports/rdma/strict_local_numa", true); + ASSERT_TRUE(helper.loadFromEnv(preset).ok()); + EXPECT_TRUE(preset.get("transports/rdma/strict_local_numa", false)); + + Config unset; + ASSERT_TRUE(helper.loadFromEnv(unset).ok()); + EXPECT_FALSE(unset.get("transports/rdma/strict_local_numa", false)); +} + +// setFromString() must still store integer env vars as integers. +TEST_F(StrictLocalNumaEnvTest, IntegerEnvVarsAreUnaffected) { + setenv("MC_IB_PORT", "1", 1); + Config config; + ConfigHelper helper; + ASSERT_TRUE(helper.loadFromEnv(config).ok()); + EXPECT_EQ(config.get("transports/rdma/device/port", -1), 1); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp index bb1c41a0ff..1f7cf6e74c 100644 --- a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp @@ -283,6 +283,34 @@ TEST(TransferEngineConfigOverrideTest, EXPECT_EQ(config.get("transports/rdma/bind_address", ""), "10.0.0.2"); } +// MOONCAKE_LOCAL_HOSTNAME is the classic Transfer Engine + store env var that +// names the local host for RPC binding and segment identity. TENT must honor +// the same env so a single MOONCAKE_LOCAL_HOSTNAME works across both engines; +// otherwise TENT's auto-discovery can pick a container/CNI IP (e.g. 10.154.0.1) +// instead of the RDMA-network IP, breaking cross-node RDMA handshake. +TEST(TransferEngineConfigOverrideTest, + LocalHostnameEnvLoadsIntoRpcServerHostname) { + EnvVarGuard guard("MOONCAKE_LOCAL_HOSTNAME", "10.0.0.2"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + EXPECT_EQ(config.get("rpc_server_hostname", ""), "10.0.0.2"); +} + +TEST(TransferEngineConfigOverrideTest, LocalHostnameEnvOverridesMcTentConf) { + EnvVarGuard conf_guard("MC_TENT_CONF", + R"({"rpc_server_hostname":"10.0.0.1"})"); + EnvVarGuard host_guard("MOONCAKE_LOCAL_HOSTNAME", "10.0.0.2"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + // Legacy env must override MC_TENT_CONF, same precedence as + // MC_RDMA_BIND_ADDRESS (env wins so per-pod injection works). + EXPECT_EQ(config.get("rpc_server_hostname", ""), "10.0.0.2"); +} + TEST(TransferEngineConfigOverrideTest, LegacyRdmaSliceAffinityLogEnvLoadsIntoTentConfig) { EnvVarGuard guard("MC_LOG_RDMA_SLICE_AFFINITY", "true"); @@ -462,6 +490,61 @@ TEST(TransferEngineConfigOverrideTest, } } +TEST(TransferEngineConfigOverrideTest, + RegisterLocalMemoryIgnoresIncompatibleCallerLocation) { +#ifdef _WIN32 + GTEST_SKIP() << "Requires local HTTP metadata server support"; +#else + // registerLocalMemory observes existing memory; the NUMA probe is the + // source of truth for transport selection. A caller-supplied location + // must not overwrite the probe with an unknown or incompatible type. + // The store's buildSegmentsLocation() emits "segments:4096:0,1" (a + // classic TE NUMA-segment encoding TENT's type system does not + // understand); blindly adopting it made getTypeEnum() return + // MTYPE_UNKNOWN and broke transport selection (warmup "Unable to find + // registered buffer"). The override must validate and keep the probe. + const auto live_port = reserveUnusedTcpPort(); + TestHttpMetadataServer metadata_server(live_port); + ASSERT_TRUE(metadata_server.start()); + const auto live_endpoint = buildHttpMetadataEndpoint(live_port); + + auto config = std::make_shared(); + config->set("metadata_type", "http"); + config->set("metadata_servers", live_endpoint); + config->set("local_segment_name", kSegmentName); + config->set("rpc_server_hostname", kLoopbackHostname); + config->set("rpc_server_port", "0"); + configureTcpOnlyTransports(*config); + + TransferEngineImpl engine(config); + ASSERT_TRUE(engine.available()); + + // Register a host buffer with an incompatible caller location. The + // probe classifies host memory as "cpu:N"; "segments:..." is an unknown + // type to TENT and must be ignored in favor of the probe. + constexpr size_t kBufSize = 4096; + std::vector buf(kBufSize, 0); + MemoryOptions options; + options.location = "segments:4096:0,1"; + std::vector addrs = {buf.data()}; + std::vector sizes = {kBufSize}; + ASSERT_TRUE(engine.registerLocalMemory(addrs, sizes, options).ok()); + + SegmentInfo info; + ASSERT_TRUE(engine.getSegmentInfo(LOCAL_SEGMENT_ID, info).ok()); + ASSERT_EQ(info.buffers.size(), 1u); + // The buffer location must be the probed "cpu:..." form, NOT the + // caller-supplied "segments:..." string. + const auto& loc = info.buffers[0].location; + EXPECT_NE(loc.find("cpu"), std::string::npos) + << "expected probed cpu location, got '" << loc << "'"; + EXPECT_EQ(loc.find("segments"), std::string::npos) + << "caller 'segments:' must not leak into buffer location"; + + engine.unregisterLocalMemory(buf.data(), kBufSize); +#endif +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index f42a974f4f..15bdda6e85 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -188,6 +188,7 @@ TEST(TransportSelectorTest, TransportTypeNameMapping) { EXPECT_STREQ(transportTypeName(SUNRISE_LINK), "sunrise_link"); EXPECT_STREQ(transportTypeName(UB), "ub"); EXPECT_STREQ(transportTypeName(MPCOMM), "mpcomm"); + EXPECT_STREQ(transportTypeName(HP_TCP), "hp_tcp"); } TEST(TransportSelectorTest, ParseTransportType) { @@ -203,6 +204,7 @@ TEST(TransportSelectorTest, ParseTransportType) { EXPECT_EQ(parseTransportType("sunrise_link"), SUNRISE_LINK); EXPECT_EQ(parseTransportType("ub"), UB); EXPECT_EQ(parseTransportType("mpcomm"), MPCOMM); + EXPECT_EQ(parseTransportType("hp_tcp"), HP_TCP); EXPECT_EQ(parseTransportType("unknown"), UNSPEC); } @@ -212,7 +214,7 @@ TEST(TransportSelectorTest, UbTransportNameRoundTrips) { EXPECT_EQ(parseTransportType(name), UB); } -TEST(TransportTypeTest, WireValuesRemainStableWithUbAppended) { +TEST(TransportTypeTest, WireValuesRemainStableWithHpTcpAppended) { EXPECT_EQ(static_cast(UNSPEC), 0); EXPECT_EQ(static_cast(RDMA), 1); EXPECT_EQ(static_cast(MNNVL), 2); @@ -226,7 +228,8 @@ TEST(TransportTypeTest, WireValuesRemainStableWithUbAppended) { EXPECT_EQ(static_cast(TPU), 10); EXPECT_EQ(static_cast(UB), 11); EXPECT_EQ(static_cast(MPCOMM), 12); - EXPECT_EQ(static_cast(kNumTransportTypes), 13); + EXPECT_EQ(static_cast(HP_TCP), 13); + EXPECT_EQ(static_cast(kNumTransportTypes), 14); } // MPComm is appended after UB, so it takes wire value 12. The same integer is @@ -245,6 +248,14 @@ TEST(TransportTypeTest, MpcommWireValueMatchesCApiAndRoundTrips) { EXPECT_EQ(parseTransportType("mpcomm"), MPCOMM); } +TEST(TransportTypeTest, HpTcpWireValueMatchesCApiAndRoundTrips) { + EXPECT_EQ(static_cast(HP_TCP), 13); + EXPECT_EQ(TRANSPORT_HP_TCP, static_cast(HP_TCP)); + EXPECT_EQ(c_to_transport_hint(TRANSPORT_HP_TCP), HP_TCP); + EXPECT_STREQ(transportTypeName(HP_TCP), "hp_tcp"); + EXPECT_EQ(parseTransportType("hp_tcp"), HP_TCP); +} + // Topology::NicType is serialized as an integer. These values are therefore a // wire-compatibility contract, not merely an implementation detail. TEST(TopologyTest, NicTypeWireValuesRemainStableWithUbAppended) { @@ -307,6 +318,81 @@ TEST(TopologyTest, UbDeviceAttributesRoundTripThroughJson) { EXPECT_EQ(round_tripped->device_attrs, ub.device_attrs); } +TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { + UbBootstrapDesc source; + source.protocol_version = 1; + source.segment_name = "peer-segment"; + source.local_nic_path = "local/ub-device-0/eid-2"; + source.peer_nic_path = "peer/ub-device-1/eid-3"; + source.local_device_name = "ub-device-0"; + source.local_device_id = 7; + source.local_eid_index = 2; + source.local_eid = "e1:02:03:04:05:06:07:08"; + source.jetty_ids = {11, 12}; + source.jetty_uasids = {21, 22}; + source.endpoint_generation = 41; + source.segment_generation = 42; + source.capabilities = {"read", "write"}; + source.reply_msg = "ok"; + + const auto parsed = json(source).get(); + EXPECT_EQ(parsed.protocol_version, source.protocol_version); + EXPECT_EQ(parsed.segment_name, source.segment_name); + EXPECT_EQ(parsed.local_nic_path, source.local_nic_path); + EXPECT_EQ(parsed.peer_nic_path, source.peer_nic_path); + EXPECT_EQ(parsed.local_device_name, source.local_device_name); + EXPECT_EQ(parsed.local_device_id, source.local_device_id); + EXPECT_EQ(parsed.local_eid_index, source.local_eid_index); + EXPECT_EQ(parsed.local_eid, source.local_eid); + EXPECT_EQ(parsed.jetty_ids, source.jetty_ids); + 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.capabilities, source.capabilities); + EXPECT_EQ(parsed.reply_msg, source.reply_msg); +} + +TEST(ControlPlaneTest, UbBootstrapJsonDefaultsOptionalFields) { + const auto parsed = json{{"protocol_version", 1}}.get(); + EXPECT_EQ(parsed.protocol_version, 1u); + EXPECT_TRUE(parsed.segment_name.empty()); + EXPECT_TRUE(parsed.local_nic_path.empty()); + EXPECT_TRUE(parsed.peer_nic_path.empty()); + EXPECT_TRUE(parsed.local_device_name.empty()); + EXPECT_EQ(parsed.local_device_id, -1); + EXPECT_EQ(parsed.local_eid_index, -1); + EXPECT_TRUE(parsed.local_eid.empty()); + EXPECT_TRUE(parsed.jetty_ids.empty()); + EXPECT_TRUE(parsed.jetty_uasids.empty()); + EXPECT_EQ(parsed.endpoint_generation, 0u); + EXPECT_EQ(parsed.segment_generation, 0u); + EXPECT_TRUE(parsed.capabilities.empty()); + EXPECT_TRUE(parsed.reply_msg.empty()); +} + +TEST(ControlPlaneTest, UbBootstrapJsonRejectsMissingOrUnknownVersion) { + EXPECT_THROW((void)json::object().get(), + std::invalid_argument); + const json unknown_version{{"protocol_version", 2}}; + EXPECT_THROW((void)unknown_version.get(), + std::invalid_argument); +} + +TEST(ControlPlaneTest, UbBootstrapRpcIdIsAppendedWithoutRenumbering) { + EXPECT_EQ(static_cast(GetSegmentDesc), 1); + EXPECT_EQ(static_cast(BootstrapRdma), 2); + EXPECT_EQ(static_cast(SendData), 3); + EXPECT_EQ(static_cast(RecvData), 4); + EXPECT_EQ(static_cast(Notify), 5); + EXPECT_EQ(static_cast(Probe), 6); + EXPECT_EQ(static_cast(Delegate), 7); + EXPECT_EQ(static_cast(Pin), 8); + EXPECT_EQ(static_cast(Unpin), 9); + EXPECT_EQ(static_cast(SubscribeSegmentUpdate), 10); + EXPECT_EQ(static_cast(NotifySegmentUpdated), 11); + EXPECT_EQ(static_cast(BootstrapUb), 12); +} + // --------------------------------------------------------------------------- // Test legacy mode // --------------------------------------------------------------------------- @@ -613,44 +699,84 @@ TEST(TransportSelectorTest, RocmMemoryTypeSupported) { TEST(TransportSelectorTest, ConfigBasedPolicySelection) { auto conf = std::make_shared(); + ASSERT_TRUE( + conf->load( + R"({"policy":[{"name":"hp_tcp_memory","segment_type":"memory","transports":["hp_tcp"]}]})") + .ok()); + TransportSelector selector(conf); - // Set up a custom policy via JSON config - conf->set("policy", json::array()); - auto policies = conf->getArray("policy"); + std::array, kSupportedTransportTypes> + transports{}; + transports[HP_TCP] = std::make_shared(HP_TCP); + static_cast(transports[HP_TCP].get())->setDramToDram(true); + + const std::vector buffer_transports = {HP_TCP}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + EXPECT_EQ(selector.select(ctx, transports).transport, HP_TCP); +} + +TEST(TransportSelectorTest, PolicyCanPreferUb) { + auto conf = std::make_shared(); json policy; - policy["name"] = "test_memory_policy"; + policy["name"] = "ub-preferred"; policy["segment_type"] = "memory"; - policy["transports"] = {"tcp", "rdma"}; // Prefer TCP over RDMA - - // We can't easily modify the config's internal JSON structure, - // so this test verifies the selector at least loads without error + policy["transports"] = {"ub", "rdma"}; + conf->set("policy", json::array({policy})); TransportSelector selector(conf); - - // Default behavior should still work std::array, kSupportedTransportTypes> transports{}; + transports[UB] = std::make_shared(UB); transports[RDMA] = std::make_shared(RDMA); - transports[TCP] = std::make_shared(TCP); + static_cast(transports[UB].get())->setDramToDram(true); + static_cast(transports[RDMA].get())->setDramToDram(true); - auto* rdma = static_cast(transports[RDMA].get()); - rdma->setDramToDram(true); - auto* tcp = static_cast(transports[TCP].get()); - tcp->setDramToDram(true); + std::vector buffer_transports = {RDMA, UB}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "ub-preferred"; - std::vector buffer_transports = {RDMA, TCP}; + EXPECT_EQ(selector.select(ctx, transports).transport, UB); +} + +TEST(TransportSelectorTest, PolicyFallsBackWhenUbIsIncapable) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "ub-with-rdma-fallback"; + policy["segment_type"] = "memory"; + policy["transports"] = {"ub", "rdma"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[UB] = std::make_shared(UB); + transports[RDMA] = std::make_shared(RDMA); + // UB is installed but deliberately lacks dram_to_dram capability. + static_cast(transports[RDMA].get())->setDramToDram(true); + std::vector buffer_transports = {UB, RDMA}; SelectionContext ctx; ctx.segment_type = SegmentType::Memory; ctx.same_machine = false; ctx.local_memory_type = MTYPE_CPU; ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "ub-with-rdma-fallback"; - auto result = selector.select(ctx, transports); - // With default policies, should use buffer_transports order (RDMA first) - EXPECT_EQ(result.transport, RDMA); + EXPECT_EQ(selector.select(ctx, transports).transport, RDMA); } // --------------------------------------------------------------------------- @@ -719,6 +845,36 @@ TEST(TransportSelectorTest, HintIsPrependedToCandidateList) { EXPECT_EQ(r1.transport, RDMA); } +TEST(TransportSelectorTest, UbHintIsPrependedAndThenFallsBackByIndex) { + auto conf = std::make_shared(); + TransportSelector selector(conf); + + std::array, kSupportedTransportTypes> + transports{}; + transports[UB] = std::make_shared(UB); + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[UB].get())->setDramToDram(true); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA, UB}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/0, /*hint=*/UB).transport, + UB); + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/1, /*hint=*/UB).transport, + RDMA); + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/2, /*hint=*/UB).transport, + UNSPEC); +} + TEST(TransportSelectorTest, HintNotInMatchingPolicyReturnsUnspec) { // Selector mode: the matching policy's transports list is the // authorization whitelist. A hint that is not in it produces UNSPEC diff --git a/mooncake-transfer-engine/tent/tests/ub_core_test.cpp b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp new file mode 100644 index 0000000000..4581130273 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp @@ -0,0 +1,364 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include + +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" +#include "tent/transport/ub/device_selection.h" + +namespace mooncake::tent::ub { +namespace { + +TEST(UbParamsTest, ParsesStrictValuesAndAllowsRetryDisable) { + Config config; + config.set("transports/ub/max_retries", 0); + config.set("transports/ub/worker_count", 2); + config.set("transports/ub/device_filter", + std::vector{"ub:fake0:eid0"}); + UbParams params; + ASSERT_TRUE(UbParams::FromConfig(config, params).ok()); + EXPECT_EQ(params.max_retries, 0U); + EXPECT_EQ(params.worker_count, 2U); + ASSERT_EQ(params.device_filter.size(), 1U); + + config.set("transports/ub/worker_count", -1); + EXPECT_TRUE(UbParams::FromConfig(config, params).IsInvalidArgument()); +} + +Request makeRequest(size_t length) { + Request request{}; + request.opcode = Request::WRITE; + request.length = length; + return request; +} + +UbPostPath makePath(Topology::NicID local_id = 0, int remote_id = 1, + uint64_t generation = 1) { + return UbPostPath{local_id, 42, remote_id, generation}; +} + +TEST(UbSliceTest, RetryRejectsStaleCompletionAndNotifiesTerminalOnce) { + int callback_count = 0; + TransferStatus callback_status{INITIAL, 0}; + auto task = UbTask::create( + makeRequest(128), + [&](const TransferStatus& status) { + ++callback_count; + callback_status = status; + }, + 1); + auto slice = task->addSlice(UbSliceSpec{nullptr, 0, 128, 0, 1}, 2); + ASSERT_NE(slice, nullptr); + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(slice->markQueued(3)); + + auto first_attempt = slice->beginAttempt(makePath(0, 1, 10), 4); + ASSERT_TRUE(first_attempt.has_value()); + auto first_completion = slice->completionToken(*first_attempt); + ASSERT_TRUE(first_completion.has_value()); + ASSERT_TRUE(first_completion->markPosted(5)); + EXPECT_EQ(first_completion->resolve(TIMEOUT, 0, true, 6), + UbAttemptResolution::kRetryScheduled); + EXPECT_EQ(task->transferStatus().s, PENDING); + EXPECT_EQ(slice->snapshot().retry_count, 1U); + + ASSERT_TRUE(slice->markQueued(7)); + auto second_attempt = slice->beginAttempt(makePath(0, 2, 11), 8); + ASSERT_TRUE(second_attempt.has_value()); + auto second_completion = slice->completionToken(*second_attempt); + ASSERT_TRUE(second_completion.has_value()); + ASSERT_TRUE(second_completion->markPosted(9)); + + // A completion from the retired endpoint generation cannot resolve the + // replacement attempt. + EXPECT_EQ(first_completion->resolve(COMPLETED, 128, false, 10), + UbAttemptResolution::kIgnored); + EXPECT_EQ(second_completion->resolve(COMPLETED, 128, false, 11), + UbAttemptResolution::kTerminal); + EXPECT_EQ(second_completion->resolve(COMPLETED, 128, false, 12), + UbAttemptResolution::kIgnored); + + EXPECT_EQ(callback_count, 1); + EXPECT_EQ(callback_status.s, COMPLETED); + EXPECT_EQ(callback_status.transferred_bytes, 128U); + const auto task_snapshot = task->snapshot(); + EXPECT_EQ(task_snapshot.remaining_slices, 0U); + EXPECT_NE(task_snapshot.terminal_ns, 0U); +} + +TEST(UbSliceTest, CancellationStopsQueuedWorkAndDrainsPostedWork) { + int callback_count = 0; + TransferStatus callback_status{INITIAL, 0}; + auto task = + UbTask::create(makeRequest(128), [&](const TransferStatus& status) { + ++callback_count; + callback_status = status; + }); + auto posted = task->addSlice(UbSliceSpec{nullptr, 0, 64, 0, 1}); + auto queued = task->addSlice(UbSliceSpec{nullptr, 64, 64, 64, 1}); + ASSERT_NE(posted, nullptr); + ASSERT_NE(queued, nullptr); + ASSERT_TRUE(task->seal()); + + ASSERT_TRUE(posted->markQueued(1)); + auto attempt = posted->beginAttempt(makePath(), 2); + ASSERT_TRUE(attempt.has_value()); + auto completion = posted->completionToken(*attempt); + ASSERT_TRUE(completion.has_value()); + ASSERT_TRUE(completion->markPosted(3)); + ASSERT_TRUE(queued->markQueued(4)); + + EXPECT_EQ(task->requestCancellation(5), 1U); + EXPECT_EQ(task->transferStatus().s, PENDING); + EXPECT_EQ(posted->snapshot().state, UbSliceState::kPosted); + EXPECT_EQ(queued->snapshot().state, UbSliceState::kCanceled); + + // Posted work is not fabricated as immediately canceled. Its real + // completion drains first, then the task reaches its single terminal + // state (CANCELED because another slice never reached the device). + EXPECT_EQ(completion->resolve(COMPLETED, 64, false, 6), + UbAttemptResolution::kTerminal); + EXPECT_EQ(completion->resolve(FAILED, 0, false, 7), + UbAttemptResolution::kIgnored); + EXPECT_EQ(callback_count, 1); + EXPECT_EQ(callback_status.s, CANCELED); + EXPECT_EQ(callback_status.transferred_bytes, 64U); +} + +TEST(UbSliceTest, ConcurrentDuplicateCompletionsChooseOneTerminalWinner) { + std::atomic callback_count{0}; + auto task = UbTask::create(makeRequest(64), [&](const TransferStatus&) { + callback_count.fetch_add(1, std::memory_order_relaxed); + }); + auto slice = task->addSlice(UbSliceSpec{nullptr, 0, 64, 0, 0}); + ASSERT_NE(slice, nullptr); + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(slice->markQueued()); + auto attempt = slice->beginAttempt(makePath()); + ASSERT_TRUE(attempt.has_value()); + auto completion = slice->completionToken(*attempt); + ASSERT_TRUE(completion.has_value()); + ASSERT_TRUE(completion->markPosted()); + + std::atomic terminal_winners{0}; + std::vector threads; + for (int i = 0; i < 8; ++i) { + threads.emplace_back([&] { + if (completion->resolve(COMPLETED, 64, false) == + UbAttemptResolution::kTerminal) { + terminal_winners.fetch_add(1, std::memory_order_relaxed); + } + }); + } + for (auto& thread : threads) thread.join(); + + EXPECT_EQ(terminal_winners.load(std::memory_order_relaxed), 1); + EXPECT_EQ(callback_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(task->transferStatus().s, COMPLETED); + EXPECT_EQ(task->transferStatus().transferred_bytes, 64U); +} + +TEST(UbQuotaTest, EnforcesBothLevelsAndReleasesIdempotently) { + QuotaManager quota(/*default_device_limits=*/{100, 2}, + /*default_path_limits=*/{60, 1}); + const auto first_path = makePath(0, 1, 1); + const auto second_path = makePath(0, 2, 1); + + auto reservation = quota.tryAcquire(first_path, 60); + ASSERT_TRUE(reservation.has_value()); + EXPECT_FALSE(quota.tryAcquire(first_path, 1).has_value()); + // The second path has room, but the physical-device byte cap is shared. + EXPECT_FALSE(quota.tryAcquire(second_path, 50).has_value()); + + auto copied_token = *reservation; + copied_token.path = second_path; + copied_token.bytes = 1; + copied_token.wrs = 99; + EXPECT_TRUE(quota.release(copied_token)); + EXPECT_FALSE(quota.release(*reservation)); + + const auto device = quota.deviceStats(0); + const auto path = quota.pathStats(first_path); + const auto aggregate = quota.aggregateStats(); + EXPECT_EQ(device.usage, QuotaUsage{}); + EXPECT_EQ(path.usage, QuotaUsage{}); + EXPECT_EQ(aggregate.usage, QuotaUsage{}); + EXPECT_EQ(aggregate.active_reservations, 0U); + EXPECT_EQ(aggregate.duplicate_release_attempts, 1U); + EXPECT_TRUE(quota.tryAcquire(first_path, 60).has_value()); +} + +TEST(UbRailMonitorTest, PausesOnErrorWindowAndRecoversAfterCooldown) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/100, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(1), -1.0); + monitor.recordSuccess(path, 100, 10, 50); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(50), 10'000'000'000.0); + + monitor.recordError(path, 100); + monitor.recordTimeout(path, 110); + EXPECT_FALSE(monitor.available(path, 159)); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(159), 0.0); + + EXPECT_TRUE(monitor.available(path, 160)); + const auto stats = monitor.stats(path, 160); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.completion_errors, 2U); + EXPECT_EQ(stats.timeouts, 1U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(160), 10'000'000'000.0); +} + +TEST(UbRailMonitorTest, ErrorsOutsideWindowDoNotPauseRail) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/10, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + monitor.recordError(path, 1); + monitor.recordError(path, 11); + EXPECT_TRUE(monitor.available(path, 11)); + EXPECT_EQ(monitor.stats(path, 11).errors_in_window, 1U); +} + +TEST(UbRailMonitorTest, OutOfOrderErrorsUseLatestEventForCooldown) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/100, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + monitor.recordError(path, 120); + monitor.recordError(path, 100); + + auto stats = monitor.stats(path, 120); + EXPECT_TRUE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 2U); + EXPECT_EQ(stats.pause_started_ns, 120U); + EXPECT_EQ(stats.cooldown_until_ns, 170U); + EXPECT_FALSE(monitor.available(path, 169)); + EXPECT_TRUE(monitor.available(path, 170)); + + stats = monitor.stats(path, 170); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); + + // Neither an older query nor an error from the recovered epoch may + // rewind the rail or resurrect its completed pause. + EXPECT_TRUE(monitor.available(path, 130)); + monitor.recordError(path, 110); + stats = monitor.stats(path, 130); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.completion_errors, 3U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); +} + +TEST(UbRailMonitorTest, OutOfOrderExpiredErrorDoesNotTriggerPause) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/50, + /*cooldown_ns=*/25, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + monitor.recordError(path, 200); + monitor.recordError(path, 100); + + const auto stats = monitor.stats(path, 200); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 1U); + EXPECT_EQ(stats.completion_errors, 2U); + EXPECT_EQ(stats.last_error_ns, 200U); + EXPECT_EQ(stats.pauses, 0U); + EXPECT_EQ(stats.recoveries, 0U); +} + +DeviceInfo makeDevice(std::string native_name, std::string topology_name = "") { + DeviceInfo info; + info.native_device_name = std::move(native_name); + info.topology_name = topology_name.empty() + ? ("ub:" + info.native_device_name + ":eid0") + : std::move(topology_name); + info.active = true; + return info; +} + +TEST(UbDeviceSelectionTest, DetectsBondingNamesCaseInsensitively) { + EXPECT_TRUE(isBondingDeviceName("bonding_dev_0")); + EXPECT_TRUE(isBondingDeviceName("Bonding_Dev_0")); + EXPECT_TRUE(isBondingDeviceName("ub:bonding_dev_0:eid0")); + EXPECT_TRUE(isBondingDeviceName("foo_bond_bar")); + EXPECT_FALSE(isBondingDeviceName("udmac1d1e2")); + EXPECT_FALSE(isBondingDeviceName("ub:udmac1d1e2:eid0")); +} + +TEST(UbDeviceSelectionTest, PrefersBondingWhenPresentAndFilterEmpty) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("bonding_dev_0"), + makeDevice("udmac0d0e1"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/false); + ASSERT_EQ(selected.size(), 1U); + EXPECT_EQ(selected[0].native_device_name, "bonding_dev_0"); +} + +TEST(UbDeviceSelectionTest, KeepsAllDevicesWhenNoBondingPresent) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("udmac0d0e1"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/false); + ASSERT_EQ(selected.size(), 2U); + EXPECT_EQ(selected[0].native_device_name, "udmac1d1e2"); + EXPECT_EQ(selected[1].native_device_name, "udmac0d0e1"); +} + +TEST(UbDeviceSelectionTest, ExplicitFilterSkipsAutoPrefer) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("bonding_dev_0"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/true); + ASSERT_EQ(selected.size(), 2U); +} + +TEST(UbDeviceSelectionTest, EmptyInputStaysEmpty) { + const auto selected = + preferBondingDevicesIfPresent({}, /*explicit_filter=*/false); + EXPECT_TRUE(selected.empty()); +} + +} // namespace +} // namespace mooncake::tent::ub 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 new file mode 100644 index 0000000000..3db656fe58 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -0,0 +1,957 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/segment_registry.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/endpoint_store.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/topology_attrs.h" +#include "tent/transport/ub/ub_transport.h" +#include "tent/transport/ub/workers.h" + +namespace mooncake::tent::ub { +namespace { + +class FakeContext final : public Context { + public: + explicit FakeContext(DeviceInfo info) : info_(std::move(info)) {} + bool valid() const noexcept override { return valid_; } + const DeviceInfo& deviceInfo() const noexcept override { return info_; } + int asyncFd() const noexcept override { return -1; } + void close() { valid_ = false; } + + private: + DeviceInfo info_; + bool valid_{true}; +}; + +class FakeJfc final : public Jfc { + public: + bool valid() const noexcept override { return valid_; } + int eventFd() const noexcept override { return -1; } + + void push(Completion completion) { + std::lock_guard lock(mutex_); + completions_.push_back(completion); + } + void poll(size_t maximum, std::vector& output) { + std::lock_guard lock(mutex_); + while (!completions_.empty() && output.size() < maximum) { + output.push_back(completions_.front()); + completions_.pop_front(); + } + } + void close() { valid_ = false; } + + private: + std::mutex mutex_; + std::deque completions_; + bool valid_{true}; +}; + +SegmentDescriptor makeDescriptor(uint64_t address, uint64_t length) { + return SegmentDescriptor{ + SegmentDescriptor::kSchemaVersion, 1, 16, + std::to_string(address) + ":" + std::to_string(length)}; +} + +bool parseDescriptor(const SegmentDescriptor& descriptor, uint64_t& address, + uint64_t& length) { + const auto colon = descriptor.hex.find(':'); + if (descriptor.schema_version != SegmentDescriptor::kSchemaVersion || + descriptor.urma_api_version != 1 || descriptor.urma_abi_size != 16 || + colon == std::string::npos) { + return false; + } + try { + address = std::stoull(descriptor.hex.substr(0, colon)); + length = std::stoull(descriptor.hex.substr(colon + 1)); + return length != 0; + } catch (...) { + return false; + } +} + +class FakeLocalSegment final : public LocalSegment { + public: + FakeLocalSegment(uint64_t address, uint64_t length) + : address_(address), + length_(length), + descriptor_(makeDescriptor(address, length)) {} + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() { valid_ = false; } + + private: + uint64_t address_; + uint64_t length_; + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeRemoteSegment final : public RemoteSegment { + public: + FakeRemoteSegment(uint64_t address, uint64_t length, + SegmentDescriptor descriptor) + : address_(address), + length_(length), + descriptor_(std::move(descriptor)) {} + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() { valid_ = false; } + + private: + uint64_t address_; + uint64_t length_; + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeJetty final : public Jetty { + public: + FakeJetty(uint32_t id, std::shared_ptr jfc) + : id_(id), jfc_(std::move(jfc)) {} + bool valid() const noexcept override { return valid_; } + uint32_t id() const noexcept override { return id_; } + uint32_t uasid() const noexcept override { return 0; } + std::shared_ptr jfc() const { return jfc_; } + bool bound() const { return bound_; } + void bind() { bound_ = true; } + void unbind() { bound_ = false; } + void close() { + valid_ = false; + bound_ = false; + } + + private: + uint32_t id_; + std::shared_ptr jfc_; + bool valid_{true}; + bool bound_{false}; +}; + +class FakeUrmaAdapter final : public UrmaAdapter { + public: + explicit FakeUrmaAdapter(DeviceInfo device) : device_(std::move(device)) {} + + bool available() const noexcept override { return true; } + uint32_t nativeApiVersion() const noexcept override { return 1; } + size_t nativeSegmentDescriptorSize() const noexcept override { return 16; } + Status initialize() override { + initialized_ = true; + return Status::OK(); + } + Status shutdown() override { + initialized_ = false; + return Status::OK(); + } + Status discoverDevices(std::vector& devices) override { + if (!initialized_) return Status::InvalidArgument("not initialized"); + devices = {device_}; + return Status::OK(); + } + Status openContext(const DeviceInfo& device, ContextPtr& context) override { + if (!initialized_ || device.topology_name != device_.topology_name) { + return Status::DeviceNotFound("fake device missing"); + } + context = std::make_shared(device); + return Status::OK(); + } + Status closeContext(ContextPtr& context) override { + if (auto fake = std::dynamic_pointer_cast(context)) { + fake->close(); + } + context.reset(); + return Status::OK(); + } + Status createJfc(const ContextPtr&, const JfcOptions&, + JfcPtr& jfc) override { + jfc = std::make_shared(); + return Status::OK(); + } + Status deleteJfc(JfcPtr& jfc) override { + if (auto fake = std::dynamic_pointer_cast(jfc)) fake->close(); + jfc.reset(); + return Status::OK(); + } + Status registerLocalSegment(const ContextPtr&, uint64_t address, + size_t length, const SegmentOptions& options, + LocalSegmentPtr& segment) override { + if (length == 0) return Status::InvalidArgument("empty segment"); + last_registered_access_.store(options.access, + std::memory_order_release); + segment = std::make_shared(address, length); + return Status::OK(); + } + Status unregisterLocalSegment(LocalSegmentPtr& segment) override { + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + return Status::OK(); + } + Status importRemoteSegment(const ContextPtr&, + const SegmentDescriptor& descriptor, + const SegmentOptions&, + RemoteSegmentPtr& segment) override { + uint64_t address = 0; + uint64_t length = 0; + if (!parseDescriptor(descriptor, address, length)) { + return Status::InvalidArgument("bad descriptor"); + } + segment = + std::make_shared(address, length, descriptor); + return Status::OK(); + } + Status unimportRemoteSegment(RemoteSegmentPtr& segment) override { + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + return Status::OK(); + } + Status createJetty(const ContextPtr&, const JfcPtr& jfc, + const JettyOptions&, JettyPtr& jetty) override { + auto fake_jfc = std::dynamic_pointer_cast(jfc); + if (!fake_jfc) return Status::InvalidArgument("bad JFC"); + jetty = std::make_shared(next_jetty_id_++, fake_jfc); + return Status::OK(); + } + Status deleteJetty(JettyPtr& jetty) override { + if (auto fake = std::dynamic_pointer_cast(jetty)) { + fake->close(); + } + jetty.reset(); + return Status::OK(); + } + Status bindJetty(const JettyPtr& jetty, const RemoteJettyInfo&) override { + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid()) + return Status::InvalidArgument("bad Jetty"); + fake->bind(); + return Status::OK(); + } + Status unbindJetty(const JettyPtr& jetty) override { + if (auto fake = std::dynamic_pointer_cast(jetty)) { + fake->unbind(); + } + return Status::OK(); + } + Status resetJetty(const JettyPtr&) override { + if (fail_resets_.load(std::memory_order_acquire)) { + return Status::RdmaError("injected Jetty reset failure"); + } + return Status::OK(); + } + Status quiesceJetty(const JettyPtr& jetty, uint32_t timeout_ms, + std::vector& completions) override { + completions.clear(); + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid() || timeout_ms == 0) { + return Status::InvalidArgument("bad Jetty quiesce"); + } + if (fail_next_quiesce_.exchange(false, std::memory_order_acq_rel)) { + return Status::RdmaError("injected quiesce failure"); + } + std::lock_guard lock(pending_mutex_); + if (!quiesce_drops_completions_.load(std::memory_order_acquire)) { + auto it = pending_.begin(); + while (it != pending_.end()) { + if (it->jetty_id != fake->id()) { + ++it; + continue; + } + completions.push_back( + Completion{CompletionCategory::ENDPOINT_ERROR, 0, + it->request.token, 0, fake->id()}); + it = pending_.erase(it); + } + } + quiesce_calls_.fetch_add(1, std::memory_order_relaxed); + return Status::OK(); + } + Status post(const JettyPtr& jetty, const std::vector& requests, + size_t& posted_count) override { + posted_count = 0; + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid() || !fake->bound()) { + return Status::InvalidArgument("unbound Jetty"); + } + for (const auto& request : requests) { + if (!request.local_segment || !request.remote_segment || + request.token == 0 || request.length == 0) { + return Status::InvalidArgument("bad WR"); + } + const auto category = next_completion_.exchange( + CompletionCategory::SUCCESS, std::memory_order_acq_rel); + if (hold_next_completion_.exchange(false, + std::memory_order_acq_rel)) { + std::lock_guard lock(pending_mutex_); + pending_.push_back(Pending{fake->id(), request}); + ++posted_count; + continue; + } + if (category == CompletionCategory::SUCCESS) { + if (request.operation == Operation::WRITE) { + std::memcpy( + reinterpret_cast(request.remote_address), + reinterpret_cast(request.local_address), + request.length); + } else { + std::memcpy( + reinterpret_cast(request.local_address), + reinterpret_cast(request.remote_address), + request.length); + } + } + fake->jfc()->push( + Completion{category, 0, request.token, + category == CompletionCategory::SUCCESS + ? static_cast(request.length) + : 0, + fake->id()}); + ++posted_count; + } + return Status::OK(); + } + Status poll(const JfcPtr& jfc, size_t maximum, + std::vector& completions) override { + completions.clear(); + auto fake = std::dynamic_pointer_cast(jfc); + if (!fake || !fake->valid()) return Status::InvalidArgument("bad JFC"); + fake->poll(maximum, completions); + return Status::OK(); + } + + void failNextCompletion(CompletionCategory category) { + next_completion_.store(category, std::memory_order_release); + } + void holdNextCompletion() { + hold_next_completion_.store(true, std::memory_order_release); + } + void failNextQuiesce() { + fail_next_quiesce_.store(true, std::memory_order_release); + } + void dropQuiesceCompletions() { + quiesce_drops_completions_.store(true, std::memory_order_release); + } + void failResets() { fail_resets_.store(true, std::memory_order_release); } + void allowResets() { fail_resets_.store(false, std::memory_order_release); } + size_t pendingCount() const { + std::lock_guard lock(pending_mutex_); + return pending_.size(); + } + uint64_t quiesceCalls() const { + return quiesce_calls_.load(std::memory_order_relaxed); + } + uint32_t lastRegisteredAccess() const { + return last_registered_access_.load(std::memory_order_acquire); + } + + private: + struct Pending { + uint32_t jetty_id{0}; + WorkRequest request; + }; + + DeviceInfo device_; + bool initialized_{false}; + uint32_t next_jetty_id_{1}; + std::atomic next_completion_{ + CompletionCategory::SUCCESS}; + std::atomic hold_next_completion_{false}; + std::atomic fail_next_quiesce_{false}; + std::atomic quiesce_drops_completions_{false}; + std::atomic fail_resets_{false}; + mutable std::mutex pending_mutex_; + std::vector pending_; + std::atomic quiesce_calls_{0}; + std::atomic last_registered_access_{0}; +}; + +class NullRegistry final : public SegmentRegistry { + public: + Status getSegmentDesc(SegmentDescRef&, const std::string&) override { + return Status::InvalidEntry("not used"); + } + Status putSegmentDesc(SegmentDescRef&) override { return Status::OK(); } + Status deleteSegmentDesc(const std::string&) override { + return Status::OK(); + } +}; + +DeviceInfo fakeDevice() { + DeviceInfo info; + info.topology_name = "ub:fake0:eid0"; + info.native_device_name = "fake0"; + info.native_device_path = "/fake/fake0"; + info.eid_index = 0; + info.eid = "0001:0000:0000:0000:0000:0000:0000:0000"; + info.active = true; + info.capabilities.max_jfc = 4; + info.capabilities.max_jetty = 64; + return info; +} + +std::shared_ptr fakeTopology(bool discovery_active = true) { + auto topology = std::make_shared(); + Topology::NicEntry nic{.name = "ub:fake0:eid0", + .pci_bus_id = "0000:00:00.0", + .type = Topology::NIC_UB, + .numa_node = 0}; + auto device = fakeDevice(); + device.active = discovery_active; + encodeTopologyDeviceAttributes(device, 0, nic.device_attrs); + topology->nic_list_.push_back(std::move(nic)); + Topology::MemEntry memory; + memory.name = kWildcardLocation; + memory.type = Topology::MEM_HOST; + memory.numa_node = -1; + memory.device_list[0].push_back(0); + topology->mem_list_.push_back(std::move(memory)); + return topology; +} + +TEST(UbNativeDataPathTest, LocalBuffersRequestProviderLocalOnlyAccess) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + UbBufferManager buffers(adapter, {context}); + + std::array storage{}; + BufferDesc descriptor{}; + descriptor.addr = reinterpret_cast(storage.data()); + descriptor.length = storage.size(); + descriptor.location = kWildcardLocation; + MemoryOptions options; + options.perm = kLocalReadWrite; + + ASSERT_TRUE(buffers.addBuffer(descriptor, options).ok()); + EXPECT_EQ(adapter->lastRegisteredAccess(), SEGMENT_ACCESS_LOCAL_ONLY); + + EXPECT_TRUE(buffers.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, + MissingCompletionIsFencedBeforeRetryAndEventuallyCompletes) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + std::vector contexts{context}; + auto topology = fakeTopology(); + UbBufferManager buffers(adapter, contexts); + + std::array source{}; + std::array target{}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i + 1); + } + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = kWildcardLocation; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = kWildcardLocation; + MemoryOptions options; + options.perm = kGlobalReadWrite; + ASSERT_TRUE(buffers.addBuffer(source_desc, options).ok()); + ASSERT_TRUE(buffers.addBuffer(target_desc, options).ok()); + + SegmentManager manager(std::make_unique()); + ASSERT_TRUE(manager + .updateLocal([&](SegmentDesc& segment) { + segment.name = "local"; + segment.type = SegmentType::Memory; + segment.detail = MemorySegmentDesc{}; + auto& memory = + std::get(segment.detail); + memory.topology = *topology; + memory.buffers = {source_desc, target_desc}; + return Status::OK(); + }) + .ok()); + + EndpointStore endpoints(adapter, 16, 1); + RailMonitor rails; + QuotaManager quota; + UbParams params; + params.worker_count = 1; + params.poller_count = 1; + params.slice_size = 16; + params.max_retries = 1; + params.slice_timeout_ms = 20; + + EndpointResolver resolver = [&](const EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + UbEndpointKey key{request.local_context->topologyId(), + request.remote_segment_id, request.remote_topology_id, + "local@ub:fake0:eid0"}; + auto status = + endpoints.getOrCreate(key, request.local_context, endpoint); + if (!status.ok() || endpoint->ready()) return status; + UbBootstrapDesc peer; + peer.local_eid = fakeDevice().eid; + peer.endpoint_generation = 100; + peer.jetty_ids = {777}; + return endpoint->bind(peer); + }; + UbWorkers workers(adapter, contexts, topology, &manager, &buffers, &rails, + "a, params, std::move(resolver), + [&](const std::shared_ptr& endpoint) { + (void)endpoints.retire(endpoint); + }); + ASSERT_TRUE(workers.start().ok()); + adapter->holdNextCompletion(); + + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(target.data()); + request.length = source.size(); + auto task = UbTask::create(request); + for (size_t offset = 0; offset < request.length; offset += 16) { + ASSERT_NE(task->addSlice(UbSliceSpec{ + source.data() + offset, + reinterpret_cast(target.data() + offset), 16, + offset, 1}), + nullptr); + } + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(workers.submit(task).ok()); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (task->transferStatus().s == PENDING && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(task->transferStatus().s, COMPLETED); + EXPECT_EQ(task->transferStatus().transferred_bytes, source.size()); + EXPECT_EQ(source, target); + EXPECT_GT(rails.aggregateBandwidth(), 0.0); + EXPECT_GE(rails.stats(UbPostPath{0, LOCAL_SEGMENT_ID, 0, 1}).timeouts, 1U); + EXPECT_EQ(adapter->pendingCount(), 0U); + EXPECT_GE(adapter->quiesceCalls(), 1U); + + EXPECT_TRUE(workers.stop().ok()); + EXPECT_TRUE(endpoints.clear().ok()); + EXPECT_TRUE(buffers.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, + ProviderLostCompletionIsReclaimedAfterSuccessfulFence) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + std::vector contexts{context}; + auto topology = fakeTopology(); + UbBufferManager buffers(adapter, contexts); + + std::array source{}; + std::array target{}; + std::array recovered{}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i + 1); + } + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = kWildcardLocation; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = kWildcardLocation; + BufferDesc recovered_desc{}; + recovered_desc.addr = reinterpret_cast(recovered.data()); + recovered_desc.length = recovered.size(); + recovered_desc.location = kWildcardLocation; + MemoryOptions options; + options.perm = kGlobalReadWrite; + ASSERT_TRUE(buffers.addBuffer(source_desc, options).ok()); + ASSERT_TRUE(buffers.addBuffer(target_desc, options).ok()); + ASSERT_TRUE(buffers.addBuffer(recovered_desc, options).ok()); + + SegmentManager manager(std::make_unique()); + ASSERT_TRUE( + manager + .updateLocal([&](SegmentDesc& segment) { + segment.name = "local"; + segment.type = SegmentType::Memory; + segment.detail = MemorySegmentDesc{}; + auto& memory = std::get(segment.detail); + memory.topology = *topology; + memory.buffers = {source_desc, target_desc, recovered_desc}; + return Status::OK(); + }) + .ok()); + + EndpointStore endpoints(adapter, 16, 1); + RailMonitor rails; + QuotaManager quota; + UbParams params; + params.worker_count = 1; + params.poller_count = 1; + params.slice_size = 16; + params.max_retries = 1; + params.slice_timeout_ms = 20; + + std::shared_ptr fenced_endpoint; + EndpointResolver resolver = [&](const EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + UbEndpointKey key{request.local_context->topologyId(), + request.remote_segment_id, request.remote_topology_id, + "local@ub:fake0:eid0"}; + auto status = + endpoints.getOrCreate(key, request.local_context, endpoint); + if (!status.ok()) return status; + if (!fenced_endpoint) fenced_endpoint = endpoint; + if (endpoint->ready()) return status; + UbBootstrapDesc peer; + peer.local_eid = fakeDevice().eid; + peer.endpoint_generation = 100; + peer.jetty_ids = {777}; + return endpoint->bind(peer); + }; + UbWorkers workers(adapter, contexts, topology, &manager, &buffers, &rails, + "a, params, std::move(resolver), + [&](const std::shared_ptr& endpoint) { + (void)endpoints.retire(endpoint); + }); + ASSERT_TRUE(workers.start().ok()); + adapter->holdNextCompletion(); + adapter->dropQuiesceCompletions(); + + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(target.data()); + request.length = source.size(); + auto task = UbTask::create(request); + for (size_t offset = 0; offset < request.length; offset += 16) { + ASSERT_NE(task->addSlice(UbSliceSpec{ + source.data() + offset, + reinterpret_cast(target.data() + offset), 16, + offset, 1}), + nullptr); + } + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(workers.submit(task).ok()); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (task->transferStatus().s == PENDING && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(task->transferStatus().s, COMPLETED); + EXPECT_EQ(task->transferStatus().transferred_bytes, source.size()); + EXPECT_EQ(source, target); + EXPECT_GE(rails.stats(UbPostPath{0, LOCAL_SEGMENT_ID, 0, 1}).timeouts, 1U); + EXPECT_GE(adapter->quiesceCalls(), 1U); + + EXPECT_EQ(workers.inflightCount(), 0U); + EXPECT_EQ(quota.activeReservationCount(), 0U); + ASSERT_NE(fenced_endpoint, nullptr); + EXPECT_EQ(fenced_endpoint->outstandingWrs(), 0U); + EXPECT_EQ(fenced_endpoint->outstandingBytes(), 0U); + + Request followup{}; + followup.opcode = Request::WRITE; + followup.source = source.data(); + followup.target_id = LOCAL_SEGMENT_ID; + followup.target_offset = reinterpret_cast(recovered.data()); + followup.length = source.size(); + auto followup_task = UbTask::create(followup); + for (size_t offset = 0; offset < followup.length; offset += 16) { + ASSERT_NE(followup_task->addSlice(UbSliceSpec{ + source.data() + offset, + reinterpret_cast(recovered.data() + offset), 16, + offset, 1}), + nullptr); + } + ASSERT_TRUE(followup_task->seal()); + ASSERT_TRUE(workers.submit(followup_task).ok()); + + const auto followup_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (followup_task->transferStatus().s == PENDING && + std::chrono::steady_clock::now() < followup_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(followup_task->transferStatus().s, COMPLETED); + EXPECT_EQ(followup_task->transferStatus().transferred_bytes, source.size()); + EXPECT_EQ(recovered, source); + + EXPECT_TRUE(workers.stop().ok()); + EXPECT_TRUE(endpoints.clear().ok()); + EXPECT_TRUE(buffers.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, EndpointStoreNeverReusesRetiredGeneration) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + EndpointStore store(adapter, 2, 1); + UbEndpointKey key{0, 9, 0, "peer@ub:fake0:eid0"}; + std::array, 8> concurrent; + std::array statuses; + std::vector threads; + for (size_t i = 0; i < concurrent.size(); ++i) { + threads.emplace_back([&, i] { + statuses[i] = store.getOrCreate(key, context, concurrent[i]); + }); + } + for (auto& thread : threads) thread.join(); + for (size_t i = 0; i < concurrent.size(); ++i) { + ASSERT_TRUE(statuses[i].ok()); + EXPECT_EQ(concurrent[i], concurrent[0]); + } + auto first = concurrent[0]; + const uint64_t old_generation = first->generation(); + EXPECT_TRUE(store.retire(key, old_generation)); + std::shared_ptr replacement; + ASSERT_TRUE(store.getOrCreate(key, context, replacement).ok()); + EXPECT_GT(replacement->generation(), old_generation); + EXPECT_FALSE(store.retire(key, old_generation)); + EXPECT_EQ(store.get(key), replacement); + EXPECT_TRUE(store.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, + EndpointStoreDoesNotReplaceEndpointWhenCleanupCannotComplete) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + EndpointStore store(adapter, 1, 1); + UbEndpointKey key{0, 9, 0, "peer@ub:fake0:eid0"}; + std::shared_ptr endpoint; + ASSERT_TRUE(store.getOrCreate(key, context, endpoint).ok()); + UbBootstrapDesc peer; + peer.local_eid = fakeDevice().eid; + peer.endpoint_generation = 100; + peer.jetty_ids = {777}; + ASSERT_TRUE(endpoint->bind(peer).ok()); + + adapter->failResets(); + std::vector completions; + EXPECT_FALSE(endpoint->quiesce(20, completions).ok()); + + std::shared_ptr replacement; + const auto status = store.getOrCreate(key, context, replacement); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(replacement, nullptr); + EXPECT_EQ(store.size(), 0U); + + adapter->allowResets(); + EXPECT_TRUE(store.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { + auto adapter = std::make_shared(fakeDevice()); + // A serialized discovery snapshot is informational. Current local + // discovery and runtime path health remain the scheduling authorities. + auto topology = fakeTopology(false); + auto control = std::make_shared("p2p", "", nullptr); + uint16_t port = 0; + ASSERT_TRUE(control->start(port).ok()); + const std::string segment_name = + "127.0.0.1:" + std::to_string(static_cast(port)); + ASSERT_TRUE(control->segmentManager() + .updateLocal([&](SegmentDesc& segment) { + segment.name = segment_name; + segment.rpc_server_addr = segment_name; + segment.machine_id = "fake-machine"; + segment.type = SegmentType::Memory; + segment.detail = MemorySegmentDesc{}; + std::get(segment.detail).topology = + *topology; + return Status::OK(); + }) + .ok()); + + auto config = std::make_shared(); + config->set("transports/ub/enable", true); + config->set("transports/ub/worker_count", 1); + config->set("transports/ub/poller_count", 1); + config->set("transports/ub/jfc_per_context", 1); + config->set("transports/ub/jetty_per_endpoint", 1); + config->set("transports/ub/max_endpoints", 16); + config->set("transports/ub/slice_size", 16); + config->set("transports/ub/max_slices_per_task", 8); + config->set("transports/ub/max_retries", 1); + config->set("transports/ub/slice_timeout_ms", 1000); + config->set("transports/ub/endpoint_cooldown_ms", 100); + + UbTransport transport(adapter); + std::string mutable_segment_name = segment_name; + ASSERT_TRUE( + transport.install(mutable_segment_name, control, topology, config) + .ok()); + EXPECT_STREQ(transport.getName(), "ub"); + EXPECT_TRUE(transport.supportsCancellation()); + EXPECT_FALSE(transport.supportNotification()); + EXPECT_TRUE(transport.capabilities().dram_to_dram); + + std::array source{}; + std::array target{}; + std::array oversized_source{}; + std::array oversized_target{}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(100 - i); + } + const auto expected = source; + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = kWildcardLocation; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = kWildcardLocation; + BufferDesc oversized_source_desc{}; + oversized_source_desc.addr = + reinterpret_cast(oversized_source.data()); + oversized_source_desc.length = oversized_source.size(); + oversized_source_desc.location = kWildcardLocation; + BufferDesc oversized_target_desc{}; + oversized_target_desc.addr = + reinterpret_cast(oversized_target.data()); + oversized_target_desc.length = oversized_target.size(); + oversized_target_desc.location = kWildcardLocation; + MemoryOptions options; + options.perm = kGlobalReadWrite; + std::vector descriptors{ + source_desc, target_desc, oversized_source_desc, oversized_target_desc}; + ASSERT_TRUE(transport.addMemoryBuffer(descriptors, options).ok()); + ASSERT_TRUE(control->segmentManager() + .updateLocal([&](SegmentDesc& segment) { + std::get(segment.detail).buffers = + descriptors; + return Status::OK(); + }) + .ok()); + + Transport::SubBatchRef rejected_batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(rejected_batch, 1).ok()); + Request oversized_request{}; + oversized_request.opcode = Request::WRITE; + oversized_request.source = oversized_source.data(); + oversized_request.target_id = LOCAL_SEGMENT_ID; + oversized_request.target_offset = + reinterpret_cast(oversized_target.data()); + oversized_request.length = oversized_source.size(); + EXPECT_TRUE( + transport.submitTransferTasks(rejected_batch, {oversized_request}) + .IsInvalidArgument()); + EXPECT_TRUE(transport.freeSubBatch(rejected_batch).ok()); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok()); + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(target.data()); + request.length = source.size(); + ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); + + TransferStatus transfer{PENDING, 0}; + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (transfer.s == PENDING && + std::chrono::steady_clock::now() < deadline) { + ASSERT_TRUE(transport.getTransferStatus(batch, 0, transfer).ok()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(transfer.s, COMPLETED); + EXPECT_EQ(transfer.transferred_bytes, source.size()); + EXPECT_EQ(source, target); + EXPECT_GT(transport.getEstimatedBandwidth(), 0.0); + + source.fill(0); + Request read_request = request; + read_request.opcode = Request::READ; + ASSERT_TRUE(transport.submitTransferTasks(batch, {read_request}).ok()); + transfer = TransferStatus{PENDING, 0}; + const auto read_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (transfer.s == PENDING && + std::chrono::steady_clock::now() < read_deadline) { + ASSERT_TRUE(transport.getTransferStatus(batch, 1, transfer).ok()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(transfer.s, COMPLETED); + EXPECT_EQ(transfer.transferred_bytes, source.size()); + EXPECT_EQ(source, expected); + + ASSERT_TRUE(transport.freeSubBatch(batch).ok()); + EXPECT_EQ(batch, nullptr); + + // A failed device fence must make uninstall retryable without destroying + // pollers, registered memory, or the held WR token. + Transport::SubBatchRef draining_batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(draining_batch, 1).ok()); + adapter->holdNextCompletion(); + ASSERT_TRUE(transport.submitTransferTasks(draining_batch, {request}).ok()); + const auto posted_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(1); + while (adapter->pendingCount() == 0 && + std::chrono::steady_clock::now() < posted_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(adapter->pendingCount(), 1U); + adapter->failNextQuiesce(); + EXPECT_FALSE(transport.uninstall().ok()); + EXPECT_EQ(adapter->pendingCount(), 1U); + EXPECT_TRUE(transport.uninstall().ok()); + EXPECT_EQ(adapter->pendingCount(), 0U); + EXPECT_TRUE(transport.freeSubBatch(draining_batch).ok()); + EXPECT_TRUE(transport.uninstall().ok()); +} + +} // namespace +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tests/endpoint_store_test.cpp b/mooncake-transfer-engine/tests/endpoint_store_test.cpp index 5f3b121656..8225a9ac93 100644 --- a/mooncake-transfer-engine/tests/endpoint_store_test.cpp +++ b/mooncake-transfer-engine/tests/endpoint_store_test.cpp @@ -20,6 +20,8 @@ #include #include +#include +#include #include "transport/rdma_transport/endpoint_store.h" #include "transport/rdma_transport/rdma_context.h" @@ -175,4 +177,38 @@ TEST_F(EndpointStoreTest, ReclaimDoesNotRequireActiveMap) { EXPECT_EQ(store.waitingListSize(), 0u); } +TEST_F(EndpointStoreTest, + StaleRawPointerLookupAndDeleteStressDoesNotDereference) { + SIEVEEndpointStore store(4); + auto sentinel = makeQuiescentEndpoint(*ctx_); + std::vector stale_ptrs; + stale_ptrs.reserve(1000); + + for (size_t i = 0; i < 1000; ++i) { + auto endpoint = makeQuiescentEndpoint(*ctx_); + const std::string peer_nic_path = "peer@" + std::to_string(i); + store.testOnlyInsertEndpoint(peer_nic_path, endpoint); + stale_ptrs.push_back(endpoint.get()); + EXPECT_EQ(endpoint, store.getEndpointByPtr(endpoint.get())); + EXPECT_EQ(0, store.deleteEndpointByPtr(endpoint.get())); + } + + EXPECT_EQ(store.getSize(), 0u); + EXPECT_EQ(store.waitingListSize(), 1000u); + store.testOnlyInsertEndpoint("sentinel@peer", sentinel); + store.reclaimEndpoint(); + EXPECT_EQ(store.waitingListSize(), 0u); + EXPECT_EQ(store.getSize(), 1u); + + // The endpoints were live in the store, deleted by raw pointer, then + // reclaimed. The raw pointers are now stale while the sentinel keeps the + // active map non-empty. Store APIs must compare pointer identity only; + // dereferencing would be caught by ASan builds. + for (auto* stale_ptr : stale_ptrs) { + EXPECT_EQ(nullptr, store.getEndpointByPtr(stale_ptr)); + EXPECT_EQ(-1, store.deleteEndpointByPtr(stale_ptr)); + } + EXPECT_EQ(sentinel, store.getEndpointByPtr(sentinel.get())); +} + } // namespace diff --git a/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp b/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp index 5bfb27c7ad..5b94739f26 100644 --- a/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp @@ -30,9 +30,12 @@ #include "error.h" #include "transfer_metadata.h" #include "transfer_metadata_plugin.h" +#include "topology.h" +#include "transport/rdma_transport/endpoint_store.h" #include "transport/rdma_transport/rdma_context.h" #include "transport/rdma_transport/rdma_endpoint.h" #include "transport/rdma_transport/rdma_transport.h" +#include "transport/rdma_transport/worker_pool.h" #if defined(__has_feature) #define MC_HAS_FEATURE(x) __has_feature(x) @@ -59,6 +62,16 @@ class RdmaTransportTestPeer { transport.local_server_name_ = local_server_name; transport.rdma_server_name_ = local_server_name; } + + static void bindTopology(RdmaTransport &transport, + std::shared_ptr topology) { + transport.local_topology_ = std::move(topology); + } + + static void addContext(RdmaTransport &transport, + std::shared_ptr context) { + transport.context_list_.push_back(std::move(context)); + } }; class RdmaContextTestPeer { @@ -72,11 +85,68 @@ class RdmaContextTestPeer { context.lid_ = 0; } + static void bindFakeEndpointStore(RdmaContext &context, ibv_cq *cq) { + bindCompletionQueue(context, cq); + context.endpoint_store_ = std::make_shared(8); + context.transfer_worker_count_ = 1; + context.endpoint_lifecycle_locks_.clear(); + context.endpoint_lifecycle_locks_.push_back( + std::make_unique()); + context.active_.store(true, std::memory_order_release); + } + + static std::unique_lock tryLockEndpointLifecycle( + RdmaContext &context, const std::string &peer_nic_path) { + const int owner_thread = context.postingThreadForPeer(peer_nic_path); + if (owner_thread < 0 || static_cast(owner_thread) >= + context.endpoint_lifecycle_locks_.size()) { + return std::unique_lock(); + } + return std::unique_lock( + *context.endpoint_lifecycle_locks_[owner_thread], std::try_to_lock); + } + + static void insertEndpointForTest(RdmaContext &context, + const std::string &peer_nic_path, + std::shared_ptr endpoint) { + auto store = std::dynamic_pointer_cast( + context.endpoint_store_); + ASSERT_NE(store, nullptr); + store->testOnlyInsertEndpoint(peer_nic_path, std::move(endpoint)); + } + static void clearCompletionQueues(RdmaContext &context) { context.cq_list_.clear(); } }; +class WorkerPoolTestPeer { + public: + static void stopWorkers(WorkerPool &pool) { + if (!pool.workers_running_.load(std::memory_order_acquire)) return; + pool.workers_running_.store(false, std::memory_order_release); + pool.cond_var_.notify_all(); + for (auto &entry : pool.worker_thread_) { + if (entry.joinable()) entry.join(); + } + } + + static void queueSlice(WorkerPool &pool, int thread_id, + const std::string &peer_nic_path, + Transport::Slice *slice) { + pool.collective_slice_queue_[thread_id][peer_nic_path].push_back(slice); + } + + static void performPostSend(WorkerPool &pool, int thread_id) { + pool.performPostSend(thread_id); + } + + static void clearQueuedSlices(WorkerPool &pool) { + for (auto &queue : pool.collective_slice_queue_) queue.clear(); + for (auto &queue : pool.worker_slice_queue_) queue.clear(); + } +}; + class RdmaEndPointTestPeer { public: static void setStatus(RdmaEndPoint &endpoint, RdmaEndPoint::Status status) { @@ -103,6 +173,114 @@ class RdmaEndPointTestPeer { namespace { +class InProcessRdmaTransport : public RdmaTransport { + public: + int sendHandshake(const std::string &peer_server_name, + const HandShakeDesc &local_desc, + HandShakeDesc &peer_desc) override { + (void)peer_server_name; + if (!local_desc.ready_ack) { + auto lifecycle_lock = RdmaContextTestPeer::tryLockEndpointLifecycle( + *local_context_, local_desc.peer_nic_path); + if (!lifecycle_lock.owns_lock()) { + lifecycle_gate_was_free_.store(false, + std::memory_order_release); + return ERR_ENDPOINT; + } + lifecycle_lock.unlock(); + + std::unique_lock lock(barrier_->mutex); + ++barrier_->arrivals; + barrier_->cv.notify_all(); + if (!barrier_->cv.wait_for(lock, std::chrono::seconds(5), [&] { + return barrier_->arrivals >= 2; + })) { + return ERR_ENDPOINT; + } + } + + return peer_->onSetupRdmaConnections(local_desc, peer_desc); + } + + bool lifecycleGateWasFree() const { + return lifecycle_gate_was_free_.load(std::memory_order_acquire); + } + + struct Barrier { + std::mutex mutex; + std::condition_variable cv; + int arrivals = 0; + }; + + RdmaContext *local_context_ = nullptr; + InProcessRdmaTransport *peer_ = nullptr; + std::shared_ptr barrier_; + + private: + std::atomic lifecycle_gate_was_free_{true}; +}; + +struct FakeRdmaPeer { + std::shared_ptr metadata; + std::unique_ptr transport; + std::shared_ptr topology; + std::shared_ptr context; + ibv_cq cq = {}; + std::unique_ptr worker_pool; + std::string server_name; + std::string device_name; +}; + +void initFakePeer(FakeRdmaPeer &peer, const std::string &server_name, + const std::string &device_name, + std::shared_ptr barrier) { + peer.server_name = server_name; + peer.device_name = device_name; + peer.metadata = std::make_shared(P2PHANDSHAKE); + peer.transport = std::make_unique(); + RdmaTransportTestPeer::bindMetadata(*peer.transport, peer.metadata, + server_name); + + peer.topology = std::make_shared(); + ASSERT_EQ( + peer.topology->parse("{\"cpu:0\": [[\"" + device_name + "\"], []]}"), + 0); + RdmaTransportTestPeer::bindTopology(*peer.transport, peer.topology); + + peer.context = + std::make_shared(*peer.transport, peer.device_name); + RdmaContextTestPeer::bindFakeEndpointStore(*peer.context, &peer.cq); + RdmaTransportTestPeer::addContext(*peer.transport, peer.context); + + peer.transport->local_context_ = peer.context.get(); + peer.transport->barrier_ = std::move(barrier); + + peer.worker_pool = std::make_unique(*peer.context); + WorkerPoolTestPeer::stopWorkers(*peer.worker_pool); +} + +void queueHandshakeSlice(FakeRdmaPeer &peer, const std::string &peer_nic_path, + Transport::Slice &slice, + Transport::TransferTask &task) { + slice = Transport::Slice(); + task = Transport::TransferTask(); + slice.task = &task; + slice.status = Transport::Slice::PENDING; + slice.peer_nic_path = peer_nic_path; + slice.rdma.retry_cnt = 0; + slice.rdma.max_retry_cnt = 8; + WorkerPoolTestPeer::queueSlice(*peer.worker_pool, 0, peer_nic_path, &slice); +} + +void installZeroQpEndpoint(FakeRdmaPeer &peer, + const std::string &peer_nic_path) { + auto endpoint = std::make_shared(*peer.context); + ASSERT_EQ(endpoint->construct(&peer.cq, 0), 0); + endpoint->setPeerNicPath(peer_nic_path); + RdmaContextTestPeer::insertEndpointForTest(*peer.context, peer_nic_path, + std::move(endpoint)); +} + class RdmaEndPointStateTest : public ::testing::Test { protected: void SetUp() override { @@ -278,4 +456,60 @@ TEST_F(RdmaEndPointStateTest, EXPECT_FALSE(endpoint_->connected()); } +TEST(RdmaEndpointLifecycleGateTest, + BidirectionalActiveHandshakeDoesNotHoldLifecycleGateDuringRpc) { + auto barrier = std::make_shared(); + FakeRdmaPeer peer_a; + FakeRdmaPeer peer_b; + + ASSERT_NO_FATAL_FAILURE( + initFakePeer(peer_a, "rdma-gate-a:10000", "mlx5_gate_a", barrier)); + ASSERT_NO_FATAL_FAILURE( + initFakePeer(peer_b, "rdma-gate-b:10000", "mlx5_gate_b", barrier)); + peer_a.transport->peer_ = peer_b.transport.get(); + peer_b.transport->peer_ = peer_a.transport.get(); + + const std::string a_to_b = + MakeNicPath(peer_b.server_name, peer_b.device_name); + const std::string b_to_a = + MakeNicPath(peer_a.server_name, peer_a.device_name); + + ASSERT_NO_FATAL_FAILURE(installZeroQpEndpoint(peer_a, a_to_b)); + ASSERT_NO_FATAL_FAILURE(installZeroQpEndpoint(peer_b, b_to_a)); + + Transport::Slice slice_a; + Transport::Slice slice_b; + Transport::TransferTask task_a; + Transport::TransferTask task_b; + queueHandshakeSlice(peer_a, a_to_b, slice_a, task_a); + queueHandshakeSlice(peer_b, b_to_a, slice_b, task_b); + + auto active_a = std::async(std::launch::async, [&] { + WorkerPoolTestPeer::performPostSend(*peer_a.worker_pool, 0); + }); + auto active_b = std::async(std::launch::async, [&] { + WorkerPoolTestPeer::performPostSend(*peer_b.worker_pool, 0); + }); + + ASSERT_EQ(active_a.wait_for(std::chrono::seconds(5)), + std::future_status::ready); + ASSERT_EQ(active_b.wait_for(std::chrono::seconds(5)), + std::future_status::ready); + active_a.get(); + active_b.get(); + + EXPECT_TRUE(peer_a.transport->lifecycleGateWasFree()); + EXPECT_TRUE(peer_b.transport->lifecycleGateWasFree()); + + auto endpoint_a = peer_a.context->findEndpoint(a_to_b); + auto endpoint_b = peer_b.context->findEndpoint(b_to_a); + ASSERT_NE(endpoint_a, nullptr); + ASSERT_NE(endpoint_b, nullptr); + EXPECT_TRUE(endpoint_a->readyToSend()); + EXPECT_TRUE(endpoint_b->readyToSend()); + + WorkerPoolTestPeer::clearQueuedSlices(*peer_a.worker_pool); + WorkerPoolTestPeer::clearQueuedSlices(*peer_b.worker_pool); +} + } // namespace