From bb39b2ebb814614b3f9f2bd28d3028dc3b43683b Mon Sep 17 00:00:00 2001 From: Aoi Date: Mon, 31 Aug 2026 14:09:12 +0800 Subject: [PATCH 01/17] [CI] Rebalance PR and nightly test coverage (#3669) Rebalance CI coverage so pull requests get comprehensive pre-merge unit and formatting feedback while resource-intensive end-to-end workloads run nightly. --- .github/workflows/ci.yml | 528 ++++++++++++++------ .github/workflows/integration-test.yml | 79 +-- .github/workflows/nightly.yml | 109 +++- .github/workflows/release-efa-cuda13.yaml | 1 + .github/workflows/release-efa-non-cuda.yaml | 1 + .github/workflows/release-efa.yaml | 1 + mooncake-pg/tests/test_pg_elastic.py | 246 +++++---- 7 files changed, 680 insertions(+), 285 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9fe555ffe7..095abf2f7c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,8 +29,9 @@ jobs: github.event_name == 'workflow_dispatch') && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) uses: ./.github/workflows/_build-wheel.yaml with: python-versions: '["3.10", "3.12"]' @@ -38,17 +39,21 @@ jobs: version-override: 0.0.0.dev0 test-wheel-ubuntu: - needs: [spell-check, clang-format, build-wheel] + needs: [build-wheel] if: >- needs.build-wheel.result == 'success' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) strategy: matrix: - ubuntu-version: [ubuntu-22.04, ubuntu-24.04] - python-version: ['3.10', '3.12'] + include: + - ubuntu-version: ubuntu-22.04 + python-version: '3.12' + - ubuntu-version: ubuntu-24.04 + python-version: '3.10' runs-on: ${{ matrix.ubuntu-version }} steps: - uses: actions/checkout@v4 @@ -96,122 +101,357 @@ jobs: bash scripts/test_installation.sh shell: bash - - name: Run tests with ssd + - name: Run all PG CPU tests with latest PyTorch + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' + env: + MC_FORCE_TCP: "true" run: | - # Reserve port 50052 (mooncake_client RPC port) so the kernel never - # auto-allocates it as ephemeral source port for other outbound - # connections in the test suite. Without this, a random Python test - # connection can pick src_port=50052, leave a TIME_WAIT on - # :50052 for 60s, and block mooncake_client's bind to - # 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes - # TIME_WAIT+bind conflict for same-IP or loopback). - sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052 source test_env/bin/activate - MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh - rm -rf /tmp/mooncake_test_ssd - deactivate + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + python -m unittest discover -s mooncake-pg/tests -k CPU -v shell: bash - - name: Start Mooncake Master + - name: Run Store API and RPC smoke tests + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' + env: + MOONCAKE_MASTER: "127.0.0.1:50051" + MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" + MOONCAKE_PROTOCOL: "tcp" + LOCAL_HOSTNAME: "127.0.0.1" run: | source test_env/bin/activate + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + python -m pip install safetensors + + reserved_ports=$(sysctl -n net.ipv4.ip_local_reserved_ports) + sudo sysctl -w \ + "net.ipv4.ip_local_reserved_ports=${reserved_ports:+$reserved_ports,}50052" + mkdir -p /tmp/mooncake_storage mooncake_master \ --default_kv_lease_ttl=500 \ --eviction_high_watermark_ratio=0.95 \ - --cluster_id=ci_test_cluster \ + --cluster_id=ci_store_api_smoke \ --port 50051 \ - --enable_http_metadata_server=true & - sleep 3 + --enable_http_metadata_server=true \ + >"$RUNNER_TEMP/mooncake-master.log" 2>&1 & + master_pid=$! + rpc_server_pid="" + + cleanup() { + if [ -n "$rpc_server_pid" ]; then + kill "$rpc_server_pid" 2>/dev/null || true + wait "$rpc_server_pid" 2>/dev/null || true + fi + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true + } + trap cleanup EXIT + + master_ready=false + for _ in {1..50}; do + if ! kill -0 "$master_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/mooncake-master.log" + echo "::error::mooncake_master exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :50051' | grep -q . && \ + ss -H -ltn 'sport = :8080' | grep -q .; then + master_ready=true + break + fi + sleep 0.1 + done + if [ "$master_ready" != true ]; then + cat "$RUNNER_TEMP/mooncake-master.log" + echo "::error::mooncake_master did not become ready within 5 seconds" + exit 1 + fi + + python scripts/test_tensor_api.py -n 1 + python scripts/test_async_store.py + python scripts/test_copy_move_api.py + python -m unittest mooncake-wheel.tests.test_safetensor_functions + python scripts/test_drain_http_api.py --timeout-sec 90 + + python -u mooncake-transfer-engine/tests/rpc_communicator_test.py \ + server --url 127.0.0.1:9004 --data-size 1 \ + >"$RUNNER_TEMP/rpc-server.log" 2>&1 & + rpc_server_pid=$! + rpc_ready=false + for _ in {1..50}; do + if ! kill -0 "$rpc_server_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :9004' | grep -q .; then + rpc_ready=true + break + fi + sleep 0.1 + done + if [ "$rpc_ready" != true ]; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server did not become ready within 5 seconds" + exit 1 + fi + + client_rc=0 + timeout 10 python -u \ + mooncake-transfer-engine/tests/rpc_communicator_test.py \ + client --url 127.0.0.1:9004 --threads 2 --data-size 1 \ + >"$RUNNER_TEMP/rpc-client.log" 2>&1 || \ + client_rc=$? + cat "$RUNNER_TEMP/rpc-client.log" + if [ "$client_rc" -ne 0 ] && [ "$client_rc" -ne 124 ]; then + echo "::error::RPC communicator client failed with exit code $client_rc" + exit "$client_rc" + fi + if ! grep -q '^bandwidth:' "$RUNNER_TEMP/rpc-client.log"; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator did not complete a successful transfer" + exit 1 + fi shell: bash - - name: Run Python Tensor API Performance Test (CI check) - env: - MOONCAKE_MASTER: "127.0.0.1:50051" - MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" - MOONCAKE_PROTOCOL: "tcp" - LOCAL_HOSTNAME: "127.0.0.1" + - name: Run SSD offload and promotion end-to-end tests + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' run: | source test_env/bin/activate - python scripts/test_tensor_api.py -n 1 + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + + metadata_pid="" + master_pid="" + cleanup() { + status=$? + trap - EXIT + if [ -n "$master_pid" ]; then + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true + fi + if [ -n "$metadata_pid" ]; then + kill "$metadata_pid" 2>/dev/null || true + wait "$metadata_pid" 2>/dev/null || true + fi + if [ "$status" -ne 0 ]; then + cat "$RUNNER_TEMP/mooncake-ssd-master.log" 2>/dev/null || true + cat "$RUNNER_TEMP/mooncake-promotion-master.log" 2>/dev/null || true + cat "$RUNNER_TEMP/mooncake-metadata.log" 2>/dev/null || true + fi + rm -rf /tmp/mooncake_ci_ssd_offload \ + /tmp/mooncake_ci_promotion + exit "$status" + } + trap cleanup EXIT + + wait_for_port() { + local pid=$1 + local port=$2 + local label=$3 + for _ in {1..50}; do + if ! kill -0 "$pid" 2>/dev/null; then + echo "::error::$label exited before listening on port $port" + return 1 + fi + if ss -H -ltn "sport = :$port" | grep -q .; then + return 0 + fi + sleep 0.1 + done + echo "::error::$label did not listen on port $port within 5 seconds" + return 1 + } + + stop_master() { + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true + master_pid="" + } + + mooncake_http_metadata_server --port 8080 \ + >"$RUNNER_TEMP/mooncake-metadata.log" 2>&1 & + metadata_pid=$! + wait_for_port "$metadata_pid" 8080 "metadata server" + + mkdir -p /tmp/mooncake_ci_ssd_offload + mooncake_master \ + --default_kv_lease_ttl=500 \ + --root_fs_dir=/tmp/mooncake_ci_ssd_offload \ + >"$RUNNER_TEMP/mooncake-ssd-master.log" 2>&1 & + master_pid=$! + wait_for_port "$master_pid" 50051 "SSD offload master" + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \ + DEFAULT_KV_LEASE_TTL=500 \ + python mooncake-wheel/tests/test_ssd_offload_in_evict.py + stop_master + + mkdir -p /tmp/mooncake_ci_promotion + mooncake_master \ + --default_kv_lease_ttl=500 \ + --root_fs_dir=/tmp/mooncake_ci_promotion \ + --enable_offload=true \ + --offload_on_evict=true \ + --promotion_on_hit=true \ + --promotion_admission_threshold=1 \ + --promotion_max_per_heartbeat=16 \ + >"$RUNNER_TEMP/mooncake-promotion-master.log" 2>&1 & + master_pid=$! + wait_for_port "$master_pid" 50051 "promotion master" + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \ + DEFAULT_KV_LEASE_TTL=500 \ + MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/tmp/mooncake_ci_promotion \ + MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS=2 \ + MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT=10 \ + MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES=10485760 \ + python mooncake-wheel/tests/test_promotion_on_hit.py + stop_master shell: bash - - name: Run Python Async API Test (CI check) - env: - MOONCAKE_MASTER: "127.0.0.1:50051" - MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" - MOONCAKE_PROTOCOL: "tcp" - LOCAL_HOSTNAME: "127.0.0.1" + unit-tests: + name: CTest unit tests (Python 3.12) + needs: [check-paths] + if: *run-ci-for-source-changes + runs-on: ubuntu-22.04 + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + + - name: Install and start etcd run: | - source test_env/bin/activate - python scripts/test_async_store.py + wget -q https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz + tar xzf etcd-v3.6.1-linux-amd64.tar.gz + sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ + etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & + sleep 3 + ETCDCTL_API=3 etcdctl --endpoints=http://127.0.0.1:2379 endpoint health shell: bash - - name: Test Mooncake Copy/Move API - env: - MOONCAKE_MASTER: "127.0.0.1:50051" - MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" - MOONCAKE_PROTOCOL: "tcp" - LOCAL_HOSTNAME: "127.0.0.1" + - name: Free up disk space + uses: ./.github/actions/free-disk-space + + - name: Install CUDA Toolkit + uses: Jimver/cuda-toolkit@v0.2.24 + with: + cuda: '12.8.1' + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc"]' + + - name: Install build utilities and dependencies run: | - source test_env/bin/activate - python scripts/test_copy_move_api.py + sudo apt-get update + sudo apt-get install -y ninja-build + sudo bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" shell: bash - - name: Run Python Drain HTTP E2E Test (CI check) - env: - MOONCAKE_MASTER: "127.0.0.1:50051" - MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" - MOONCAKE_PROTOCOL: "tcp" - LOCAL_HOSTNAME: "127.0.0.1" + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project with unit tests run: | - source test_env/bin/activate - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - python scripts/test_drain_http_api.py --timeout-sec 90 + cmake -S . -B build -G Ninja \ + -DUSE_HTTP=ON \ + -DUSE_CXL=ON \ + -DUSE_UB=ON \ + -DUSE_ETCD=ON \ + -DUSE_CUDA=ON \ + -DWITH_P2P_STORE=ON \ + -DSTORE_USE_ETCD=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UNIT_TESTS=ON \ + -DENABLE_SCCACHE=ON shell: bash - - name: Run RPC Communicator Bandwidth Test + - name: Configure CUDA driver runtime run: | - source test_env/bin/activate - python mooncake-transfer-engine/tests/rpc_communicator_test.py server --url 127.0.0.1:9004 --data-size 1 & - SERVER_PID=$! - sleep 5 - timeout 10 python mooncake-transfer-engine/tests/rpc_communicator_test.py client --url 127.0.0.1:9004 --threads 2 --data-size 1 || true - kill $SERVER_PID 2>/dev/null || true - wait $SERVER_PID 2>/dev/null || true - - - name: Test Mooncake PyTorch Backend (CPU Only) - env: - MC_FORCE_TCP: "true" + cuda_driver_library=$(sed -n \ + 's/^CUDA_cuda_driver_LIBRARY:FILEPATH=//p' build/CMakeCache.txt) + if [ -z "$cuda_driver_library" ] || [ ! -f "$cuda_driver_library" ]; then + echo "::error::CMake did not resolve the CUDA driver library" + exit 1 + fi + + cuda_driver_dir=$(dirname "$cuda_driver_library") + if [ ! -e "$cuda_driver_dir/libcuda.so.1" ]; then + sudo ln -s "$(basename "$cuda_driver_library")" \ + "$cuda_driver_dir/libcuda.so.1" + fi + echo "LIBRARY_PATH=$cuda_driver_dir:${LIBRARY_PATH:-}" >> "$GITHUB_ENV" + echo "LD_LIBRARY_PATH=$cuda_driver_dir:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV" + shell: bash + + - name: Build project and unit tests run: | - source test_env/bin/activate - python mooncake-pg/tests/test_pg_collectives.py + cmake --build build --parallel 128 + sudo -E cmake --install build shell: bash - - name: Test PyTorch >= 2.13 Single-Buffer Collectives (CPU Only) + - name: Run Mooncake Store Rust unit tests env: - MC_FORCE_TCP: "true" + MOONCAKE_BUILD_DIR: ${{ github.workspace }}/build + MOONCAKE_STORE_LIB_DIR: ${{ github.workspace }}/build/mooncake-store/src + MOONCAKE_STORE_INCLUDE_DIR: ${{ github.workspace }}/mooncake-store/include run: | - source test_env/bin/activate - python mooncake-pg/tests/test_pg_collectives.py \ - TestMooncakePGCollectivesCPU.test_all_gather_into_tensor \ - TestMooncakePGCollectivesCPU.test_reduce_scatter_sum + export LD_LIBRARY_PATH="$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd:/usr/local/lib:${LD_LIBRARY_PATH:-}" + cargo test --manifest-path mooncake-store/rust/Cargo.toml --lib shell: bash - - name: Test Safetensor Functions + - name: Start Metadata Server run: | - source test_env/bin/activate - pip install safetensors - python -m unittest mooncake-wheel.tests.test_safetensor_functions + cd mooncake-transfer-engine/example/http-metadata-server-python + pip install aiohttp + python ./bootstrap_server.py & + sleep 2 + shell: bash + + - name: Run all CTest unit tests + id: ctest + run: | + reserved_ports=$(sysctl -n net.ipv4.ip_local_reserved_ports) + sudo sysctl -w "net.ipv4.ip_local_reserved_ports=${reserved_ports:+$reserved_ports,}50052" + mkdir -p build/test-results + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \ + DEFAULT_KV_LEASE_TTL=500 \ + ctest --test-dir build --parallel "$(nproc)" --output-on-failure \ + --output-junit "$GITHUB_WORKSPACE/build/test-results/ctest.xml" shell: bash + - name: Preserve CTest diagnostics + if: ${{ always() && steps.ctest.outcome == 'failure' }} + uses: ./.github/actions/ctest-diagnostics + with: + summary-title: PR CTest failure + artifact-name: ctest-diagnostics-pr + junit-report: build/test-results/ctest.xml + failed-tests-log: build/Testing/Temporary/LastTestsFailed.log + last-test-log: build/Testing/Temporary/LastTest.log + build-flags: - needs: [spell-check, clang-format, check-paths] + needs: [check-paths] if: *run-ci-for-source-changes runs-on: ubuntu-22.04 - strategy: - matrix: - python-version: ['3.10', '3.12'] env: CI: "true" SCCACHE_GHA_ENABLED: "true" @@ -223,10 +463,10 @@ jobs: with: persist-credentials: false - - name: Set up Python ${{ matrix.python-version }} + - name: Set up Python 3.12 uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python-version }} + python-version: '3.12' - name: Free up disk space uses: ./.github/actions/free-disk-space @@ -294,8 +534,9 @@ jobs: if: &run-ci >- (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) runs-on: ubuntu-22.04 steps: - name: Checkout Actions Repository @@ -391,6 +632,52 @@ jobs: ./scripts/code_format.sh --check --changed-lines --base "${BASE_REF}" shell: bash + python-lint: + name: Check Python with Ruff + if: *run-ci + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install Ruff + run: python -m pip install --disable-pip-version-check ruff==0.6.9 + + - name: Check changed Python files + run: | + if [ "${{ github.event_name }}" = "pull_request" ]; then + base_ref="origin/${{ github.base_ref }}" + elif [ "${{ github.event.before }}" = \ + "0000000000000000000000000000000000000000" ]; then + base_ref="origin/${{ github.event.repository.default_branch }}" + elif [ -n "${{ github.event.before }}" ]; then + base_ref="${{ github.event.before }}" + else + base_ref="HEAD^" + fi + + mapfile -d '' python_files < <( + git diff --name-only --diff-filter=ACMR -z "$base_ref"...HEAD -- \ + '*.py' \ + ':(exclude)extern/**' \ + ':(exclude)FAST25-release/**' + ) + if [ "${#python_files[@]}" -eq 0 ]; then + echo "No changed Python files to check." + exit 0 + fi + + ruff check "${python_files[@]}" + ruff format --check "${python_files[@]}" + shell: bash + docs-check: name: Check Sphinx docs build @@ -436,7 +723,6 @@ jobs: outputs: should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }} should-run-tent: ${{ steps.dispatch-override.outputs.tent || steps.filter.outputs.tent }} - should-run-reshard: ${{ steps.dispatch-override.outputs.reshard || steps.filter.outputs.reshard }} steps: # workflow_dispatch has no PR/push diff context — skip paths-filter and default to true - name: Default to true for workflow_dispatch @@ -445,7 +731,6 @@ jobs: run: | echo "src=true" >> $GITHUB_OUTPUT echo "tent=true" >> $GITHUB_OUTPUT - echo "reshard=true" >> $GITHUB_OUTPUT - uses: actions/checkout@v4 if: github.event_name != 'workflow_dispatch' with: @@ -470,49 +755,6 @@ jobs: - 'CMakeLists.txt' - 'dependencies.sh' - '.github/workflows/ci.yml' - reshard: - - 'mooncake-reshard/**' - - 'mooncake-wheel/pyproject.toml' - - 'requirements.txt' - - 'scripts/check_reshard_types.sh' - - '.github/workflows/ci.yml' - - reshard-test: - name: Test Reshard (Python ${{ matrix.python-version }}) - needs: [check-paths] - if: >- - (needs.check-paths.outputs.should-run-reshard == 'true' || - github.event_name == 'workflow_dispatch') && - (github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) - runs-on: ubuntu-22.04 - strategy: - fail-fast: false - matrix: - python-version: ['3.10', '3.12'] - steps: - - uses: actions/checkout@v4 - with: - persist-credentials: false - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Install Reshard test dependencies - run: | - python -m pip install --disable-pip-version-check --upgrade pip - python -m pip install --disable-pip-version-check \ - -r requirements.txt pytest==8.3.5 - - - name: Test Reshard - run: | - PYTHONPATH=mooncake-reshard/python \ - python -m pytest -q mooncake-reshard/tests - shell: bash build-wheel-cu13: needs: [spell-check, clang-format, check-paths] @@ -520,7 +762,7 @@ jobs: uses: ./.github/workflows/_build-wheel.yaml with: variant: cuda13 - python-versions: '["3.10", "3.12"]' + python-versions: '["3.12"]' artifact-prefix: mooncake-wheel-cu130 version-override: 0.0.0.dev0 @@ -537,17 +779,12 @@ jobs: github.event_name == 'workflow_dispatch') && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) uses: ./.github/workflows/ci_rocm.yml secrets: inherit - integration-test: - needs: [build-wheel-cu13, check-paths] - if: needs.check-paths.outputs.should-run-downstream == 'true' - uses: ./.github/workflows/integration-test.yml - secrets: inherit - tent-ci: needs: [spell-check, clang-format, check-paths] if: >- @@ -555,8 +792,9 @@ jobs: github.event_name == 'workflow_dispatch') && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) runs-on: ubuntu-22.04 strategy: fail-fast: false @@ -718,18 +956,18 @@ jobs: needs: - spell-check - clang-format + - python-lint - docs-check - reshard-type-check - release-gate-tests - - reshard-test - build-wheel + - unit-tests - build-flags - test-wheel-ubuntu - build-wheel-cu13 - build-wheel-efa - build-wheel-rocm - tent-ci - - integration-test runs-on: ubuntu-latest steps: - name: Check required job results diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 262136d8c1..3888e68237 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -11,6 +11,14 @@ on: description: 'PR number (passed from parent workflow for workflow_dispatch)' required: false type: string + source_run_id: + description: 'Workflow run containing the wheel artifact' + required: false + type: string + artifact_name: + description: 'Exact wheel artifact name' + required: false + type: string testpypi_version: description: 'Exact TestPyPI CUDA 13 package version to test instead of a CI artifact' required: false @@ -23,14 +31,14 @@ jobs: env: tone_user_name: ${{ secrets.TONE_USER_NAME }} steps: - - name: Require T-one credentials for the pre-release gate - if: ${{ inputs.testpypi_version != '' }} + - name: Require T-one credentials for gated integration + if: ${{ inputs.testpypi_version != '' || inputs.source_run_id != '' || inputs.artifact_name != '' }} env: TONE_USER_NAME: ${{ secrets.TONE_USER_NAME }} TONE_USER_TOKEN: ${{ secrets.TONE_USER_TOKEN }} run: | if [ -z "$TONE_USER_NAME" ] || [ -z "$TONE_USER_TOKEN" ]; then - echo "TONE_USER_NAME and TONE_USER_TOKEN are required for pre-release validation" + echo "TONE_USER_NAME and TONE_USER_TOKEN are required for gated integration" exit 1 fi @@ -72,6 +80,8 @@ jobs: # Priority: explicit inputs > PR event context > push SHA SHA="${{ inputs.pr_sha || github.event.pull_request.head.sha || github.sha }}" PR_ID="${{ inputs.pr_number || github.event.pull_request.number }}" + SOURCE_RUN_ID="${{ inputs.source_run_id }}" + ARTIFACT_NAME="${{ inputs.artifact_name }}" if [ "${{ github.event_name }}" = "push" ]; then SHA="${{ github.sha }}" @@ -85,46 +95,53 @@ jobs: echo "Attempt $attempt: Fetching artifact..." echo "Target SHA=${SHA}" artifact_id="" - run_id="" - if curl -L -fs -o runs.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "https://api.github.com/repos/${{ github.repository }}/actions/runs?head_sha=${SHA}&per_page=100"; then - if jq empty runs.json >/dev/null 2>&1; then - run_id=$(jq -r '.workflow_runs[] | select((.path == ".github/workflows/ci.yml") or (.name == "Build & Test (Linux)")) | .id' runs.json | head -n 1) + run_id="$SOURCE_RUN_ID" + if [ -z "$run_id" ]; then + if curl -L -fs -o runs.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "https://api.github.com/repos/${{ github.repository }}/actions/runs?head_sha=${SHA}&per_page=100"; then + if jq empty runs.json >/dev/null 2>&1; then + run_id=$(jq -r '.workflow_runs[] | select((.path == ".github/workflows/ci.yml") or (.name == "Build & Test (Linux)")) | .id' runs.json | head -n 1) + else + echo "Failed to download workflow run list. Retrying..." + fi else - echo "Failed to download workflow run list. Retrying..." + echo "Failed to fetch workflow runs. Retrying..." fi - if [ -n "$run_id" ]; then - echo "Matched workflow run id $run_id" - if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "https://api.github.com/repos/${{ github.repository }}/actions/runs/${run_id}/artifacts?per_page=100"; then - if jq empty artifact.json >/dev/null 2>&1; then - artifact_id=$(jq -r '.artifacts[] | select(.name | contains("py312") ) | select(.name | contains("mooncake") ) | select(.name | contains("cu130") ) | .id' artifact.json | head -n 1) - if [ -z "$artifact_id" ]; then - echo "Available artifacts in workflow run $run_id:" - jq -r '.artifacts[].name' artifact.json || true - fi + fi + + if [ -n "$run_id" ]; then + echo "Matched workflow run id $run_id" + if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "https://api.github.com/repos/${{ github.repository }}/actions/runs/${run_id}/artifacts?per_page=100"; then + if jq empty artifact.json >/dev/null 2>&1; then + if [ -n "$ARTIFACT_NAME" ]; then + artifact_id=$(jq -r --arg name "$ARTIFACT_NAME" '.artifacts[] | select(.name == $name) | .id' artifact.json | head -n 1) else - echo "Failed to download artifact list. Retrying..." + artifact_id=$(jq -r '.artifacts[] | select(.name | contains("py312") ) | select(.name | contains("mooncake") ) | select(.name | contains("cu130") ) | .id' artifact.json | head -n 1) + fi + if [ -z "$artifact_id" ]; then + echo "Available artifacts in workflow run $run_id:" + jq -r '.artifacts[].name' artifact.json || true fi else - echo "Failed to fetch artifacts for workflow run $run_id. Retrying..." + echo "Failed to download artifact list. Retrying..." fi else - echo "Failed to find Build & Test workflow run for SHA $SHA. Retrying..." + echo "Failed to fetch artifacts for workflow run $run_id. Retrying..." + fi + else + echo "Failed to find Build & Test workflow run for SHA $SHA. Retrying..." + if jq empty runs.json >/dev/null 2>&1; then echo "Available workflow runs for SHA:" jq -r '.workflow_runs[] | "\(.id) \(.name) \(.path) \(.status) \(.conclusion)"' runs.json || true fi - if [ -n "$artifact_id" ]; then - echo "Successfully fetched expected artifact id $artifact_id" - break - else - echo "Failed to fetch expected artifact. Retrying..." - if [ $attempt -lt $max_attempts ]; then - sleep $((attempt * 60 < 600 ? attempt * 60 : 600)) - fi - fi + fi + + if [ -n "$artifact_id" ]; then + echo "Successfully fetched expected artifact id $artifact_id" + break else - echo "Failed to fetch workflow runs. Retrying..." + echo "Failed to fetch expected artifact. Retrying..." if [ $attempt -lt $max_attempts ]; then - sleep $((attempt * 60)) + sleep $((attempt * 60 < 600 ? attempt * 60 : 600)) fi fi attempt=$((attempt + 1)) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index c9f7327898..55a4675c32 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -74,14 +74,25 @@ jobs: - variant: non-cuda architecture: x86_64 artifact-prefix: nightly-non-cuda-x86 + - variant: non-cuda + architecture: arm64 + artifact-prefix: nightly-non-cuda-arm64 uses: ./.github/workflows/_build-wheel.yaml with: variant: ${{ matrix.variant }} architecture: ${{ matrix.architecture }} - python-versions: '["3.10", "3.12"]' artifact-prefix: ${{ matrix.artifact-prefix }} version-override: ${{ needs.version-stamp.outputs.nightly_version }} + tone-sglang-integration: + needs: build-wheels + uses: ./.github/workflows/integration-test.yml + with: + pr_sha: ${{ github.sha }} + source_run_id: ${{ github.run_id }} + artifact_name: nightly-cuda13-x86-py312 + secrets: inherit + build-musa: runs-on: ubuntu-22.04 container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64 @@ -286,6 +297,7 @@ jobs: cd mooncake-transfer-engine/example/http-metadata-server-python pip install aiohttp python ./bootstrap_server.py & + echo "NIGHTLY_METADATA_SERVER_PID=$!" >> "$GITHUB_ENV" sleep 2 - name: Run CTest unit tests @@ -356,13 +368,103 @@ jobs: MC_METADATA_SERVER: http://127.0.0.1:8080/metadata RUN_TESTS_METADATA_SERVER_MODE: external DEFAULT_KV_LEASE_TTL: "500" - TEST_SSD_OFFLOAD_IN_EVICT: "1" - TEST_PROMOTION_ON_HIT: "1" TEST_CXL: "1" run: | export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib bash scripts/run_tests.sh + - name: Run Python API end-to-end tests + env: + MOONCAKE_MASTER: "127.0.0.1:50051" + MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata" + MOONCAKE_PROTOCOL: "tcp" + LOCAL_HOSTNAME: "127.0.0.1" + run: | + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + kill "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null || true + for _ in {1..50}; do + if ! kill -0 "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null; then + break + fi + sleep 0.1 + done + if kill -0 "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null; then + kill -KILL "$NIGHTLY_METADATA_SERVER_PID" + fi + mkdir -p /tmp/mooncake_storage + mooncake_master \ + --default_kv_lease_ttl=500 \ + --eviction_high_watermark_ratio=0.95 \ + --cluster_id=nightly_test_cluster \ + --port 50051 \ + --enable_http_metadata_server=true & + master_pid=$! + trap 'kill "$master_pid" 2>/dev/null || true; wait "$master_pid" 2>/dev/null || true' EXIT + sleep 3 + python scripts/test_tensor_api.py -n 1 + python scripts/test_async_store.py + python scripts/test_copy_move_api.py + python scripts/test_drain_http_api.py --timeout-sec 90 + + - name: Run RPC Communicator Bandwidth Test + run: | + rpc_server_pid="" + cleanup() { + status=$? + trap - EXIT + if [ -n "$rpc_server_pid" ]; then + kill "$rpc_server_pid" 2>/dev/null || true + wait "$rpc_server_pid" 2>/dev/null || true + fi + if [ "$status" -ne 0 ]; then + cat "$RUNNER_TEMP/rpc-server.log" 2>/dev/null || true + fi + exit "$status" + } + trap cleanup EXIT + + python -u mooncake-transfer-engine/tests/rpc_communicator_test.py \ + server --url 127.0.0.1:9004 --data-size 1 \ + >"$RUNNER_TEMP/rpc-server.log" 2>&1 & + rpc_server_pid=$! + + rpc_ready=false + for _ in {1..50}; do + if ! kill -0 "$rpc_server_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :9004' | grep -q .; then + rpc_ready=true + break + fi + sleep 0.1 + done + if [ "$rpc_ready" != true ]; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server did not become ready within 5 seconds" + exit 1 + fi + + client_rc=0 + timeout 10 python -u \ + mooncake-transfer-engine/tests/rpc_communicator_test.py \ + client --url 127.0.0.1:9004 --threads 2 --data-size 1 \ + >"$RUNNER_TEMP/rpc-client.log" 2>&1 || \ + client_rc=$? + cat "$RUNNER_TEMP/rpc-client.log" + if [ "$client_rc" -ne 0 ] && [ "$client_rc" -ne 124 ]; then + echo "::error::RPC communicator client failed with exit code $client_rc" + exit "$client_rc" + fi + if ! grep -q '^bandwidth:' "$RUNNER_TEMP/rpc-client.log"; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator did not complete a successful transfer" + exit 1 + fi + shell: bash + nightly-coverage: runs-on: ubuntu-22.04 env: @@ -527,6 +629,7 @@ jobs: needs: - version-stamp - build-wheels + - tone-sglang-integration - build-musa - build-docker - ascend-test diff --git a/.github/workflows/release-efa-cuda13.yaml b/.github/workflows/release-efa-cuda13.yaml index dcbfe7434f..1b0809b4d3 100644 --- a/.github/workflows/release-efa-cuda13.yaml +++ b/.github/workflows/release-efa-cuda13.yaml @@ -8,6 +8,7 @@ on: # Publishes mooncake-transfer-engine-efa-cuda13, the CUDA 13-aware AWS EFA wheel. jobs: build: + if: ${{ !contains(github.ref_name, '-') }} permissions: contents: write uses: ./.github/workflows/_build-efa-wheel.yaml diff --git a/.github/workflows/release-efa-non-cuda.yaml b/.github/workflows/release-efa-non-cuda.yaml index 6ea27b00d6..6c6f8edc2b 100644 --- a/.github/workflows/release-efa-non-cuda.yaml +++ b/.github/workflows/release-efa-non-cuda.yaml @@ -8,6 +8,7 @@ on: # Publishes mooncake-transfer-engine-efa-non-cuda, the CPU/DRAM AWS EFA wheel. jobs: build: + if: ${{ !contains(github.ref_name, '-') }} permissions: contents: write uses: ./.github/workflows/_build-efa-wheel.yaml diff --git a/.github/workflows/release-efa.yaml b/.github/workflows/release-efa.yaml index c10d46bc89..a16aa5a7e8 100644 --- a/.github/workflows/release-efa.yaml +++ b/.github/workflows/release-efa.yaml @@ -8,6 +8,7 @@ on: # Publishes mooncake-transfer-engine-efa, the CUDA 12-aware AWS EFA wheel. jobs: build: + if: ${{ !contains(github.ref_name, '-') }} permissions: contents: write uses: ./.github/workflows/_build-efa-wheel.yaml diff --git a/mooncake-pg/tests/test_pg_elastic.py b/mooncake-pg/tests/test_pg_elastic.py index db7d35e0f2..320c0f0152 100644 --- a/mooncake-pg/tests/test_pg_elastic.py +++ b/mooncake-pg/tests/test_pg_elastic.py @@ -69,8 +69,9 @@ def _extension_worker( description=f"rank {ctx.proc_rank} waiting for joiner ready", ) resp = pg.recover_ranks(backend, join_ranks) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" # The Coordinator has committed the joiner as active, so the visible # rank-space extent now covers the full group. @@ -81,14 +82,18 @@ def _extension_worker( ) # Final collective - final_tensor = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) + final_tensor = torch.tensor( + [ctx.proc_rank + 1], dtype=torch.int32, device=device + ) dist.all_reduce(final_tensor, op=dist.ReduceOp.SUM) - ctx.record_result({ - "role": "original", - "rank": ctx.proc_rank, - "baseline": baseline, - }) + ctx.record_result( + { + "role": "original", + "rank": ctx.proc_rank, + "baseline": baseline, + } + ) else: # Extension rank if not extend_event.wait(timeout=30.0): @@ -134,13 +139,17 @@ def _extension_worker( ) # Final collective - final_tensor = torch.tensor([extension_rank + 1], dtype=torch.int32, device=device) + final_tensor = torch.tensor( + [extension_rank + 1], dtype=torch.int32, device=device + ) dist.all_reduce(final_tensor, op=dist.ReduceOp.SUM) - ctx.record_result({ - "role": "extension", - "rank": extension_rank, - }) + ctx.record_result( + { + "role": "extension", + "rank": extension_rank, + } + ) def _extension_p2p_worker( @@ -174,8 +183,9 @@ def _extension_p2p_worker( description=f"rank {ctx.proc_rank} waiting for joiner ready", ) resp = pg.recover_ranks(backend, join_ranks) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" else: if not extend_event.wait(timeout=30.0): raise TimeoutError("timed out waiting for extend_event") @@ -212,9 +222,7 @@ def _extension_p2p_worker( numel = 1024 if ctx.proc_rank == src_rank: - send_tensor = torch.full( - (numel,), src_rank, dtype=torch.int32, device=device - ) + send_tensor = torch.full((numel,), src_rank, dtype=torch.int32, device=device) works = dist.batch_isend_irecv( [dist.P2POp(op=dist.isend, tensor=send_tensor, peer=dst_rank)] ) @@ -304,8 +312,16 @@ def _extension_worker_with_subgroups( backend=ctx.backend_name, pg_options=pg.MooncakeBackendOptions(4), ) - a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 0 else None - b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 1 else None + a_backend = ( + get_mooncake_backend(group_a, device_type=ctx.device_type) + if ctx.proc_rank == 0 + else None + ) + b_backend = ( + get_mooncake_backend(group_b, device_type=ctx.device_type) + if ctx.proc_rank == 1 + else None + ) c_backend = get_mooncake_backend(group_c, device_type=ctx.device_type) # Pre-activation: WORLD primary ranks sum to 1+2=3 @@ -331,8 +347,9 @@ def _extension_worker_with_subgroups( description=f"rank {ctx.proc_rank} waiting for joiners", ) resp = pg.recover_ranks(world_backend, join_ranks) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {ctx.proc_rank}: recover_ranks(world) should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {ctx.proc_rank}: recover_ranks(world) should apply, got {resp.status}" # group_a: rank 0 waits for joiner (local rank 1 = global rank 2) if ctx.proc_rank == 0: @@ -343,8 +360,9 @@ def _extension_worker_with_subgroups( description="rank 0 waiting for group_a joiner", ) resp = pg.recover_ranks(a_backend, [1]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank 0: recover_ranks(group_a) should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank 0: recover_ranks(group_a) should apply, got {resp.status}" # group_b: rank 1 waits for joiner (local rank 1 = global rank 3) if ctx.proc_rank == 1: @@ -355,8 +373,9 @@ def _extension_worker_with_subgroups( description="rank 1 waiting for group_b joiner", ) resp = pg.recover_ranks(b_backend, [1]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank 1: recover_ranks(group_b) should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank 1: recover_ranks(group_b) should apply, got {resp.status}" # group_c: both primaries wait for both joiners (local ranks 2,3) wait_until( @@ -366,8 +385,9 @@ def _extension_worker_with_subgroups( description=f"rank {ctx.proc_rank} waiting for group_c joiners", ) resp = pg.recover_ranks(c_backend, [2, 3]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {ctx.proc_rank}: recover_ranks(group_c) should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {ctx.proc_rank}: recover_ranks(group_c) should apply, got {resp.status}" # Post-activation: WORLD all 4 ranks → 1+2+3+4=10 t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) @@ -380,20 +400,26 @@ def _extension_worker_with_subgroups( ta = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(ta, op=dist.ReduceOp.SUM, group=group_a) if int(ta.cpu().item()) != 4: - raise AssertionError(f"group_a post: expected 4, got {int(ta.cpu().item())}") + raise AssertionError( + f"group_a post: expected 4, got {int(ta.cpu().item())}" + ) # group_b: ranks [1,3], values [2,4], sum=6 if ctx.proc_rank == 1: tb = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tb, op=dist.ReduceOp.SUM, group=group_b) if int(tb.cpu().item()) != 6: - raise AssertionError(f"group_b post: expected 6, got {int(tb.cpu().item())}") + raise AssertionError( + f"group_b post: expected 6, got {int(tb.cpu().item())}" + ) # group_c: all 4 ranks, sum=10 tc = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tc, op=dist.ReduceOp.SUM, group=group_c) if int(tc.cpu().item()) != 10: - raise AssertionError(f"group_c post: expected 10, got {int(tc.cpu().item())}") + raise AssertionError( + f"group_c post: expected 10, got {int(tc.cpu().item())}" + ) ctx.record_result({"role": "extension_subgroups", "rank": ctx.proc_rank}) else: @@ -428,8 +454,16 @@ def _extension_worker_with_subgroups( backend=ctx.backend_name, pg_options=pg.MooncakeBackendOptions(4, True), ) - a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 2 else None - b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 3 else None + a_backend = ( + get_mooncake_backend(group_a, device_type=ctx.device_type) + if ctx.proc_rank == 2 + else None + ) + b_backend = ( + get_mooncake_backend(group_b, device_type=ctx.device_type) + if ctx.proc_rank == 3 + else None + ) c_backend = get_mooncake_backend(group_c, device_type=ctx.device_type) # Before any join_group call, WORLD collectives are local-only on each @@ -439,8 +473,7 @@ def _extension_worker_with_subgroups( local_value = int(t.cpu().item()) if local_value != ctx.proc_rank + 1: raise AssertionError( - f"WORLD local-only: expected {ctx.proc_rank + 1}, " - f"got {local_value}" + f"WORLD local-only: expected {ctx.proc_rank + 1}, " f"got {local_value}" ) # Join groups in same order primaries created them @@ -462,20 +495,26 @@ def _extension_worker_with_subgroups( ta = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(ta, op=dist.ReduceOp.SUM, group=group_a) if int(ta.cpu().item()) != 4: - raise AssertionError(f"group_a post: expected 4, got {int(ta.cpu().item())}") + raise AssertionError( + f"group_a post: expected 4, got {int(ta.cpu().item())}" + ) # group_b: ranks [1,3], sum=6 if ctx.proc_rank == 3: tb = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tb, op=dist.ReduceOp.SUM, group=group_b) if int(tb.cpu().item()) != 6: - raise AssertionError(f"group_b post: expected 6, got {int(tb.cpu().item())}") + raise AssertionError( + f"group_b post: expected 6, got {int(tb.cpu().item())}" + ) # group_c: all 4 ranks, sum=10 tc = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tc, op=dist.ReduceOp.SUM, group=group_c) if int(tc.cpu().item()) != 10: - raise AssertionError(f"group_c post: expected 10, got {int(tc.cpu().item())}") + raise AssertionError( + f"group_c post: expected 10, got {int(tc.cpu().item())}" + ) ctx.record_result({"role": "extension_subgroups", "rank": ctx.proc_rank}) @@ -570,8 +609,9 @@ def _allgather_reduce_scatter_extension_worker( description=f"rank {ctx.proc_rank} waiting for joiner", ) resp = pg.recover_ranks(backend, join_ranks) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" # Post-activation: all 4 ranks active. _run_allgather_reduce_scatter(device, ctx.world_size, ctx.proc_rank) @@ -602,6 +642,7 @@ def _allgather_reduce_scatter_extension_worker( def _allgather_reduce_scatter_recovery_worker( ctx: MooncakePGWorkerContext, + pre_failure_barrier: mp.Barrier, broken_exited: mp.Event, start_recovery: mp.Event, ) -> None: @@ -619,17 +660,26 @@ def _allgather_reduce_scatter_recovery_worker( device = ctx.init_group(rank=logical_rank) backend = ctx.get_backend() - # Pre-failure: all 4 ranks active. + # Pre-failure: all 4 ranks active. Wait for every rank's collectives to + # return before killing one process so the injected failure cannot race + # with completion of the healthy round. _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank) + ctx.synchronize() + pre_failure_barrier.wait(timeout=30.0) if logical_rank == broken_rank: ctx.record_result({"role": "broken"}) broken_exited.set() os._exit(0) - # Survivors: 3 active ranks, max_group_size=4 → overflow path. + # First use the basic collective path to detect and deactivate the + # departed rank. The assertions below target the three-rank buffer + # layout, not failure detection inside all-gather/reduce-scatter. broken_exited.wait() + probe = torch.tensor([logical_rank], dtype=torch.int32, device=device) + dist.all_reduce(probe, op=dist.ReduceOp.SUM) + # Survivors: 3 active ranks, max_group_size=4 → overflow path. _run_allgather_reduce_scatter(device, ctx.world_size - 1, logical_rank) if logical_rank == 0: @@ -642,8 +692,9 @@ def _allgather_reduce_scatter_recovery_worker( description=f"rank {logical_rank} waiting for replacement", ) resp = pg.recover_ranks(backend, [broken_rank]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" # Post-recovery: all 4 ranks active again. _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank) @@ -715,13 +766,14 @@ def _replacement_recovery_worker( # view used by that collective. ctx.synchronize() resp = pg.deactivate_ranks(backend, [logical_rank]) - assert resp.status == pg.ProposalStatus.Applied, \ - "graceful self-deactivation should apply, " \ + assert resp.status == pg.ProposalStatus.Applied, ( + "graceful self-deactivation should apply, " f"got {resp.status}: {resp.reject_reason}" + ) dist.destroy_process_group() ctx.record_result({"role": "gracefully_removed"}) - # Do not set broken_exited here. The parent sets the event only + # Do not set broken_exited here. The parent sets the event only # after it observes that this process has exited. return @@ -752,8 +804,9 @@ def _replacement_recovery_worker( # All ranks call recover_ranks to include replacement resp = pg.recover_ranks(backend, [BROKEN_RANK]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" # Final collective with all 4 ranks tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device) @@ -789,9 +842,7 @@ def _run_p2p_ping_pong( raise AssertionError("recovery P2P test expects exactly two ranks") peer = 1 - logical_rank - send_tensor = torch.tensor( - [logical_rank], dtype=torch.int64, device=device - ) + send_tensor = torch.tensor([logical_rank], dtype=torch.int64, device=device) recv_tensor = torch.empty_like(send_tensor) works = dist.batch_isend_irecv( [ @@ -803,9 +854,7 @@ def _run_p2p_ping_pong( work.wait() if not all(pg.get_local_success(work) for work in works): - raise AssertionError( - f"rank {logical_rank}: P2P with rank {peer} failed" - ) + raise AssertionError(f"rank {logical_rank}: P2P with rank {peer} failed") received = int(recv_tensor.cpu().item()) if received != peer: raise AssertionError( @@ -830,9 +879,10 @@ def _recovery_p2p_worker( if logical_rank == BROKEN_RANK: if graceful_group_destroy: resp = pg.deactivate_ranks(backend, [logical_rank]) - assert resp.status == pg.ProposalStatus.Applied, \ - "graceful self-deactivation should apply, " \ + assert resp.status == pg.ProposalStatus.Applied, ( + "graceful self-deactivation should apply, " f"got {resp.status}: {resp.reject_reason}" + ) dist.destroy_process_group() ctx.record_result({"role": "gracefully_removed"}) @@ -869,8 +919,9 @@ def _recovery_p2p_worker( description=f"rank {logical_rank} waiting for replacement", ) resp = pg.recover_ranks(backend, [BROKEN_RANK]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" role = "survivor" else: @@ -894,12 +945,8 @@ def _run_rejoin_operation( verify: bool, ): if operation == "collective": - tensor = torch.tensor( - [ctx.rank + 1], dtype=torch.int32, device=device - ) - works = [ - dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) - ] + tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device) + works = [dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True)] for work in works: work.wait() @@ -909,16 +956,12 @@ def _run_rejoin_operation( expected = ctx.world_size * (ctx.world_size + 1) // 2 actual = int(tensor.cpu().item()) if actual != expected: - raise AssertionError( - f"collective expected {expected}, got {actual}" - ) + raise AssertionError(f"collective expected {expected}, got {actual}") return works if operation == "p2p": peer = 1 - ctx.rank - send_tensor = torch.tensor( - [ctx.rank], dtype=torch.int64, device=device - ) + send_tensor = torch.tensor([ctx.rank], dtype=torch.int64, device=device) recv_tensor = torch.empty_like(send_tensor) works = dist.batch_isend_irecv( [ @@ -957,9 +1000,7 @@ def _run_fault_round( return fault_started.set() - works = _run_rejoin_operation( - ctx, device, operation, verify=False - ) + works = _run_rejoin_operation(ctx, device, operation, verify=False) if all(pg.get_local_success(work) for work in works): raise AssertionError(f"injected {operation} unexpectedly succeeded") @@ -991,9 +1032,7 @@ def _inplace_rejoin_worker( epoch_before_failure = pg.get_current_epoch(backend) # 2. Fail rank 1's data plane. Auto sync must make it inactive. - _run_fault_round( - ctx, device, backend, operation, fault_library, fault_started - ) + _run_fault_round(ctx, device, backend, operation, fault_library, fault_started) # 3. At an application-selected safe point, the failed rank declares # itself ready to rejoin. @@ -1051,8 +1090,9 @@ def _manual_deactivate_recovery_worker( work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) work.wait() assert int(tensor.cpu().item()) == expected_all - assert pg.get_local_success(work), \ - f"rank {ctx.rank}: round 1 should succeed locally" + assert pg.get_local_success( + work + ), f"rank {ctx.rank}: round 1 should succeed locally" failed_ranks_hint = pg.get_failed_ranks_hint(work) assert failed_ranks_hint.tolist() == [0] * ctx.world_size @@ -1074,8 +1114,9 @@ def _manual_deactivate_recovery_worker( work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) work.wait() assert int(tensor.cpu().item()) == expected_reduced - assert not pg.get_local_success(work), \ - f"rank {ctx.rank}: round 2 should detect broken rank" + assert not pg.get_local_success( + work + ), f"rank {ctx.rank}: round 2 should detect broken rank" failed_ranks_hint = pg.get_failed_ranks_hint(work) expected_failed_ranks_hint = [0] * ctx.world_size @@ -1094,8 +1135,9 @@ def _manual_deactivate_recovery_worker( pg.SyncAfterFailureStatus.Reconciled, pg.SyncAfterFailureStatus.NoPending, ), f"rank {logical_rank}: sync_after_failure failed: {sync_resp.reject_reason}" - assert not pg.get_peer_state(backend, [BROKEN_RANK])[0], \ - f"rank {logical_rank}: failed rank remained locally activatable after sync" + assert not pg.get_peer_state(backend, [BROKEN_RANK])[ + 0 + ], f"rank {logical_rank}: failed rank remained locally activatable after sync" # Survivors deactivate the dead rank before issuing new collectives. resp = pg.deactivate_ranks(backend, [BROKEN_RANK]) @@ -1116,8 +1158,9 @@ def _manual_deactivate_recovery_worker( work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) work.wait() assert int(tensor.cpu().item()) == expected_reduced - assert pg.get_local_success(work), \ - f"rank {ctx.rank}: round 3 should succeed after manual deactivate" + assert pg.get_local_success( + work + ), f"rank {ctx.rank}: round 3 should succeed after manual deactivate" if logical_rank == 0: start_recovery.set() @@ -1130,12 +1173,11 @@ def _manual_deactivate_recovery_worker( ) resp = pg.recover_ranks(backend, [BROKEN_RANK]) - assert resp.status == pg.ProposalStatus.Applied, \ - f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" + assert ( + resp.status == pg.ProposalStatus.Applied + ), f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" - tensor = torch.tensor( - [logical_rank + 1], dtype=torch.int32, device=device - ) + tensor = torch.tensor([logical_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) assert int(tensor.cpu().item()) == expected_all @@ -1154,9 +1196,7 @@ def _manual_deactivate_recovery_worker( backend = ctx.get_backend() pg.join_group(backend) - tensor = torch.tensor( - [logical_rank + 1], dtype=torch.int32, device=device - ) + tensor = torch.tensor([logical_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) ctx.record_result({"role": "replacement"}) @@ -1212,7 +1252,7 @@ def test_recovery(self) -> None: self.assertGreaterEqual(len(broken_rows), 1) def test_recovery_after_graceful_group_destroy(self) -> None: - """A rank can destroy its group normally and later be replaced. """ + """A rank can destroy its group normally and later be replaced.""" spawn_ctx = mp.get_context("spawn") removed_process_exited = spawn_ctx.Event() start_recovery = spawn_ctx.Event() @@ -1234,9 +1274,7 @@ def test_recovery_after_graceful_group_destroy(self) -> None: survivor_rows = [r for r in rows if r.get("role") == "survivor"] replacement_rows = [r for r in rows if r.get("role") == "replacement"] - removed_rows = [ - r for r in rows if r.get("role") == "gracefully_removed" - ] + removed_rows = [r for r in rows if r.get("role") == "gracefully_removed"] self.assertEqual(len(survivor_rows), self.world_size - 1) self.assertEqual(len(replacement_rows), 1) @@ -1257,9 +1295,7 @@ def _run_recovery_p2p(self, graceful_group_destroy: bool) -> None: nprocs=recovery_world_size + 1, timeout_s=75.0, process_exit_events=( - {BROKEN_RANK: broken_exited} - if graceful_group_destroy - else None + {BROKEN_RANK: broken_exited} if graceful_group_destroy else None ), ) @@ -1405,11 +1441,13 @@ def test_allgather_reduce_scatter_recovery(self) -> None: and max_group_size=4, so the buggy code would access slot 3 of a size-3 buffer. """ spawn_ctx = mp.get_context("spawn") + pre_failure_barrier = spawn_ctx.Barrier(self.world_size) broken_exited = spawn_ctx.Event() start_recovery = spawn_ctx.Event() rows = self.spawn_backend_and_collect( _allgather_reduce_scatter_recovery_worker, + pre_failure_barrier, broken_exited, start_recovery, nprocs=self.world_size + 1, @@ -1459,15 +1497,11 @@ def test_manual_evict_recovery(self) -> None: self.assertGreaterEqual(len(broken_rows), 1) -class TestMooncakePGElasticCPU( - _ElasticMixin, MooncakePGCPUBackendTestCase -): +class TestMooncakePGElasticCPU(_ElasticMixin, MooncakePGCPUBackendTestCase): pass -class TestMooncakePGElasticCUDA( - _ElasticMixin, MooncakePGCUDABackendTestCase -): +class TestMooncakePGElasticCUDA(_ElasticMixin, MooncakePGCUDABackendTestCase): @classmethod def configure_for_cuda_device_count(cls, device_count: int) -> None: if device_count < 2: From 9fd23e2c30ff1548a2919b2af524f3fdabf2e840 Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 31 Aug 2026 14:09:57 +0800 Subject: [PATCH 02/17] [TENT] Add task failure reason counter and in-flight gauges (#3680) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Add task failure reason counter and in-flight gauges Add three metrics for production diagnosis, built on the cached label cells: - tent_task_failures_total{transport, reason}: task-level failures with reason submit/poll/timeout/canceled. TaskInfo.failure_stage marks where the first failure originated (set at the submit-rejection and poll sites, never overwritten, so a poll failure followed by a rejected failover resubmit still counts as poll). Additive to the legacy read/write_failures_total counters, and the first failure metric that records TIMEOUT/CANCELED outcomes. - tent_inflight_attempts{transport}: gauge incremented on attempt submit and decremented on attempt finish — how much work is sitting in the engine right now. - tent_registered_buffer_bytes{transport}: gauge maintained on register/unregisterLocalMemory from desc.transports. * [TENT] Keep state gauges symmetric across setEnabled transitions Gauge updates are paired add/sub operations (attempt start/finish, register/unregister). Gating them on the runtime enable flag let a transition between the two halves of a pair permanently corrupt the gauge (stuck positive or negative). Gauges now update whenever initialized and ignore setEnabled: they track engine state, not samples. Counters and histograms keep the runtime gate. Add tests toggling setEnabled between paired updates for both gauges. --- .../tent/include/tent/metrics/tent_metrics.h | 55 +++++++ .../tent/runtime/transfer_engine_impl.h | 16 +- .../tent/src/metrics/tent_metrics.cpp | 98 +++++++++++ .../tent/src/runtime/transfer_engine_impl.cpp | 69 +++++++- .../tent/tests/metrics_recording_test.cpp | 155 ++++++++++++++++++ 5 files changed, 389 insertions(+), 4 deletions(-) diff --git a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h index dab9c1829b..9b6b352915 100644 --- a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h +++ b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h @@ -124,6 +124,31 @@ class TentMetrics { // are parked until shutdown and the last reclaim error is in the log. void recordBatchQuarantined(); + // Why a task failed. "submit" failures were observed synchronously at + // submission (the transport rejected the request); "poll" failures were + // observed by polling after the request had been accepted; timeout and + // canceled come from the terminal status itself. + enum class TaskFailureReason { Submit, Poll, Timeout, Canceled }; + + // Record a task-level failure with its reason. Additive to the legacy + // read/write_failures_total counters (which stay label-compatible); also + // the only failure metric that records TIMEOUT/CANCELED outcomes. + void recordTaskFailure(TransportType tp, TaskFailureReason reason); + + // In-flight transport attempts (gauge): incremented when an attempt is + // submitted, decremented when it finishes. Reflects how much work is + // sitting in the engine right now — a persistently high or stuck value + // is the signature of a stalled pipeline. Gauge updates are paired + // add/sub operations, so they execute whenever initialized and ignore + // setEnabled(): skipping one half of a pair across an enable/disable + // transition would permanently corrupt the gauge. + void recordInflightAttemptStarted(TransportType tp); + void recordInflightAttemptFinished(TransportType tp); + + // Registered buffer bytes per transport (gauge), maintained on + // register/unregisterLocalMemory. Same pairing rule as above. + void recordRegisteredBufferBytes(TransportType tp, int64_t delta); + enum class Stage { QueueWait, Dispatch, @@ -199,6 +224,8 @@ class TentMetrics { "to"}; static inline const std::array kAttemptLabels{"transport", "operation"}; + static inline const std::array kTaskFailureLabels{ + "transport", "reason"}; // Hot-path metrics record through pre-resolved label cells (see // cached_metric.h). Label domain sizes: TransportType has @@ -207,6 +234,7 @@ class TentMetrics { static constexpr size_t kTransportDomain = static_cast(kNumTransportTypes); static constexpr size_t kAttemptDomain = kTransportDomain * 2; + static constexpr size_t kTaskFailureDomain = kTransportDomain * 4; // Stable slot indices into the cached-cell label domain. static size_t transportSlot(TransportType tp) { @@ -215,6 +243,9 @@ class TentMetrics { static size_t attemptSlot(TransportType tp, Request::OpCode op) { return static_cast(tp) * 2 + (op == Request::READ ? 0u : 1u); } + static size_t taskFailureSlot(TransportType tp, TaskFailureReason reason) { + return static_cast(tp) * 4 + static_cast(reason); + } // Per-transport counters (label: transport). Values are int64_t. metrics::CachedDynamicCounter<1> read_bytes_total_{ @@ -248,6 +279,13 @@ class TentMetrics { "tent_transport_attempt_failures_total", "Physical transport attempts that terminated with FAILED", kAttemptLabels, kAttemptDomain}; + // Task-level failures with a reason label, so Grafana can alert on the + // failure signature (submit vs poll vs timeout vs canceled) without log + // scraping. Additive to the legacy read/write_failures_total counters. + metrics::CachedDynamicCounter<2> task_failures_total_{ + "tent_task_failures_total", + "Task-level failures by terminal transport and failure reason", + kTaskFailureLabels, kTaskFailureDomain}; metrics::CachedDynamicCounter<1> deadline_infeasible_total_{ "tent_deadline_infeasible_total", "Transfers whose deadline was already in the past at submit", @@ -259,6 +297,18 @@ class TentMetrics { "Batches abandoned by lazyFreeBatch after repeated failed reclaim " "attempts; reclaimed only at engine teardown"}; + // Gauges, backed by cached cells and updated with add/sub. Serialized as + // Prometheus gauges (current value per label), NOT via the counters_ + // vector. + metrics::CachedDynamicCounter<1> inflight_attempts_{ + "tent_inflight_attempts", + "In-flight transport attempts (submitted, not yet finished)", + kTransportLabel, kTransportDomain}; + metrics::CachedDynamicCounter<1> registered_buffer_bytes_{ + "tent_registered_buffer_bytes", "Registered local buffer bytes", + kTransportLabel, kTransportDomain}; + std::vector*> gauges_; + // Latency histograms use microseconds (us) as unit // Default buckets: 100us, 500us, 1ms, 5ms, 10ms, 50ms, 100ms, 500ms, 1s static inline const std::vector kLatencyBuckets{ @@ -321,6 +371,11 @@ class TentMetrics { // getPrometheusMetrics / getJsonMetrics) via the same templated helpers. std::vector*> histograms_; + // Serialize a gauge backed by a cached-cell counter: emits the current + // value per label with # TYPE ... gauge. + void serializeGaugePrometheus(metrics::CachedDynamicCounter<1>& gauge, + std::string& out) const; + // Helper to register all metrics to the vectors void registerMetrics(); 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 cb40724db5..681e98a7a8 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 @@ -84,6 +84,14 @@ struct TaskInfo { std::chrono::steady_clock::time_point attempt_post_time{}; TransportType attempt_type{UNSPEC}; bool attempt_active{false}; + // Failure attribution for tent_task_failures_total (first failure wins): + // -1 = none, 0 = submit-stage, 1 = poll-stage. Set where the failure + // originates and never overwritten, so a poll-observed failure stays + // "poll" even when a later failover resubmit is synchronously rejected. + // Invariant for integrators: mark submit failures before any recovery + // attempt, so a task that recovers and later fails at poll still + // attributes its root cause to submit. + int8_t failure_stage{-1}; TaskInfo() = default; @@ -105,7 +113,8 @@ struct TaskInfo { post_time(other.post_time), attempt_post_time(other.attempt_post_time), attempt_type(other.attempt_type), - attempt_active(other.attempt_active) {} + attempt_active(other.attempt_active), + failure_stage(other.failure_stage) {} TaskInfo(TaskInfo&& other) noexcept : type(other.type), @@ -125,7 +134,8 @@ struct TaskInfo { post_time(other.post_time), attempt_post_time(other.attempt_post_time), attempt_type(other.attempt_type), - attempt_active(other.attempt_active) {} + attempt_active(other.attempt_active), + failure_stage(other.failure_stage) {} TaskInfo& operator=(const TaskInfo& other) { if (this != &other) { @@ -149,6 +159,7 @@ struct TaskInfo { attempt_post_time = other.attempt_post_time; attempt_type = other.attempt_type; attempt_active = other.attempt_active; + failure_stage = other.failure_stage; } return *this; } @@ -175,6 +186,7 @@ struct TaskInfo { attempt_post_time = other.attempt_post_time; attempt_type = other.attempt_type; attempt_active = other.attempt_active; + failure_stage = other.failure_stage; } return *this; } diff --git a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp index 6f1b4d3b34..82db575d28 100644 --- a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp +++ b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp @@ -35,6 +35,20 @@ namespace { const char* operationName(Request::OpCode operation) { return operation == Request::READ ? "read" : "write"; } + +const char* taskFailureReasonName(TentMetrics::TaskFailureReason reason) { + switch (reason) { + case TentMetrics::TaskFailureReason::Submit: + return "submit"; + case TentMetrics::TaskFailureReason::Poll: + return "poll"; + case TentMetrics::TaskFailureReason::Timeout: + return "timeout"; + case TentMetrics::TaskFailureReason::Canceled: + return "canceled"; + } + return "unknown"; +} } // namespace Status TentMetrics::initialize(const MetricsConfig& config) { @@ -195,6 +209,7 @@ void TentMetrics::shutdown() { // Clear metric vectors counters_.clear(); histograms_.clear(); + gauges_.clear(); // Reset bound port so httpPort() returns 0 after shutdown, not a stale // port from a previous initialization. Without this, a re-initialize @@ -221,10 +236,16 @@ void TentMetrics::registerMetrics() { &failover_total_, &transport_attempts_total_, &transport_attempt_failures_total_, + &task_failures_total_, &deadline_infeasible_total_, &quarantined_batches_total_, }; + gauges_ = { + &inflight_attempts_, + ®istered_buffer_bytes_, + }; + // Register the N=1 per-transport histograms for unified serialization. // The N=2 transport_attempt_latency_ histogram is serialized separately // (it can't share this N=1-typed vector); see getPrometheusMetrics / @@ -301,6 +322,43 @@ void TentMetrics::recordBatchQuarantined() { quarantined_batches_total_.inc(); } +void TentMetrics::recordTaskFailure(TransportType tp, + TaskFailureReason reason) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + task_failures_total_.incCached(taskFailureSlot(tp, reason), [tp, reason] { + return std::array{transportTypeName(tp), + taskFailureReasonName(reason)}; + }); +} + +// Gauge updates below deliberately ignore runtime_enabled_: gauges track +// engine state through paired add/sub operations (attempt start/finish, +// register/unregister), and skipping one half of a pair across a +// setEnabled() transition would permanently corrupt the value. Counters and +// histograms are samples and may be dropped while disabled; gauges are not. +void TentMetrics::recordInflightAttemptStarted(TransportType tp) { + if (!initialized_) return; + inflight_attempts_.incCached(transportSlot(tp), [tp] { + return std::array{transportTypeName(tp)}; + }); +} + +void TentMetrics::recordInflightAttemptFinished(TransportType tp) { + if (!initialized_) return; + inflight_attempts_.incCached( + transportSlot(tp), + [tp] { return std::array{transportTypeName(tp)}; }, -1); +} + +void TentMetrics::recordRegisteredBufferBytes(TransportType tp, int64_t delta) { + if (!initialized_) return; + registered_buffer_bytes_.incCached( + transportSlot(tp), + [tp] { return std::array{transportTypeName(tp)}; }, + delta); +} + void TentMetrics::recordStageLatency(Stage stage, TransportType tp, double latency_us) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) @@ -419,6 +477,11 @@ std::string TentMetrics::getPrometheusMetrics() { // operation) goes through the same helper, instantiated for N=2. serializeHistogramPrometheus(transport_attempt_latency_, result); + // Gauges: current value per label. + for (auto* gauge : gauges_) { + serializeGaugePrometheus(*gauge, result); + } + return result; } catch (const std::exception& e) { LOG(ERROR) << "Failed to serialize Prometheus metrics: " << e.what(); @@ -587,6 +650,29 @@ void TentMetrics::serializeHistogramPrometheus( } } +void TentMetrics::serializeGaugePrometheus( + metrics::CachedDynamicCounter<1>& gauge, std::string& out) const { + auto cells = gauge.copy(); + if (cells.empty()) return; + const std::string& name = gauge.str_name(); + out.append("# HELP ") + .append(name) + .append(" ") + .append(gauge.help()) + .append("\n"); + out.append("# TYPE ").append(name).append(" gauge\n"); + for (auto& e : cells) { + out.append(name) + .append("{") + .append(gauge.labels_name()[0]) + .append("=\"") + .append(e->label[0]) + .append("\"} ") + .append(std::to_string(e->value.load(std::memory_order_relaxed))) + .append("\n"); + } +} + std::string TentMetrics::getJsonMetrics() { if (!initialized_) return "{}"; @@ -618,6 +704,14 @@ std::string TentMetrics::getJsonMetrics() { root[quarantined_batches_total_.str_name()] = static_cast(quarantined_batches_total_.value()); + // Gauges: aggregate across transport labels (total in-flight attempts + // / total registered bytes). Per-transport breakdown is via the + // Prometheus endpoint. + root[inflight_attempts_.str_name()] = + sumCounterValues(&inflight_attempts_); + root[registered_buffer_bytes_.str_name()] = + sumCounterValues(®istered_buffer_bytes_); + // Histograms: sum bucket counts across all transport labels. The // templated helper also reads back the histogram's sum counter so the // JSON endpoint emits "sum" alongside "count" (and stays in sync with @@ -711,6 +805,10 @@ void TentMetrics::recordTransportAttemptFinished(TransportType, Request::OpCode, void TentMetrics::recordDeadlineMLU(TransportType, double) {} void TentMetrics::recordDeadlineInfeasible(TransportType) {} void TentMetrics::recordBatchQuarantined() {} +void TentMetrics::recordTaskFailure(TransportType, TaskFailureReason) {} +void TentMetrics::recordInflightAttemptStarted(TransportType) {} +void TentMetrics::recordInflightAttemptFinished(TransportType) {} +void TentMetrics::recordRegisteredBufferBytes(TransportType, int64_t) {} void TentMetrics::recordStageLatency(Stage, TransportType, double) {} std::string TentMetrics::getPrometheusMetrics() { 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 07ffd9ba83..739b770f0d 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -843,6 +843,14 @@ Status TransferEngineImpl::registerLocalMemory(std::vector addr_list, auto s = transport_list_[type]->addMemoryBuffer(descs, options); if (!s.ok()) LOG(WARNING) << s.ToString(); } + // desc.transports lists the transports that actually registered + // the buffer (each transport appends itself on success). + for (auto& desc : descs) { + for (auto type : desc.transports) { + TentMetrics::instance().recordRegisteredBufferBytes( + type, static_cast(desc.length)); + } + } return Status::OK(); }); if (!status.ok()) return status; @@ -862,6 +870,10 @@ Status TransferEngineImpl::unregisterLocalMemory(void* addr, size_t size) { auto status = transport_list_[type]->removeMemoryBuffer(desc); if (!status.ok()) LOG(WARNING) << status.ToString(); } + for (auto type : desc.transports) { + TentMetrics::instance().recordRegisteredBufferBytes( + type, -static_cast(desc.length)); + } return Status::OK(); }); if (!status.ok()) return status; @@ -886,6 +898,10 @@ Status TransferEngineImpl::unregisterLocalMemory( auto s = transport_list_[type]->removeMemoryBuffer(desc); if (!s.ok()) LOG(WARNING) << s.ToString(); } + for (auto type : desc.transports) { + TentMetrics::instance().recordRegisteredBufferBytes( + type, -static_cast(desc.length)); + } return Status::OK(); }); if (!status.ok()) return status; @@ -1879,11 +1895,16 @@ Status TransferEngineImpl::commitPreparedSubmit( auto status = transport->submitTransferTasks(sub_batch, requests); if (!status.ok()) { auto attempt_end = std::chrono::steady_clock::now(); + // failure_stage must be marked in this first segment, before + // any recovery/failover attempt on the failure: a task that + // recovers and later fails at poll must still attribute its + // root cause to submit. for (const auto physical_task_id : group) { for (const auto public_task_id : public_tasks_by_physical_owner.at(physical_task_id)) { - finishTransportAttempt(batch->task_list[public_task_id], - FAILED, attempt_end); + auto& task = batch->task_list[public_task_id]; + if (task.failure_stage < 0) task.failure_stage = 0; + finishTransportAttempt(task, FAILED, attempt_end); } } // Recover by failing over to the remaining candidate @@ -2106,6 +2127,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { startTransportAttempt(task, task.type, std::chrono::steady_clock::now()); auto status = transport->submitTransferTasks(sub_batch, {task.request}); if (!status.ok()) { + if (task.failure_stage < 0) task.failure_stage = 0; finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); // Submit-stage failover: walk the remaining candidate transports // before giving up (the queued path dispatches one owner task at a @@ -2410,6 +2432,7 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { startTransportAttempt(task, type, std::chrono::steady_clock::now()); auto status = transport->submitTransferTasks(sub_batch, {task.request}); if (!status.ok()) { + if (task.failure_stage < 0) task.failure_stage = 0; finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); } return status; @@ -2472,6 +2495,14 @@ void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id, bool allow_failover) { auto& task = batch->task_list[task_id]; task.status = task_status.s; + // First-failure attribution: a terminal failure observed by polling marks + // the poll stage. Submit-stage failures were already marked at their + // origin and are not overwritten (a poll failure followed by a rejected + // failover resubmit still counts as poll). + if (task_status.s == FAILED || task_status.s == TIMEOUT || + 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) return; @@ -2850,6 +2881,35 @@ void TransferEngineImpl::recordTaskCompletionMetrics( #if TENT_METRICS_ENABLED if (prev_status == PENDING && new_status != PENDING && !task.derived) { auto end_time = std::chrono::steady_clock::now(); + // Failure reason: task.failure_stage marks where the first failure + // originated (0 = submit, 1 = poll), set at the failure site — a + // poll-observed failure stays "poll" even when the failover resubmit + // is synchronously rejected. TIMEOUT/CANCELED come from the status + // itself and were previously not recorded anywhere. + if (new_status == FAILED || new_status == TIMEOUT || + new_status == CANCELED) { + TentMetrics::TaskFailureReason reason; + if (new_status == TIMEOUT) { + reason = TentMetrics::TaskFailureReason::Timeout; + } else if (new_status == CANCELED) { + reason = TentMetrics::TaskFailureReason::Canceled; + } else if (task.failure_stage == 0) { + reason = TentMetrics::TaskFailureReason::Submit; + } else if (task.failure_stage == 1) { + reason = TentMetrics::TaskFailureReason::Poll; + } else { + // Unmarked path: an attempt still in flight means the + // failure was observed by polling. + reason = task.attempt_active + ? TentMetrics::TaskFailureReason::Poll + : TentMetrics::TaskFailureReason::Submit; + } + // Prefer the attempt transport: task.type may already be UNSPEC + // after a submit-stage failure. + TransportType failure_tp = + task.attempt_type != UNSPEC ? task.attempt_type : task.type; + TentMetrics::instance().recordTaskFailure(failure_tp, reason); + } finishTransportAttempt(task, new_status, end_time); auto start_time = task.start_time; if (start_time.time_since_epoch().count() > 0) { @@ -2951,6 +3011,7 @@ void TransferEngineImpl::startTransportAttempt( #if TENT_METRICS_ENABLED TentMetrics::instance().recordTransportAttemptStarted(type, task.request.opcode); + TentMetrics::instance().recordInflightAttemptStarted(type); #else (void)type; #endif @@ -2962,6 +3023,10 @@ void TransferEngineImpl::finishTransportAttempt( if (!task.attempt_active) return; task.attempt_active = false; #if TENT_METRICS_ENABLED + // Decrement before the early return below so the in-flight gauge stays + // symmetric with recordInflightAttemptStarted() for every attempt that + // actually started. + TentMetrics::instance().recordInflightAttemptFinished(task.attempt_type); auto post_time = task.attempt_post_time; if (post_time.time_since_epoch().count() == 0) return; double latency_us = diff --git a/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp index 11c639c0db..1e715134b2 100644 --- a/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp +++ b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp @@ -1152,6 +1152,161 @@ TEST_F(MetricsRecordingTest, ConcurrentRecordsAcrossTransportsAreExact) { 0); } +// Task failure with reason="submit": a synchronously rejected submit fails +// the attempt at the submit site, so recordTaskCompletionMetrics sees +// attempt_active == false. The counter must attribute the failure to the +// attempt transport (rdma) even though task.type is UNSPEC by then — the +// attribution leak the reason counter must not reproduce. +TEST_F(MetricsRecordingTest, SubmitFailureRecordsTaskFailureWithSubmitReason) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA, /*force_fail=*/false, + /*force_submit_fail=*/true); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xCD); + auto before_register = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kLen).ok()); + auto after_register = MetricsSnapshot(TentMetrics::instance()); + EXPECT_EQ(after_register.series( + "tent_registered_buffer_bytes{transport=\"rdma\"}") - + before_register.series( + "tent_registered_buffer_bytes{transport=\"rdma\"}"), + static_cast(kLen)); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), TransferStatusEnum::FAILED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ( + after.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"submit" + "\"}") - + before.series("tent_task_failures_total{transport=\"rdma\",reason=" + "\"submit\"}"), + 1); + EXPECT_EQ( + after.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"poll\"}") - + before.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"" + "poll\"}"), + 0); + // In-flight gauge returns to zero after the terminal failure. + EXPECT_EQ(after.counter("tent_inflight_attempts") - + before.counter("tent_inflight_attempts"), + 0); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kLen).ok()); + auto after_unregister = MetricsSnapshot(TentMetrics::instance()); + EXPECT_EQ(after_unregister.series( + "tent_registered_buffer_bytes{transport=\"rdma\"}"), + 0); +} + +// Task failure with reason="poll": the transport accepted the request and +// reported FAILED from getTransferStatus, so the attempt is still active when +// recordTaskCompletionMetrics runs. +TEST_F(MetricsRecordingTest, PollFailureRecordsTaskFailureWithPollReason) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA, /*force_fail=*/true, + /*force_submit_fail=*/false); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xAB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kLen).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), TransferStatusEnum::FAILED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ( + after.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"poll\"}") - + before.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"" + "poll\"}"), + 1); + EXPECT_EQ( + after.series( + "tent_task_failures_total{transport=\"rdma\",reason=\"submit" + "\"}") - + before.series("tent_task_failures_total{transport=\"rdma\",reason=" + "\"submit\"}"), + 0); + EXPECT_EQ(after.counter("tent_inflight_attempts") - + before.counter("tent_inflight_attempts"), + 0); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kLen).ok()); +} + +// Gauges must stay symmetric across setEnabled() transitions: the runtime +// switch governs statistical sampling, not state tracking. Skipping one half +// of a paired add/sub would permanently corrupt the gauge (review feedback). +TEST_F(MetricsRecordingTest, InflightGaugeSurvivesEnableToggle) { + auto& m = TentMetrics::instance(); + auto before = MetricsSnapshot(m); + + // Enabled at start, disabled before finish. + TentMetrics::setEnabled(true); + m.recordInflightAttemptStarted(RDMA); + TentMetrics::setEnabled(false); + m.recordInflightAttemptFinished(RDMA); + + // Disabled at start, enabled before finish (reverse straddle). + m.recordInflightAttemptStarted(RDMA); + TentMetrics::setEnabled(true); + m.recordInflightAttemptFinished(RDMA); + + auto after = MetricsSnapshot(m); + EXPECT_EQ(after.counter("tent_inflight_attempts") - + before.counter("tent_inflight_attempts"), + 0); + EXPECT_EQ(after.series("tent_inflight_attempts{transport=\"rdma\"}") - + before.series("tent_inflight_attempts{transport=\"rdma\"}"), + 0); +} + +TEST_F(MetricsRecordingTest, RegisteredBytesGaugeSurvivesEnableToggle) { + auto& m = TentMetrics::instance(); + auto before = MetricsSnapshot(m); + + TentMetrics::setEnabled(true); + m.recordRegisteredBufferBytes(RDMA, 4096); + TentMetrics::setEnabled(false); + m.recordRegisteredBufferBytes(RDMA, -4096); + TentMetrics::setEnabled(true); + + auto after = MetricsSnapshot(m); + EXPECT_EQ(after.counter("tent_registered_buffer_bytes") - + before.counter("tent_registered_buffer_bytes"), + 0); + EXPECT_EQ( + after.series("tent_registered_buffer_bytes{transport=\"rdma\"}") - + before.series("tent_registered_buffer_bytes{transport=\"rdma\"}"), + 0); +} + // --------------------------------------------------------------------------- // L2 HTTP integration: scrape the real /metrics, /metrics/json, /health // endpoints via coro_http_client and assert on status + body. Validates the From 71342bac44fd9fa534c4a33ea7593e6d9446133f Mon Sep 17 00:00:00 2001 From: mjwtom Date: Mon, 31 Aug 2026 14:12:52 +0800 Subject: [PATCH 03/17] [Store] Expose build version over HTTP and metrics (#3617) (#3619) Expose build version information over HTTP and Prometheus so operators can verify which binary a running process is using without shell access to the host. - Master admin server: new `GET /version` handler, registered alongside `/health` so it stays available in standby. - Real client HTTP server: same `/version` payload on the client HTTP port; it does not touch `client_`, so it is unaffected by `MC_STORE_CLIENT_METRIC` and by client initialization state. - `mooncake_client`: set the gflags version string so `--version` works, and log the display version at startup. - Master and client metrics: new `mooncake_build_info` gauge following the Prometheus "info" pattern, so dashboards and alerts can group or filter by the running build. Both surfaces share the metric name and are told apart by the scrape target's job/instance labels; on the client the version labels are added on top of the caller supplied labels. Both endpoints and the metric expose `version` (used for RPC handshake compatibility) and `display_version` (release plus short git hash). Docs: document the endpoint in the master admin endpoint table, the client endpoint table, the HTTP service reference, and the three client HTTP server switches that gate it. Tests: add `/version` coverage for the real client HTTP server and for the master admin server in both serving and standby states, plus build info metric serialization tests on the master and client sides. Co-authored-by: majingwei --- .../source/api-reference/http/http-service.md | 23 ++++++ .../api-reference/python/mooncake-store.md | 2 +- .../mooncake-store-deployment-guide.md | 20 ++++-- docs/source/getting_started/observability.md | 14 ++++ mooncake-store/include/client_metric.h | 6 ++ mooncake-store/include/master_admin_service.h | 2 + .../include/master_metric_manager.h | 10 +++ mooncake-store/src/client_metric.cpp | 21 ++++++ mooncake-store/src/master.cpp | 3 + mooncake-store/src/master_admin_service.cpp | 19 +++++ mooncake-store/src/master_metric_manager.cpp | 12 +++- mooncake-store/src/real_client.cpp | 10 +++ mooncake-store/src/real_client_main.cpp | 5 ++ mooncake-store/tests/client_metrics_test.cpp | 71 +++++++++++++++++++ mooncake-store/tests/health_check_test.cpp | 22 ++++++ .../tests/master_admin_server_test.cpp | 30 ++++++++ mooncake-store/tests/master_metrics_test.cpp | 37 ++++++++++ 17 files changed, 301 insertions(+), 6 deletions(-) diff --git a/docs/source/api-reference/http/http-service.md b/docs/source/api-reference/http/http-service.md index fd34b638c6..0050adcd50 100644 --- a/docs/source/api-reference/http/http-service.md +++ b/docs/source/api-reference/http/http-service.md @@ -260,6 +260,29 @@ Basic health check endpoint for service availability verification. curl http://localhost:8080/health ``` +#### `/version` +Report the master version. Always available, including while the master is in +standby. + +**Method**: `GET` +**Content-Type**: `application/json; charset=utf-8` +**Response**: JSON object with: +- `version` (string): Store version used for RPC handshake compatibility +- `display_version` (string): Human-readable release plus short git hash + +**Example**: +```bash +curl http://localhost:8080/version +``` + +```json +{"version":"2.0.0","display_version":"0.3.12.post1 (git: f9e8311f)"} +``` + +Real clients expose the same `/version` payload on their own client HTTP port +when `enable_client_http_server` is on. See +[Client Metrics Endpoint](../../getting_started/observability.md#client-metrics-endpoint). + ## Store REST API Endpoints The following endpoints are served by the Python store REST service, which wraps diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md index 151f8f7f6b..6a9f3cc9b4 100644 --- a/docs/source/api-reference/python/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -1150,7 +1150,7 @@ the positional overload. data is stored under `MOONCAKE_DFS_ROOT_DIR`, but this separate directory is still validated during FileStorage initialization. - `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`. -- `enable_client_http_server` (bool): Enable the client-local `/health`, `/metrics`, and `/metrics/summary` HTTP endpoints. Defaults to `False`. +- `enable_client_http_server` (bool): Enable the client-local `/health`, `/metrics`, `/metrics/summary`, and `/version` HTTP endpoints. Defaults to `False`. - `client_http_port` (int): Port for the client-local HTTP endpoints. Defaults to `9300`. **Store segment pinned memory:** CUDA-enabled builds can register Store-managed diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index bafa12a4ec..5d8419a9d7 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -962,7 +962,7 @@ Arguments of `MooncakeDistributedStore.setup(...)`: | `enable_ssd_offload` | bool | `false` | *(advanced)* Initialize client-side `FileStorage`; required for SSD offload and descriptor-based DFS | | `ssd_offload_path` | str | empty | *(advanced)* FileStorage path; with the distributed backend, DFS data uses `MOONCAKE_DFS_ROOT_DIR` | | `tenant_id` | str | `default` | *(advanced)* Tenant identifier | -| `enable_client_http_server` | bool | `false` | Enable the client-side HTTP `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `enable_client_http_server` | bool | `false` | Enable the client-side HTTP `/health`, `/metrics`, `/metrics/summary`, and `/version` endpoints | | `client_http_port` | int | `9300` | Client-side HTTP endpoint port, used only when `enable_client_http_server=true` | ```{note} @@ -993,7 +993,7 @@ The store service CLI only accepts `--config`, `-D/--define`, `--port`, and `--m | `MOONCAKE_OFFLOAD_ENABLED` | `enable_ssd_offload` | `false` | Initialize client-side `FileStorage`; required for SSD offload and descriptor-based DFS | | `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `ssd_offload_path` | empty | FileStorage path; DFS shard data uses `MOONCAKE_DFS_ROOT_DIR` with the distributed backend | | `MOONCAKE_TENANT_ID` | `tenant_id` | `default` | Tenant identifier | -| `MOONCAKE_ENABLE_CLIENT_HTTP_SERVER` | `enable_client_http_server` | `false` | Enable client-side `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `MOONCAKE_ENABLE_CLIENT_HTTP_SERVER` | `enable_client_http_server` | `false` | Enable client-side `/health`, `/metrics`, `/metrics/summary`, and `/version` endpoints | | `MOONCAKE_CLIENT_HTTP_PORT` | `client_http_port` | `9300` | Client-side HTTP endpoint port | | `MOONCAKE_CONFIG_PATH` | — | unset | Path to a JSON config file (takes precedence over the variables above) | @@ -1062,9 +1062,12 @@ mooncake_client \ | `--tenant_id` | `default` | Tenant identifier | | `--enable_offload` | `false` | Enable client-side SSD offload | | `--start_offload_rpc_server` | `true` | Start the offload RPC server for dummy clients | -| `--enable_http_server` | `false` | Enable client-side `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `--enable_http_server` | `false` | Enable client-side `/health`, `/metrics`, `/metrics/summary`, and `/version` endpoints | | `--http_port` | `9300` | Client-side HTTP endpoint port | +`mooncake_client --version` prints the release version plus the short git hash, +and the same value is logged at startup. + ### Client HTTP Health and Metrics Endpoint Each real client can expose its own lightweight HTTP endpoint independently of the master admin HTTP server and the Python store REST API. This endpoint is disabled by default for programmatic clients and `mooncake_store_service`; enable it explicitly when you want to scrape client-local metrics: @@ -1090,9 +1093,18 @@ For `mooncake_store_service`, use `MOONCAKE_ENABLE_CLIENT_HTTP_SERVER=true` and | `GET /health` | Client health check | | `GET /metrics` | Prometheus-format client metrics | | `GET /metrics/summary` | Human-readable client metrics summary | +| `GET /version` | Client version as JSON (`version` for RPC handshake compatibility, `display_version` for release plus short git hash) | + +```bash +curl http://:9300/version +``` + +```json +{"version":"2.0.0","display_version":"0.3.12.post1 (git: f9e8311f)"} +``` ```{note} -`MC_STORE_CLIENT_METRIC` controls whether client metrics are collected. If the client HTTP server is enabled but `MC_STORE_CLIENT_METRIC=0`, `/metrics` and `/metrics/summary` return HTTP 503 with `metrics not available`. +`MC_STORE_CLIENT_METRIC` controls whether client metrics are collected. If the client HTTP server is enabled but `MC_STORE_CLIENT_METRIC=0`, `/metrics` and `/metrics/summary` return HTTP 503 with `metrics not available`. `/health` and `/version` are unaffected. ``` ### Engine Runtime Tuning (`MC_*`) diff --git a/docs/source/getting_started/observability.md b/docs/source/getting_started/observability.md index ffc6b007bb..769574995a 100644 --- a/docs/source/getting_started/observability.md +++ b/docs/source/getting_started/observability.md @@ -122,6 +122,7 @@ The admin HTTP server runs on `metrics_port` (default: **9003**) and exposes the | `GET /metrics` | `text/plain; version=0.0.4` | All metrics in Prometheus exposition format | | `GET /metrics/summary` | `text/plain; version=0.0.4` | Human-readable summary (same content as the periodic log) | | `GET /health` | `application/json` | Health check with role, HA state, and service readiness | +| `GET /version` | `application/json` | Master version (`version` for RPC compatibility, `display_version` for release + git hash) | | `GET /role` | `text/plain` | Current HA role (`leader` / `standby`) | | `GET /ha_status` | `text/plain` | Current HA runtime state (`serving` / `starting` / etc.) | @@ -150,6 +151,9 @@ curl http://:9003/metrics/summary # Check health curl http://:9003/health + +# Check version +curl http://:9003/version ``` ### Configuration @@ -198,13 +202,23 @@ For `mooncake.mooncake_store_service`, set | `GET /health` | `application/json` | Client health check | | `GET /metrics` | `text/plain; version=0.0.4` | Prometheus-format client metrics | | `GET /metrics/summary` | `text/plain` | Human-readable client metrics summary | +| `GET /version` | `application/json` | Client version (`version` for RPC compatibility, `display_version` for release + git hash) | ```bash curl http://:9300/health curl http://:9300/metrics curl http://:9300/metrics/summary +curl http://:9300/version +``` + +```json +{"version":"2.0.0","display_version":"0.3.12.post1 (git: f9e8311f)"} ``` +`/version` does not depend on client metric collection or on a fully +initialized client, so it stays available whenever the client HTTP server is +running. + Set `MC_STORE_CLIENT_METRIC=0` to disable client metric collection. If the client HTTP server remains enabled while metrics are disabled, `/metrics` and `/metrics/summary` return HTTP 503 with `metrics not available`. diff --git a/mooncake-store/include/client_metric.h b/mooncake-store/include/client_metric.h index 87247f1770..9784600d3a 100644 --- a/mooncake-store/include/client_metric.h +++ b/mooncake-store/include/client_metric.h @@ -658,6 +658,12 @@ struct ClientMetric { MasterClientMetric master_client_metric; TransferOperationMetric transfer_operation_metric; SsdMetric ssd_metric; + // Prometheus "info" pattern: the value carries no meaning and is always 1, + // the version strings ride along as static labels next to any caller + // supplied labels so a scrape can attribute samples to a concrete build. + // Shares the `mooncake_build_info` name with the master-side metric; the + // two are told apart by the scrape target's job/instance labels. + ylt::metric::gauge_t build_info; /** * @brief Creates a ClientMetric instance based on environment variables diff --git a/mooncake-store/include/master_admin_service.h b/mooncake-store/include/master_admin_service.h index 752b1f7a74..051598320f 100644 --- a/mooncake-store/include/master_admin_service.h +++ b/mooncake-store/include/master_admin_service.h @@ -68,6 +68,8 @@ class MasterAdminServer { coro_http::coro_http_response& resp); void HandleHealth(coro_http::coro_http_request& req, coro_http::coro_http_response& resp); + void HandleVersion(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); void HandleRole(coro_http::coro_http_request& req, coro_http::coro_http_response& resp); void HandleHaStatus(coro_http::coro_http_request& req, diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 899596eb22..60e909fb52 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -755,6 +755,16 @@ class MasterMetricManager { ylt::metric::counter_t fetch_tasks_failures_; ylt::metric::counter_t mark_task_to_complete_requests_; ylt::metric::counter_t mark_task_to_complete_failures_; + + // Build Info Metric + // Prometheus "info" pattern: the value carries no meaning and is always 1, + // the version strings are exposed as labels so dashboards and alerts can + // group or filter by the running build. The label values are compile-time + // constants, so this is a single-series gauge with static labels rather + // than a dynamic-label metric. + // Shares the `mooncake_build_info` name with the client-side metric; the + // two are told apart by the scrape target's job/instance labels. + ylt::metric::gauge_t build_info_; }; } // namespace mooncake diff --git a/mooncake-store/src/client_metric.cpp b/mooncake-store/src/client_metric.cpp index 4fbaef026f..449417ab39 100644 --- a/mooncake-store/src/client_metric.cpp +++ b/mooncake-store/src/client_metric.cpp @@ -7,11 +7,21 @@ #include "bool_parser.h" #include "integer_parser.h" +#include "version.h" namespace mooncake { namespace { +// Build info is exposed as an info-style metric, so the version strings live in +// labels. Caller supplied labels are preserved to keep instance identification. +std::map WithBuildInfoLabels( + std::map labels) { + labels["version"] = GetMooncakeStoreVersion(); + labels["display_version"] = MOONCAKE_DISPLAY_VERSION; + return labels; +} + bool parseMetricsEnabled() { const char* metric_env = std::getenv("MC_STORE_CLIENT_METRIC"); if (!metric_env) { @@ -71,10 +81,16 @@ ClientMetric::ClientMetric(uint64_t interval_seconds, master_client_metric(labels), transfer_operation_metric(labels), ssd_metric(labels), + build_info("mooncake_build_info", + "Build version of the running client; the value is always 1 " + "and the version strings are carried by the labels", + WithBuildInfoLabels(labels)), should_stop_metrics_thread_(false), metrics_interval_seconds_(interval_seconds), bandwidth_reporting_enabled_(bandwidth_reporting_enabled), master_rpc_metrics_enabled_(master_rpc_metrics_enabled) { + // Set once: the compiled-in version never changes at runtime. + build_info.update(1); last_report_snapshot_ = TransferSnapshot{ static_cast(transfer_metric.total_read_bytes.value()), static_cast(transfer_metric.total_write_bytes.value()), @@ -116,11 +132,16 @@ void ClientMetric::serialize(std::string& str) { } transfer_operation_metric.serialize(str); ssd_metric.serialize(str); + build_info.serialize(str); } std::string ClientMetric::summary_metrics() { std::stringstream ss; ss << "Client Metrics Summary\n"; + // Identify the build inline so one /metrics/summary request is enough to + // tell which binary produced the numbers below. + ss << "Version: " << GetMooncakeStoreVersion() << " (" + << MOONCAKE_DISPLAY_VERSION << ")\n"; ss << transfer_metric.summary_metrics(bandwidth_reporting_enabled_); ss << "\n"; if (master_rpc_metrics_enabled_) { diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index bf7c98a980..052c44fe82 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -1412,6 +1412,9 @@ int main(int argc, char* argv[]) { google::SetLogSymlink(google::GLOG_INFO, "mooncake_master"); } + LOG(INFO) << "Mooncake master version: " + << mooncake::MOONCAKE_DISPLAY_VERSION; + // Initialize the master configuration mooncake::MasterConfig master_config; std::string conf_path = FLAGS_config_path; diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp index cc85bcb219..462358b54d 100644 --- a/mooncake-store/src/master_admin_service.cpp +++ b/mooncake-store/src/master_admin_service.cpp @@ -21,6 +21,7 @@ #include "master_metric_manager.h" #include "rpc_service.h" #include "types.h" +#include "version.h" namespace mooncake { @@ -501,6 +502,20 @@ void MasterAdminServer::HandleHealth(coro_http::coro_http_request&, WriteJsonResponse(resp, coro_http::status_type::ok, payload); } +struct HttpVersionResponse { + std::string version; + std::string display_version; +}; +YLT_REFL(HttpVersionResponse, version, display_version); + +void MasterAdminServer::HandleVersion(coro_http::coro_http_request&, + coro_http::coro_http_response& resp) { + WriteJsonResponse( + resp, coro_http::status_type::ok, + HttpVersionResponse{.version = GetMooncakeStoreVersion(), + .display_version = MOONCAKE_DISPLAY_VERSION}); +} + struct HttpLeaderResponse { bool present{false}; std::optional leader_address; @@ -1218,6 +1233,10 @@ void MasterAdminServer::RegisterHandler() { "/health", [this](coro_http_request& req, coro_http_response& resp) { HandleHealth(req, resp); }); + http_server_.set_http_handler( + "/version", [this](coro_http_request& req, coro_http_response& resp) { + HandleVersion(req, resp); + }); http_server_.set_http_handler( "/role", [this](coro_http_request& req, coro_http_response& resp) { HandleRole(req, resp); diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index 0125bcfbcf..4339a74fe6 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -9,6 +9,7 @@ #include "segment.h" #include "utils.h" +#include "version.h" namespace mooncake { @@ -483,9 +484,17 @@ MasterMetricManager::MasterMetricManager() "Total number of MarkTaskToComplete requests received"), mark_task_to_complete_failures_( "master_update_task_failures_total", - "Total number of failed MarkTaskToComplete requests") { + "Total number of failed MarkTaskToComplete requests"), + build_info_("mooncake_build_info", + "Build version of the running master; the value is always 1 " + "and the version strings are carried by the labels", + {{"version", GetMooncakeStoreVersion()}, + {"display_version", MOONCAKE_DISPLAY_VERSION}}) { // Update all metrics once to ensure zero values are serialized update_metrics_for_zero_output(); + // Info-style metric: emit a single series for the build this binary was + // compiled from. Set once here because the value never changes at runtime. + build_info_.update(1); } // --- Metric Interface Methods --- @@ -1987,6 +1996,7 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(promotion_candidate_dropped_limit_); serialize_metric(tenant_quota_reject_total_); serialize_metric(tenant_evict_bytes_total_); + serialize_metric(build_info_); // Serialize Snapshot Metrics serialize_metric(snapshot_duration_ms_); diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 16d8d0fe84..e764f49184 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -37,6 +37,7 @@ #include "device/cuda_ipc_buffer.h" #include "shm_helper.h" #include "memory_location.h" +#include "version.h" #ifdef USE_NOF #include "spdk/spdk_wrapper.h" #endif @@ -1926,6 +1927,15 @@ int RealClient::start_http_server(int port) { resp.set_status_and_content(status_type::ok, std::move(*result)); }); + http_server_->set_http_handler( + "/version", [](coro_http_request &req, coro_http_response &resp) { + std::string body = "{\"version\":\"" + GetMooncakeStoreVersion() + + "\",\"display_version\":\"" + + std::string(MOONCAKE_DISPLAY_VERSION) + "\"}"; + resp.add_header("Content-Type", "application/json"); + resp.set_status_and_content(status_type::ok, std::move(body)); + }); + auto ec = http_server_->async_start(); if (ec.hasResult()) { LOG(WARNING) << "Failed to start HTTP server on port " << port; diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 19712ee578..a0b627fb5e 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -6,6 +6,7 @@ #include "common.h" #include "config.h" #include "real_client.h" +#include "version.h" using namespace mooncake; @@ -110,11 +111,15 @@ int main(int argc, char *argv[]) { // spawning threads, leading to missing signal processing. mooncake::ResourceTracker::getInstance(); + gflags::SetVersionString(mooncake::MOONCAKE_DISPLAY_VERSION); gflags::ParseCommandLineFlags(&argc, &argv, true); if (!FLAGS_log_dir.empty()) { google::InitGoogleLogging(argv[0]); } + LOG(INFO) << "Mooncake real client version: " + << mooncake::MOONCAKE_DISPLAY_VERSION; + size_t global_segment_size = string_to_byte_size(FLAGS_global_segment_size); size_t local_buffer_size = string_to_byte_size(FLAGS_local_buffer_size); #ifdef USE_ASCEND_DIRECT diff --git a/mooncake-store/tests/client_metrics_test.cpp b/mooncake-store/tests/client_metrics_test.cpp index ee446a4f74..9f777b17f5 100644 --- a/mooncake-store/tests/client_metrics_test.cpp +++ b/mooncake-store/tests/client_metrics_test.cpp @@ -14,6 +14,7 @@ #include "real_client.h" #include "test_server_helpers.h" #include "utils.h" +#include "version.h" namespace mooncake::test { namespace { @@ -205,6 +206,10 @@ TEST_F(ClientMetricsTest, ClientMetricsSummaryTest) { EXPECT_TRUE(summary.find("ExistKey: count=1") != std::string::npos); EXPECT_TRUE(summary.find("get_buffer: count=1") != std::string::npos); EXPECT_TRUE(summary.find("put_batch: count=1") != std::string::npos); + // The build is named inline so the summary alone identifies the binary. + EXPECT_TRUE(summary.find("Version: " + GetMooncakeStoreVersion()) != + std::string::npos); + EXPECT_TRUE(summary.find(MOONCAKE_DISPLAY_VERSION) != std::string::npos); std::cout << "Full Client Metrics Summary:\n" << summary << std::endl; } @@ -409,6 +414,62 @@ TEST_F(ClientMetricsTest, SerializeWithoutDynamicLabels) { } } +TEST_F(ClientMetricsTest, BuildInfoMetricIsSerialized) { + ClientMetric metrics(0); + + std::string serialized; + metrics.serialize(serialized); + + ASSERT_NE(serialized.find("mooncake_build_info{"), std::string::npos) + << "build info metric missing from serialized output"; + // Label order inside the series is not asserted: only the presence of both + // labels with the compiled-in values matters for scraping and grouping. + EXPECT_NE(serialized.find("version=\"" + GetMooncakeStoreVersion() + "\""), + std::string::npos) + << "RPC handshake version missing from build info labels"; + EXPECT_NE(serialized.find("display_version=\"" + + std::string(MOONCAKE_DISPLAY_VERSION) + "\""), + std::string::npos) + << "display version missing from build info labels"; + + // Info-style metric: the value is fixed at 1, so exactly one series is + // emitted and its value follows the closing brace of the label set. + EXPECT_EQ(CountOccurrences(serialized, "mooncake_build_info{"), 1u); + const auto metric_pos = serialized.find("mooncake_build_info{"); + const auto brace_end = serialized.find('}', metric_pos); + ASSERT_NE(brace_end, std::string::npos); + const auto line_end = serialized.find('\n', brace_end); + const std::string value_part = + serialized.substr(brace_end + 1, line_end == std::string::npos + ? std::string::npos + : line_end - brace_end - 1); + EXPECT_NE(value_part.find('1'), std::string::npos) + << "build info value should be 1, got:" << value_part; +} + +TEST_F(ClientMetricsTest, BuildInfoMetricKeepsCallerLabels) { + // Caller supplied labels identify the instance; the version labels are + // added on top of them rather than replacing them. + std::map static_labels = { + {"instance_id", "12345"}, {"cluster_id", "cluster1"}}; + ClientMetric metrics(0, static_labels); + + std::string serialized; + metrics.serialize(serialized); + + const auto metric_pos = serialized.find("mooncake_build_info{"); + ASSERT_NE(metric_pos, std::string::npos); + const auto brace_end = serialized.find('}', metric_pos); + ASSERT_NE(brace_end, std::string::npos); + const std::string labels = + serialized.substr(metric_pos, brace_end - metric_pos); + + EXPECT_NE(labels.find("instance_id=\"12345\""), std::string::npos); + EXPECT_NE(labels.find("cluster_id=\"cluster1\""), std::string::npos); + EXPECT_NE(labels.find("version=\"" + GetMooncakeStoreVersion() + "\""), + std::string::npos); +} + TEST_F(ClientMetricsTest, HttpMetricsEndpointsReturnData) { std::unordered_set used_ports; int master_rpc_port = GetTestPort(used_ports); @@ -544,6 +605,16 @@ TEST_F(ClientMetricsTest, HttpMetricsEndpointReturns503WhenMetricsDisabled) { EXPECT_EQ(metrics.status, 503); EXPECT_NE(metrics.body.find("metrics not available"), std::string::npos); + // `/version` is served without touching the metric collector, so disabling + // metrics must not take it down along with `/metrics`. + auto version = + FetchUrl("http://127.0.0.1:" + std::to_string(http_port) + "/version"); + EXPECT_EQ(version.status, 200); + EXPECT_NE( + version.body.find("\"version\":\"" + GetMooncakeStoreVersion() + "\""), + std::string::npos); + EXPECT_NE(version.body.find("display_version"), std::string::npos); + EXPECT_EQ(client->tearDownAll(), 0); } diff --git a/mooncake-store/tests/health_check_test.cpp b/mooncake-store/tests/health_check_test.cpp index 7a059d7355..c755cc2f5d 100644 --- a/mooncake-store/tests/health_check_test.cpp +++ b/mooncake-store/tests/health_check_test.cpp @@ -14,6 +14,7 @@ #include "real_client.h" #include "test_server_helpers.h" #include "default_config.h" +#include "version.h" DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma"); @@ -241,5 +242,26 @@ TEST_F(HealthCheckTest, MetricsEndpointsReturnCorrectData) { master_.Stop(); } +// Test 8: HTTP /version returns 200 with version information +TEST_F(HealthCheckTest, VersionEndpointReturnsVersion) { + int http_port = getFreeTcpPort(); + FLAGS_http_port = http_port; + ASSERT_EQ(StartMasterAndSetupClient(18930), 0) << "setup_real failed"; + + std::this_thread::sleep_for(std::chrono::seconds(2)); + + auto resp = fetch_url(http_port, "/version"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("\"version\""), std::string::npos); + EXPECT_NE(resp.body.find("\"display_version\""), std::string::npos); + EXPECT_NE(resp.body.find(GetMooncakeStoreVersion()), std::string::npos) + << "store version missing from /version response"; + EXPECT_NE(resp.body.find(MOONCAKE_DISPLAY_VERSION), std::string::npos) + << "display version missing from /version response"; + + py_client_->tearDownAll(); + master_.Stop(); +} + } // namespace testing } // namespace mooncake diff --git a/mooncake-store/tests/master_admin_server_test.cpp b/mooncake-store/tests/master_admin_server_test.cpp index 3b2b9b2b2c..898cf4440a 100644 --- a/mooncake-store/tests/master_admin_server_test.cpp +++ b/mooncake-store/tests/master_admin_server_test.cpp @@ -21,6 +21,7 @@ #include "tenant_quota_policy_store.h" #include "types.h" #include "utils.h" +#include "version.h" #include @@ -208,6 +209,35 @@ TEST_F(MasterAdminServerTest, HealthEndpointReturns200InServing) { admin.Stop(); } +TEST_F(MasterAdminServerTest, VersionEndpointReturnsVersion) { + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + + auto resp = HttpGet(port, "/version"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("\"version\""), std::string::npos); + EXPECT_NE(resp.body.find("\"display_version\""), std::string::npos); + EXPECT_NE(resp.body.find(GetMooncakeStoreVersion()), std::string::npos); + EXPECT_NE(resp.body.find(MOONCAKE_DISPLAY_VERSION), std::string::npos); + + admin.Stop(); +} + +TEST_F(MasterAdminServerTest, VersionEndpointAvailableInStandby) { + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kStandby); + + auto resp = HttpGet(port, "/version"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find(GetMooncakeStoreVersion()), std::string::npos); + + admin.Stop(); +} + TEST_F(MasterAdminServerTest, HealthEndpointIncludesLeaderInfoWhenSet) { int port = getFreeTcpPort(); MasterAdminServer admin(static_cast(port), false); diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp index 16a98b95a4..da557a88ff 100644 --- a/mooncake-store/tests/master_metrics_test.cpp +++ b/mooncake-store/tests/master_metrics_test.cpp @@ -18,6 +18,7 @@ #include "types.h" #include "master_config.h" #include "master_metric_manager.h" +#include "version.h" namespace mooncake::test { @@ -1162,6 +1163,42 @@ TEST_F(MasterMetricsTest, SsdOffloadCacheHitAndTotalConsistent) { service_.Remove(ssd_only_key, "default"); } +// Build info is an "info"-style metric: the value is always 1 and the version +// strings are carried by labels, so serialization must emit a single series +// containing both version labels. +TEST_F(MasterMetricsTest, BuildInfoMetricIsSerialized) { + auto& metrics = MasterMetricManager::instance(); + + const std::string serialized = metrics.serialize_metrics(); + + ASSERT_NE(serialized.find("mooncake_build_info"), std::string::npos) + << "build info metric missing from serialized output"; + // Label order inside the series is not asserted: only the presence of both + // labels with the compiled-in values matters for scraping and grouping. + EXPECT_NE(serialized.find("version=\"" + GetMooncakeStoreVersion() + "\""), + std::string::npos) + << "RPC handshake version missing from build info labels"; + EXPECT_NE(serialized.find("display_version=\"" + + std::string(MOONCAKE_DISPLAY_VERSION) + "\""), + std::string::npos) + << "display version missing from build info labels"; + + // The series value is fixed at 1; extract the number after the closing + // brace of the build info series to confirm it is exposed as such. + const auto metric_pos = serialized.find("mooncake_build_info{"); + ASSERT_NE(metric_pos, std::string::npos) + << "build info metric has no labelled series"; + const auto brace_end = serialized.find('}', metric_pos); + ASSERT_NE(brace_end, std::string::npos); + const auto line_end = serialized.find('\n', brace_end); + const std::string value_part = + serialized.substr(brace_end + 1, line_end == std::string::npos + ? std::string::npos + : line_end - brace_end - 1); + EXPECT_NE(value_part.find('1'), std::string::npos) + << "build info value should be 1, got:" << value_part; +} + } // namespace mooncake::test int main(int argc, char** argv) { From 7b1ed37b42e3eac8422faf19c9075a3efef2c217 Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 31 Aug 2026 14:13:51 +0800 Subject: [PATCH 04/17] [TENT] Drop ineffective unit tests and bind coverage to production paths (#3757) Tautological intent/failover tests and TCP tests that never constructed TcpTransport could not catch runtime regressions. Route default RDMA promotion through DecidePromotionHeadOnly and test TCP install plus batch worst-failure aggregation against the real engine. Signed-off-by: staryxchen Co-authored-by: Cursor --- docs/source/design/tent/testing.md | 21 ++ .../hip_bandwidth_bench.cpp | 2 +- .../tent/transport/rdma/promotion_policy.h | 16 +- .../include/tent/transport/rdma/workers.h | 5 +- .../tent/transport/tcp/tcp_transport.h | 12 + .../tent/src/transport/rdma/workers.cpp | 20 +- .../tent/src/transport/tcp/tcp_transport.cpp | 2 +- .../tent/tests/CMakeLists.txt | 16 +- .../tent/tests/engine_failover_e2e_test.cpp | 65 ++++ .../tent/tests/failover_test.cpp | 308 ------------------ .../tent/tests/intent_type_test.cpp | 95 ------ .../tent/tests/promotion_policy_test.cpp | 15 +- .../tent/tests/tcp_transport_test.cpp | 183 +++-------- .../tent/tests/thread_local_storage_test.cpp | 40 --- 14 files changed, 169 insertions(+), 631 deletions(-) rename mooncake-transfer-engine/tent/{tests => benchmark}/hip_bandwidth_bench.cpp (99%) delete mode 100644 mooncake-transfer-engine/tent/tests/failover_test.cpp delete mode 100644 mooncake-transfer-engine/tent/tests/intent_type_test.cpp diff --git a/docs/source/design/tent/testing.md b/docs/source/design/tent/testing.md index c5544a66b0..595ac43aec 100644 --- a/docs/source/design/tent/testing.md +++ b/docs/source/design/tent/testing.md @@ -179,3 +179,24 @@ and CUDA stubs would bypass the fakes. Concurrency tests that touch shared runtime maps should also be run under ThreadSanitizer (`-fsanitize=thread`). Which binaries those are belongs in the test sources, not here. + +## What belongs in this suite + +Tests should call production functions or the public engine API. Do not +reimplement a state machine in the test file and assert against that +copy. Enum assignment, `std::atomic` store/load, and `Config::get/set` +without going through `install()` / `TransferEngineImpl` do not lock +runtime behavior. + +Binaries that are not ctest: + +- `tent_metrics_example` — HTTP metrics demo +- `deadline_promotion_bench` — hot-path microbenchmark +- `tent/benchmark/hip_bandwidth_bench.cpp` — standalone `hipcc` HIP + bandwidth sweep; CMake does not build it + +Hardware data-path tests (`tent_rdma_transport_test` roundtrip, +`tent_nvlink_transport_test`, Sunrise, MPComm) skip or are omitted from +ctest when the device is missing. That skip is expected; a green +`cuda-off` run does not mean those paths were exercised. `tent_nvlink_transport_test` +is built when CUDA is on but is not registered with ctest. diff --git a/mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp b/mooncake-transfer-engine/tent/benchmark/hip_bandwidth_bench.cpp similarity index 99% rename from mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp rename to mooncake-transfer-engine/tent/benchmark/hip_bandwidth_bench.cpp index a265a80863..7d95bae52e 100644 --- a/mooncake-transfer-engine/tent/tests/hip_bandwidth_bench.cpp +++ b/mooncake-transfer-engine/tent/benchmark/hip_bandwidth_bench.cpp @@ -4,7 +4,7 @@ // Measures H2D, D2H, D2D (same GPU), and GPU-to-GPU (P2P via XGMI) // bandwidth across a sweep of transfer sizes. // -// Build: +// Build (from this directory): // hipcc -O3 -o hip_bandwidth_bench hip_bandwidth_bench.cpp // // Run: diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h index a5d3a92549..ecdcbf043a 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h @@ -37,10 +37,10 @@ struct PromotionDecision { bool promoted_any() const { return !promote_indices.empty(); } }; -// Historical policy (as implemented in Workers::promoteTimedOutRequests today): -// the pass drains the whole queue, inspects ONLY the head entry, and if the -// head has timed out it promotes EVERY entry in the queue. This is the behavior -// #2528 flags as unintended: +// Default policy (Workers::promoteTimedOutRequests when +// transports/rdma/priority_promotion_per_entry is false): the pass drains the +// whole queue, inspects ONLY the head entry, and if the head has timed out it +// promotes EVERY entry in the queue. Issue #2528 flags this as unintended: // * decision is head-only but applied to the whole queue -> freshly enqueued // entries that are not starving get promoted alongside the starving head; // * if the head has NOT timed out, later timed-out entries are not promoted. @@ -68,10 +68,10 @@ inline PromotionDecision DecidePromotionHeadOnly( return d; } -// Per-entry policy: promote exactly the entries that have themselves timed out, -// leaving freshly enqueued entries in place. This is the behavior #2528 -// proposes; kept here next to the historical policy so a test can contrast the -// two and a follow-up fix can switch Workers over to it. +// Opt-in per-entry policy (transports/rdma/priority_promotion_per_entry=true): +// promote exactly the entries that have themselves timed out, leaving freshly +// enqueued entries in place. Workers::promoteTimedOutRequests selects this or +// DecidePromotionHeadOnly from the config flag. inline PromotionDecision DecidePromotionPerEntry( const std::vector& enqueue_ts, uint64_t current_ts, uint64_t promotion_timeout_ns) { 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 cc8a00681f..9868100007 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -262,9 +262,8 @@ class Workers { uint64_t slice_timeout_ns_; uint64_t priority_promotion_timeout_ns_; // Timeout for priority promotion // Opt-in (issue #2528): when true, a promotion pass promotes exactly the - // entries that have themselves timed out, instead of promoting the whole - // queue whenever only the head has timed out. Default false keeps the - // historical "flush the tier" behavior. + // entries that have themselves timed out (DecidePromotionPerEntry). + // Default false keeps DecidePromotionHeadOnly ("flush the tier"). bool priority_promotion_per_entry_ = false; std::unique_ptr device_selector_; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h index 3d5f2ebedd..6fd1962b81 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h @@ -15,7 +15,9 @@ #ifndef TCP_TRANSPORT_H_ #define TCP_TRANSPORT_H_ +#include #include +#include #include #include #include @@ -36,6 +38,14 @@ struct TcpParams { size_t max_concurrent_tasks = 16; // worker thread pool size }; +// Next sleep used by doTransferWithRetry after a failed attempt. Kept as a +// pure function so the cap is unit-testable without standing up a peer. +inline uint64_t nextTcpRetryDelay(uint64_t delay_ms, uint64_t cap) { + return std::min(delay_ms * 2, cap); +} + +class TcpTransportTestPeer; + struct TcpTask { Request request; BatchID progress_batch_id{0}; @@ -64,6 +74,8 @@ struct TcpSubBatch : public Transport::SubBatch { }; class TcpTransport : public Transport { + friend class TcpTransportTestPeer; + public: TcpTransport(); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 9a3c33a93e..13e973af86 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -629,18 +629,6 @@ void Workers::promoteTimedOutRequests(WorkerContext& worker) { } }; - if (!priority_promotion_per_entry_) { - auto* slice = drained.front().first; - const bool head_timed_out = slice && slice->enqueue_ts > 0 && - current_ts >= slice->enqueue_ts && - (current_ts - slice->enqueue_ts) >= - priority_promotion_timeout_ns_; - for (auto& slice_list : drained) { - requeue(head_timed_out ? to : from, slice_list); - } - return head_timed_out; - } - std::vector enqueue_ts; enqueue_ts.reserve(drained.size()); for (auto& slice_list : drained) { @@ -648,8 +636,12 @@ void Workers::promoteTimedOutRequests(WorkerContext& worker) { enqueue_ts.push_back(slice ? slice->enqueue_ts : 0); } - PromotionDecision decision = DecidePromotionPerEntry( - enqueue_ts, current_ts, priority_promotion_timeout_ns_); + PromotionDecision decision = + priority_promotion_per_entry_ + ? DecidePromotionPerEntry(enqueue_ts, current_ts, + priority_promotion_timeout_ns_) + : DecidePromotionHeadOnly(enqueue_ts, current_ts, + priority_promotion_timeout_ns_); if (!decision.promoted_any()) { for (auto& slice_list : drained) requeue(from, slice_list); diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp index 973173ca72..99d4352af9 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp @@ -241,7 +241,7 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { std::this_thread::sleep_for(std::chrono::milliseconds( std::min(static_cast(100), delay_ms - i))); } - delay_ms = std::min(delay_ms * 2, params_.retry_max_delay_ms); + delay_ms = nextTcpRetryDelay(delay_ms, params_.retry_max_delay_ms); } if (task->request.opcode == Request::WRITE) { diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index e560009065..875a111470 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -1,4 +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 Metrics Example add_executable(tent_metrics_example tent_metrics_example.cpp) @@ -246,13 +248,6 @@ if(USE_CUDA) PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) endif() -add_executable(tent_failover_test failover_test.cpp) -target_link_libraries(tent_failover_test PRIVATE gtest gtest_main - tent_link_group) -target_include_directories(tent_failover_test - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_failover_test COMMAND tent_failover_test) - add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp) target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main tent_link_group) @@ -424,13 +419,6 @@ target_include_directories(tent_transport_hint_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_transport_hint_test COMMAND tent_transport_hint_test) -add_executable(tent_intent_type_test intent_type_test.cpp) -target_link_libraries(tent_intent_type_test PRIVATE gtest gtest_main - tent_common) -target_include_directories(tent_intent_type_test - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_intent_type_test COMMAND tent_intent_type_test) - # ProgressWorker skeleton test: covers default-off behavior, event-driven # progress without poll-failover, and freeBatch races (issue #2116). add_executable(tent_progress_worker_test progress_worker_test.cpp) 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 ca417609ec..e5149dcdbc 100644 --- a/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp +++ b/mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp @@ -557,6 +557,71 @@ TEST(EngineFailoverE2E, ProgressBatchKeepsOverallPendingWithMixedOutcomes) { EXPECT_TRUE(engine.unregisterLocalMemory(pending_buf.data(), kBufLen).ok()); } +// A terminal TIMEOUT plus a COMPLETED sibling must report TIMEOUT overall, +// not a generic FAILED. The old aggregation collapsed every non-success +// terminal status into FAILED. +TEST(EngineFailoverE2E, OverallStatusUsesWorstFailureNotGenericFailed) { + auto cfg = makeMinimalP2PConfig(); + cfg->set("enable_auto_failover_on_poll", false); + cfg->set("max_failover_attempts", 0); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + constexpr size_t kBufLen = 4096; + std::vector timeout_buf(kBufLen, 0xB1); + std::vector completed_buf(kBufLen, 0xB2); + const uint64_t timeout_addr = + reinterpret_cast(timeout_buf.data()); + + auto fake_rdma = std::make_shared( + RDMA, [timeout_addr](const Request& req) { + if (reinterpret_cast(req.source) == timeout_addr) { + return TransferStatus{TransferStatusEnum::TIMEOUT, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, req.length}; + }); + auto fake_tcp = std::make_shared(TCP); + + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok()); + ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(RDMA, fake_rdma); + engine.swapTransportForTest(TCP, fake_tcp); + + ASSERT_TRUE(engine.registerLocalMemory(timeout_buf.data(), kBufLen).ok()); + ASSERT_TRUE(engine.registerLocalMemory(completed_buf.data(), kBufLen).ok()); + + BatchID batch_id = engine.allocateBatch(2); + ASSERT_NE(batch_id, (BatchID)0); + + Request timeout_req; + timeout_req.opcode = Request::WRITE; + timeout_req.source = timeout_buf.data(); + timeout_req.target_id = LOCAL_SEGMENT_ID; + timeout_req.target_offset = timeout_addr; + timeout_req.length = kBufLen; + + Request completed_req; + completed_req.opcode = Request::WRITE; + completed_req.source = completed_buf.data(); + completed_req.target_id = LOCAL_SEGMENT_ID; + completed_req.target_offset = + reinterpret_cast(completed_buf.data()); + completed_req.length = kBufLen; + + ASSERT_TRUE( + engine.submitTransfer(batch_id, {timeout_req, completed_req}).ok()); + + TransferStatus overall_status{}; + ASSERT_TRUE(engine.getTransferStatus(batch_id, overall_status).ok()); + EXPECT_EQ(overall_status.s, TransferStatusEnum::TIMEOUT); + + EXPECT_TRUE(engine.freeBatch(batch_id).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(timeout_buf.data(), kBufLen).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(completed_buf.data(), kBufLen).ok()); +} + TEST(EngineFailoverE2E, WaitTransferCompletionUsesProgressBatchWhenPollDisabled) { auto cfg = makeMinimalP2PConfig(); diff --git a/mooncake-transfer-engine/tent/tests/failover_test.cpp b/mooncake-transfer-engine/tent/tests/failover_test.cpp deleted file mode 100644 index 6c5b9441e0..0000000000 --- a/mooncake-transfer-engine/tent/tests/failover_test.cpp +++ /dev/null @@ -1,308 +0,0 @@ -// 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 "tent/common/config.h" -#include "tent/common/types.h" -#include "tent/runtime/transfer_engine_impl.h" - -namespace mooncake { -namespace tent { -namespace { - -// --------------------------------------------------------------------------- -// TaskInfo failover_count field tests -// --------------------------------------------------------------------------- - -TEST(TaskInfoTest, DefaultFailoverCount) { - TaskInfo task; - EXPECT_EQ(task.failover_count, 0); -} - -TEST(TaskInfoTest, FailoverCountIncrement) { - TaskInfo task; - ++task.failover_count; - EXPECT_EQ(task.failover_count, 1); - ++task.failover_count; - ++task.failover_count; - EXPECT_EQ(task.failover_count, 3); -} - -TEST(TaskInfoTest, FailoverCountLimitCheck) { - // Simulate the limit check pattern from resubmitTransferTask - constexpr int kMaxAttempts = 3; - TaskInfo task; - - // Attempts 1..3 should pass - for (int i = 0; i < kMaxAttempts; ++i) { - ++task.failover_count; - EXPECT_LE(task.failover_count, kMaxAttempts) - << "Attempt " << task.failover_count << " should be within limit"; - } - - // Attempt 4 should exceed - ++task.failover_count; - EXPECT_GT(task.failover_count, kMaxAttempts); -} - -// --------------------------------------------------------------------------- -// TaskInfo preserves other fields after failover-related mutations -// --------------------------------------------------------------------------- - -TEST(TaskInfoTest, FailoverFieldsIndependent) { - TaskInfo task; - task.type = RDMA; - task.xport_priority = 0; - task.status = TransferStatusEnum::PENDING; - task.failover_count = 0; - - // Simulate failover: increment priority and failover_count, change type - task.xport_priority++; - task.failover_count++; - task.type = TCP; - task.status = TransferStatusEnum::PENDING; // Reset to PENDING - - EXPECT_EQ(task.xport_priority, 1); - EXPECT_EQ(task.failover_count, 1); - EXPECT_EQ(task.type, TCP); - EXPECT_EQ(task.status, TransferStatusEnum::PENDING); -} - -// --------------------------------------------------------------------------- -// Config loading for max_failover_attempts -// --------------------------------------------------------------------------- - -TEST(FailoverConfigTest, DefaultValue) { - auto conf = std::make_shared(); - // When not set, should return the default - EXPECT_EQ(conf->get("max_failover_attempts", 3), 3); -} - -TEST(FailoverConfigTest, CustomValue) { - auto conf = std::make_shared(); - conf->set("max_failover_attempts", 5); - EXPECT_EQ(conf->get("max_failover_attempts", 3), 5); -} - -TEST(FailoverConfigTest, ZeroDisablesFailover) { - auto conf = std::make_shared(); - conf->set("max_failover_attempts", 0); - int max_attempts = conf->get("max_failover_attempts", 3); - EXPECT_EQ(max_attempts, 0); - - // With max=0, the very first ++failover_count > 0 check should fail - TaskInfo task; - ++task.failover_count; - EXPECT_GT(task.failover_count, max_attempts); -} - -TEST(FailoverConfigTest, AutoFailoverOnPollDefaultEnabled) { - auto conf = std::make_shared(); - EXPECT_TRUE(conf->get("enable_auto_failover_on_poll", true)); -} - -TEST(FailoverConfigTest, AutoFailoverOnPollCanBeDisabled) { - auto conf = std::make_shared(); - conf->set("enable_auto_failover_on_poll", false); - EXPECT_FALSE(conf->get("enable_auto_failover_on_poll", true)); -} - -// --------------------------------------------------------------------------- -// TransportType name coverage (tests the static helper indirectly via -// the enum values — the function itself is file-local in the .cpp, so we -// verify the enum values are well-defined and usable in switch) -// --------------------------------------------------------------------------- - -TEST(TransportTypeTest, UnspecIsZeroSlot) { - // UNSPEC must be value 0 so that zero-initialized tent_request / - // tent_memory_options structs (the C ABI) default to UNSPEC - // rather than silently pinning to whichever transport happened - // to occupy slot 0. - EXPECT_EQ(static_cast(UNSPEC), 0); - EXPECT_EQ(kSupportedTransportTypes, static_cast(kNumTransportTypes)); - EXPECT_GT(kSupportedTransportTypes, static_cast(UNSPEC)); -} - -// --------------------------------------------------------------------------- -// Failover state machine simulation (no real transport needed) -// --------------------------------------------------------------------------- - -TEST(FailoverStateMachineTest, SimulateFullFailoverSequence) { - // Simulate the sequence: submit on RDMA → fail → failover to TCP → succeed - TaskInfo task; - task.type = RDMA; - task.xport_priority = 0; - task.status = TransferStatusEnum::PENDING; - task.failover_count = 0; - - // Step 1: RDMA reports FAILED - task.status = TransferStatusEnum::FAILED; - EXPECT_EQ(task.status, TransferStatusEnum::FAILED); - - // Step 2: resubmitTransferTask logic (simulated) - constexpr int kMaxAttempts = 3; - ++task.failover_count; - EXPECT_LE(task.failover_count, kMaxAttempts); // Within limit - - task.xport_priority++; - // Simulate resolveTransport returning TCP - task.type = TCP; - task.status = TransferStatusEnum::PENDING; - - EXPECT_EQ(task.type, TCP); - EXPECT_EQ(task.xport_priority, 1); - EXPECT_EQ(task.failover_count, 1); - EXPECT_EQ(task.status, TransferStatusEnum::PENDING); - - // Step 3: TCP succeeds - task.status = TransferStatusEnum::COMPLETED; - EXPECT_EQ(task.status, TransferStatusEnum::COMPLETED); -} - -TEST(FailoverStateMachineTest, ExhaustAllTransports) { - // Simulate exhausting all failover attempts - constexpr int kMaxAttempts = 3; - TaskInfo task; - task.type = RDMA; - task.xport_priority = 0; - - for (int i = 0; i < kMaxAttempts; ++i) { - task.status = TransferStatusEnum::FAILED; - ++task.failover_count; - EXPECT_LE(task.failover_count, kMaxAttempts); - task.xport_priority++; - task.status = TransferStatusEnum::PENDING; - } - - // Next failure should exceed the limit - task.status = TransferStatusEnum::FAILED; - ++task.failover_count; - EXPECT_GT(task.failover_count, kMaxAttempts); - - // Task stays FAILED — no more failover - EXPECT_EQ(task.status, TransferStatusEnum::FAILED); - EXPECT_EQ(task.failover_count, kMaxAttempts + 1); -} - -TEST(FailoverStateMachineTest, StagingBypassesPriorityIncrement) { - // When task.staging is true, xport_priority should NOT increment - // (mirroring resubmitTransferTask logic) - TaskInfo task; - task.type = TCP; - task.xport_priority = 0; - task.staging = true; - - // Simulate resubmit logic - if (task.staging) - task.staging = false; - else - task.xport_priority++; - - EXPECT_FALSE(task.staging); - EXPECT_EQ(task.xport_priority, 0); // Should NOT have incremented -} - -// --------------------------------------------------------------------------- -// Batch status aggregation correctness -// --------------------------------------------------------------------------- - -// Verifies the invariant: a batch with one permanently-FAILED task and -// another still-PENDING task must report PENDING (not FAILED) overall, -// because the PENDING task may still complete. The old code latched -// overall_status to FAILED as soon as any task was terminal. -TEST(BatchStatusAggregationTest, PendingTaskPreventsEarlyBatchFailure) { - // Two tasks in a batch scenario (simulated): - // Task A: permanently FAILED (exhausted failover budget) - // Task B: still PENDING (retrying on secondary transport) - // - // Expected overall status: PENDING (not FAILED), because Task B is - // still in-flight. Only once Task B reaches a terminal state should - // the batch become terminal. - - struct MockTask { - TransferStatusEnum status{PENDING}; - bool derived{false}; - size_t length{1024}; - }; - - auto aggregateStatus = - [](const std::vector& tasks) -> TransferStatusEnum { - size_t success_tasks = 0; - size_t failed_tasks = 0; - size_t total_tasks = 0; - for (auto& t : tasks) { - if (t.derived) continue; - total_tasks++; - if (t.status == COMPLETED) - success_tasks++; - else if (t.status != PENDING) - failed_tasks++; - } - if (success_tasks == total_tasks) return COMPLETED; - if (success_tasks + failed_tasks == total_tasks) return FAILED; - return PENDING; - }; - - // Case 1: one FAILED + one PENDING → overall PENDING - { - std::vector tasks = {{FAILED, false, 1024}, - {PENDING, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), PENDING); - } - - // Case 2: one FAILED + one COMPLETED → overall FAILED - { - std::vector tasks = {{FAILED, false, 1024}, - {COMPLETED, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), FAILED); - } - - // Case 3: all COMPLETED → overall COMPLETED - { - std::vector tasks = {{COMPLETED, false, 1024}, - {COMPLETED, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), COMPLETED); - } - - // Case 4: all PENDING → overall PENDING - { - std::vector tasks = {{PENDING, false, 1024}, - {PENDING, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), PENDING); - } - - // Case 5: derived tasks are skipped - { - std::vector tasks = {{FAILED, false, 1024}, - {FAILED, true, 1024}, // derived: skip - {PENDING, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), PENDING); - } - - // Case 6: all non-derived are terminal with at least one failure - { - std::vector tasks = {{COMPLETED, false, 1024}, - {FAILED, true, 1024}, // derived: skip - {FAILED, false, 1024}}; - EXPECT_EQ(aggregateStatus(tasks), FAILED); - } -} - -} // namespace -} // namespace tent -} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/intent_type_test.cpp b/mooncake-transfer-engine/tent/tests/intent_type_test.cpp deleted file mode 100644 index ac56a2852e..0000000000 --- a/mooncake-transfer-engine/tent/tests/intent_type_test.cpp +++ /dev/null @@ -1,95 +0,0 @@ -// 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. -// -// Unit tests for IntentType enum and its integration with Request. - -#include - -#include - -#include "tent/common/types.h" - -namespace mooncake { -namespace tent { -namespace { - -TEST(IntentTypeTest, DefaultIsUnspec) { - Request r{}; - EXPECT_EQ(r.intent_type, IntentType::INTENT_UNSPEC); -} - -TEST(IntentTypeTest, AllValuesAssignable) { - Request r{}; - r.intent_type = IntentType::FOREGROUND_GET; - EXPECT_EQ(r.intent_type, IntentType::FOREGROUND_GET); - r.intent_type = IntentType::BACKGROUND_PREFETCH; - EXPECT_EQ(r.intent_type, IntentType::BACKGROUND_PREFETCH); - r.intent_type = IntentType::MIGRATION; - EXPECT_EQ(r.intent_type, IntentType::MIGRATION); - r.intent_type = IntentType::CHECKPOINT; - EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); - r.intent_type = IntentType::WEIGHT_LOADING; - EXPECT_EQ(r.intent_type, IntentType::WEIGHT_LOADING); - r.intent_type = IntentType::STAGING_INTERNAL; - EXPECT_EQ(r.intent_type, IntentType::STAGING_INTERNAL); -} - -TEST(IntentTypeTest, IntegerValues) { - EXPECT_EQ(static_cast(IntentType::INTENT_UNSPEC), 0); - EXPECT_EQ(static_cast(IntentType::FOREGROUND_GET), 1); - EXPECT_EQ(static_cast(IntentType::BACKGROUND_PREFETCH), 2); - EXPECT_EQ(static_cast(IntentType::MIGRATION), 3); - EXPECT_EQ(static_cast(IntentType::CHECKPOINT), 4); - EXPECT_EQ(static_cast(IntentType::WEIGHT_LOADING), 5); - EXPECT_EQ(static_cast(IntentType::STAGING_INTERNAL), 6); -} - -TEST(IntentTypeTest, DoesNotAffectOtherFields) { - Request r{}; - r.opcode = Request::READ; - r.priority = PRIO_LOW; - r.deadline_ns = 12345; - r.transport_hint = RDMA; - r.intent_type = IntentType::CHECKPOINT; - - EXPECT_EQ(r.opcode, Request::READ); - EXPECT_EQ(r.priority, PRIO_LOW); - EXPECT_EQ(r.deadline_ns, 12345u); - EXPECT_EQ(r.transport_hint, RDMA); - EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); -} - -TEST(IntentTypeTest, CopyPreservesIntentType) { - Request r{}; - r.intent_type = IntentType::WEIGHT_LOADING; - Request copy = r; - EXPECT_EQ(copy.intent_type, IntentType::WEIGHT_LOADING); -} - -TEST(IntentTypeTest, VectorOfRequests) { - std::vector batch(4); - batch[0].intent_type = IntentType::FOREGROUND_GET; - batch[1].intent_type = IntentType::BACKGROUND_PREFETCH; - batch[2].intent_type = IntentType::MIGRATION; - batch[3].intent_type = IntentType::INTENT_UNSPEC; - - EXPECT_EQ(batch[0].intent_type, IntentType::FOREGROUND_GET); - EXPECT_EQ(batch[1].intent_type, IntentType::BACKGROUND_PREFETCH); - EXPECT_EQ(batch[2].intent_type, IntentType::MIGRATION); - EXPECT_EQ(batch[3].intent_type, IntentType::INTENT_UNSPEC); -} - -} // namespace -} // namespace tent -} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp index ab49faef73..7739d8a858 100644 --- a/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp +++ b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. // -// Deterministic reproduction of the promotion issues reported in #2528. The -// historical head-only policy is contrasted with a per-entry policy on the same -// inputs so the unintended behavior is unambiguous and re-runnable (no RDMA -// stack, no timing noise). +// Pins both promotion policies used by Workers::promoteTimedOutRequests +// (issue #2528). Head-only is the default; per-entry is opt-in via +// transports/rdma/priority_promotion_per_entry. Contrasted on the same inputs +// so the default "flush the tier" behavior stays unambiguous (no RDMA stack, +// no timing noise). #include "tent/transport/rdma/promotion_policy.h" @@ -44,7 +45,7 @@ TEST(PromotionPolicyTest, HeadOnlyPromotesFreshEntriesWhenHeadTimedOut) { auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); - // BUG: all three are promoted, including the two fresh (non-starving) ones. + // Default policy: all three are promoted, including the two fresh ones. EXPECT_EQ(d.promote_indices, (std::vector{0, 1, 2})); // The per-entry policy promotes only the genuinely starving head. @@ -58,7 +59,7 @@ TEST(PromotionPolicyTest, HeadOnlyMissesTimedOutTailWhenHeadFresh) { auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); - // BUG: nothing is promoted even though indices 1 and 2 are starving. + // Default policy: nothing is promoted even though indices 1 and 2 starve. EXPECT_TRUE(d.promote_indices.empty()); auto fixed = DecidePromotionPerEntry(q, kNow, kTimeout); @@ -86,7 +87,7 @@ TEST(PromotionPolicyTest, ZeroTimestampNeverTimesOut) { TEST(PromotionPolicyTest, AllTimedOutPromotesAllUnderBothPolicies) { // When every entry is starving the two policies agree — this is the case - // the historical policy was designed around. + // the default head-only policy was designed around. std::vector q = {timedOut(), timedOut(), timedOut()}; EXPECT_EQ(DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices, (std::vector{0, 1, 2})); diff --git a/mooncake-transfer-engine/tent/tests/tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/tcp_transport_test.cpp index b33d8e1bbb..4b5fb5efb8 100644 --- a/mooncake-transfer-engine/tent/tests/tcp_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/tcp_transport_test.cpp @@ -14,103 +14,70 @@ #include -#include -#include #include -#include +#include #include #include "tent/common/config.h" -#include "tent/common/types.h" +#include "tent/runtime/control_plane.h" #include "tent/transport/tcp/tcp_transport.h" namespace mooncake { namespace tent { -namespace { - -// --------------------------------------------------------------------------- -// TcpParams unit tests -// --------------------------------------------------------------------------- -TEST(TcpParamsTest, DefaultValues) { - TcpParams params; - EXPECT_EQ(params.max_retry_count, 3); - EXPECT_EQ(params.retry_base_delay_ms, 100ULL); - EXPECT_EQ(params.retry_max_delay_ms, 2'000ULL); - EXPECT_EQ(params.max_concurrent_tasks, 16); -} +class TcpTransportTestPeer { + public: + static const TcpParams& params(const TcpTransport& transport) { + return transport.params_; + } +}; -// --------------------------------------------------------------------------- -// TcpTask unit tests — atomic semantics -// --------------------------------------------------------------------------- +namespace { -TEST(TcpTaskTest, DefaultConstruction) { - TcpTask task; - EXPECT_EQ(task.status_word.load(), TransferStatusEnum::PENDING); - EXPECT_EQ(task.transferred_bytes.load(), 0u); - EXPECT_EQ(task.target_addr, 0u); +std::shared_ptr makeP2PMetadata() { + return std::make_shared("p2p", "", nullptr); } -TEST(TcpTaskTest, AtomicStatusTransitions) { - TcpTask task; - EXPECT_EQ(task.status_word.load(std::memory_order_acquire), - TransferStatusEnum::PENDING); +TEST(TcpTransportConfigTest, InstallWithoutConfigKeepsStructDefaults) { + TcpTransport transport; + auto metadata = makeP2PMetadata(); + std::string name = "tcp-config-default"; + ASSERT_TRUE(transport.install(name, metadata, nullptr).ok()); - task.status_word.store(TransferStatusEnum::COMPLETED, - std::memory_order_release); - EXPECT_EQ(task.status_word.load(std::memory_order_acquire), - TransferStatusEnum::COMPLETED); + const TcpParams defaults; + const auto& params = TcpTransportTestPeer::params(transport); + EXPECT_EQ(params.max_retry_count, defaults.max_retry_count); + EXPECT_EQ(params.retry_base_delay_ms, defaults.retry_base_delay_ms); + EXPECT_EQ(params.retry_max_delay_ms, defaults.retry_max_delay_ms); + EXPECT_EQ(params.max_concurrent_tasks, defaults.max_concurrent_tasks); - task.status_word.store(TransferStatusEnum::FAILED, - std::memory_order_release); - EXPECT_EQ(task.status_word.load(std::memory_order_acquire), - TransferStatusEnum::FAILED); + EXPECT_TRUE(transport.uninstall().ok()); } -TEST(TcpTaskTest, AtomicTransferredBytes) { - TcpTask task; - constexpr size_t kPayload = 1024 * 1024; - task.transferred_bytes.store(kPayload, std::memory_order_release); - EXPECT_EQ(task.transferred_bytes.load(std::memory_order_acquire), kPayload); -} - -TEST(TcpTaskTest, MoveConstruction) { - TcpTask a; - a.status_word.store(TransferStatusEnum::COMPLETED, - std::memory_order_relaxed); - a.transferred_bytes.store(4096, std::memory_order_relaxed); - a.target_addr = 0xdeadbeef; - - TcpTask b(std::move(a)); - EXPECT_EQ(b.status_word.load(), TransferStatusEnum::COMPLETED); - EXPECT_EQ(b.transferred_bytes.load(), 4096u); - EXPECT_EQ(b.target_addr, 0xdeadbeef); -} - -// --------------------------------------------------------------------------- -// TcpSubBatch unit tests -// --------------------------------------------------------------------------- +TEST(TcpTransportConfigTest, InstallReadsConfigKeys) { + auto conf = std::make_shared(); + conf->set("transports/tcp/max_retry_count", 5); + conf->set("transports/tcp/retry_base_delay_ms", 200ULL); + conf->set("transports/tcp/retry_max_delay_ms", 4000ULL); + conf->set("transports/tcp/max_concurrent_tasks", 32); -TEST(TcpSubBatchTest, EmptyBatch) { - TcpSubBatch batch; - batch.max_size = 64; - EXPECT_EQ(batch.size(), 0u); -} + TcpTransport transport; + auto metadata = makeP2PMetadata(); + std::string name = "tcp-config-override"; + ASSERT_TRUE(transport.install(name, metadata, nullptr, conf).ok()); -TEST(TcpSubBatchTest, EmplaceAndSize) { - TcpSubBatch batch; - batch.max_size = 64; - batch.task_list.reserve(batch.max_size); + const auto& params = TcpTransportTestPeer::params(transport); + EXPECT_EQ(params.max_retry_count, 5u); + EXPECT_EQ(params.retry_base_delay_ms, 200ULL); + EXPECT_EQ(params.retry_max_delay_ms, 4000ULL); + EXPECT_EQ(params.max_concurrent_tasks, 32u); - batch.task_list.emplace_back(); - batch.task_list.emplace_back(); - batch.task_list.emplace_back(); - EXPECT_EQ(batch.size(), 3u); + EXPECT_TRUE(transport.uninstall().ok()); } TEST(TcpSubBatchTest, PointerStabilityAfterReserve) { - // Verify that pointers taken after reserve remain valid when more - // tasks are emplaced (important for async dispatch correctness). + // allocateSubBatch reserves task_list to max_size so submitTransferTasks + // can take stable TcpTask* after emplace. TcpSubBatch batch; batch.max_size = 8; batch.task_list.reserve(batch.max_size); @@ -119,78 +86,15 @@ TEST(TcpSubBatchTest, PointerStabilityAfterReserve) { TcpTask* first = &batch.task_list[0]; first->target_addr = 42; - // Add more tasks (within reserved capacity) for (int i = 1; i < 8; ++i) { batch.task_list.emplace_back(); } - // First pointer must still be valid EXPECT_EQ(first->target_addr, 42u); EXPECT_EQ(&batch.task_list[0], first); } -// --------------------------------------------------------------------------- -// TcpTransport config loading test -// --------------------------------------------------------------------------- - -TEST(TcpTransportConfigTest, ConfigOverridesDefaults) { - auto conf = std::make_shared(); - conf->set("transports/tcp/max_retry_count", 5); - conf->set("transports/tcp/retry_base_delay_ms", 200ULL); - conf->set("transports/tcp/retry_max_delay_ms", 4000ULL); - conf->set("transports/tcp/max_concurrent_tasks", 32); - - EXPECT_EQ(conf->get("transports/tcp/max_retry_count", 0), 5); - EXPECT_EQ(conf->get("transports/tcp/retry_base_delay_ms", 0ULL), 200ULL); - EXPECT_EQ(conf->get("transports/tcp/retry_max_delay_ms", 0ULL), 4000ULL); - EXPECT_EQ(conf->get("transports/tcp/max_concurrent_tasks", 0), 32); -} - -TEST(TcpTransportConfigTest, MissingConfigUsesDefaults) { - auto conf = std::make_shared(); - - TcpParams defaults; - EXPECT_EQ( - conf->get("transports/tcp/max_retry_count", defaults.max_retry_count), - 3); - EXPECT_EQ(conf->get("transports/tcp/max_concurrent_tasks", - defaults.max_concurrent_tasks), - 16); -} - -// --------------------------------------------------------------------------- -// Cross-thread visibility test (verifies atomic correctness) -// --------------------------------------------------------------------------- - -TEST(TcpTaskTest, CrossThreadVisibility) { - TcpTask task; - std::atomic writer_done{false}; - - std::thread writer([&]() { - task.transferred_bytes.store(8192, std::memory_order_release); - task.status_word.store(TransferStatusEnum::COMPLETED, - std::memory_order_release); - writer_done.store(true, std::memory_order_release); - }); - - // Spin until writer signals done - while (!writer_done.load(std::memory_order_acquire)) { - std::this_thread::yield(); - } - - EXPECT_EQ(task.status_word.load(std::memory_order_acquire), - TransferStatusEnum::COMPLETED); - EXPECT_EQ(task.transferred_bytes.load(std::memory_order_acquire), 8192u); - - writer.join(); -} - -// --------------------------------------------------------------------------- -// Exponential backoff calculation test -// --------------------------------------------------------------------------- - TEST(TcpRetryBackoffTest, ExponentialGrowthWithCap) { - // Simulate the backoff logic from doTransferWithRetry const uint64_t base = 100; const uint64_t cap = 2000; uint64_t delay = base; @@ -198,16 +102,15 @@ TEST(TcpRetryBackoffTest, ExponentialGrowthWithCap) { std::vector delays; for (int attempt = 0; attempt < 6; ++attempt) { delays.push_back(delay); - delay = std::min(delay * 2, cap); + delay = nextTcpRetryDelay(delay, cap); } - // 100 → 200 → 400 → 800 → 1600 → 2000 (capped) EXPECT_EQ(delays[0], 100u); EXPECT_EQ(delays[1], 200u); EXPECT_EQ(delays[2], 400u); EXPECT_EQ(delays[3], 800u); EXPECT_EQ(delays[4], 1600u); - EXPECT_EQ(delays[5], 2000u); // capped at retry_max_delay_ms + EXPECT_EQ(delays[5], 2000u); } } // namespace diff --git a/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp index cb6921a461..6611043c2b 100644 --- a/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp +++ b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp @@ -22,8 +22,6 @@ #include #include -#include -#include #include #include #include @@ -192,43 +190,5 @@ TEST(ThreadLocalStorageTest, OrphanedValuesSweptOnInstanceChurn) { } } -// Informational: hot-path cost of get(). The remote-desc cache consults this -// on every transfer submit, so the common case must stay a thread_local -// access plus a compare. -TEST(ThreadLocalStorageTest, HotPathMicrobench) { - ThreadLocalStorage storage; - storage.get().value = 1; - constexpr uint64_t kOps = 20'000'000; - volatile int sink = 0; - auto t0 = std::chrono::steady_clock::now(); - for (uint64_t i = 0; i < kOps; ++i) { - sink += storage.get().value; - } - auto t1 = std::chrono::steady_clock::now(); - double ns = - (double)std::chrono::duration_cast(t1 - t0) - .count() / - (double)kOps; - printf("get_hot_path_ns_per_op %.2f\n", ns); - (void)sink; -#if defined(__SANITIZE_THREAD__) || defined(__SANITIZE_ADDRESS__) - constexpr bool kSanitized = true; -#elif defined(__has_feature) -#if __has_feature(thread_sanitizer) || __has_feature(address_sanitizer) - constexpr bool kSanitized = true; -#else - constexpr bool kSanitized = false; -#endif -#else - constexpr bool kSanitized = false; -#endif - // Wall-clock assertions flake under sanitizers (TSAN alone is ~14x); - // elsewhere keep a loose ceiling that still catches syscall- or - // contention-class regressions on the hot path. - if (!kSanitized) { - EXPECT_LT(ns, 100.0); - } -} - } // namespace tent } // namespace mooncake From 619a48fcf673b431918ffb9223cc9e18b90526bc Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 31 Aug 2026 14:28:10 +0800 Subject: [PATCH 05/17] [Bugfix][TENT] Offload TCP SendData/RecvData and raise default RPC threads (#3767) TCP bulk copies ran inline on the RPC io_context, serializing concurrent transfers and stalling Probe/Bootstrap. Offload the handlers and default rpc_server_threads higher when TCP is enabled so attachments can be read in parallel. Signed-off-by: staryxchen Co-authored-by: Cursor --- mooncake-transfer-engine/tent/include/tent/rpc/rpc.h | 7 ++++--- mooncake-transfer-engine/tent/src/common/config.cpp | 8 +++++--- .../tent/src/runtime/control_plane.cpp | 10 ++++++++-- .../tent/src/runtime/transfer_engine_impl.cpp | 12 +++++++++++- .../tent/tests/tcp_datapath_roundtrip_test.cpp | 5 +++-- 5 files changed, 31 insertions(+), 11 deletions(-) diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 8616865454..23a96aa7f5 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -78,9 +78,10 @@ class CoroRpcAgent { bool offload = false); // threads: number of io_context worker threads (default 1, the - // historical behavior). The TCP data-path handlers do full-payload - // blocking copies inline, so a single thread caps TCP throughput; - // sourced from the rpc_server_threads config key. + // historical RDMA-only behavior). TCP SendData/RecvData copies are + // offloaded off this pool, but attachments are still read here; when + // TCP is enabled the engine defaults rpc_server_threads higher so + // concurrent bulk transfers are not serialized on one io_context. Status start(uint16_t &port, bool ipv6 = false, size_t threads = 1); Status stop(); diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index e7d98a12c1..cec776d2ac 100644 --- a/mooncake-transfer-engine/tent/src/common/config.cpp +++ b/mooncake-transfer-engine/tent/src/common/config.cpp @@ -179,9 +179,11 @@ Status ConfigHelper::loadFromEnv(Config& config) { // MC_CUSTOM_TOPO_JSON works under MC_USE_TENT. Inline // topology/priority_matrix in MC_TENT_CONF still takes precedence. setConfig(config, "MC_CUSTOM_TOPO_JSON", "topology/custom_json_path"); - // TENT RPC server io_context threads. The TCP data-path handlers do - // full-payload blocking copies inline, so deployments pushing bulk data - // over the TENT TCP transport raise this above the default of 1. + // TENT RPC server io_context threads. TCP SendData/RecvData copies are + // offloaded onto the blocking executor; this pool still reads the RPC + // attachments. When TCP is enabled and this key is unset, the engine + // defaults to several threads so concurrent bulk transfers are not + // serialized on one io_context. setConfig(config, "MC_TENT_RPC_THREADS", "rpc_server_threads"); return status; } diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 4c3de85d8c..1abf18251e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -267,16 +267,22 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onBootstrapRdma(request, response); }); + // SendData/RecvData copy the full TCP payload. Running them inline on the + // 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( SendData, [this](const std::string_view& request, std::string& response) { onSendData(request, response); - }); + }, + /*offload=*/true); rpc_server_->registerFunction( RecvData, [this](const std::string_view& request, std::string& response) { onRecvData(request, response); - }); + }, + /*offload=*/true); rpc_server_->registerFunction( Notify, [this](const std::string_view& request, std::string& response) { onNotify(request, response); 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 739b770f0d..4e15b63d02 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -343,8 +343,18 @@ Status TransferEngineImpl::construct() { hostname_ = conf_->get("rpc_server_hostname", ""); local_segment_name_ = conf_->get("local_segment_name", ""); CHECK_STATUS(getRpcServerPortFromConfig(*conf_, 0, port_)); + // TCP SendData/RecvData copies are offloaded, but the RPC io_context still + // reads the full attachment. One thread serializes concurrent bulk TCP. + // Leave RDMA-only at 1; when TCP is on and the user did not set + // rpc_server_threads, use several so attachments can be read in parallel. size_t rpc_server_threads = 1; - CHECK_STATUS(getRpcServerThreadsFromConfig(*conf_, 1, rpc_server_threads)); + const size_t rpc_threads_default = + conf_->get("transports/tcp/enable", false) + ? std::min( + 8, std::max(4, std::thread::hardware_concurrency())) + : 1; + CHECK_STATUS(getRpcServerThreadsFromConfig(*conf_, rpc_threads_default, + rpc_server_threads)); merge_requests_ = conf_->get("merge_requests", true); max_failover_attempts_ = conf_->get("max_failover_attempts", 3); enable_auto_failover_on_poll_ = diff --git a/mooncake-transfer-engine/tent/tests/tcp_datapath_roundtrip_test.cpp b/mooncake-transfer-engine/tent/tests/tcp_datapath_roundtrip_test.cpp index f54a0b4437..5dce2520f8 100644 --- a/mooncake-transfer-engine/tent/tests/tcp_datapath_roundtrip_test.cpp +++ b/mooncake-transfer-engine/tent/tests/tcp_datapath_roundtrip_test.cpp @@ -236,8 +236,9 @@ TEST(TcpDataPathRoundtripTest, WriteThenReadAcrossProcesses) { runWriteThenReadAcrossProcesses(1); } -// Same round trip with a multi-threaded RPC server, as used by bulk-TCP -// deployments (MC_TENT_RPC_THREADS / rpc_server_threads). +// Same round trip with a multi-threaded RPC server. SendData/RecvData are +// offloaded, so this also covers concurrent bulk copies overlapping other +// RPCs (MC_TENT_RPC_THREADS / rpc_server_threads). TEST(TcpDataPathRoundtripTest, WriteThenReadAcrossProcessesMultiThreadedRpc) { runWriteThenReadAcrossProcesses(4); } From fbce6fdc8a007b3661ec5dc5422c07330a57c681 Mon Sep 17 00:00:00 2001 From: Miguel <60073809+migarci2@users.noreply.github.com> Date: Mon, 31 Aug 2026 08:39:41 +0200 Subject: [PATCH 06/17] [CI/Build] Run etcd wrapper Go tests with CTest (#3753) Register the wrapper module with CTest when unit tests are enabled so Go test failures propagate through the existing test path. Signed-off-by: Miguel Garcia --- mooncake-common/etcd/CMakeLists.txt | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mooncake-common/etcd/CMakeLists.txt b/mooncake-common/etcd/CMakeLists.txt index 570e8bbf17..0e3d8ceb80 100644 --- a/mooncake-common/etcd/CMakeLists.txt +++ b/mooncake-common/etcd/CMakeLists.txt @@ -14,6 +14,13 @@ add_custom_target( DEPENDS ${ETCD_WRAPPER_LIB} ) +if(BUILD_UNIT_TESTS) + add_test( + NAME etcd_wrapper_go_test + COMMAND go test ./... + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) +endif() + install( FILES ${ETCD_WRAPPER_LIB} DESTINATION lib From 4251eacedc1ad244e5b1a9860bc91b7a7a7f33bb Mon Sep 17 00:00:00 2001 From: Posedge_Lin Date: Mon, 31 Aug 2026 00:07:05 -0700 Subject: [PATCH 07/17] [TransferEngine] Add per-target tebench metrics (#3779) --- docs/source/performance/mooncake/tebench.md | 16 ++ .../benchmark/CMakeLists.txt | 45 ++++-- .../benchmark/bench_runner.h | 2 + mooncake-transfer-engine/benchmark/main.cpp | 33 ++++ .../benchmark/target_metrics.cpp | 144 ++++++++++++++++++ .../benchmark/target_metrics.h | 75 +++++++++ .../benchmark/te_backend.h | 4 + .../benchmark/tent_backend.h | 4 + .../benchmark/tests/target_metrics_test.cpp | 66 ++++++++ mooncake-transfer-engine/benchmark/utils.cpp | 5 + mooncake-transfer-engine/benchmark/utils.h | 1 + 11 files changed, 381 insertions(+), 14 deletions(-) create mode 100644 mooncake-transfer-engine/benchmark/target_metrics.cpp create mode 100644 mooncake-transfer-engine/benchmark/target_metrics.h create mode 100644 mooncake-transfer-engine/benchmark/tests/target_metrics_test.cpp 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-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..42e5b7aeb7 100644 --- a/mooncake-transfer-engine/benchmark/bench_runner.h +++ b/mooncake-transfer-engine/benchmark/bench_runner.h @@ -56,6 +56,8 @@ 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, diff --git a/mooncake-transfer-engine/benchmark/main.cpp b/mooncake-transfer-engine/benchmark/main.cpp index 19f9bbf37a..ce1f81ba17 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; @@ -208,10 +216,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 +248,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) { 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.h b/mooncake-transfer-engine/benchmark/tent_backend.h index 4f6cba5ffc..adbc50fa84 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, 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..ffb39eb857 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."); @@ -132,6 +135,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 +175,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; From 88230cb08b3c15d41b4176aa9821ac8ff4648ee0 Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 31 Aug 2026 15:49:03 +0800 Subject: [PATCH 08/17] [CI] Stop run-e2e-ci from retriggering Build & Test (#3756) ci.yml listened for every labeled event, and auto-labeler already applies run-ci, so adding run-e2e-ci cancelled in-progress PR CI and reran every job. Keep Build & Test on open/push and move same-SHA retrigger to a workflow that only reacts to a human-applied run-ci label. Signed-off-by: staryxchen Co-authored-by: Cursor --- .github/workflows/ci-on-label.yml | 53 +++++++++++++++++++++++++++++++ .github/workflows/ci.yml | 6 +++- 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/ci-on-label.yml 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..a6cdb72374 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: From 6e1dd41dcec4edbff8f9210fbfa5daf068287170 Mon Sep 17 00:00:00 2001 From: Aoi Date: Mon, 31 Aug 2026 16:43:12 +0800 Subject: [PATCH 09/17] [Store] Introduce stateful region resource drivers (#3703) * [Store] Introduce stateful region resource drivers * [Store] Refine region driver recovery boundaries * [Store] Validate CacheLib slab count * [Store] Clarify region driver ownership contracts --- docs/source/api-reference/cpp/index.md | 2 +- mooncake-store/include/allocator.h | 64 +++- mooncake-store/include/placement/target.h | 32 ++ mooncake-store/include/replica.h | 9 +- mooncake-store/include/segment.h | 4 +- mooncake-store/include/segment/region.h | 24 ++ .../include/segment/region_driver.h | 120 +++++++ mooncake-store/src/CMakeLists.txt | 2 +- mooncake-store/src/allocator.cpp | 166 ++++++--- mooncake-store/src/master_service.cpp | 36 +- mooncake-store/src/segment.cpp | 87 ++--- mooncake-store/src/segment/region_driver.cpp | 336 ++++++++++++++++++ mooncake-store/tests/CMakeLists.txt | 1 + .../tests/allocation_strategy_test.cpp | 14 +- .../tests/buffer_allocator_test.cpp | 229 ++++++------ .../tests/ha/master_service_ha_test.cpp | 8 + mooncake-store/tests/region_driver_test.cpp | 181 ++++++++++ 17 files changed, 1034 insertions(+), 281 deletions(-) create mode 100644 mooncake-store/include/placement/target.h create mode 100644 mooncake-store/include/segment/region.h create mode 100644 mooncake-store/include/segment/region_driver.h create mode 100644 mooncake-store/src/segment/region_driver.cpp create mode 100644 mooncake-store/tests/region_driver_test.cpp 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/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/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/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 2c2d5f9089..95d6464b46 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -19,6 +19,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 @@ -106,7 +107,6 @@ 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 "") 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/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/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/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 66372aedf2..1f9c4d88af 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) 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/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/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 From 73471dc2803f0de7e42484cfa1e2f44c09dc1de5 Mon Sep 17 00:00:00 2001 From: Aoi Date: Mon, 31 Aug 2026 17:37:41 +0800 Subject: [PATCH 10/17] [CI/Build] Keep non-CUDA wheels CUDA-free (#3745) * [CI/Build] Keep non-CUDA wheels CUDA-free * [CI/Build] Run wheel smoke test under bash for CUDA scan The non-CUDA CUDA-dependency scan uses process substitution (< <(...)), which sh (dash) rejects at parse time, failing every wheel build variant regardless of VARIANT_FLAG. Run the smoke test step under bash. --- .github/workflows/_build-wheel.yaml | 20 +++++++++++++++----- mooncake-store/src/CMakeLists.txt | 29 +++++++++++++++-------------- 2 files changed, 30 insertions(+), 19 deletions(-) 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/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 95d6464b46..786e8e0227 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -299,12 +299,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 +414,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 +494,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 +550,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() From 964b8cf0c8208a3fa974d1afe7b10e53e8f05070 Mon Sep 17 00:00:00 2001 From: XiaokunDing Date: Tue, 1 Sep 2026 10:34:35 +0800 Subject: [PATCH 11/17] Fix data copy while not on same device (#3476) * Fix data copy while not on same device * Format code * Format code * Fix code address comment by @staryxchen --------- Co-authored-by: shawnding --- .../tent/include/tent/platform/cuda.h | 3 ++ .../tent/include/tent/platform/rocm.h | 3 ++ .../src/platform/ascend/ascend_allocator.cpp | 6 +++ .../tent/src/platform/cuda/cuda_allocator.cpp | 14 ++++++- .../tent/src/platform/cuda/cuda_probe.cpp | 20 ++++++++++ .../tent/src/platform/rocm/rocm_allocator.cpp | 13 ++++++- .../tent/src/platform/rocm/rocm_probe.cpp | 14 +++++++ .../tent/src/runtime/control_plane.cpp | 38 +++++++++++++++---- 8 files changed, 101 insertions(+), 10 deletions(-) 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/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/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 1abf18251e..cef9348498 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -102,8 +102,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 +126,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) { @@ -452,7 +454,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 +479,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 +502,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"); } } From 27e00201944fa0b121b5b6bd25e5b38a3c9834d0 Mon Sep 17 00:00:00 2001 From: Colors-111 <70190328+Colors-111@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:07:00 +0800 Subject: [PATCH 12/17] [TENT]: Map MOONCAKE_LOCAL_HOSTNAME to TENT rpc_server_hostname (#3790) * Map MOONCAKE_LOCAL_HOSTNAME to TENT rpc_server_hostname * code format --------- Co-authored-by: ruanzhao --- .../tent/src/common/config.cpp | 1 + .../transfer_engine_config_override_test.cpp | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index cec776d2ac..fa997e2c0d 100644 --- a/mooncake-transfer-engine/tent/src/common/config.cpp +++ b/mooncake-transfer-engine/tent/src/common/config.cpp @@ -135,6 +135,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"); 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..62f15cd17d 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"); From 40aefcc6f8c389533af76eb90a1daccaf8f31957 Mon Sep 17 00:00:00 2001 From: SongOf <46475785+SongOf@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:08:30 +0800 Subject: [PATCH 13/17] [TENT] Prefer the LAG-effective port speed from ibv_query_port_speed (#3777) * [TENT] Prefer the LAG-effective port speed from ibv_query_port_speed * [TENT] Let tests inject verbs into RdmaContext and cover the effective-speed path * [TENT] Hold the last effective speed over transient query failures and count them --------- Co-authored-by: maxlisongsong --- docs/source/design/tent/slice-spraying.md | 10 +- mooncake-common/include/ib_link_speed.h | 12 + mooncake-common/tests/ib_link_speed_test.cpp | 17 ++ .../include/tent/transport/rdma/context.h | 34 ++- .../include/tent/transport/rdma/ibv_loader.h | 7 +- .../tent/src/transport/rdma/context.cpp | 42 +++- .../tent/src/transport/rdma/ibv_loader.cpp | 17 +- .../tent/src/transport/rdma/workers.cpp | 5 +- .../tent/tests/rdma_transport_test.cpp | 220 ++++++++++++++++++ 9 files changed, 356 insertions(+), 8 deletions(-) diff --git a/docs/source/design/tent/slice-spraying.md b/docs/source/design/tent/slice-spraying.md index 5c8db04587..1e7caec198 100644 --- a/docs/source/design/tent/slice-spraying.md +++ b/docs/source/design/tent/slice-spraying.md @@ -270,7 +270,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/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-transfer-engine/tent/include/tent/transport/rdma/context.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h index 434694fdd1..10921a2bcd 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 @@ -131,6 +134,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 +156,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 +190,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 +210,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/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/src/transport/rdma/context.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp index 9a861cc9a7..1e1d85fd44 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,49 @@ int RdmaContext::refreshPortAttributes() { return -1; } recordPortSpeed(port_attr); + queryEffectiveSpeed(); 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/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/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 13e973af86..380a17c731 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -944,7 +944,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); } diff --git a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp index d4b70a89eb..8f6d9aae2c 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 { @@ -89,6 +92,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 +503,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); From 82a90a9e7cf8cff90f1b3e1f27026c91370a7160 Mon Sep 17 00:00:00 2001 From: Icedcoco <102317026+Icedcoco@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:22:06 +0800 Subject: [PATCH 14/17] [Store] Add batch OpLog snapshot coordinator (#3794) * [Store] Add batch OpLog snapshot coordinator * [Bugfix][Store] Fix batch snapshot coordinator races --------- Co-authored-by: Yuchen Kou --- .../batch_oplog_snapshot_coordinator.h | 126 +++++ .../batch_oplog_snapshot_coordinator.h | 5 + mooncake-store/include/hot_standby_service.h | 19 + mooncake-store/src/CMakeLists.txt | 1 + .../batch_oplog_snapshot_coordinator.cpp | 462 ++++++++++++++++++ mooncake-store/src/hot_standby_service.cpp | 75 +++ mooncake-store/tests/CMakeLists.txt | 2 + .../snapshot/batch_oplog/coordinator_test.cpp | 306 ++++++++++++ 8 files changed, 996 insertions(+) create mode 100644 mooncake-store/include/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.h create mode 100644 mooncake-store/include/ha/snapshot/batch_oplog_snapshot_coordinator.h create mode 100644 mooncake-store/src/ha/snapshot/batch_oplog/batch_oplog_snapshot_coordinator.cpp create mode 100644 mooncake-store/tests/ha/snapshot/batch_oplog/coordinator_test.cpp 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/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 786e8e0227..b539334ab8 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -44,6 +44,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 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/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 1f9c4d88af..8a7b4a1a22 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -189,6 +189,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/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 From a14168c9bad901d702fa798455ff0281bddec6a2 Mon Sep 17 00:00:00 2001 From: Schatten Date: Tue, 1 Sep 2026 11:25:25 +0800 Subject: [PATCH 15/17] [Store] Structure ClientAutoPortConfig environment settings (#3785) Signed-off-by: Schatten --- .../include/environment_variables.h | 6 + .../include/client_auto_port_config.h | 13 ++ mooncake-store/src/CMakeLists.txt | 1 + .../src/config/client_auto_port_config.cpp | 27 ++++ mooncake-store/src/real_client.cpp | 30 ++-- mooncake-store/tests/CMakeLists.txt | 1 + .../tests/client_auto_port_config_test.cpp | 143 ++++++++++++++++++ 7 files changed, 203 insertions(+), 18 deletions(-) create mode 100644 mooncake-store/include/client_auto_port_config.h create mode 100644 mooncake-store/src/config/client_auto_port_config.cpp create mode 100644 mooncake-store/tests/client_auto_port_config_test.cpp diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 552cd2b6bb..90739e9ec6 100644 --- a/mooncake-common/include/environment_variables.h +++ b/mooncake-common/include/environment_variables.h @@ -38,6 +38,12 @@ 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); +}; + #undef MC_DEFINE_ENV_VAR } // namespace mooncake 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/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index b539334ab8..4eb689df05 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -91,6 +91,7 @@ set(MOONCAKE_STORE_CLIENT_SOURCES shm_helper.cpp file_storage.cpp config/file_storage_config.cpp + config/client_auto_port_config.cpp device/accelerator_device.cpp device/accelerator_registry.cpp device/runtime_accelerator.cpp 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/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/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 8a7b4a1a22..ed3c035f97 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -144,6 +144,7 @@ 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(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) 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 From 4956149bcba2af5dcb99d51a0e7365dc07025f3e Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Tue, 1 Sep 2026 11:56:46 +0800 Subject: [PATCH 16/17] [Store] Remove redundant per-file deletion delay (#3758) * perf(store): remove redundant per-file deletion delay * test(store): make concurrent remove test deterministic --- mooncake-store/src/storage_backend.cpp | 11 +- mooncake-store/tests/storage_backend_test.cpp | 122 ++++++++++++++++++ 2 files changed, 125 insertions(+), 8 deletions(-) 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/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; From 9ab8e1b278a2dcb7b668bf1c1233f404679e61db Mon Sep 17 00:00:00 2001 From: Yuchen Kou Date: Tue, 1 Sep 2026 12:35:00 +0800 Subject: [PATCH 17/17] [Store] Enable fenced batch OpLog writer --- mooncake-store/include/master_service.h | 3 +- mooncake-store/src/master_service.cpp | 32 +++++++-- .../tests/ha/master_service_ha_test.cpp | 71 +++++++++++++++++++ 3 files changed, 100 insertions(+), 6 deletions(-) diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index f3c7589513..1cf01d93e0 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -2917,7 +2917,8 @@ class MasterService { std::string SerializeMetadataForOpLogFromReplicaDescriptors( const ObjectMetadata& metadata, const std::vector& replicas) const; - ErrorCode InitializeBatchOpLogWriter(std::shared_ptr backend); + ErrorCode InitializeBatchOpLogWriter(std::shared_ptr backend, + bool require_fenced_writer); tl::expected AppendOpLogVisibleBeforeDurable( OpType type, const std::string& tenant_id, const std::string& key, const std::string& payload); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 68f26b5ddd..69553eb6fd 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -456,7 +456,8 @@ MasterService::MasterService(const MasterServiceConfig& config) toString(connect_err))); } auto backend = std::make_shared(); - ErrorCode err = InitializeBatchOpLogWriter(std::move(backend)); + ErrorCode err = InitializeBatchOpLogWriter( + std::move(backend), /*require_fenced_writer=*/true); if (err != ErrorCode::OK) { throw std::runtime_error(fmt::format( "failed to create HA batch-record OpLog writer: {}", @@ -708,7 +709,9 @@ MasterService::~MasterService() { ErrorCode MasterService::SetBatchOpLogBackendForTesting( std::shared_ptr backend) { - return InitializeBatchOpLogWriter(std::move(backend)); + // Explicit test injection keeps the zero-view fixture API. A configured + // view still exercises the production fenced path. + return InitializeBatchOpLogWriter(std::move(backend), view_version_ > 0); } void MasterService::SetBatchOpLogWriterFactoryForTesting( @@ -13210,10 +13213,15 @@ std::string MasterService::SerializeMetadataForOpLogFromReplicaDescriptors( } ErrorCode MasterService::InitializeBatchOpLogWriter( - std::shared_ptr backend) { + std::shared_ptr backend, bool require_fenced_writer) { if (!backend || !backend->SupportsTxn()) { return ErrorCode::INVALID_PARAMS; } + if (require_fenced_writer && view_version_ == 0) { + LOG(ERROR) << "Fenced batch OpLog writer requires a non-zero acquired " + "producer view"; + return ErrorCode::INVALID_PARAMS; + } auto storage = std::make_unique(cluster_id_, *backend); DurablePrefix durable_prefix; @@ -13221,14 +13229,28 @@ ErrorCode MasterService::InitializeBatchOpLogWriter( if (err != ErrorCode::OK) { return err; } + if (require_fenced_writer) { + err = storage->ClaimProducerView(view_version_); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to claim producer view for batch OpLog " + << "writer: " << toString(err); + return err; + } + } OrderedOpLogWriterConfig writer_config; writer_config.max_entries_per_batch = oplog_batch_max_entries_; writer_config.initial_durable_prefix = durable_prefix; OpLogBatchStorage* storage_ptr = storage.get(); + const ViewVersionId producer_view_version = view_version_; OrderedOpLogWriter::WriteBatchFn write_batch = - [storage_ptr](const OpLogBatchRecord& batch, - const DurablePrefix& expected_prefix) { + [storage_ptr, require_fenced_writer, producer_view_version]( + const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + if (require_fenced_writer) { + return storage_ptr->WriteBatchAndAdvancePrefix( + batch, expected_prefix, producer_view_version); + } return storage_ptr->WriteBatchAndAdvancePrefix(batch, expected_prefix); }; diff --git a/mooncake-store/tests/ha/master_service_ha_test.cpp b/mooncake-store/tests/ha/master_service_ha_test.cpp index 3eaffb4bb1..8f925150bd 100644 --- a/mooncake-store/tests/ha/master_service_ha_test.cpp +++ b/mooncake-store/tests/ha/master_service_ha_test.cpp @@ -331,6 +331,13 @@ class MasterServiceHATest : public ::testing::Test { return service.batch_oplog_storage_ != nullptr; } + static std::optional + GetWriterTerminalStateForTesting(const MasterService& service) { + return service.ordered_oplog_writer_ + ? service.ordered_oplog_writer_->GetTerminalState() + : std::nullopt; + } + Segment MakeSegment(std::string name = "test_segment", size_t base = kDefaultSegmentBase, size_t size = kDefaultSegmentSize) const { @@ -1903,6 +1910,70 @@ TEST_F(MasterServiceHATest, OplogExplicitEnableCreatesWriter) { EXPECT_TRUE(HasBatchOpLogStorage(service)); } +TEST_F(MasterServiceHATest, FencedWriterClaimsConfiguredProducerView) { + constexpr ViewVersionId kProducerView = 7; + const std::string cluster_id = "fenced_writer_claim"; + auto backend = std::make_shared(); + auto config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_view_version(kProducerView) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + + MasterService service(config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + std::string producer_view; + ASSERT_EQ(ErrorCode::OK, + backend->Get(BuildProducerViewKey(cluster_id), producer_view)); + EXPECT_EQ(std::to_string(kProducerView), producer_view); + + ASSERT_TRUE(AppendVisibleForTesting(service, OpType::PUT_END, "default", + "fenced_writer_key", {}) + .has_value()); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildProducerViewKey(cluster_id), "8")); + ASSERT_TRUE(AppendVisibleForTesting(service, OpType::PUT_END, "default", + "stale_writer_key", {}) + .has_value()); + std::optional terminal_state; + for (int i = 0; i < 100; ++i) { + terminal_state = GetWriterTerminalStateForTesting(service); + if (terminal_state.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_TRUE(terminal_state.has_value()); + EXPECT_EQ(ErrorCode::ETCD_TRANSACTION_FAIL, terminal_state->error); + EXPECT_EQ(OrderedOpLogWriterTerminalReason::kFenced, + terminal_state->reason); +} + +TEST_F(MasterServiceHATest, FencedWriterRejectsContendedProducerViewClaim) { + const std::string cluster_id = "fenced_writer_contention"; + auto backend = std::make_shared(); + ASSERT_EQ(ErrorCode::OK, + backend->Put(BuildProducerViewKey(cluster_id), "8")); + auto config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_view_version(7) + .set_cluster_id(cluster_id) + .build(); + + MasterService service(config); + EXPECT_EQ(ErrorCode::ETCD_TRANSACTION_FAIL, + service.SetBatchOpLogBackendForTesting(backend)); + EXPECT_FALSE(HasOpLogWriter(service)); +} + TEST_F(MasterServiceHATest, OplogDoesNotStartWithUnsupportedHABackend) { auto config = MasterServiceConfig::builder() .set_enable_ha(true)