diff --git a/.github/workflows/all_test.yml b/.github/workflows/all_test.yml index 3f36d2d..f082c74 100644 --- a/.github/workflows/all_test.yml +++ b/.github/workflows/all_test.yml @@ -15,6 +15,31 @@ jobs: ci-2-virt-node: runs-on: ubuntu-latest steps: + - name: Reclaim runner disk before CI + run: | + set -euo pipefail + df -h / /mnt + + # ubuntu-latest includes large SDKs that this Rust/Python workflow never uses. + sudo rm -rf \ + /opt/ghc \ + /opt/hostedtoolcache/CodeQL \ + /usr/local/.ghcup \ + /usr/local/lib/android \ + /usr/share/dotnet \ + /usr/share/swift + + sudo apt-get clean + sudo rm -rf /var/lib/apt/lists/* + + if command -v docker >/dev/null 2>&1; then + sudo systemctl start docker + docker system prune --all --force --volumes + docker system df + fi + + df -h / /mnt + - name: Check out repository uses: actions/checkout@v6 with: @@ -34,6 +59,8 @@ jobs: protobuf-compiler \ python3-venv \ rsync + sudo apt-get clean + sudo rm -rf /var/lib/apt/lists/* - name: Verify Docker availability run: | @@ -97,9 +124,16 @@ jobs: "scale": "n1_kvowner_dram_20gib", "case_config": True, "scene_config": { - "kv_test_rounds": "p2p_only", + "kv_test_rounds": "p2p_only,p2p_only_ssd", }, }, + "ci_top_attention_cargo_kv_unit": { + "subject": "rust", + "runtime_contract": "rust_self_managed", + "scale": "n1_kvowner_dram_20gib", + "case_config": True, + "scene_config": {}, + }, "ci_top_attention_log_mgmt": { "subject": "rust", "runtime_contract": "rust_self_managed", @@ -181,7 +215,8 @@ jobs: # Scene selection: # - ci_top_attention_doc_page_build validates doc build through the prebuilt Docker image. - # - ci_top_attention_bin_kvtest keeps the Rust kv_test entry under the testbed scene contract. + # - ci_top_attention_bin_kvtest covers the memory and native SSD P2P rounds. + # - ci_top_attention_cargo_kv_unit runs the Rust KV crate unit tests. # - ci_top_attention_log_mgmt keeps log rolling/sharding coverage under the same CI testbed contract. # - ci_top_attention_ctrl_c_kv keeps runtime child-retirement Ctrl-C coverage in this workflow. # - ci_top_attention_ctrl_c_mq keeps MQ Ctrl-C exit coverage in this workflow. @@ -214,6 +249,7 @@ jobs: # - Keep the original per-scene scales from ci_test_list.yaml. # - ci_top_attention_doc_page_build stays on n1_kvowner_dram_3gib. # - ci_top_attention_bin_kvtest stays on n1_kvowner_dram_20gib. + # - ci_top_attention_cargo_kv_unit stays on n1_kvowner_dram_20gib. # - ci_top_attention_log_mgmt stays on n1_kvowner_dram_20gib. # - ci_top_attention_ctrl_c_kv stays on n1_kvowner_dram_3gib. # - ci_top_attention_ctrl_c_mq stays on n1_kvowner_dram_20gib. @@ -234,7 +270,9 @@ jobs: python3 fluxon_test_stack/ci_2_virt_node.py \ --suite-path "$RUNNER_TEMP/ci_test_list.ci.yaml" \ --print-generated \ - --testbed-hostworkdir "$RUNNER_TEMP/fluxon_deploy" + --testbed-hostworkdir "$RUNNER_TEMP/fluxon_deploy" \ + --cleanup-pack-runtime-after-success \ + --cleanup-successful-case-artifacts - name: Print ci_2_virt_node failure summary if: ${{ failure() }} @@ -282,7 +320,7 @@ jobs: PY - name: Normalize ci_2_virt_node debug artifact permissions - if: ${{ always() }} + if: ${{ failure() }} run: | chmod -R a+rX .dever || true chmod -R a+rX fluxon_release || true @@ -300,7 +338,7 @@ jobs: -exec chmod -R a+rX {} + || true - name: Upload ci_2_virt_node debug artifacts - if: ${{ always() }} + if: ${{ failure() }} uses: actions/upload-artifact@v4 with: name: ci-2-virt-node-debug-${{ github.sha }} @@ -313,3 +351,9 @@ jobs: .dever/**/pack_release_runtime/project-data/**/instances/**/logs/** .dever/**/pack_release_runtime/project-data/**/instances/**/release/** .dever/**/pack_release_runtime/project-data/**/assemblies/**/profile/** + + - name: Report final runner disk usage + if: ${{ always() }} + run: | + df -h / /mnt + docker system df || true diff --git a/AGENTS.md b/AGENTS.md index 1b1367b..634c5f6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,8 @@ Keep this document concise. - Core user, developer, and design docs are in-repo under fluxon_doc_cn/ and fluxon_doc_en/ - Detailed bilingual doc writing rules are indexed at `fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md` and `fluxon_doc_cn/dev_doc/开发者 - 3 - 文档写作规约.md` +- The bilingual technical copy-editing workflow and one-shot example are at `fluxon_doc_en/dev_doc/Developer - 5 - Technical Documentation Copy Editing.md` and `fluxon_doc_cn/dev_doc/开发者 - 5 - 技术文档审校.md` +- Bilingual code review rules, finding levels, and review templates are at `fluxon_doc_en/dev_doc/Developer - 6 - Code Review Guidelines.md` and `fluxon_doc_cn/dev_doc/开发者 - 6 - Code Review 规约.md` - teststack has two steps: start testbed and testrunner - teststack has UI support; testrunner should own the UI authority and API surface, and the UI should run as a long-lived service that reuses the ops interfaces underneath - All Python code in this project must be compatible with Python >=3.10 @@ -9,6 +11,9 @@ Keep this document concise. - Git operations are limited to basic `stage`, `unstage`, `commit`, and `push`. Do not use other Git operations. - Prefer contraction over compatibility by default. Do not add compatibility layers, deprecated paths, or aliases unless the task explicitly requires them. - Prefer one canonical name for one concept. Avoid synonym parameters, duplicated entrypoints, and parallel config surfaces. +- When a change crosses module boundaries or moves resource cleanup, identify one final-release owner. Keep public layers on contracts, composition layers on ordering, and internal modules on their own state and cleanup; fix the invariant at its owner without adding outer field mutation or a duplicate close path. +- When a change adds lifecycle dependencies or background work, write the dependency order and implement shutdown as admission stop, scoped wake or cancel, quiescence, dependent release, then dependency release. Keep independent branches parallel and avoid a global lock or coordinator unless the dependency graph requires it. +- When shutdown intent and cleanup completion can diverge, assign each state a scope and sole writer, keep transitions monotonic, make close repeatable, and treat successful close as a completion barrier. Do not add multiple sources of truth for the same lifecycle fact. - Do not use environment variables for ordinary parameter passing. Prefer configuration files first, then explicit command-line arguments. - Prefer convention over configuration. When one canonical path or default wiring is sufficient, do not add extra config knobs. - Minimize multi-path config delivery. Do not pass the same config through parallel channels such as env vars, CLI flags, and files at the same time. diff --git a/AGENTS_CN.md b/AGENTS_CN.md index 38f0a50..383a2e1 100644 --- a/AGENTS_CN.md +++ b/AGENTS_CN.md @@ -1,6 +1,8 @@ 保持本文档简洁。 - 核心用户文档、开发文档和设计文档都在仓库内的 `fluxon_doc_cn/` 和 `fluxon_doc_en/` 下 - 详细的中英文文档写作规约索引见 `fluxon_doc_cn/dev_doc/开发者 - 3 - 文档写作规约.md` 和 `fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md` +- 中英文技术文档审校流程和 one-shot 示例见 `fluxon_doc_cn/dev_doc/开发者 - 5 - 技术文档审校.md` 和 `fluxon_doc_en/dev_doc/Developer - 5 - Technical Documentation Copy Editing.md` +- 中英文 Code Review 规约、finding 等级和 Review 模板见 `fluxon_doc_cn/dev_doc/开发者 - 6 - Code Review 规约.md` 和 `fluxon_doc_en/dev_doc/Developer - 6 - Code Review Guidelines.md` - `teststack` 有两个步骤：`start testbed` 和 `testrunner` - `teststack` 支持 UI；`testrunner` 应负责 UI 的 authority 和 API surface，但 UI 应作为常驻服务运行，并复用下层的 ops 接口 - 本项目所有 Python 代码都必须兼容 Python `>= 3.10` @@ -9,6 +11,9 @@ - Git 操作仅限基础的 `stage`、`unstage`、`commit` 和 `push`。不要使用其他 Git 操作 - 默认优先收束而不是兼容。除非任务明确要求，否则不要添加兼容层、废弃路径或别名 - 一个概念优先只保留一个正式名字。避免同义参数、重复入口和并行配置面 +- 当改动跨越模块边界或移动资源清理职责时，必须指定唯一 final-release owner。公共层只依赖契约，组合层只安排顺序，内部模块维护自己的状态和 cleanup；在不变量 owner 处修复，不增加外层字段操作或第二条 close 路径 +- 当改动新增生命周期依赖或后台任务时，必须写明依赖顺序，并按“关闭入口、定向 wake/cancel、建立静默点、释放依赖方、最后释放被依赖方”执行关闭。互不依赖的分支保持并行，依赖图不需要时不引入全局锁或全局 coordinator +- 当停止意图和清理完成可能分离时，必须为每个状态指定作用域和唯一 writer，保持状态单调、close 可重复，并把 close 成功作为完成屏障。不要为同一个生命周期事实增加多个 source of truth - 不要用环境变量传递普通参数。优先使用配置文件，其次使用显式命令行参数 - 优先约定而不是配置。当一个标准路径或默认连接方式已经足够时，不要增加额外配置项 - 最小化多路径配置传递。不要同时通过环境变量、CLI 参数和文件等并行通道传递同一份配置 diff --git a/README.md b/README.md index de50828..1e85399 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,12 @@ The `TCP Benchmark` shows that Fluxon outperforms `MooncakeStore` and `Redis` on ![](./pics/kv_benchmark_chart.png) +Fluxon KV can also use an owner's local SSD as a runtime backing layer for DRAM replicas. The chart below shows a single-node H100 SSD-pressure experiment measured through CUDA event completion. At c16, Fluxon's hit payload throughput for 4/8/16 MiB payloads was 3.83×/5.61×/6.73× that of the faster of the two Mooncake topologies. + +![GPU-ready KV throughput and hit rate across Fluxon and Mooncake SSD configurations](./pics/kv_ssd_gpu_sweep.png) + +See the Chinese deep dive, [Fluxon KV SSD Storage: Using Local SSD as a Backing Layer for In-Memory Replicas](https://tele-ai.github.io/Fluxon/cn/blog/blog_2_Fluxon_KV_SSD%E5%AD%98%E5%82%A8%EF%BC%9A%E6%8A%8A%E6%9C%AC%E5%9C%B0SSD%E6%8E%A5%E6%88%90%E5%86%85%E5%AD%98%E5%89%AF%E6%9C%AC%E7%9A%84%E5%9B%9E%E5%A1%AB%E5%B1%82), for the full setup, hit rates, and scope. + ### Fluxon FS Benchmark The benchmark results show that small-file reads and large-file writes already outperform `Alluxio`, large-file read performance is broadly on par, and small-file write performance still has further room to improve. @@ -285,6 +291,9 @@ Contributions are welcome. Before you start, please read the developer docs on G - [Developer - 2 - Package middleware and images](https://tele-ai.github.io/Fluxon/dev_doc/Developer---2---Package-Middleware-and-Images/) - [Developer - 3 - Documentation Writing Rules](https://tele-ai.github.io/Fluxon/dev_doc/Developer---3---Documentation-Writing-Rules/) - [Developer - 4 - Publish a release](https://tele-ai.github.io/Fluxon/dev_doc/Developer---4---Publish-a-Release/) +- [Developer - 5 - Technical Documentation Copy Editing](https://tele-ai.github.io/Fluxon/dev_doc/Developer---5---Technical-Documentation-Copy-Editing/) +- [Developer - 6 - Code Review Guidelines](https://tele-ai.github.io/Fluxon/dev_doc/Developer---6---Code-Review-Guidelines/) +- [Developer - 7 - Event Subscription and Full Snapshot Guidelines](https://tele-ai.github.io/Fluxon/dev_doc/Developer---7---Event-Subscription-and-Full-Snapshot-Guidelines/) diff --git a/README_CN.md b/README_CN.md index 62d3d9e..4f96bd9 100644 --- a/README_CN.md +++ b/README_CN.md @@ -127,6 +127,12 @@ Fluxon FS 是一款面向 AI 数据与模型文件、兼容 `S3` 的高性能文 ![](./pics/kv_benchmark_chart.png) +Fluxon KV 还可将 owner 本地 SSD 作为 DRAM 副本的运行期回填层。下图为单机 H100 的 SSD-pressure 实验，以 CUDA event 完成为计数边界；c16 下，Fluxon 在 4/8/16 MiB 组中的 hit payload 吞吐相对两种 Mooncake 拓扑中更快的一种分别达到 3.83×/5.61×/6.73×。 + +![Fluxon 与 Mooncake 不同 SSD 配置下的 GPU 可用 KV 吞吐和命中率](./pics/kv_ssd_gpu_sweep.png) + +完整实验设置、命中率和适用边界见博客：[Fluxon KV SSD 存储：把本地 SSD 接成内存副本的回填层](https://tele-ai.github.io/Fluxon/cn/blog/blog_2_Fluxon_KV_SSD%E5%AD%98%E5%82%A8%EF%BC%9A%E6%8A%8A%E6%9C%AC%E5%9C%B0SSD%E6%8E%A5%E6%88%90%E5%86%85%E5%AD%98%E5%89%AF%E6%9C%AC%E7%9A%84%E5%9B%9E%E5%A1%AB%E5%B1%82)。 + ### Fluxon FS 基准测试 测试结果显示，小文件读取和大文件写入性能已显著优于 `Alluxio`，大文件读取性能基本持平，小文件写入性能仍有进一步优化的空间。 @@ -287,6 +293,9 @@ ui - [开发者 - 2 - 打包中间件和镜像](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---2---%E6%89%93%E5%8C%85%E4%B8%AD%E9%97%B4%E4%BB%B6%E5%92%8C%E9%95%9C%E5%83%8F/) - [开发者 - 3 - 文档写作规约](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---3---%E6%96%87%E6%A1%A3%E5%86%99%E4%BD%9C%E8%A7%84%E7%BA%A6/) - [开发者 - 4 - 发布 Release](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---4---%E5%8F%91%E5%B8%83-Release/) +- [开发者 - 5 - 技术文档审校](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---5---%E6%8A%80%E6%9C%AF%E6%96%87%E6%A1%A3%E5%AE%A1%E6%A0%A1/) +- [开发者 - 6 - Code Review 规约](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---6---Code-Review-%E8%A7%84%E7%BA%A6/) +- [开发者 - 7 - 事件订阅与全量快照规约](https://tele-ai.github.io/Fluxon/cn/dev_doc/%E5%BC%80%E5%8F%91%E8%80%85---7---%E4%BA%8B%E4%BB%B6%E8%AE%A2%E9%98%85%E4%B8%8E%E5%85%A8%E9%87%8F%E5%BF%AB%E7%85%A7%E8%A7%84%E7%BA%A6/) diff --git a/deployment/gen_k8s_daemonset.py b/deployment/gen_k8s_daemonset.py index 1499a85..cc0ad5d 100644 --- a/deployment/gen_k8s_daemonset.py +++ b/deployment/gen_k8s_daemonset.py @@ -349,7 +349,7 @@ def _mirror_daemonset_yaml_outdir(*, src_outdir: str, dst_outdir: str) -> None: # # Causal chain: # - Operators often run the generator from the repo workdir, while fluxon-deployer reads - # manifests from a hostworkdir-mounted directory (e.g. /opt/store_team_dev/fluxon_deployer). + # manifests from a hostworkdir-mounted directory (e.g. /tmp/fluxon-example/deploy). # - If the two directories drift, deployer applies stale manifests and can crash-loop. # - Therefore we provide an explicit opt-in "mirror_outdir" that: # - cleans the destination directory (same semantics as the primary outdir) diff --git a/deployment/tests/test_build_doc_site_contract.py b/deployment/tests/test_build_doc_site_contract.py index 7236d99..9418ee5 100644 --- a/deployment/tests/test_build_doc_site_contract.py +++ b/deployment/tests/test_build_doc_site_contract.py @@ -99,6 +99,16 @@ def test_rewrite_homepage_target_path_for_chinese_home(self) -> None: "../fluxon_rs/rust-toolchain.toml", ) + def test_rewrite_variant_repo_image_for_quartz_route(self) -> None: + for language in ("en", "cn"): + with self.subTest(language=language): + self.assertEqual( + _DOC_SITE.rewrite_variant_target_path( + "../../pics/kv_ssd_cpu_sweep.png", language=language + ), + "../pics/kv_ssd_cpu_sweep.png", + ) + def test_stage_cn_only_design_into_en_root_when_english_design_missing(self) -> None: tmpdir = Path(tempfile.mkdtemp(prefix="fluxon_doc_site_contract_")) old_stage_root = _DOC_SITE.STAGE_DOCS_ROOT diff --git a/deployment/tests/test_start_test_bed_bootstrap_log.py b/deployment/tests/test_start_test_bed_bootstrap_log.py index 2bd6b00..36eaf28 100644 --- a/deployment/tests/test_start_test_bed_bootstrap_log.py +++ b/deployment/tests/test_start_test_bed_bootstrap_log.py @@ -1021,11 +1021,11 @@ def test_bare_then_apply_success_path_does_not_run_post_apply_stop() -> None: gitops_config_path: ./gitops.yaml bootstrap_phases: - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - etcd - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - fluxon_core_controller deploy_workloads: @@ -1047,7 +1047,7 @@ def test_bare_then_apply_success_path_does_not_run_post_apply_stop() -> None: FLUXON_RELEASE_PYLIB_SRC_TAR: pylib_src.tar.gz FLUXON_RELEASE_WHEEL: fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl cluster_nodes: - - hostname: infra44-ThinkStation-PX + - hostname: example-node-a ip: 127.0.0.1 hostworkdir: /tmp/fluxon_testbed ssh_user: tester @@ -1056,28 +1056,28 @@ def test_bare_then_apply_success_path_does_not_run_post_apply_stop() -> None: atomic_groups: fluxon_core_controller: phase: 1 - nodes: ["infra44-ThinkStation-PX"] + nodes: ["example-node-a"] services: ["master", "owner", "ops_controller", "ops_agent"] bootstrap_bare_services: ["etcd"] service: etcd: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] master: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] owner: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] ops_controller: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] ops_agent: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] fluxon_fs_agent: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] """.strip() + "\n", encoding="utf-8", @@ -1252,11 +1252,11 @@ def test_bare_only_stops_after_controller_ready() -> None: gitops_config_path: ./gitops.yaml bootstrap_phases: - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - etcd - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - fluxon_core_controller - mode: coverage_bare @@ -1282,7 +1282,7 @@ def test_bare_only_stops_after_controller_ready() -> None: FLUXON_RELEASE_PYLIB_SRC_TAR: pylib_src.tar.gz FLUXON_RELEASE_WHEEL: fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl cluster_nodes: - - hostname: infra44-ThinkStation-PX + - hostname: example-node-a ip: 127.0.0.1 hostworkdir: /tmp/fluxon_testbed ssh_user: tester @@ -1291,31 +1291,31 @@ def test_bare_only_stops_after_controller_ready() -> None: atomic_groups: fluxon_core_controller: phase: 1 - nodes: ["infra44-ThinkStation-PX"] + nodes: ["example-node-a"] services: ["master", "owner", "ops_controller", "ops_agent"] bootstrap_bare_services: ["etcd"] service: etcd: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] greptime: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] master: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] owner: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] ops_controller: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] ops_agent: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] fluxon_fs_agent: node_bind: - node: ["infra44-ThinkStation-PX"] + node: ["example-node-a"] """.strip() + "\n", encoding="utf-8", @@ -1414,7 +1414,7 @@ def _fail_deploy_or_wait(**_: object) -> dict[str, str]: { "launches": [ { - "node": "infra44-ThinkStation-PX", + "node": "example-node-a", "selection_name": "etcd", } ] @@ -1422,7 +1422,7 @@ def _fail_deploy_or_wait(**_: object) -> dict[str, str]: { "launches": [ { - "node": "infra44-ThinkStation-PX", + "node": "example-node-a", "selection_name": "fluxon_core_controller", } ] @@ -1440,7 +1440,7 @@ def test_initial_controller_handover_uses_local_bare_stop() -> None: with tempfile.TemporaryDirectory(prefix="test_start_test_bed_controller_handover_stop_") as td: workdir = Path(td) local_node_cfg = { - "hostname": "infra44-ThinkStation-PX-a", + "hostname": "example-node-a-a", "ip": "127.0.0.1", "hostworkdir": str(workdir / "hostworkdir"), "execution_mode": "local", @@ -1452,15 +1452,15 @@ def test_initial_controller_handover_uses_local_bare_stop() -> None: "atomic_groups": { "fluxon_core_controller": { "phase": 1, - "nodes": ["infra44-ThinkStation-PX-a"], + "nodes": ["example-node-a-a"], "services": ["master", "owner", "ops_controller", "ops_agent"], } }, "service": { - "master": {"node_bind": {"node": ["infra44-ThinkStation-PX-a"]}}, - "owner": {"node_bind": {"node": ["infra44-ThinkStation-PX-a"]}}, - "ops_controller": {"node_bind": {"node": ["infra44-ThinkStation-PX-a"]}}, - "ops_agent": {"node_bind": {"node": ["infra44-ThinkStation-PX-a"]}}, + "master": {"node_bind": {"node": ["example-node-a-a"]}}, + "owner": {"node_bind": {"node": ["example-node-a-a"]}}, + "ops_controller": {"node_bind": {"node": ["example-node-a-a"]}}, + "ops_agent": {"node_bind": {"node": ["example-node-a-a"]}}, }, } stop_calls: list[str] = [] @@ -1480,10 +1480,10 @@ def test_initial_controller_handover_uses_local_bare_stop() -> None: module._run_local_stop = original_run_local_stop assert stopped == ["master", "owner", "ops_controller", "ops_agent"], stopped assert stop_calls == [ - "infra44-ThinkStation-PX-a:master", - "infra44-ThinkStation-PX-a:owner", - "infra44-ThinkStation-PX-a:ops_controller", - "infra44-ThinkStation-PX-a:ops_agent", + "example-node-a-a:master", + "example-node-a-a:owner", + "example-node-a-a:ops_controller", + "example-node-a-a:ops_agent", ], stop_calls print("PASS: test_initial_controller_handover_uses_local_bare_stop") diff --git a/examples/dataloader_video_benchmark/README.md b/examples/dataloader_video_benchmark/README.md new file mode 100644 index 0000000..19598e2 --- /dev/null +++ b/examples/dataloader_video_benchmark/README.md @@ -0,0 +1,109 @@ +# Dataloader Video Benchmark + +`dataloader_video_benchmark` compares the original `decord.VideoReader` +path with the FluxonFS video reader path at the dataloader boundary. + +It is an example benchmark, not a training dataset implementation. It replays +decode cases from CSV, timeout logs, or parquet metadata and reports decode +latency, QPS, output bandwidth, FluxonFS range-read counters, and reader-pool +stats. + +## Entrypoints + +- `benchmark.py` + - runs original and Fluxon video decode backends +- `analyze_results.py` + - summarizes benchmark CSV output +- `test_benchmark_contract.py` + - local contract test for the example script + +## Decode Case CSV + +The simplest input is a CSV with these columns: + +```csv +idx,video_path,height,width,start_idx,end_idx,num_frames +1,/abs/path/sample.mp4,480,832,0,63,16 +``` + +`video_path` must be an absolute local path for the original backend. For the +Fluxon backend, the same path must be under `--fluxon-remote-root`; the example +converts it to `export_name + relpath` before calling FluxonFS. + +## Original Baseline + +```bash +python3 examples/dataloader_video_benchmark/benchmark.py \ + --backend original \ + --case-csv /tmp/cases.csv \ + --rounds 100 \ + --warmup-rounds 5 \ + --workers 4 \ + --prefetch-factor 2 \ + --num-threads 2 \ + --output-csv /tmp/original.csv \ + --output-json /tmp/original.summary.json +``` + +The original backend runs: + +```text +decord.VideoReader(video_path).get_batch(indices).asnumpy() +``` + +## Fluxon Backend + +Start a FluxonFS master and agent first, then run: + +```bash +python3 examples/dataloader_video_benchmark/benchmark.py \ + --backend fluxon \ + --case-csv /tmp/cases.csv \ + --rounds 100 \ + --warmup-rounds 5 \ + --workers 4 \ + --prefetch-factor 2 \ + --decode-batch-size 3 \ + --num-threads 2 \ + --fluxon-kv-config /tmp/external.yaml \ + --fluxon-remote-root /tmp/video_root \ + --fluxon-agent-instance-key fs-agent-1 \ + --fluxon-reader-cache-size 32 \ + --fluxon-request-username bench \ + --fluxon-request-password bench \ + --output-csv /tmp/fluxon.csv \ + --output-json /tmp/fluxon.summary.json +``` + +`--decode-batch-size` controls the benchmark's sample-batch window. The Fluxon +backend passes that window to `FluxonFsVideoReaderPool.read_many_numpy_with_stats`. +Requests with the same video reader key are merged into one native decode call +and split back into per-sample rows. + +For the current synthetic benchmark shape, `--decode-batch-size 3` is the +recommended first value: it coalesces the three clips per video while keeping +enough scheduling granularity for multi-worker runs. + +## Analyze Results + +```bash +python3 examples/dataloader_video_benchmark/analyze_results.py \ + --input-csv /tmp/fluxon.csv \ + --output-json /tmp/fluxon.analysis.json \ + --output-md /tmp/fluxon.analysis.md \ + --title fluxon_batch3_w4 +``` + +The analyzer reports: + +- backend QPS, frames/s, and output MiB/s +- p50/p95/p99 decode latency +- pairwise original-vs-Fluxon speedup when both backends are present +- FluxonFS remote read bytes, range-read calls, page cache hit rate, and reader + pool hit rate + +## Local Contract Test + +```bash +python3 examples/dataloader_video_benchmark/test_benchmark_contract.py +``` diff --git a/examples/dataloader_video_benchmark/analyze_results.py b/examples/dataloader_video_benchmark/analyze_results.py new file mode 100644 index 0000000..486243b --- /dev/null +++ b/examples/dataloader_video_benchmark/analyze_results.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python3 +"""Analyze original-vs-Fluxon dataloader benchmark CSV results.""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional + + +LATENCY_FIELDS = ("elapsed_s", "open_s", "decode_s", "materialize_s") +FLUXON_VIDEO_COUNTER_FIELDS = ( + "fluxon_video_read_at_calls", + "fluxon_video_read_at_requested_bytes", + "fluxon_video_read_at_returned_bytes", + "fluxon_video_page_cache_hits", + "fluxon_video_page_cache_misses", + "fluxon_video_remote_read_calls", + "fluxon_video_remote_read_bytes", + "fluxon_video_decode_calls", + "fluxon_video_decode_frames_requested", + "fluxon_video_decode_errors", +) + + +def load_rows(path: Path) -> list[dict[str, str]]: + with path.open(newline="") as f: + return list(csv.DictReader(f)) + + +def fval(row: Mapping[str, str], key: str, default: float = 0.0) -> float: + raw = row.get(key, "") + if raw == "": + return default + try: + return float(raw) + except ValueError: + return default + + +def ival(row: Mapping[str, str], key: str, default: int = 0) -> int: + raw = row.get(key, "") + if raw == "": + return default + try: + return int(float(raw)) + except ValueError: + return default + + +def percentile(values: list[float], q: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + pos = (len(ordered) - 1) * q + lo = int(pos) + hi = min(lo + 1, len(ordered) - 1) + frac = pos - lo + return ordered[lo] * (1.0 - frac) + ordered[hi] * frac + + +def numeric_summary(values: Iterable[float]) -> dict[str, float]: + data = [float(v) for v in values] + if not data: + return { + "count": 0, + "min": 0.0, + "mean": 0.0, + "p50": 0.0, + "p95": 0.0, + "p99": 0.0, + "max": 0.0, + } + return { + "count": len(data), + "min": min(data), + "mean": statistics.fmean(data), + "p50": percentile(data, 0.50), + "p95": percentile(data, 0.95), + "p99": percentile(data, 0.99), + "max": max(data), + } + + +def analyze(rows: list[dict[str, str]]) -> dict[str, Any]: + out: dict[str, Any] = { + "total_rows": len(rows), + "backends": {}, + "pairwise": {}, + "fluxon_video": {}, + } + + backends = sorted({row.get("backend", "") for row in rows if row.get("backend", "")}) + for backend in backends: + backend_rows = [row for row in rows if row.get("backend") == backend] + ok_rows = [row for row in backend_rows if row.get("status") == "ok"] + status_counts = Counter(row.get("status", "unknown") for row in backend_rows) + frames = sum(ival(row, "num_frames") for row in ok_rows) + output_bytes = sum(ival(row, "nbytes") for row in ok_rows) + elapsed_sum = sum(fval(row, "elapsed_s") for row in ok_rows) + backend_summary: dict[str, Any] = { + "rows": len(backend_rows), + "ok": len(ok_rows), + "error": len(backend_rows) - len(ok_rows), + "status_counts": dict(sorted(status_counts.items())), + "frames": frames, + "output_bytes": output_bytes, + "elapsed_sum_s": elapsed_sum, + "qps": len(ok_rows) / elapsed_sum if elapsed_sum > 0 else 0.0, + "frames_per_s": frames / elapsed_sum if elapsed_sum > 0 else 0.0, + "output_mib_per_s": (output_bytes / (1024 * 1024)) / elapsed_sum if elapsed_sum > 0 else 0.0, + } + for field in LATENCY_FIELDS: + backend_summary[field] = numeric_summary(fval(row, field) for row in ok_rows) + backend_summary["batch_wall_s"] = numeric_summary(fval(row, "batch_wall_s") for row in ok_rows) + backend_summary["decode_batch_input_size"] = numeric_summary( + ival(row, "decode_batch_input_size") for row in ok_rows + ) + out["backends"][backend] = backend_summary + + out["pairwise"] = pairwise_analysis(rows) + out["fluxon_video"] = fluxon_video_analysis(rows) + return out + + +def pairwise_analysis(rows: list[dict[str, str]]) -> dict[str, Any]: + by_sample: dict[str, dict[str, dict[str, str]]] = defaultdict(dict) + for row in rows: + sample_id = row.get("sample_id", "") + backend = row.get("backend", "") + if sample_id and backend: + by_sample[sample_id][backend] = row + + comparable = [] + mismatches = [] + for sample_id, sample_rows in sorted(by_sample.items()): + original = sample_rows.get("original") + fluxon = sample_rows.get("fluxon") + if original is None or fluxon is None: + continue + if original.get("status") != "ok" or fluxon.get("status") != "ok": + continue + comparable.append((sample_id, original, fluxon)) + for field in ("shape", "dtype", "fingerprint"): + lhs = original.get(field, "") + rhs = fluxon.get(field, "") + if lhs and rhs and lhs != rhs: + mismatches.append( + { + "sample_id": sample_id, + "field": field, + "original": lhs, + "fluxon": rhs, + } + ) + + speedups = [] + deltas = [] + for _sample_id, original, fluxon in comparable: + original_elapsed = fval(original, "elapsed_s") + fluxon_elapsed = fval(fluxon, "elapsed_s") + if original_elapsed > 0 and fluxon_elapsed > 0: + speedups.append(original_elapsed / fluxon_elapsed) + deltas.append(original_elapsed - fluxon_elapsed) + + return { + "comparable_samples": len(comparable), + "speedup": numeric_summary(speedups), + "elapsed_delta_s": numeric_summary(deltas), + "mismatch_count": len(mismatches), + "mismatches": mismatches[:50], + } + + +def fluxon_video_analysis(rows: list[dict[str, str]]) -> dict[str, Any]: + fluxon_ok = [row for row in rows if row.get("backend") == "fluxon" and row.get("status") == "ok"] + counters = { + field: sum(ival(row, field) for row in fluxon_ok) + for field in FLUXON_VIDEO_COUNTER_FIELDS + } + page_hits = counters["fluxon_video_page_cache_hits"] + page_misses = counters["fluxon_video_page_cache_misses"] + page_lookups = page_hits + page_misses + remote_read_bytes = counters["fluxon_video_remote_read_bytes"] + output_bytes = sum(ival(row, "nbytes") for row in fluxon_ok) + elapsed_sum = sum(fval(row, "elapsed_s") for row in fluxon_ok) + + out: dict[str, Any] = { + "rows": len(fluxon_ok), + "counters": counters, + "reader_pool": fluxon_reader_pool_analysis(fluxon_ok), + "page_cache_hit_rate": page_hits / page_lookups if page_lookups > 0 else 0.0, + "remote_read_mib": remote_read_bytes / (1024 * 1024), + "remote_read_mib_per_s": (remote_read_bytes / (1024 * 1024)) / elapsed_sum if elapsed_sum > 0 else 0.0, + "compressed_to_output_byte_ratio": remote_read_bytes / output_bytes if output_bytes > 0 else 0.0, + "remote_read_bytes_per_row": numeric_summary( + ival(row, "fluxon_video_remote_read_bytes") for row in fluxon_ok + ), + "read_at_calls_per_row": numeric_summary( + ival(row, "fluxon_video_read_at_calls") for row in fluxon_ok + ), + } + return out + + +def fluxon_reader_pool_analysis(rows: list[dict[str, str]]) -> dict[str, Any]: + if not rows: + return { + "reader_cache_hit_rate": 0.0, + "reader_cache_hits": 0, + "reader_cache_misses": 0, + "open_count": 0, + "close_count": 0, + "evict_count": 0, + "max_current_readers": 0, + "max_active_readers": 0, + } + row_hits = sum(ival(row, "fluxon_video_pool_reader_cache_hit") for row in rows) + snapshot_hit_max = max(ival(row, "fluxon_video_pool_reader_cache_hits") for row in rows) + snapshot_miss_max = max(ival(row, "fluxon_video_pool_reader_cache_misses") for row in rows) + return { + "reader_cache_hit_rate": row_hits / len(rows), + "reader_cache_hits": snapshot_hit_max, + "reader_cache_misses": snapshot_miss_max, + "open_count": max(ival(row, "fluxon_video_pool_open_count") for row in rows), + "close_count": max(ival(row, "fluxon_video_pool_close_count") for row in rows), + "evict_count": max(ival(row, "fluxon_video_pool_evict_count") for row in rows), + "max_current_readers": max(ival(row, "fluxon_video_pool_current_readers") for row in rows), + "max_active_readers": max(ival(row, "fluxon_video_pool_active_readers") for row in rows), + } + + +def render_markdown(analysis: Mapping[str, Any], *, title: str) -> str: + lines = [f"# {title}", ""] + lines.append(f"- total_rows: {analysis.get('total_rows', 0)}") + + lines.extend(["", "## Backend Summary", ""]) + lines.append("| backend | rows | ok | error | p50 elapsed s | p95 elapsed s | qps | frames/s | output MiB/s |") + lines.append("| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |") + backends = analysis.get("backends", {}) + if isinstance(backends, Mapping): + for backend, data in backends.items(): + if not isinstance(data, Mapping): + continue + elapsed = data.get("elapsed_s", {}) + if not isinstance(elapsed, Mapping): + elapsed = {} + lines.append( + "| " + f"{backend} | {data.get('rows', 0)} | {data.get('ok', 0)} | {data.get('error', 0)} | " + f"{float(elapsed.get('p50', 0.0)):.6f} | {float(elapsed.get('p95', 0.0)):.6f} | " + f"{float(data.get('qps', 0.0)):.2f} | {float(data.get('frames_per_s', 0.0)):.2f} | " + f"{float(data.get('output_mib_per_s', 0.0)):.2f} |" + ) + + pairwise = analysis.get("pairwise", {}) + if isinstance(pairwise, Mapping): + speed = pairwise.get("speedup", {}) + delta = pairwise.get("elapsed_delta_s", {}) + lines.extend(["", "## Pairwise", ""]) + lines.append(f"- comparable_samples: {pairwise.get('comparable_samples', 0)}") + if isinstance(speed, Mapping): + lines.append( + "- speedup original/fluxon: " + f"p50={float(speed.get('p50', 0.0)):.3f}x, " + f"p95={float(speed.get('p95', 0.0)):.3f}x, " + f"mean={float(speed.get('mean', 0.0)):.3f}x" + ) + if isinstance(delta, Mapping): + lines.append( + "- elapsed_delta original-fluxon: " + f"p50={float(delta.get('p50', 0.0)):.6f}s, " + f"mean={float(delta.get('mean', 0.0)):.6f}s" + ) + lines.append(f"- mismatch_count: {pairwise.get('mismatch_count', 0)}") + + fluxon_video = analysis.get("fluxon_video", {}) + if isinstance(fluxon_video, Mapping): + counters = fluxon_video.get("counters", {}) + if not isinstance(counters, Mapping): + counters = {} + lines.extend(["", "## Fluxon VideoReader I/O", ""]) + lines.append(f"- rows: {fluxon_video.get('rows', 0)}") + lines.append(f"- page_cache_hit_rate: {float(fluxon_video.get('page_cache_hit_rate', 0.0)):.4f}") + lines.append(f"- remote_read_mib: {float(fluxon_video.get('remote_read_mib', 0.0)):.3f}") + lines.append(f"- remote_read_mib_per_s: {float(fluxon_video.get('remote_read_mib_per_s', 0.0)):.3f}") + lines.append( + "- compressed_to_output_byte_ratio: " + f"{float(fluxon_video.get('compressed_to_output_byte_ratio', 0.0)):.4f}" + ) + reader_pool = fluxon_video.get("reader_pool", {}) + if isinstance(reader_pool, Mapping): + lines.extend(["", "## Fluxon VideoReader Pool", ""]) + lines.append( + f"- reader_cache_hit_rate: {float(reader_pool.get('reader_cache_hit_rate', 0.0)):.4f}" + ) + lines.append(f"- reader_cache_hits: {reader_pool.get('reader_cache_hits', 0)}") + lines.append(f"- reader_cache_misses: {reader_pool.get('reader_cache_misses', 0)}") + lines.append(f"- open_count: {reader_pool.get('open_count', 0)}") + lines.append(f"- close_count: {reader_pool.get('close_count', 0)}") + lines.append(f"- evict_count: {reader_pool.get('evict_count', 0)}") + lines.append(f"- max_current_readers: {reader_pool.get('max_current_readers', 0)}") + lines.append(f"- max_active_readers: {reader_pool.get('max_active_readers', 0)}") + for field in FLUXON_VIDEO_COUNTER_FIELDS: + lines.append(f"- {field}: {counters.get(field, 0)}") + + lines.append("") + return "\n".join(lines) + + +def write_outputs( + analysis: Mapping[str, Any], + *, + output_json: Optional[Path], + output_md: Optional[Path], + title: str, +) -> None: + if output_json is not None: + output_json.parent.mkdir(parents=True, exist_ok=True) + output_json.write_text(json.dumps(analysis, indent=2, sort_keys=True), encoding="utf-8") + if output_md is not None: + output_md.parent.mkdir(parents=True, exist_ok=True) + output_md.write_text(render_markdown(analysis, title=title), encoding="utf-8") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input-csv", required=True) + parser.add_argument("--output-json", default="") + parser.add_argument("--output-md", default="") + parser.add_argument("--title", default="Dataloader VideoReader Benchmark Analysis") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + rows = load_rows(Path(args.input_csv)) + analysis = analyze(rows) + output_json = Path(args.output_json) if args.output_json else None + output_md = Path(args.output_md) if args.output_md else None + write_outputs( + analysis, + output_json=output_json, + output_md=output_md, + title=args.title, + ) + print(render_markdown(analysis, title=args.title), end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/dataloader_video_benchmark/benchmark.py b/examples/dataloader_video_benchmark/benchmark.py new file mode 100644 index 0000000..7d95103 --- /dev/null +++ b/examples/dataloader_video_benchmark/benchmark.py @@ -0,0 +1,1125 @@ +#!/usr/bin/env python3 +"""Compare original decord video loading with FluxonFS VideoReader. + +This is an integration benchmark entrypoint for the dataloader boundary. It +replays decode cases from CSV/log/parquet metadata and runs the same frame +selection through: + +* original: decord.VideoReader(video_path).get_batch(...).asnumpy() +* fluxon: FluxonFsPatcher.open_video_reader_pool(...).read_frames_numpy_with_stats(...) + +The Fluxon path uses export_name + relpath and does not materialize a local +file path for FFmpeg. It requires fluxon_pyo3 built with the +fluxon_fs_video_ffmpeg feature. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import hashlib +import json +import os +import random +import re +import statistics +import sys +import threading +import time +import traceback +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional + + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +DEFAULT_MANIFEST = "/mayukuo/flow-lite/runs/dualpointer/cached_metadata/parquets_list.txt" + +TIMEOUT_RE = re.compile( + r"Timeout loading data at index (?P\d+):.*?" + r"'(?P/[^']+)'\s*,\s*" + r"(?P\d+)\s*,\s*" + r"(?P\d+)\s*,\s*" + r"(?P\d+)\s*,\s*" + r"(?P\d+)\s*,\s*" + r"(?P\d+)", +) + + +class RawDefaultsHelpFormatter( + argparse.ArgumentDefaultsHelpFormatter, + argparse.RawDescriptionHelpFormatter, +): + pass + + +@dataclass(frozen=True) +class DecodeCase: + source: str + idx: int + video_path: str + height: int + width: int + start_idx: int + end_idx: int + num_frames: int + + @property + def key(self) -> tuple[Any, ...]: + return ( + self.video_path, + self.height, + self.width, + self.start_idx, + self.end_idx, + self.num_frames, + ) + + +@dataclass(frozen=True) +class BenchmarkJob: + sample_id: str + round_id: int + case_position: int + backend: str + case: DecodeCase + + +class OriginalDecordBackend: + name = "original" + + def __init__(self, *, num_threads: int, fingerprint_bytes: int) -> None: + from decord import VideoReader, cpu # type: ignore + + self._video_reader_cls = VideoReader + self._cpu = cpu + self._num_threads = num_threads + self._fingerprint_bytes = fingerprint_bytes + + def decode(self, case: DecodeCase) -> dict[str, Any]: + indices = frame_indices(case) + + t0 = time.perf_counter() + reader = self._video_reader_cls( + case.video_path, + width=case.width, + height=case.height, + num_threads=self._num_threads, + ctx=self._cpu(0), + ) + t1 = time.perf_counter() + batch = reader.get_batch(indices) + t2 = time.perf_counter() + array = batch.asnumpy() + t3 = time.perf_counter() + + row = { + "status": "ok", + "open_s": t1 - t0, + "decode_s": t2 - t1, + "materialize_s": t3 - t2, + "elapsed_s": t3 - t0, + } + row.update(array_result_fields(array, self._fingerprint_bytes)) + return row + + def decode_batch(self, cases: list[DecodeCase]) -> list[dict[str, Any]]: + return [self.decode(case) for case in cases] + + +class FluxonVideoBackend: + name = "fluxon" + + def __init__( + self, + *, + runtime: "FluxonBenchmarkRuntime", + num_threads: int, + fingerprint_bytes: int, + ) -> None: + self._runtime = runtime + self._num_threads = num_threads + self._fingerprint_bytes = fingerprint_bytes + + def decode(self, case: DecodeCase) -> dict[str, Any]: + return self.decode_batch([case])[0] + + def decode_batch(self, cases: list[DecodeCase]) -> list[dict[str, Any]]: + from fluxon_py.fluxon_fs.video import FluxonFsVideoReadRequest + + requests: list[FluxonFsVideoReadRequest] = [] + relpaths: list[str] = [] + frame_counts: list[int] = [] + for case in cases: + indices = tuple(frame_indices(case)) + relpath = self._runtime.relpath_for_case(case) + requests.append( + FluxonFsVideoReadRequest( + export_name=self._runtime.export_name, + relpath=relpath, + height=case.height, + width=case.width, + num_threads=self._num_threads, + indices=indices, + ) + ) + relpaths.append(relpath) + frame_counts.append(len(indices)) + + t0 = time.perf_counter() + results = self._runtime.video_reader_pool.read_many_numpy_with_stats(requests) + t1 = time.perf_counter() + batch_elapsed_s = t1 - t0 + weights = [max(1, int(getattr(result.array, "nbytes", 0))) for result in results] + elapsed_shares = split_float_by_weights(batch_elapsed_s, weights) + + rows: list[dict[str, Any]] = [] + for case, relpath, frame_count, result, elapsed_s in zip( + cases, + relpaths, + frame_counts, + results, + elapsed_shares, + ): + row = { + "status": "ok", + "fluxon_relpath": relpath, + "open_s": 0.0, + "decode_s": elapsed_s, + "materialize_s": 0.0, + "elapsed_s": elapsed_s, + "batch_wall_s": batch_elapsed_s, + "decode_batch_input_size": len(cases), + "decode_batch_input_frames": sum(frame_counts), + "decode_batch_sample_frames": frame_count, + "fluxon_video_pool_reader_cache_hit": int(result.reader_cache_hit), + } + row.update(array_result_fields(result.array, self._fingerprint_bytes)) + row.update({f"fluxon_video_{key}": value for key, value in result.reader_stats.items()}) + row.update({f"fluxon_video_pool_{key}": value for key, value in result.pool_stats.items()}) + row.update({f"fluxon_video_batch_{key}": value for key, value in result.batch_stats.items()}) + rows.append(row) + return rows + + +class FluxonBenchmarkRuntime: + """Owns the FluxonFS client used by the benchmark.""" + + def __init__( + self, + *, + patcher: Any, + stores: list[Any], + video_reader_pool: Any, + export_name: str, + remote_root_abs: Path, + close_timeout_s: float, + ) -> None: + self.patcher = patcher + self._stores = stores + self.video_reader_pool = video_reader_pool + self.export_name = export_name + self.remote_root_abs = remote_root_abs + self.close_timeout_s = close_timeout_s + + @classmethod + def open(cls, args: argparse.Namespace) -> "FluxonBenchmarkRuntime": + import yaml + from fluxon_py.config import FluxonKvClientConfig + from fluxon_py.fluxon_fs.patcher import FluxonFsPatcher + from fluxon_py.kvclient import new_store + + kv_cfg = load_yaml_mapping(args.fluxon_kv_config, "fluxon kv config") + suffix = f"{os.getpid()}_{int(time.time() * 1000)}" + client_key = args.fluxon_client_instance_key or f"dataloader_perf_client_{suffix}" + agent_node_id = resolve_fluxon_agent_node_id(args) + + client_store = new_fluxon_store( + new_store_fn=new_store, + config_cls=FluxonKvClientConfig, + base_config=kv_cfg, + instance_key=client_key, + ) + + remote_root_abs = require_abs_existing_dir( + args.fluxon_remote_root, + "fluxon remote root", + ) + cache_yaml = yaml.safe_dump( + { + "stale_window_ms": int(args.fluxon_stale_window_ms), + "rules": [], + "exports": { + args.fluxon_export_name: { + "remote_root_dir_abs": str(remote_root_abs), + "nodes": [str(agent_node_id)], + "cache_max_bytes": int(args.fluxon_cache_max_bytes), + "metadata_cache_ttl_ms": int(args.fluxon_metadata_cache_ttl_ms), + } + }, + }, + sort_keys=False, + ) + + patcher = FluxonFsPatcher(client_store) + patcher.set_cache_config_yaml(cache_yaml) + patcher.wait_cache_config_loaded() + if args.fluxon_request_username: + patcher.set_request_identity(args.fluxon_request_username, args.fluxon_request_password) + video_reader_pool = patcher.open_video_reader_pool( + max_readers=int(args.fluxon_reader_cache_size), + ) + + return cls( + patcher=patcher, + stores=[client_store], + video_reader_pool=video_reader_pool, + export_name=str(args.fluxon_export_name), + remote_root_abs=remote_root_abs, + close_timeout_s=float(args.fluxon_close_timeout_s), + ) + + def relpath_for_case(self, case: DecodeCase) -> str: + video_path = Path(case.video_path) + if not video_path.is_absolute(): + raise ValueError(f"Fluxon backend requires absolute video_path, got {case.video_path!r}") + video_abs = video_path.resolve(strict=False) + try: + rel = video_abs.relative_to(self.remote_root_abs) + except ValueError as exc: + raise ValueError( + f"video_path is outside fluxon remote root: path={video_abs} root={self.remote_root_abs}" + ) from exc + rel_s = rel.as_posix() + if not rel_s or rel_s == ".": + raise ValueError(f"invalid empty Fluxon relpath for {video_abs}") + return rel_s + + def close(self) -> None: + self.video_reader_pool.close() + for store in self._stores: + close_store_with_timeout(store, timeout_s=self.close_timeout_s) + + def __enter__(self) -> "FluxonBenchmarkRuntime": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.close() + + +def unwrap_result(result: Any, ctx: str) -> Any: + if result.is_ok(): + return result.unwrap() + raise RuntimeError(f"{ctx} failed: {result.unwrap_error()}") + + +def resolve_fluxon_agent_node_id(args: argparse.Namespace) -> str: + agent_instance_key = str(args.fluxon_agent_instance_key).strip() + agent_node_id = str(args.fluxon_agent_node_id).strip() + if agent_instance_key and agent_node_id and agent_instance_key != agent_node_id: + raise ValueError( + "--fluxon-agent-instance-key and --fluxon-agent-node-id must match when both are set" + ) + node_id = agent_instance_key or agent_node_id + if not node_id: + raise ValueError( + "--fluxon-agent-instance-key is required for the Fluxon backend; " + "start fluxon_py.runtime.start_fs_agent with the same instance_key first" + ) + return node_id + + +def close_store_with_timeout(store: Any, *, timeout_s: float) -> None: + result_holder: dict[str, Any] = {} + + def close_worker() -> None: + try: + close_res = store.close() + if close_res.is_ok(): + result_holder["ok"] = close_res.unwrap() + else: + result_holder["error"] = close_res.unwrap_error() + except BaseException as exc: + result_holder["exception"] = exc + result_holder["traceback"] = traceback.format_exc(limit=6) + + worker = threading.Thread(target=close_worker, daemon=True) + worker.start() + worker.join(max(0.0, float(timeout_s))) + if worker.is_alive(): + print( + f"[dataloader_perf][WARNING] store.close exceeded {timeout_s:.3f}s; continuing", + flush=True, + ) + return + if "error" in result_holder: + print(f"[dataloader_perf][WARNING] store.close failed: {result_holder['error']}", flush=True) + if "exception" in result_holder: + print( + "[dataloader_perf][WARNING] store.close raised: " + f"{type(result_holder['exception']).__name__}: {result_holder['exception']}", + flush=True, + ) + + +def new_fluxon_store( + *, + new_store_fn: Any, + config_cls: Any, + base_config: Mapping[str, Any], + instance_key: str, +) -> Any: + cfg = copy.deepcopy(dict(base_config)) + cfg["instance_key"] = instance_key + result = new_store_fn(config_cls(cfg)) + return unwrap_result(result, f"new_store({instance_key})") + + +def frame_indices(case: DecodeCase) -> list[int]: + if case.num_frames <= 0: + raise ValueError(f"num_frames must be positive, got {case.num_frames}") + if case.num_frames == 1: + return [int(case.start_idx)] + span = case.end_idx - case.start_idx + return [ + int(case.start_idx + (span * i) / (case.num_frames - 1)) + for i in range(case.num_frames) + ] + + +def split_float_by_weights(value: float, weights: list[int]) -> list[float]: + if not weights: + return [] + positive_weights = [max(0, int(weight)) for weight in weights] + total_weight = sum(positive_weights) + if total_weight <= 0: + return [value / len(weights) for _ in weights] + shares = [(value * weight) / total_weight for weight in positive_weights] + if shares: + shares[-1] += value - sum(shares) + return shares + + +def array_result_fields(array: Any, fingerprint_bytes: int) -> dict[str, Any]: + shape = tuple(int(x) for x in getattr(array, "shape", ())) + dtype = str(getattr(array, "dtype", "unknown")) + nbytes = int(getattr(array, "nbytes", 0)) + out = { + "shape": "x".join(str(x) for x in shape), + "dtype": dtype, + "nbytes": nbytes, + } + if fingerprint_bytes > 0: + out["fingerprint"] = array_fingerprint(array, fingerprint_bytes) + return out + + +def array_fingerprint(array: Any, fingerprint_bytes: int) -> str: + if fingerprint_bytes <= 0: + return "" + data = array.tobytes(order="C") + sample_n = min(len(data), fingerprint_bytes) + h = hashlib.blake2b(digest_size=16) + h.update(str(getattr(array, "shape", "")).encode("utf-8")) + h.update(str(getattr(array, "dtype", "")).encode("utf-8")) + h.update(data[:sample_n]) + if len(data) > sample_n: + h.update(data[-sample_n:]) + return h.hexdigest() + + +def parse_case_csv(path: str) -> list[DecodeCase]: + cases: list[DecodeCase] = [] + csv_path = Path(path) + if not csv_path.exists(): + raise FileNotFoundError(path) + with csv_path.open(newline="") as f: + for row in csv.DictReader(f): + video_path = row.get("video_path") or row.get("path") + if not video_path: + continue + cases.append( + DecodeCase( + source=row.get("source") or f"csv:{path}", + idx=int(row.get("idx") or len(cases)), + video_path=video_path, + height=int(row["height"]), + width=int(row["width"]), + start_idx=int(row["start_idx"]), + end_idx=int(row["end_idx"]), + num_frames=int(row["num_frames"]), + ) + ) + return cases + + +def parse_timeout_log(path: str) -> list[DecodeCase]: + cases: list[DecodeCase] = [] + log_path = Path(path) + if not log_path.exists(): + raise FileNotFoundError(path) + for line in log_path.read_text(errors="replace").splitlines(): + match = TIMEOUT_RE.search(line) + if match is None: + continue + cases.append( + DecodeCase( + source=f"log:{path}", + idx=int(match.group("idx")), + video_path=match.group("path"), + height=int(match.group("height")), + width=int(match.group("width")), + start_idx=int(match.group("start")), + end_idx=int(match.group("end")), + num_frames=int(match.group("num")), + ) + ) + return cases + + +def read_manifest(path: str) -> list[str]: + manifest = Path(path) + if not manifest.exists(): + return [] + return [ + line.strip() + for line in manifest.read_text().splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + +def case_from_parquet_row(dataset: Any, idx: int, source: str) -> DecodeCase: + row = dataset.take([idx]).to_pydict() + video_path = row["video_path"][0] + height = int(row["height"][0]) + width = int(row["width"][0]) + caption_start = int(row["caption_start"][0]) + caption_end = int(row["caption_end"][0]) + real_num_frames = int(row["real_num_frames"][0]) + dst_num_frames = int(row["dst_num_frames"][0]) + + span = min(real_num_frames, caption_end - caption_start) + max_start = caption_end - span + if max_start <= caption_start: + window_start = caption_start + else: + window_start = random.randint(caption_start, max_start) + window_end = window_start + span - 1 + + return DecodeCase( + source=source, + idx=idx, + video_path=video_path, + height=height, + width=width, + start_idx=window_start, + end_idx=window_end, + num_frames=dst_num_frames, + ) + + +def sample_parquet_cases( + parquet_paths: list[str], + *, + count: int, + seed: int, +) -> list[DecodeCase]: + if count <= 0: + return [] + if not parquet_paths: + raise FileNotFoundError("No parquet paths provided for random sampling.") + + import pyarrow.dataset as ds # type: ignore + + random.seed(seed) + dataset = ds.dataset(parquet_paths, format="parquet") + total_rows = dataset.count_rows() + indices = [random.randrange(total_rows) for _ in range(count)] + return [case_from_parquet_row(dataset, idx, "parquet-random") for idx in indices] + + +def build_cases(args: argparse.Namespace) -> list[DecodeCase]: + cases: list[DecodeCase] = [] + for csv_path in args.case_csv: + parsed = parse_case_csv(csv_path) + print(f"Loaded {len(parsed)} decode cases from {csv_path}", flush=True) + cases.extend(parsed) + for log_path in args.case_log: + parsed = parse_timeout_log(log_path) + print(f"Loaded {len(parsed)} decode cases from {log_path}", flush=True) + cases.extend(parsed) + + parquet_paths: list[str] = list(args.parquet) + if not parquet_paths and args.manifest: + parquet_paths = read_manifest(args.manifest) + if args.random_cases > 0: + parsed = sample_parquet_cases( + parquet_paths, + count=args.random_cases, + seed=args.seed, + ) + print(f"Sampled {len(parsed)} decode cases from parquet metadata", flush=True) + cases.extend(parsed) + + if args.dedupe: + unique: dict[tuple[Any, ...], DecodeCase] = {} + for case in cases: + unique.setdefault(case.key, case) + cases = list(unique.values()) + + if args.shuffle: + random.Random(args.seed).shuffle(cases) + if args.limit > 0: + cases = cases[: args.limit] + + validate_cases(cases) + return cases + + +def validate_cases(cases: Iterable[DecodeCase]) -> None: + for case in cases: + if case.height <= 0 or case.width <= 0: + raise ValueError(f"case idx={case.idx} has invalid size {case.height}x{case.width}") + if case.num_frames <= 0: + raise ValueError(f"case idx={case.idx} has invalid num_frames={case.num_frames}") + if not case.video_path: + raise ValueError(f"case idx={case.idx} has empty video_path") + + +def selected_backends(args: argparse.Namespace) -> list[str]: + if args.backend == "both": + return ["original", "fluxon"] + return [args.backend] + + +def build_jobs( + cases: list[DecodeCase], + *, + backends: list[str], + rounds: int, + alternate_backend_order: bool, +) -> list[BenchmarkJob]: + jobs: list[BenchmarkJob] = [] + for round_id in range(1, rounds + 1): + for case_pos, case in enumerate(cases, start=1): + ordered = list(backends) + if alternate_backend_order and len(ordered) == 2 and (round_id + case_pos) % 2 == 1: + ordered.reverse() + sample_id = f"r{round_id:04d}_c{case_pos:06d}" + for backend in ordered: + jobs.append( + BenchmarkJob( + sample_id=sample_id, + round_id=round_id, + case_position=case_pos, + backend=backend, + case=case, + ) + ) + return jobs + + +def build_job_row(job: BenchmarkJob, started: str) -> dict[str, Any]: + return { + "sample_id": job.sample_id, + "round_id": job.round_id, + "case_position": job.case_position, + "backend": job.backend, + "started": started, + "source": job.case.source, + "idx": job.case.idx, + "video_path": job.case.video_path, + "height": job.case.height, + "width": job.case.width, + "start_idx": job.case.start_idx, + "end_idx": job.case.end_idx, + "num_frames": job.case.num_frames, + } + + +def run_job(job: BenchmarkJob, backend: Any) -> dict[str, Any]: + started = time.strftime("%Y-%m-%d %H:%M:%S") + t0 = time.perf_counter() + row = build_job_row(job, started) + try: + row.update(backend.decode(job.case)) + except BaseException as exc: + row.update( + { + "status": "error", + "elapsed_s": time.perf_counter() - t0, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(limit=8), + } + ) + return row + + +def run_backend_job_batch(jobs: list[BenchmarkJob], backend: Any) -> list[dict[str, Any]]: + started = time.strftime("%Y-%m-%d %H:%M:%S") + rows = [build_job_row(job, started) for job in jobs] + t0 = time.perf_counter() + try: + result_rows = backend.decode_batch([job.case for job in jobs]) + if len(result_rows) != len(rows): + raise RuntimeError( + "backend returned unexpected batch result count: " + f"expected={len(rows)} actual={len(result_rows)}" + ) + for row, result in zip(rows, result_rows): + row.update(result) + except BaseException as exc: + elapsed_s = time.perf_counter() - t0 + elapsed_shares = split_float_by_weights(elapsed_s, [1 for _ in rows]) + for row, elapsed_share in zip(rows, elapsed_shares): + row.update( + { + "status": "error", + "elapsed_s": elapsed_share, + "batch_wall_s": elapsed_s, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(limit=8), + } + ) + return rows + + +def run_job_window(jobs: list[BenchmarkJob], backends: Mapping[str, Any]) -> list[dict[str, Any]]: + grouped: dict[str, list[BenchmarkJob]] = {} + for job in jobs: + grouped.setdefault(job.backend, []).append(job) + + rows: list[dict[str, Any]] = [] + for backend_name, backend_jobs in grouped.items(): + rows.extend(run_backend_job_batch(backend_jobs, backends[backend_name])) + return rows + + +def build_job_windows(jobs: list[BenchmarkJob], decode_batch_size: int) -> list[list[BenchmarkJob]]: + if decode_batch_size <= 1: + return [[job] for job in jobs] + return [ + jobs[start : start + decode_batch_size] + for start in range(0, len(jobs), decode_batch_size) + ] + + +def run_jobs( + jobs: list[BenchmarkJob], + *, + backends: Mapping[str, Any], + workers: int, + prefetch_factor: int, + decode_batch_size: int, + label: str, +) -> tuple[list[dict[str, Any]], float]: + rows: list[dict[str, Any]] = [] + started = time.perf_counter() + if not jobs: + return rows, 0.0 + + job_windows = build_job_windows(jobs, decode_batch_size) + print( + f"{label}: jobs={len(jobs)} windows={len(job_windows)} workers={workers} " + f"prefetch_factor={prefetch_factor} decode_batch_size={decode_batch_size}", + flush=True, + ) + if workers <= 1: + for window in job_windows: + for row in run_job_window(window, backends): + rows.append(row) + print_progress(label, len(rows), len(jobs), row) + return rows, time.perf_counter() - started + + max_inflight = max(workers, workers * prefetch_factor) + window_iter = iter(job_windows) + with ThreadPoolExecutor(max_workers=workers) as executor: + pending: dict[Any, list[BenchmarkJob]] = {} + + def submit_next() -> bool: + try: + window = next(window_iter) + except StopIteration: + return False + pending[executor.submit(run_job_window, window, backends)] = window + return True + + for _ in range(min(len(job_windows), max_inflight)): + submit_next() + + while pending: + completed, _ = wait(pending.keys(), return_when=FIRST_COMPLETED) + for future in completed: + pending.pop(future) + for row in future.result(): + rows.append(row) + print_progress(label, len(rows), len(jobs), row) + submit_next() + + rows.sort(key=lambda r: (int(r["round_id"]), int(r["case_position"]), str(r["backend"]))) + return rows, time.perf_counter() - started + + +def print_progress(label: str, done: int, total: int, row: Mapping[str, Any]) -> None: + print( + f"[{label} {done}/{total}] " + f"backend={row.get('backend')} status={row.get('status')} " + f"elapsed={float(row.get('elapsed_s', 0.0)):.3f}s " + f"idx={row.get('idx')} path={row.get('video_path')}", + flush=True, + ) + + +def summarize(rows: list[dict[str, Any]], *, wall_s: float) -> dict[str, Any]: + summary: dict[str, Any] = { + "wall_s": wall_s, + "total_rows": len(rows), + "backends": {}, + } + + for backend in sorted({str(r.get("backend")) for r in rows}): + backend_rows = [r for r in rows if r.get("backend") == backend] + ok_rows = [r for r in backend_rows if r.get("status") == "ok"] + elapsed = [float(r.get("elapsed_s", 0.0)) for r in ok_rows] + frames = sum(int(r.get("num_frames", 0)) for r in ok_rows) + backend_summary = { + "rows": len(backend_rows), + "ok": len(ok_rows), + "error": len(backend_rows) - len(ok_rows), + "frames": frames, + } + if elapsed: + elapsed_sum = sum(elapsed) + backend_summary.update( + { + "elapsed_sum_s": elapsed_sum, + "decode_rows_per_s": len(ok_rows) / elapsed_sum if elapsed_sum > 0 else 0.0, + "elapsed_min_s": min(elapsed), + "elapsed_mean_s": statistics.fmean(elapsed), + "elapsed_p50_s": percentile(elapsed, 0.50), + "elapsed_p95_s": percentile(elapsed, 0.95), + "elapsed_p99_s": percentile(elapsed, 0.99), + "elapsed_max_s": max(elapsed), + "decode_frames_per_s": frames / elapsed_sum if elapsed_sum > 0 else 0.0, + } + ) + summary["backends"][backend] = backend_summary + + if {"original", "fluxon"}.issubset(summary["backends"].keys()): + original_elapsed = [ + float(r["elapsed_s"]) + for r in rows + if r.get("backend") == "original" and r.get("status") == "ok" + ] + fluxon_elapsed = [ + float(r["elapsed_s"]) + for r in rows + if r.get("backend") == "fluxon" and r.get("status") == "ok" + ] + if original_elapsed and fluxon_elapsed: + original_p50 = percentile(original_elapsed, 0.50) + fluxon_p50 = percentile(fluxon_elapsed, 0.50) + summary["fluxon_vs_original"] = { + "p50_speedup": original_p50 / fluxon_p50 if fluxon_p50 > 0 else 0.0, + "original_p50_s": original_p50, + "fluxon_p50_s": fluxon_p50, + } + return summary + + +def percentile(values: list[float], q: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + pos = (len(ordered) - 1) * q + lo = int(pos) + hi = min(lo + 1, len(ordered) - 1) + frac = pos - lo + return ordered[lo] * (1.0 - frac) + ordered[hi] * frac + + +def print_summary(summary: Mapping[str, Any]) -> None: + print("Summary:", flush=True) + print(f" wall_s: {float(summary.get('wall_s', 0.0)):.3f}", flush=True) + backends = summary.get("backends", {}) + if isinstance(backends, Mapping): + for backend, data in backends.items(): + if not isinstance(data, Mapping): + continue + print( + f" {backend}: rows={data.get('rows')} ok={data.get('ok')} error={data.get('error')}", + flush=True, + ) + if "elapsed_p50_s" in data: + print( + " elapsed_s: " + f"mean={float(data['elapsed_mean_s']):.3f}, " + f"p50={float(data['elapsed_p50_s']):.3f}, " + f"p95={float(data['elapsed_p95_s']):.3f}, " + f"max={float(data['elapsed_max_s']):.3f}, " + f"qps={float(data.get('decode_rows_per_s', 0.0)):.2f}, " + f"frames/s={float(data['decode_frames_per_s']):.2f}", + flush=True, + ) + speed = summary.get("fluxon_vs_original") + if isinstance(speed, Mapping): + print( + " fluxon_vs_original: " + f"p50_speedup={float(speed.get('p50_speedup', 0.0)):.3f}x " + f"(original={float(speed.get('original_p50_s', 0.0)):.3f}s, " + f"fluxon={float(speed.get('fluxon_p50_s', 0.0)):.3f}s)", + flush=True, + ) + + +def write_rows_csv(path: str, rows: list[dict[str, Any]]) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + fieldnames = sorted({key for row in rows for key in row.keys()}) + with output.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def write_json(path: str, payload: Mapping[str, Any]) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + + +def load_yaml_mapping(path: str, name: str) -> dict[str, Any]: + import yaml + + p = Path(path) + if not p.exists(): + raise FileNotFoundError(path) + obj = yaml.safe_load(p.read_text(encoding="utf-8")) + if not isinstance(obj, dict): + raise ValueError(f"{name} must be a YAML mapping") + return obj + + +def require_abs_existing_dir(path: str, name: str) -> Path: + p = Path(path) + if not p.is_absolute(): + raise ValueError(f"{name} must be an absolute path") + resolved = p.resolve() + if not resolved.is_dir(): + raise ValueError(f"{name} must exist and be a directory: {resolved}") + return resolved + + +def validate_args(args: argparse.Namespace) -> None: + if args.rounds <= 0: + raise ValueError("--rounds must be positive") + if args.warmup_rounds < 0: + raise ValueError("--warmup-rounds must be non-negative") + if args.workers <= 0: + raise ValueError("--workers must be positive") + if args.prefetch_factor <= 0: + raise ValueError("--prefetch-factor must be positive") + if args.decode_batch_size <= 0: + raise ValueError("--decode-batch-size must be positive") + if args.num_threads <= 0: + raise ValueError("--num-threads must be positive") + if args.fingerprint_bytes < 0: + raise ValueError("--fingerprint-bytes must be non-negative") + if args.backend in ("fluxon", "both"): + if args.fluxon_reader_cache_size <= 0: + raise ValueError("--fluxon-reader-cache-size must be positive") + if not args.fluxon_kv_config: + raise ValueError("--fluxon-kv-config is required for the Fluxon backend") + if not args.fluxon_remote_root: + raise ValueError("--fluxon-remote-root is required for the Fluxon backend") + if bool(args.fluxon_request_username) != bool(args.fluxon_request_password): + raise ValueError( + "--fluxon-request-username and --fluxon-request-password must be provided together" + ) + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=RawDefaultsHelpFormatter, + ) + + input_group = parser.add_argument_group("decode cases") + input_group.add_argument("--case-csv", action="append", default=[]) + input_group.add_argument("--case-log", action="append", default=[]) + input_group.add_argument("--manifest", default=DEFAULT_MANIFEST) + input_group.add_argument("--parquet", action="append", default=[]) + input_group.add_argument("--random-cases", type=int, default=0) + input_group.add_argument("--limit", type=int, default=0) + input_group.add_argument("--shuffle", action="store_true") + input_group.add_argument("--no-dedupe", dest="dedupe", action="store_false") + input_group.set_defaults(dedupe=True) + + bench_group = parser.add_argument_group("benchmark") + bench_group.add_argument("--backend", choices=["original", "fluxon", "both"], default="both") + bench_group.add_argument("--rounds", type=int, default=1) + bench_group.add_argument("--warmup-rounds", type=int, default=0) + bench_group.add_argument("--workers", type=int, default=1) + bench_group.add_argument("--prefetch-factor", type=int, default=2) + bench_group.add_argument("--decode-batch-size", type=int, default=1) + bench_group.add_argument("--num-threads", type=int, default=8) + bench_group.add_argument("--seed", type=int, default=1234) + bench_group.add_argument("--fingerprint-bytes", type=int, default=0) + bench_group.add_argument("--no-alternate-backend-order", dest="alternate_backend_order", action="store_false") + bench_group.add_argument("--allow-errors", action="store_true") + bench_group.set_defaults(alternate_backend_order=True) + + fluxon_group = parser.add_argument_group("fluxon backend") + fluxon_group.add_argument("--fluxon-kv-config", default="") + fluxon_group.add_argument("--fluxon-remote-root", default="") + fluxon_group.add_argument("--fluxon-export-name", default="dataloader-videos") + fluxon_group.add_argument( + "--fluxon-agent-instance-key", + default="", + help="External fluxon_py.runtime.start_fs_agent instance_key used as the static export node.", + ) + fluxon_group.add_argument("--fluxon-client-instance-key", default="") + fluxon_group.add_argument( + "--fluxon-agent-node-id", + default="", + help="Optional explicit node id; must match --fluxon-agent-instance-key when both are set.", + ) + fluxon_group.add_argument("--fluxon-cache-max-bytes", type=int, default=1 << 40) + fluxon_group.add_argument("--fluxon-reader-cache-size", type=int, default=32) + fluxon_group.add_argument("--fluxon-metadata-cache-ttl-ms", type=int, default=1000) + fluxon_group.add_argument("--fluxon-stale-window-ms", type=int, default=1000) + fluxon_group.add_argument("--fluxon-close-timeout-s", type=float, default=10.0) + fluxon_group.add_argument( + "--fluxon-request-username", + default="", + help="FluxonFS username used to sign file RPC tokens.", + ) + fluxon_group.add_argument( + "--fluxon-request-password", + default="", + help="FluxonFS password used to sign file RPC tokens.", + ) + + output_group = parser.add_argument_group("output") + output_group.add_argument( + "--output-csv", + default="runs/dualpointer/dataloader_video_benchmark/results.csv", + ) + output_group.add_argument( + "--output-json", + default="runs/dualpointer/dataloader_video_benchmark/summary.json", + ) + return parser + + +def main(argv: Optional[list[str]] = None) -> int: + parser = build_arg_parser() + args = parser.parse_args(argv) + validate_args(args) + + cases = build_cases(args) + if not cases: + raise SystemExit( + "No decode cases found. Pass --case-csv, --case-log, or --random-cases with parquet metadata." + ) + + backends_to_run = selected_backends(args) + print( + "Dataloader decode benchmark: " + f"cases={len(cases)} backend={args.backend} rounds={args.rounds} " + f"warmup_rounds={args.warmup_rounds} workers={args.workers} " + f"prefetch_factor={args.prefetch_factor} decode_batch_size={args.decode_batch_size} " + f"num_threads={args.num_threads}", + flush=True, + ) + + fluxon_runtime: Optional[FluxonBenchmarkRuntime] = None + try: + if "fluxon" in backends_to_run: + fluxon_runtime = FluxonBenchmarkRuntime.open(args) + + backend_objs: dict[str, Any] = {} + if "original" in backends_to_run: + backend_objs["original"] = OriginalDecordBackend( + num_threads=args.num_threads, + fingerprint_bytes=args.fingerprint_bytes, + ) + if "fluxon" in backends_to_run: + if fluxon_runtime is None: + raise RuntimeError("Fluxon runtime was not initialized") + backend_objs["fluxon"] = FluxonVideoBackend( + runtime=fluxon_runtime, + num_threads=args.num_threads, + fingerprint_bytes=args.fingerprint_bytes, + ) + + warmup_jobs = build_jobs( + cases, + backends=backends_to_run, + rounds=args.warmup_rounds, + alternate_backend_order=args.alternate_backend_order, + ) + if warmup_jobs: + run_jobs( + warmup_jobs, + backends=backend_objs, + workers=args.workers, + prefetch_factor=args.prefetch_factor, + decode_batch_size=args.decode_batch_size, + label="warmup", + ) + + jobs = build_jobs( + cases, + backends=backends_to_run, + rounds=args.rounds, + alternate_backend_order=args.alternate_backend_order, + ) + rows, wall_s = run_jobs( + jobs, + backends=backend_objs, + workers=args.workers, + prefetch_factor=args.prefetch_factor, + decode_batch_size=args.decode_batch_size, + label="measure", + ) + finally: + if fluxon_runtime is not None: + fluxon_runtime.close() + + summary = summarize(rows, wall_s=wall_s) + write_rows_csv(args.output_csv, rows) + write_json( + args.output_json, + { + "args": sanitized_args(args), + "summary": summary, + "rows": rows, + }, + ) + print_summary(summary) + print(f"Wrote CSV: {args.output_csv}", flush=True) + print(f"Wrote JSON: {args.output_json}", flush=True) + + if not args.allow_errors and any(row.get("status") != "ok" for row in rows): + return 1 + return 0 + + +def sanitized_args(args: argparse.Namespace) -> dict[str, Any]: + out = dict(vars(args)) + if out.get("fluxon_request_password"): + out["fluxon_request_password"] = "" + return out + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/dataloader_video_benchmark/test_benchmark_contract.py b/examples/dataloader_video_benchmark/test_benchmark_contract.py new file mode 100644 index 0000000..28db0a6 --- /dev/null +++ b/examples/dataloader_video_benchmark/test_benchmark_contract.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import csv +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + + +EXAMPLE_DIR = Path(__file__).resolve().parent +MODULE_PATH = EXAMPLE_DIR / "benchmark.py" + + +def _load_module(): + spec = importlib.util.spec_from_file_location("dataloader_video_benchmark", MODULE_PATH) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load module spec: {MODULE_PATH}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +DL = _load_module() + + +class TestDataloaderPerfContract(unittest.TestCase): + def test_parse_case_csv_and_frame_indices(self) -> None: + with tempfile.TemporaryDirectory() as tmp_s: + csv_path = Path(tmp_s) / "cases.csv" + with csv_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "idx", + "video_path", + "height", + "width", + "start_idx", + "end_idx", + "num_frames", + ], + ) + writer.writeheader() + writer.writerow( + { + "idx": "7", + "video_path": "/data/a.mp4", + "height": "480", + "width": "832", + "start_idx": "10", + "end_idx": "18", + "num_frames": "5", + } + ) + + cases = DL.parse_case_csv(str(csv_path)) + + self.assertEqual(len(cases), 1) + self.assertEqual(cases[0].idx, 7) + self.assertEqual(DL.frame_indices(cases[0]), [10, 12, 14, 16, 18]) + + def test_build_jobs_alternates_backend_order(self) -> None: + cases = [ + DL.DecodeCase( + source="test", + idx=1, + video_path="/data/a.mp4", + height=4, + width=4, + start_idx=0, + end_idx=3, + num_frames=2, + ), + DL.DecodeCase( + source="test", + idx=2, + video_path="/data/b.mp4", + height=4, + width=4, + start_idx=0, + end_idx=3, + num_frames=2, + ), + ] + + jobs = DL.build_jobs( + cases, + backends=["original", "fluxon"], + rounds=1, + alternate_backend_order=True, + ) + + self.assertEqual([job.backend for job in jobs], ["original", "fluxon", "fluxon", "original"]) + self.assertEqual(jobs[0].sample_id, jobs[1].sample_id) + self.assertEqual(jobs[2].sample_id, jobs[3].sample_id) + + def test_summarize_reports_fluxon_speedup(self) -> None: + rows = [ + { + "backend": "original", + "status": "ok", + "elapsed_s": 2.0, + "num_frames": 4, + }, + { + "backend": "fluxon", + "status": "ok", + "elapsed_s": 1.0, + "num_frames": 4, + }, + ] + + summary = DL.summarize(rows, wall_s=3.0) + + self.assertEqual(summary["backends"]["original"]["ok"], 1) + self.assertEqual(summary["backends"]["fluxon"]["ok"], 1) + self.assertEqual(summary["fluxon_vs_original"]["p50_speedup"], 2.0) + + def test_resolve_fluxon_agent_node_id_requires_external_agent(self) -> None: + args = DL.argparse.Namespace( + fluxon_agent_instance_key="fs-agent-1", + fluxon_agent_node_id="", + ) + self.assertEqual(DL.resolve_fluxon_agent_node_id(args), "fs-agent-1") + + args = DL.argparse.Namespace( + fluxon_agent_instance_key="fs-agent-1", + fluxon_agent_node_id="fs-agent-2", + ) + with self.assertRaisesRegex(ValueError, "must match"): + DL.resolve_fluxon_agent_node_id(args) + + args = DL.argparse.Namespace( + fluxon_agent_instance_key="", + fluxon_agent_node_id="", + ) + with self.assertRaisesRegex(ValueError, "start fluxon_py.runtime.start_fs_agent"): + DL.resolve_fluxon_agent_node_id(args) + + def test_sanitized_args_redacts_fluxon_password(self) -> None: + args = DL.argparse.Namespace( + fluxon_request_username="bench", + fluxon_request_password="secret", + other="value", + ) + + self.assertEqual( + DL.sanitized_args(args), + { + "fluxon_request_username": "bench", + "fluxon_request_password": "", + "other": "value", + }, + ) + + def test_run_jobs_uses_decode_batch_windows(self) -> None: + class BatchBackend: + def __init__(self) -> None: + self.calls = [] + + def decode_batch(self, cases): + self.calls.append([case.idx for case in cases]) + return [ + { + "status": "ok", + "open_s": 0.0, + "decode_s": 0.1, + "materialize_s": 0.0, + "elapsed_s": 0.1, + "shape": "1", + "dtype": "uint8", + "nbytes": 1, + } + for _case in cases + ] + + cases = [ + DL.DecodeCase( + source="test", + idx=i, + video_path=f"/data/{i}.mp4", + height=4, + width=4, + start_idx=0, + end_idx=3, + num_frames=2, + ) + for i in range(1, 4) + ] + jobs = DL.build_jobs( + cases, + backends=["fluxon"], + rounds=1, + alternate_backend_order=True, + ) + backend = BatchBackend() + + rows, _wall_s = DL.run_jobs( + jobs, + backends={"fluxon": backend}, + workers=1, + prefetch_factor=2, + decode_batch_size=2, + label="test", + ) + + self.assertEqual(backend.calls, [[1, 2], [3]]) + self.assertEqual(len(rows), 3) + self.assertTrue(all(row["status"] == "ok" for row in rows)) + + +def main() -> None: + unittest.main() + + +if __name__ == "__main__": + main() diff --git "a/fluxon_doc_cn/blog/blog_2_Fluxon_KV_SSD\345\255\230\345\202\250\357\274\232\346\212\212\346\234\254\345\234\260SSD\346\216\245\346\210\220\345\206\205\345\255\230\345\211\257\346\234\254\347\232\204\345\233\236\345\241\253\345\261\202.md" "b/fluxon_doc_cn/blog/blog_2_Fluxon_KV_SSD\345\255\230\345\202\250\357\274\232\346\212\212\346\234\254\345\234\260SSD\346\216\245\346\210\220\345\206\205\345\255\230\345\211\257\346\234\254\347\232\204\345\233\236\345\241\253\345\261\202.md" new file mode 100644 index 0000000..497bdd6 --- /dev/null +++ "b/fluxon_doc_cn/blog/blog_2_Fluxon_KV_SSD\345\255\230\345\202\250\357\274\232\346\212\212\346\234\254\345\234\260SSD\346\216\245\346\210\220\345\206\205\345\255\230\345\211\257\346\234\254\347\232\204\345\233\236\345\241\253\345\261\202.md" @@ -0,0 +1,1307 @@ +# Mooncake 6x 吞吐：让 Fluxon KV 支持 DRAM + SSD 多级缓存 + +Fluxon KV 是 Fluxon 统一 AI 数据面中的分布式键值缓存服务，面向世界模型推理中的潜变量缓存（latent cache）、大模型推理中的键值缓存（KV cache），以及跨进程、跨节点的大对象复用。应用通过 `put`、`get` 和 `delete` 存取数据，Fluxon KV 统一管理 key-version 路由、副本生命周期和数据搬运，并以 `MemHolder` 交付读取结果。 + +纯 DRAM 缓存能提供低延迟访问，但可命中工作集受内存容量约束；当工作集大于 DRAM 时，被驱逐的数据会转为 miss，owner 的本地 NVMe 容量和带宽也尚未进入读取路径。将本地 SSD 纳入缓存层次后，系统可以保留更多运行期数据，并在内存副本缺失时从 SSD 回填。对于存在数据复用且并发度足够的 AI 负载，这种分层具有提高整体命中率、减少上游重复计算或远程重取，并更充分利用本地存储与传输带宽的潜力。实际收益取决于访问分布、并发度、介质性能和回填开销。 + +本文将存储节点（owner）的本地 SSD 作为 DRAM 内存层的运行期回填层（backing store）。`put` 操作在内存副本可读后返回，系统随后在后台异步写入 SSD；`get` 操作优先查询内存，仅在没有可读内存副本时从 SSD 回填。SSD 不引入独立的公共 API，也不承担冷启动恢复。 + +下文将控制节点（master）维护的 key-version 级副本索引称为路由表（`route`），将贡献内存或 SSD 容量的节点称为存储节点（owner）。全文围绕四个核心要点展开： + +| 设计方面 | 核心要点说明 | +| --- | --- | +| 用户契约 | 继续使用 `put`、`get` 和 `delete` 操作；`get` 仍返回 `MemHolder`。 | +| 控制面与资源 | master 维护 `route`，并统一调度用于最终 value 副本和临时传输的内存分配（allocation）。owner 提供 memory segment，承载真实 bytes，并管理本地 SSD。 | +| 数据路径 | 写入的同步完成点是内存副本发布；读取优先使用内存，SSD 只在内存无可读副本时回填。 | +| 恢复边界 | SSD 仅作为运行期缓存。owner 重启后重建空 shard 和 ring，不扫描旧 shard，也不从 master 的 `route` 中恢复旧 SSD bytes。 | + +正文先说明 master 与 owner 的分工、设计取舍和当前边界，再介绍配置方式，并沿写入、读取与 SSD 回填的数据流展开实现。随后讨论 owner 本地 IO 和一致性，最后给出性能测试结果及其适用范围。 + +## 1. 架构概览：master 管理路由和内存分配，owner 管理本地 SSD + +内存 allocation 分为两类：一类承载最终 value 副本，另一类供跨 owner 传输临时使用，例如远端 SSD 回填的 source staging。master 统一调度这两类区间，因为它们的分配、保活和释放都与 `route` 及 in-flight 请求状态联动。allocation 对应的真实 bytes 仍位于 owner 注册的 memory segment。 + +SSD allocation 只负责最终 value 副本，因此只涉及 owner 本地 shard/ring 的文件区间。文件 offset、覆盖保护和 IO 调度都是本地状态，因此 owner 负责 SSD 分配和 ring 驱逐。 + +| 对象 | 资源或状态归属 | 分配与驱逐职责 | +| --- | --- | --- | +| `route` | master 持有 key-version 级副本索引。 | master 选择内存副本的放置位置并维护内存 allocation 生命周期。在内存或 SSD 副本失效后，master 收敛逻辑路由。 | +| 内存 | bytes 位于 owner 注册的 memory segment。 | master 从这些 segment 分配两类区间：内存副本区间和通信临时区间。master 同时管理对应 allocation 的保活与驱逐。 | +| SSD | bytes 位于 owner 本地的 shard 和 ring。 | owner 分配 SSD 文件区间并执行 ring 驱逐，再把 commit 或驱逐结果通知 master 更新 `route`。 | + +在 `route` 实现中，master 的 `OneKvNodesRoutes` 按 key-version 组织路由。`node_replicas` 是一张以 owner 为 key 的 map，每个 entry 用 `memory` 和 `ssd` 记录两层副本。二者共享当前 route 的 `put_id`，不能跨版本复用。SSD 文件位置、读写队列、文件 IO 对齐和覆盖保护都由 owner 管理。 + +### 1.1 角色与图示约定 + +全文时序图用角色前缀标明逻辑边界；前缀不表示独立进程。 + +| 前缀 | 表示的职责 | +| --- | --- | +| `external-*` | 集群外的公共 API 调用方或 benchmark 消费方，不贡献 Fluxon 存储容量。 | +| `master-*` | master 内的控制面，管理 `route`、in-flight 状态、master 侧 `Allocation` 保活句柄（guard）及 holding 生命周期，不中转 payload。 | +| `owner-*` | owner 侧的请求入口、内存或 SSD 存储组件，以及本地数据传输端点。多个 `owner-*` 角色可以位于同一 owner 进程，也可以分布在不同 owner。 | + +后续时序图同时涉及请求协调、bytes 位置、allocation 保活和数据搬运，四者不能混用： + +| 关系 | 文中的规范说法 | 当前实现落点 | +| --- | --- | --- | +| 请求由谁协调 | `owner-requester` 接收并推进本次公共请求。 | owner 侧 `ClientKvApi`；external 包装层在图中通常折叠。 | +| payload bytes 位于哪里 | bytes 位于 source/target owner 已注册的 memory segment，或 SSD owner 的本地 shard。 | RPC plan 中的 node、base address 和 absolute address。 | +| allocation 由谁保活 | master 持有表示 owner segment 区间的 RAII `Allocation` guard（保活句柄）。 | `InflightPutInfo`、`OneKvNodesRoutes.memory`、`InflightGetInfo` 和 `get_holding`。 | +| 数据由谁搬运 | memory get 由 requester pull；远端 SSD 回填由 SSD source push。 | owner 侧 `ClientTransferEngine`。 | + +`owner-requester` 只表示请求协调者，不代表固定 buffer，也不持有 master 侧的 `Allocation` 对象。它在不同操作中的数据面职责如下： + +| 操作 | owner 侧 range | master 侧 guard | 完成结果 | +| --- | --- | --- | --- | +| `put` | 常规路径的 source/staging bytes 位于 requester owner；side-transfer 路径位于单独的 `owner-put-source`。最终 target 由 master 选择，可能与 source 同 owner，也可能位于 `owner-put-target`。 | 本地 placement 用一个 `InflightPutAllocation::Local` guard 同时表示 source 和 target；远端 placement 用 `Remote { src, target }` 两个 guard。 | `PutDone` 把 target guard 转入 memory route；`put` 不产生 holder。 | +| `get` | route 路径的最终 target bytes 位于 requester owner。已有本地 replica 时直接复用；否则 memory source 或 SSD source 把 payload 写入该 target。 | `InflightGetInfo.allocation` 保活 target；`GetDone` 后转入 master `get_holding`。memory source 或远端 SSD source staging 的额外 guard 位于 `source_allocation`。 | owner 返回 holder metadata，external 在自己的进程中构造 `MemHolder`。 | +| `delete` | 没有 payload source、staging 或 target range。 | 不创建 payload `Allocation`。 | 只更新控制面 route，并执行既有 holder 失效流程。 | + +后续 `put` / `get` 时序从 owner 进入 master route 路径开始，折叠 external 包装 RPC。owner 返回共享内存 offset、长度和 `holder_id` 等 metadata，external 再据此构造 `MemHolder`，形成完整的公共结果。公共 `get` 操作还可能先命中 owner 本地的 `get_cached_info`，直接复用已有 `MemoryInfo`。第 4–7 节描述本地缓存未命中后进入 `GetStart` 的路径。 + +时序图默认采用常规 owner 请求入口。专用 side-transfer worker 会把请求处理者与 put source owner 分开：`PutStartReq.source_node_id` 指向与 worker 共享 mmap 的 owner，source/staging bytes 位于该 owner 的 segment。这个 fast path 不改变 master 选择 put target、`PutDone` 发布 target 的主契约。 + +### 1.2 四条主链路 + +第 3–7 节将沿这四条链路展开 allocation、传输和失败清理。 + +| 链路 | 核心流程 | 完成或收敛点 | +| --- | --- | --- | +| 写入与 SSD commit | master 分配 source/target；`PutDone` 发布内存副本；owner 随后异步写入 SSD。 | 公共 `put` 操作在内存副本发布后完成。SSD commit 只有在版本、memory 和 tomb 条件仍成立时才进入 `route`。 | +| 内存命中读取 | requester 复用本地副本，或从远端 memory owner pull 到最终 target。 | `GetDone` 把 target guard 转入 `get_holding`，external 根据 metadata 构造 `MemHolder`。 | +| SSD 回填读取 | 当前版本没有可读内存副本时，本地路径直接读入 requester target；远端路径读入 source staging 后按 chunk push。 | 两条路径复用同一个 `GetDone` 和 holder 交付链。 | +| SSD 驱逐与陈旧 route 兜底 | owner 覆盖 ring entry 后批量通知 master；读取端仍会复查本地 entry。 | 批量通知收敛正常驱逐；陈旧 source 读取失败时执行 revoke 并有界重新选路。 | + +payload 不经过 master。master 只维护逻辑路由和 allocation 生命周期；SSD shard、offset、IO 对齐与 ring 状态只在 owner 本地解析。 + +### 1.3 设计取舍与当前边界 + +引入 SSD 后，需要先明确三项职责边界： + +1. **同步语义**：SSD 不改变 `put` / `get` 操作的同步语义。 +2. **路由信息**：master 的 `route` 不记录 SSD 物理布局。 +3. **恢复能力**：SSD 不承担重启后的数据恢复。 + +在这些边界内，本设计保持现有公共 API 和内存副本语义，利用 owner 本地 SSD 扩大运行期缓存容量，并在没有可读内存副本时提供回填来源。文件布局和 IO 复杂度均封装在 owner 内部。 + +因此，当前 SSD 被定位为运行期回填层。它不向用户暴露独立的 SSD API，也不支持单 value 跨 device 条带化。SSD 文件布局、IO 调度和 ring 生命周期均由 owner 管理，不进入公共 API 或 master route。 + +| 设计选择 | 作用 | 代价 | 当前边界 | +| --- | --- | --- | --- | +| master 不保存 SSD offset | 跨节点控制面只处理 key-version、owner、payload 长度和 tomb tag 等逻辑信息。 | master 无法据此检查本地物理布局；陈旧 route 要靠节点 tomb、批量清理与读取兜底收敛。 | **route 是逻辑索引，不是恢复元数据**。 | +| `PutDone` 不等待 SSD | 同步写入只等待内存副本，SSD 背压留在 owner 本地 persist 路径。 | SSD 副本是 late commit，需要用 `put_id` 拒绝旧版本 commit，并在 master 处理 commit RPC 的窗口 pin 住本地 entry。 | **内存副本 ready 是同步提交点**。 | +| owner 本地驱逐 | ring 决定物理覆盖顺序，再批量通知 master 清理逻辑 route；发送失败的当前 batch 会重试。 | 有界队列满、进程退出或 master 持续不可用时，master 可能保留陈旧 route。 | **读取失败会 revoke 并有限重试**。 | +| `O_DIRECT + io_uring`，而非 `mmap` / PageCache | KV SSD 路径绕过 PageCache，直接提交对齐 IO。 | 实现必须协调 staging 地址、文件 offset、IO 长度、aligned buffer、ring 覆盖和 read pin。 | **文件 IO 实现封装在 owner**。 | +| 本地 SSD 回填 | SSD owner 与 requester owner 相同时复用最终 target，不申请独立 source staging。 | target capacity 必须覆盖对齐 IO；地址不满足 direct read 条件时仍会经过 scratch copy。 | **direct 条件成立时省掉本地二次 copy**。 | +| SSD owner 主动推送 | chunk ready 后由 SSD owner 直接推到 requester target，读盘和网络传输可重叠。 | requester 必须先通过 `SsdStageReadReq` 提供 target 地址和 `get_id`。 | **省去 stage-ready 后的一次 RTT**。 | +| value 级 device 选择 | 单 value 使用一个 device 上的连续文件区间，不需要跨 device 组合读结果。 | 单个大 value 不跨 device 条带化。 | **多 device 并行发生在 value 粒度**。 | +| 启动时重建空 shard 和 ring | 明确 SSD 只作为运行期缓存：不扫描 shard，也不通过 WAL 或 checkpoint 恢复 entry。 | 重启后需要由新写入重新形成 SSD 副本。 | **不支持冷启动恢复**。 | +| scratch 读路径 | 不满足 direct read 对齐时仍能正确读取真实 payload。 | 该 chunk 多一次本地 copy。 | **padding 不进入用户 holder**。 | + +## 2. 配置与观测：容量按介质分层可见 + +KV SSD 复用 owner 已有的 `large_file_paths`，不增加第二套路径参数，以保持配置项尽可能简洁。`large_limit_size` 与这些路径一一对应；运行时按实际 device 去重，并分别上报内存与 SSD 的容量和占用。 + +### 2.1 配置入口 + +配置仍位于 owner 的 `fluxonkv_spec`；`large_limit_size` 与 `large_file_paths` 等长： + +```yaml +fluxonkv_spec: + large_file_paths: [/data/fluxon_large0, /data/fluxon_large1] # Derive one KV SSD cache root from each path. + large_limit_size: ["4GB", "8GB"] # Limit the matching cache roots in the same order. +``` + +- **启用条件** + - `large_limit_size` 未配置或为 `null`：不启用 KV SSD。 + - `large_limit_size` 已配置：必须与 `large_file_paths` 长度一致，数组下标一一对应。 +- **容量值格式** + - 可以使用整数 bytes。 + - 也可以使用 `"512.9B"`、`"1.5GB"` 这类 size 字符串。 + - `MB/GB` 按二进制单位解析，等价于 `MiB/GiB`；小数结果按 bytes 向下截断。 + - 解析结果必须大于或等于 512 bytes，满足当前 `O_DIRECT` 对齐约束。 +- **external 约束** + - 不贡献存储容量的 external 不能声明 `large_limit_size`。 + - external 只继承 owner 暴露的共享 bundle 和运行时路径信息，不贡献 SSD 容量。 + +### 2.2 路径派生与设备去重 + +owner 不读取单独的 SSD 路径参数。每个 SSD cache root 都从对应的 `large_file_paths[i]` 派生： + +```text +/_cluster_kv_ssd_storage// +``` + +- **创建 root**：owner 创建派生目录，再通过 `metadata.dev()` 获取实际 device 身份。 +- **同设备去重** + - 多个 path 指向同一实际 device：只保留第一个 root。 + - 如果不去重，同一块盘上的两个目录会被误认为两条独立 IO 路径，导致调度层高估并行度。 + +### 2.3 容量、ring 与 per-device 对象 + +- **容量作用域** + - `large_limit_size[i]` 只约束 `large_file_paths[i]` 派生的 KV SSD cache root。 + - 它限制 key-value payload 在 KV SSD ring 中的用量，不代表整块 SSD 的物理容量。 + - 日志、profile、FS disk cache 等其他大文件资产不受该值约束。 +- **ring 空间不足** + + **ring 是构建在定长 shard 文件之上的循环空间分配器和本地 key-version 索引。** 每个 shard 用 `head` 为新 value 分配对齐后的连续区间；写到文件末尾时，物理 offset 回到文件开头继续复用空间。`tail` 随覆盖边界向前推进，回收最旧且已经可以覆盖的 entry。ring 同时记录 entry 的 `Writing/Committed` 状态、read pin 和 route commit pin，遇到仍在写入、读取或提交 route 的区间时会等待，不会直接覆盖。 + + ring 的作用是在固定容量内持续复用 SSD cache 空间，并把被覆盖的 key-version 交给 owner 后台通知 master 收敛逻辑 `route`。文中的 ring 与 `io_uring` 职责不同：前者管理文件区间、entry 状态和驱逐顺序，后者负责提交实际的文件读写。第 6.2 节再展开状态转换和覆盖保护。 + + 选择 ring 的直接工程动机，是让每个 shard 内的写入尽量保持追加式、对齐且连续。`head` 单向推进，空间用完后再回绕覆盖，省去了随机空洞分配、碎片合并和后台 compaction。现代 NVMe 也能处理随机 IO；这里更看重 `O_DIRECT` 连续写入、固定容量复用和回收路径的可预测性。 + + Moka 与 ring 解决的是不同层次的问题： + + | 机制 | 当前职责 | 无法替代的另一层能力 | + | --- | --- | --- | + | Moka | master 侧按 weight 管理内存副本的准入和驱逐。 | 不分配 SSD 文件 offset，也不管理 `Writing/Committed`、read pin 或 route commit pin。 | + | ring | owner 侧分配连续 SSD 区间，并按物理写入顺序回收和覆盖。 | 不提供 Moka 的访问热度与内存准入策略。 | + + 如果直接用 Moka 选择任意 SSD key 驱逐，被删除 value 的文件区间会形成不连续空洞，仍需额外实现 extent allocator、空洞合并或 compaction。Moka 可以叠加为逻辑淘汰策略，但不能替代 ring 对 SSD 物理空间的管理。 + + - 目标区域可以覆盖：ring 推进 tail 并驱逐旧 entry。 + - 目标区域仍有未完成写入、active read pin 或 route commit pin：当前写入等待空间释放。 +- **`KvSsdStorage`** + - 持有一组去重后的 device root。 + - 每个 device root 都独立持有： + - **shard 文件集合**：存放 ring 分配的连续 SSD offset。 + - **writer queue**：接收本 device 的写任务。 + - **reader queue**：接收本 device 的读任务。 + - **`UringIoEngine`**：提交本 device shard 文件上的 `read/write`；`Iovec` 只作为测试和对照路径保留。 +- **shard 到 device 的映射** + - 每个 device root 的 `large_limit_size` 会划分给该 device 的多个 shard。 + - `shard_to_device` 记录 shard 所属 device。 + - 读路径根据 committed entry 的 `shard_id` 找回对应 reader queue。 + +### 2.4 能力标记与容量观测 + +启用 KV SSD 的 owner 会在 `ClusterMember.metadata` 中发布 `kv_ssd_storage=true`。master 在 `PutStart` 时读取当前 live client members，并按下表确定 target 候选集和后续 SSD persist 决策： + +| live member 状态 | put target 候选集 | `PutDone` 后的处理 | +| --- | --- | --- | +| 没有 owner 声明 `kv_ssd_storage=true` | 所有具有已注册内存 allocator 的 live owner。 | `persist_to_ssd=false`，跳过 `SsdReplicaPersistReq`，按纯内存模式完成。 | +| 至少一个 owner 声明 `kv_ssd_storage=true` | 同时具有内存 allocator 和该标记的 live owner。 | `persist_to_ssd=true`，向选中的 target owner 异步发送 `SsdReplicaPersistReq`。 | + +master 把这次决策保存在 `InflightPutInfo` 中，`PutDone` 不重新推断。该标记只表达 owner 是否具备 KV SSD 能力，不暴露本地 shard、offset 或容量；容量和占用仍通过指标单独上报。target owner 收到 persist 请求后还会检查本地 `KvSsdStorage`，用于兜住 member capability 快照与 owner 实际状态之间的竞态。 + +owner 按资源类型分别上报 capacity/used，UI 也分别展示内存和 SSD： + +- **`memory_segment`** + - 指标：`kvcache_segment_capacity_bytes` / `kvcache_segment_used_bytes`。 + - 含义：owner 可直接承载 `MemHolder` 的内存 segment 容量和占用。 +- **`kv_ssd`** + - 指标：`kv_ssd_capacity_bytes` / `kv_ssd_used_bytes`。 + - 含义：owner 本地 SSD 介质层的容量和 ring 已用空间。 + +## 3. 写入：先内存可见，再异步进入 SSD + +**`put` 操作的同步完成点是内存副本发布。** `PutDoneResp` 返回后，用户已经可以读到该副本。SSD 对齐、写盘、ring 空间等待和 route 提交都在后台执行。 + +SSD persist 复用内存 target 所在 owner 的本地地址，因此 master 选择内存 target 时也同时决定是否生成 SSD 副本以及由哪个 owner 写盘。这里分为三层： + +| 层次 | 行为 | 分支定位 | +| --- | --- | --- | +| 能力上报 | owner 已通过 member metadata 上报 `kv_ssd_storage=true`。内存 segment 注册继续只描述 master 可分配的地址区间；SSD 文件区间仍由 owner 本地管理。 | master 已经能判断 owner 是否具备 KV SSD 能力，无需通过 persist RPC 探测。 | +| master 决策 | live SSD-capable owner 集合为空时使用纯内存 placement；集合非空时，target 必须同时具有内存 allocator 和 SSD capability。决策结果以 `persist_to_ssd` 写入 `InflightPutInfo`。 | 无 SSD owner 时不发送 persist RPC；启用 SSD placement 时不会主动选择无 SSD owner。 | +| owner 侧防御 | target owner 在处理 persist 请求时再次检查本地 `KvSsdStorage`。 | `persisted=false` 只覆盖 capability 快照过期、owner 重启或重新注册等竞态，不参与正常能力探测。 | + +一次写入分为三个阶段： + +| 阶段 | 动作 | 完成后的状态 | +| --- | --- | --- | +| 1. 填充 target | master 根据 live member capability 和已注册内存 allocator 生成 placement plan，并在 `InflightPutInfo` 中记录 `persist_to_ssd`。系统先将 payload 写入 `owner-put-source` 的 range。对于本地 placement，系统直接复用该 range；对于远端 placement，系统再将 payload 传输到 `owner-put-target`。 | payload 已到 target 地址，SSD persist 决策已经固定，尚未对其他读取发布。 | +| 2. 发布内存副本 | `owner-requester` 发送 `PutDoneReq`。master 将 target guard 转入 route 的 `memory`，并释放不再需要的 source staging。 | master 返回 `PutDoneResp`；公共 `put` 操作完成，内存副本可读。 | +| 3. 按 placement 决策处理 SSD | `persist_to_ssd=false` 时跳过 SSD RPC；为 `true` 时，master 保活 target 并发送 `SsdReplicaPersistReq`。target owner 再确认本地 `KvSsdStorage` 可用，然后对齐 payload、写入 shard，并通过 `SsdReplicaCommitReq` 提交 SSD route。 | 纯内存模式直接进入内存 cache；SSD commit 校验通过后 `ssd` 字段可用。竞态兜底、persist 失败或过期 commit 都不回滚第 2 阶段。 | + +下图按公共 `put` 操作的语义展开这三个阶段。图中省略 external 使用的 `ExternalPutStartReq`、共享内存写入和 `ExternalPutTransferEndReq` 包装层。常规路径将 `owner-put-source` 合并到 `owner-requester`；side-transfer 的分离边界见第 1.1 节。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant owner_put_target as owner-put-target + participant master_route as master-route + participant owner_ssd as owner-ssd-storage + participant owner_ring as owner-ssd-ring + + external_user->>owner_requester: put(key, value) + owner_requester->>master_route: PutStartReq + master_route->>master_route: 读取 live member 的 kv_ssd_storage capability + alt 存在 SSD-capable owner + master_route->>master_route: target = memory allocator ∩ SSD owner；persist_to_ssd = true + else 不存在 SSD-capable owner + master_route->>master_route: target = memory allocator；persist_to_ssd = false + end + master_route->>master_route: 分配并保活 InflightPutAllocation + master_route-->>owner_requester: source / target node-address plan + alt target 与 requester/source 同 owner + owner_requester->>owner_requester: payload 写入本地 target range + else target 在其他 owner + owner_requester->>owner_put_target: source/staging → target range + end + owner_requester->>master_route: PutDoneReq + master_route->>master_route: target guard 转入 memory route；远端 source guard 释放 + + par 返回公共 put + master_route-->>owner_requester: PutDoneResp + owner_requester-->>external_user: put 完成，内存副本可读 + and PutDone 后的介质处理 + alt persist_to_ssd = false + master_route->>master_route: 跳过 SSD RPC；内存副本进入驱逐 cache + else persist_to_ssd = true + master_route->>master_route: 保活 target allocation 并预留 cache capacity + master_route->>owner_ssd: SsdReplicaPersistReq(key, put_id, target_addr, len) + alt owner 本地无 KvSsdStorage（capability 状态竞态兜底） + owner_ssd-->>master_route: SsdReplicaPersistResp(persisted = false) + else owner 本地 KvSsdStorage 可用 + owner_ssd->>owner_ring: 对齐、分配 offset 并写 shard + alt 本地 persist 失败 + owner_ring-->>owner_ssd: write / ring error + owner_ssd-->>master_route: persist error + else 本地 persist 成功 + owner_ring-->>owner_ssd: Committed entry + KvSsdPersistGuard + owner_ssd->>master_route: SsdReplicaCommitReq + master_route->>master_route: 校验 key、put_id、memory 和 tomb tag + alt route 仍然有效 + master_route->>master_route: 更新对应 owner entry 的 ssd 字段 + master_route-->>owner_ssd: SsdReplicaCommitResp(OK) + else route 已删除或版本已变化 + master_route->>master_route: 忽略 late commit，不更新 route + master_route-->>owner_ssd: SsdReplicaCommitResp(OK) + end + owner_ssd->>owner_ring: 释放 route commit pin + owner_ssd-->>master_route: SsdReplicaPersistResp(persisted = true) + end + end + master_route->>master_route: 释放临时预留，内存副本加入驱逐 cache + end + end +``` + +除 SSD capability 兜底外，后台路径还需要处理三类并发风险： + +| 风险 | 机制 | 保证 | +| --- | --- | --- | +| persist 期间 target 被释放或容量超卖 | `Arc` 保活 target，capacity reservation 按 allocation capacity 扣减 cache 预算。RAII guard 覆盖成功、失败、取消和 master 关闭。 | persist 始终引用有效 allocation，容量账本与生命周期一致。 | +| 旧版本完成得更晚 | master 用 `put_id` 校验 late commit。 | 旧 SSD 副本不会进入新版本 route。 | +| 本地已写完，master 尚未处理完 commit RPC | `KvSsdPersistGuard` 在该窗口期 pin 住 ring entry。 | commit RPC 返回前，对应 SSD bytes 不会被 ring 覆盖；返回 `OK` 不代表 route 一定已发布。 | + +### 3.1 route 数据结构和 SSD commit 条件 + +`OneKvNodesRoutes` 中与写入可见性直接相关的结构如下，其他字段已省略： + +```rust +pub struct OneKvNodesRoutes { + pub put_id: PutIDForAKey, // Bind all replicas to one key version. + // ... + pub node_replicas: RwLock>, // Index both cache layers by owner. + // ... +} + +pub struct KvNodeReplicas { + pub tomb_tag: NodeTombTag, // Track this owner incarnation once. + pub memory: Option>, // Hold the readable memory replica when present. + pub ssd: Option, // Record the readable SSD replica when present. +} + +pub struct KvSsdReplicaInfo { + pub len: u64, // Preserve the real payload length. +} +``` + +`node_replicas` 中的每个 owner 只对应一个 entry： + +- **owner 身份**：来自 `HashMap` 的 `NodeID` key，不在 entry 内重复保存。 +- **节点实例**：`tomb_tag` 绑定本次 owner incarnation；节点离开或重启后，旧副本不能被新实例复用。 +- **介质状态**：`memory` 和 `ssd` 独立更新。内存驱逐只清空 `memory`，SSD ring 驱逐或读取失败只清空 `ssd`；两者都为空时才删除 entry。 + +owner 本地写盘成功后仍需提交 `SsdReplicaCommitReq`。master 只在以下条件全部满足时写入 `KvSsdReplicaInfo { len }`： + +- 当前 `kv_routes` 仍有该 key，且 route 的 `put_id` 与请求一致。 +- 对应 owner 的内存副本仍然存在。 +- owner 的 tomb tag 仍然 live。 + +`SsdReplicaCommitResp` 当前不区分“已发布”和“已忽略”。route 缺失、版本变化、memory 缺失或节点 tomb 时，master 会忽略 commit，但仍返回 `OK`。因此，owner 侧的 `persisted = true` 只表示本地 entry 已写入且 commit RPC 无协议错误，不能确认 master route 一定包含该 SSD 副本。 + +SSD persist 可能晚于同 key 的后续覆盖写。`put_id` 校验会丢弃旧版本 late commit，避免它进入新版本 route。`KvSsdPersistGuard` 从本地 commit 开始保活 entry，直到 master 处理完这次 commit RPC；无论 route 最终发布还是请求被忽略，guard 都会释放。 + +## 4. 读取：内存优先，SSD 只做回填 + +**只有当前 key-version 没有可读内存副本时，master 才选 SSD。** 以下流程从 owner 本地 `get_cached_info` 未命中后开始。进入 route 路径后，`GetStartReq` 按下表为 `get` 操作选择 source： + +| 分支 | 触发条件 | payload 路径 | 完成方式 | +| --- | --- | --- | --- | +| 内存 | route 中存在可读 `memory` 副本。 | source 在 requester owner 时，master 复用对应的 replica guard，requester 无需搬运 payload；否则由 requester 从 `owner-memory` pull 到最终 target。 | requester 发送 `GetDoneReq`。owner 返回 holder metadata，external 再构造公共 `MemHolder`。 | +| 本地 SSD | 没有可读内存副本，且 SSD source 就在 requester owner。 | shard 读到最终 target。满足 direct read 条件时直接写入；否则经过 owner 内部 scratch buffer 复制一次。 | requester 发送 `GetDoneReq`，后续 holder 交付与内存分支相同。 | +| 远端 SSD | 没有可读内存副本，且 SSD source 在其他 owner。 | `owner-ssd-source` 按 chunk 读入本地 staging，每个 chunk ready 后主动 push 到 requester target。 | SSD source owner 发送 `GetDoneReq`；requester 用 `done_*` 字段构造同一套 holder metadata，再由 external 构造公共 `MemHolder`。 | +| 未命中 | 当前 key-version 没有可读的 memory 或 SSD 副本。 | 不传输 payload。 | 返回 `KeyNotFound`。 | + +本文把“回填”限定为“把 SSD payload 写入 requester 的最终 target”。staging 是远端 SSD source owner 上的临时、对齐内存；它不会成为用户的 `MemHolder`。 + +下图展开四条 route 分支、数据面 range 的位置和 master 侧 allocation 控制关系。图中已折叠 owner 本地 `get_cached_info` 未命中之前的包装层。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant master_route as master-route + participant owner_memory as owner-memory + participant owner_ssd as owner-ssd-source + participant owner_media as owner-local-ssd + + external_user->>owner_requester: get(key) + owner_requester->>master_route: GetStartReq + master_route->>master_route: 读当前 key-version 的 node_replicas 快照 + + alt 存在可用 memory 副本 + master_route-->>owner_requester: memory source + requester target plan + alt source 就在 requester owner + owner_requester->>owner_requester: 复用本地 replica range + else source 在其他 owner + owner_requester->>owner_memory: transfer_data_no_copy (pull) + owner_memory-->>owner_requester: payload 写入最终 target + end + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_requester: GetDoneResp(holder_id) + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else 无 memory 副本，但存在可用 SSD 副本 + master_route->>master_route: 在 requester segment 分配并保活最终 target + alt SSD source 在其他 owner + master_route->>master_route: 在 SSD source segment 分配并保活对齐 staging + end + master_route-->>owner_requester: SSD source + target + stage plan + alt SSD source 就在 requester owner + owner_requester->>owner_ssd: 本地 load_and_push_kv_from_ssd 调用 + owner_ssd->>owner_media: pin entry 并读取 shard (direct or scratch) + owner_media-->>owner_ssd: IO 完成，最终 target 已写入 + owner_ssd-->>owner_requester: 本地读取完成 + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_requester: GetDoneResp(holder_id) + else SSD source 在其他 owner + owner_requester->>owner_ssd: SsdStageReadReq(get_id, stage, target, len) + loop 每个 payload chunk + owner_ssd->>owner_media: pin 并读入 source staging + owner_media-->>owner_ssd: IO 完成，staging chunk ready + owner_ssd->>owner_requester: push chunk 到 target offset + end + owner_ssd->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_ssd: GetDoneResp(holder_id) + owner_ssd-->>owner_requester: SsdStageReadResp(done_*) + end + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else 当前 key-version 没有可读副本 + master_route-->>owner_requester: KeyNotFound + owner_requester-->>external_user: miss + end +``` + +三处实现边界需要单独说明： + +- **本地 SSD 回填**：`src_addr == target_addr`，控制面不申请 source staging。地址、offset 和容量满足 direct read 条件时，`KvSsdStorage` 直接写最终 target；否则内部 scratch fallback 会增加一次本地 copy。 +- **远端 SSD staging**：master 先计算 `ssd_stage_len = align_ssd_io_len(len)`，再从 SSD source owner 已注册的 segment 中额外分配 `SSD_ALIGNMENT - 1` bytes，从 allocation range 内取得 512-byte 对齐的 `src_addr`。requester 只发送一次 `SsdStageReadReq`，后续 chunk 均由 SSD owner 主动 push。 +- **陈旧 SSD route**：读取失败后先 revoke 当前 source，再重新选路；该失败不会向用户返回半成品 holder。 + +### 4.1 target 分配与有界背压 + +`get` target 遇到内存压力时按以下顺序处理： + +1. **遍历 allocator**：先尝试所有可用 allocator，不在单个 allocator 上连续空转。 +2. **推进回收**：运行该 owner 的 replica Moka cache、`inflight_puts` 和 `inflight_gets` maintenance tasks，然后立即重试。`Cache::remove()` 会让 entry 立即不可查询，但 RAII value 可能仍在维护队列中保活 allocation。maintenance 完成后，空间才真正释放。 +3. **等待外部释放**：仍无空间时，让出执行权并等待 route 删除、owner 缓存失效 RPC 或 holder ACK。重试间隔为 2 ms，总等待上限为 5 s。 +4. **有界失败**：5 s 后仍无空间才返回带 total/free capacity 的 `NoSpace`；等待不会扩大 owner 的配置容量。 + +新 target 在传输期间尚未进入驱逐 cache，但已经占用 requester owner 的 segment。master 为可提升为内存副本的 durable target 立即预留 allocation capacity，并通过 `InflightGetInfo` 或 allocation 的 on-drop callback 保活 reservation。 + +这份预留与异步 SSD persist、lease 共用 `cache_reserved_bytes`，避免把在途 `get` 占用误算为空闲预算。`GetRevoke`、60 秒 in-flight TTL 或 `GetStartResp` 发送失败都会释放它。 + +temporary target 不进入 Moka。`GetDone` 后，master 的 `get_holding` 继续持有 target guard，直到 owner 侧最后一个 `MemoryInfo` 释放并向 master ACK，或 member-left 清理该 requester。requester owner 的非 cache 空间因此要容纳并发 target，以及尚未走完释放链的 temporary holder。 + +### 4.2 `GetDone` 与 target 生命周期 + +- **非 lease 的 durable promotion 成功**：master 先更新 route，再释放独立的临时预留，并在同一次请求处理中按相同 weight 插入 Moka，避免 route 与容量索引不同步。 +- **带 lease 的 durable promotion 成功**：target 不进入 Moka；reservation 已绑定到 allocation 的 on-drop callback，随该 leased replica 的 allocation 一起存活。 +- **无法 promotion**：route 版本变化、节点 tomb 或 route 已删除时，target 降级为 temporary holder 并释放 durable slot。非 lease 路径的独立 reservation 随 `InflightGetInfo` 释放；lease 路径的 reservation 已绑定 allocation，会保留到 temporary holder 释放。 +- **所有成功路径**：`GetDone` 都把 requester target guard 写入 master 的 `get_holding`；远端 SSD source staging 和 memory source 的额外保活引用随 `InflightGetInfo` 释放。 + +与 allocation 生命周期直接相关的字段如下： + +```rust +pub struct InflightGetInfo { + pub put_id: PutIDForAKey, // Guard the selected key version. + pub src_node_id: NodeID, // Identify the selected source owner. + pub allocation: Arc, // Hold the requester target until GetDone. + pub source_allocation: Option>, // Keep a memory source or remote SSD staging alive. + pub route: Arc, // Keep the route and its durable slot alive. + pub allocation_mode: GetAllocationMode, // Decide whether promotion is allowed. + pub source_kind: GetSourceKind, // Select memory or SSD cleanup semantics. + pub cache_capacity_reservation: Option>, // Reserve durable target capacity. + // ... Other request identity and length fields. +} +``` + +`InflightGetInfo` 的核心作用是在一次 `get` 内同时保活 target、source 和 route： + +- **`allocation`** 始终表示 requester segment 中的 target；`GetDone` 后由 `OwnerHoldingGetInfo` 继续保活。用户侧 `MemHolder` 引用这段 bytes，但不持有这个 master 侧 Rust 对象。 +- **`source_allocation`** 是 master 侧的保活引用：memory 路径用它保活内存副本，远端 SSD 路径用它保活 staging，本地 SSD 回填时为 `None`。`GetDone`、`GetRevoke` 或 in-flight TTL 清理都会释放该引用。 +- **`route` 与容量预留** 共同保证 durable promotion 期间的 route 生命周期和容量预算。非 lease 成功路径在同一次请求中更新 route、释放独立 guard 并插入 Moka；lease 路径则由 allocation callback 延续 reservation。 +- **`put_id`、`src_node_id` 和 `source_kind`** 限定失败清理范围，避免撤销新版本 route 或错误的介质副本。 + +external 公共 `get` 的 holder 交付跨越三层对象。只有 master 的 `get_holding` 持有 allocation guard，不能把三层对象统称为“holder 持有 allocation”： + +| 层级 | 实际持有 | 正常释放链 | +| --- | --- | --- | +| master | `get_holding[(requester_owner_id, holder_id)]` 中的 `OwnerHoldingGetInfo`，内部持有 requester target 的 `Arc`。 | owner 的 delete-ack batch 通过 `BatchDeleteAckReq` 到达后移除；requester member-left 时也会批量清理。 | +| requester owner | `external_get_holding[(external_client_id, holder_id)]` 中的 `Arc`。它引用 target metadata，未持有 master 进程里的 Rust guard。 | external holder 释放后收到 `ExternalDeleteAckReq` 并移除该引用；owner 侧最后一个 `MemoryInfo` 释放时把记录放入 delete-ack batch，再由 `BatchDeleteAckReq` 通知 master。`get_cached_info` 若仍有引用，master holding 会继续保留。 | +| external process | 根据共享内存 base、`ExternalMemHolderInfo.offset/len` 和 `holder_id` 构造的公共 `MemHolder`。 | 对象释放时向 requester owner 发送 `ExternalDeleteAckReq`；它不直接访问或释放 master 的 `Allocation` 对象。 | + +holder 释放只移除 `get_holding` 这条引用。`ReuseReplica` 或 durable promotion 成功后，route 的 `memory` 仍持有独立的 `Arc`；该副本要经过缓存驱逐或 route 清理才会释放。temporary target 没有这条 route 引用。 + +## 5. 读取子步骤——回填传输：本地复用 target，远端 source 主动 push + +上一节已经确定了读取 source、requester target 及其生命周期；进入实际搬运阶段后，memory 与 SSD 链路在 source 形态上分开。memory route 指向 owner 已注册 memory segment 中的现成 bytes，requester 取得 source 和 target 地址后即可发起 pull。SSD route 指向 shard 文件中的 entry，磁盘内容还不是 transfer engine 可直接使用的内存 source，必须先由 SSD owner 完成本地读盘。 + +最终 target 始终位于 requester owner 的 memory segment，并由 master 侧 guard 保活。当 SSD owner 与 requester owner 为同一节点时，本地 IO 可直接写入 target。当两者不同时，本地 IO 无法跨 owner 写入 segment。此时，SSD owner 先用本地 source staging 接收 chunk，待 chunk ready 后主动 push 到远端 target。这个方向既由 SSD IO 与跨节点传输的边界决定，也让远端路径可以重叠读盘和网络传输。 + +### 5.1 本地回填复用最终 target + +SSD owner 与 requester owner 在同一节点时，master 令 `src_addr == target_addr`，并把 target allocation 的实际 capacity 放进 `ssd_stage_len`。owner 随后调用 `load_into_addr`：满足 direct read 条件时，读盘直接写 target；否则由 `KvSsdStorage` 通过 scratch buffer 完成。两条分支都不需要控制面申请独立 source allocation。 + +```rust +if peer_id.is_none() && stage_addr == target_addr { + return store + .load_into_addr( + key, // Select the logical KV entry. + put_id, // Select the exact committed key version. + target_addr, // Read into the final range in the requester segment. + len, // Expose only the real payload length. + stage_len, // Provide the target allocation capacity for aligned IO. + ) + .await; +} +``` + +### 5.2 远端传输方向 + +两条远端路径的数据都从 source 流向 requester target，差别在 transfer 发起方： + +- **内存：requester pull**：requester 已从 `GetStartResp` 取得远端 `src_addr` 和本地 `target_addr`，由它发起从 peer source 到本地 target 的 transfer。 +- **SSD：owner push**：requester 只发一次 `SsdStageReadReq`。SSD owner 每读好一个 chunk 就 push 到 requester target，省去 stage-ready 后的一次 RTT，并让读盘与网络传输重叠。 + +`transfer_data_no_copy` 的方向位决定 peer 端地址在本次传输中是 source 还是 target。该参数只描述数据方向，不表示 peer 持有 master 侧的 `Allocation` guard。下面是删减后的接口： + +```rust +pub async fn transfer_data_no_copy( + &self, // Use this node's transfer engine. + peer_node: Option, // Select the remote peer; None means local transfer. + peer_src_or_target: bool, // True: peer address is source; false: it is target. + src_addr: u64, // Provide the absolute source address. + target_addr: u64, // Provide the absolute target address. + len: u64, // Transfer this many real payload bytes. + seg_guard: Option, // Optionally keep the local segment alive. +) -> KvResult; +``` + +方向位同时决定 peer 地址的语义和本地 segment guard： + +- **`true`**：peer 地址是 source，当前节点 pull 到本地 target；owner 侧 segment guard 守住 `target_addr`。 +- **`false`**：peer 地址是 target，当前节点从本地 source push；owner 侧 segment guard 守住 `src_addr`。 + +远端 memory get 和远端 SSD 回填分别使用这两个方向： + +```rust +// Memory get: requester owner initiates transfer with peer as source. +client_transfer_engine + .transfer_data_no_copy( + Some(memory_owner_id), // Contact the owner that holds the memory replica. + true, // Interpret the peer address as the source. + remote_src_addr, // Read from the replica range in the peer's segment. + requester_target_addr, // Write into the requester's local target. + len, // Move only the real payload bytes. + None, // Let the engine acquire the local target guard. + ) + .await?; + +// SSD refill: SSD owner pushes a ready chunk to requester target. +let chunk_target_addr = requester_target_addr + .checked_add(chunk.offset) // Place this chunk at its payload offset. + .ok_or_else(|| KvError::Api(ApiError::InvalidArgument { // Convert overflow into a KV error. + detail: "chunk target addr overflow".to_string(), // Explain which address calculation failed. + }))?; + +client_transfer_engine + .transfer_data_no_copy( + Some(requester_owner_id), // Contact the requester whose segment contains the target. + false, // Interpret the peer address as the target. + chunk.stage_addr, // Read from this SSD owner's ready staging chunk. + chunk_target_addr, // Write to the matching offset in requester target. + chunk.len, // Move this chunk's real payload length. + None, // Let the engine acquire the local source guard. + ) + .await?; +``` + +只有远端 SSD 回填会发送 `SsdStageReadReq`。请求把 key-version、`get_id`、staging 地址与容量、requester target 位置和真实 payload 长度交给 SSD owner： + +```rust +pub struct SsdStageReadReq { + pub key: String, // Identify the logical KV entry. + pub put_id: PutIDForAKey, // Reject data from a different key version. + pub get_id: u64, // Bind the stage operation to its in-flight get. + pub stage_addr: u64, // Point to aligned staging on the SSD owner. + pub stage_len: u64, // State the available aligned staging capacity. + pub target_node_id: NodeIDString, // Identify the requester whose segment contains the target. + pub target_addr: u64, // Point to the final target range in that segment. + pub len: u64, // Preserve the real payload length. +} +``` + +### 5.3 chunk pipeline 与背压 + +远端 SSD owner 把读盘和传输拆成两个并发 future： + +- **producer**：按 chunk 从 SSD 读到 source staging。 +- **consumer**：收到 ready chunk 后立刻 push 到 requester target。 + +两者通过有界 `mpsc` ready queue 连接。默认 chunk 大小为 4 MiB；producer 最多并发 4 个 SSD read，consumer 最多保留 4 个 transfer inflight，ready queue 容纳 8 个已就绪 chunk 描述。第一个 chunk 读好后即可开始网络传输。 + +#### producer / consumer 实现片段 + +```rust +let ready_queue_capacity = DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT + .saturating_mul(2) // Buffer two transfer windows between reader and sender. + .max(1); // Keep the bounded channel valid for any configured window. +let (chunk_tx, chunk_rx) = + ::tokio::sync::mpsc::channel(ready_queue_capacity); // Connect producer and consumer. + +let producer = store.load_into_addr_chunks( + key, // Select the logical KV entry. + put_id, // Select the exact committed version. + stage_addr, // Read chunks into aligned source staging. + len, // Stop after the real payload length. + stage_len, // Enforce the staging allocation capacity. + DEFAULT_READ_TRANSFER_PIPELINE_CHUNK_BYTES, // Bound each SSD read to 4 MiB. + DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT, // Allow up to four concurrent SSD reads. + chunk_tx, // Publish each ready chunk to the sender. +); + +let consumer = self.transfer_loaded_ssd_chunks( + peer_id, // Select the remote requester; None means the target is local. + target_addr, // Use the final target range in the requester segment. + chunk_rx, // Consume chunks as soon as their SSD reads complete. +); +let (producer_res, consumer_res) = ::tokio::join!( + producer, // Run SSD reads concurrently with network transfers. + consumer, // Push ready chunks without waiting for the full value. +); +match (producer_res, consumer_res) { + (Ok(()), Ok(())) => Ok(()), + (_, Err(err)) => Err(err), + (Err(err), _) => Err(err), +} +``` + +consumer 只在 transfer inflight 少于 4 时从 ready queue 取新 chunk。达到上限后，有界 queue 会把背压传给 producer；已提交的 transfer 完成后，consumer 再继续取 chunk。 + +这条 pipeline 中有三种不同的保活机制： + +| 机制 | 保护对象 | 释放时点 | +| --- | --- | --- | +| ring read pin | SSD entry 与文件 offset，防止尚未完成的 SSD read 被 tail 覆盖。 | producer 完成全部 SSD read，并把最后一个 ready chunk 送入有界 queue 后。 | +| `InflightGetInfo.source_allocation` | master 侧 guard，保活 SSD owner segment 中的整块 source staging range。 | `GetDone`、`GetRevoke` 或 in-flight TTL。 | +| transfer segment guard | 某个在途 push 使用的本地 staging 地址范围。 | 对应 chunk transfer 完成。 | + +ring read pin 只覆盖 SSD 读取阶段，不覆盖整个网络 push 窗口。producer 释放 pin 后，尚未完成的 push 由 master 侧 staging guard 和每次 transfer 的 segment guard 保护。 + +下图聚焦远端 SSD 回填内部的 pipeline；`GetStart/GetDone` 的完整流程见第 4 节。`owner-requester-target` 只表示数据传输端点，其 allocation 生命周期仍由 master 的 in-flight 状态和 `get_holding` 管理。 + +```mermaid +sequenceDiagram + autonumber + participant owner_reader as owner-ssd-reader + participant owner_media as owner-local-ssd + participant owner_queue as owner-ready-queue + participant owner_sender as owner-ssd-sender + participant owner_target as owner-requester-target + + owner_reader->>owner_media: pin committed ring entry + par read producer + loop 按 4 MiB chunk 读取 + owner_reader->>owner_media: read chunk N into source staging + owner_media-->>owner_reader: IO complete, chunk N ready + owner_reader->>owner_queue: enqueue chunk N + end + owner_reader->>owner_media: 最后一个 ready chunk 入队后 release ring pin + and transfer consumer + loop 按 payload offset 推送 + owner_queue-->>owner_sender: deliver chunk N + owner_sender->>owner_target: push chunk N + owner_target-->>owner_sender: transfer complete + opt transfer inflight 已达 4 + owner_sender->>owner_queue: pause dequeue + owner_queue-->>owner_reader: bounded-channel backpressure + end + end + end + Note over owner_sender,owner_target: 未完成的 push 继续由 master staging guard 和 transfer guard 保活 +``` + +## 6. owner 本地机制：对齐、ring 和 IO 调度 + +**SSD 介质的控制引擎留在 owner。** SSD 文件 offset、`O_DIRECT` 文件 IO 对齐、ring 覆盖保护和 `io_uring` 调度都由 owner 处理。master 仍负责分配 requester target 和跨节点 staging，并传递相应地址、容量与真实 payload 长度。用户只看到真实 payload 语义。 + +master 使用已注册的 owner segment allocator 为每次回填分配 requester target，并创建对应 guard；跨节点回填还会分配 SSD source staging。bytes 位于相应 owner 的 segment，RAII guard 位于 master。master 不记录 SSD 文件布局，用户侧仍通过 `MemHolder` 访问真实 payload。 + +### 6.1 写入：区分真实长度和对齐长度 + +**`len` 是 payload 长度，`aligned_len` 是 owner 本地 IO 长度。** target owner 从请求携带的 target 地址复制 payload，再构造 512-byte 对齐的写入 buffer。route 中的 `Arc` 在此期间保活 target 地址范围。 + +| 字段 | 含义 | 使用范围 | +| --- | --- | --- | +| `len` | 真实 payload 长度。 | master route、`SsdReplicaCommitReq`、transfer 长度和 `MemHolder.len`。 | +| `aligned_len` | `align_up(len, 512)` 得到的对齐长度。 | owner 本地 ring 分配、shard offset 推进和 `O_DIRECT` read/write。 | + +例如，1000-byte payload 的 `aligned_len` 是 1024 bytes： + +1. owner 创建 1024-byte 的对齐 `AlignedBuffer`，先置零，再把前 1000 bytes 填入 payload。 +2. ring 分配、shard offset 推进和 SSD 写入都使用 1024 bytes。 +3. `SsdIndexEntry` 同时记录 `len = 1000` 和 `aligned_len = 1024`；owner 向 master 提交的仍是 `len = 1000`。 +4. 后续读取会校验 `entry.len`，transfer 和 `MemHolder` 只暴露前 1000 bytes。末尾 24 bytes padding 留在 owner 内部。 + +### 6.2 ring：从 `Writing` 到 `Committed` + +**只有完整写入的 `Committed` entry 才能参与读取；写入、route 提交或读取期间，ring 不会覆盖对应区间。** + +本地 SSD 引擎按 value 粒度选路： + +- 用 `next_write_device` round-robin 选择一个有效 device。 +- 把整个 value 发送到选中 device 的 writer queue。 +- 当前实现不把同一个 payload 拆到多块 device 做条带化。 + +因此，多 device 并行发生在 value 粒度：不同 value 可以落到不同 device；一个 value 内部仍是某个 shard 的连续 offset。 + +owner 本地 `SsdRingBuffer` 决定 SSD 文件位置。writer task 只在本 device 的 `shard_ids` 里选择 shard，并为对齐后的 payload 分配连续文件区间。ring 按以下规则分配： + +- `aligned_len = align_up(len, 512)`，ring 按对齐后长度分配空间。 +- 每个 shard 是环形空间，`head` 前进分配新写入，`tail` 表示可以覆盖到哪里。 +- 如果当前 shard 尾部剩余空间不足，会跳到文件开头继续找连续空间。 +- 如果推进 `tail` 会覆盖未完成写入、active read pin 或 route commit pin，返回 `BlockedByBusyIo`。 + +分配完成后，状态转换、覆盖保护和 route 清理遵循以下规则： + +- **`Writing`** + - offset 已分配，SSD 写入尚未成功。 + - 默认单 buffer 路径使用 `O_DIRECT` 和 `io_uring` write。 +- **`Committed`** + - 只有实际写入字节数等于 aligned buffer 长度，entry 才从 `Writing` 转为 `Committed`。 + - `pin_read` 只接受 committed 且 offset 仍有效的 entry，未完成写入不会被 `get` 选中。 +- **ring 覆盖保护** + - route commit pin：保护本地 commit 到 master 处理完 commit RPC 的窗口。 + - read pin：保护 SSD read 使用 entry 与 file offset 的窗口；远端 push 后半段由 staging allocation 和 transfer guard 保活。 +- **tail 驱逐通知** + - tail 推进时，ring 返回本次被覆盖的全部 committed `KvSsdKey`。 + - writer 把它们送入有界失效队列，不在 SSD 写入热路径逐条调用 master。 + - owner batch task 按 256 条或 10 ms 的阈值聚合 RPC；第 7 节展开 master 的条件清理和读取兜底。 + +下图展示 owner 本地 `Writing → Committed` 与 master 发布 SSD route 之间的时序。`owner-memory-target` 只有前 `len` bytes 是真实 payload，其地址范围由 master route 中的 `Arc` 保活。 + +```mermaid +sequenceDiagram + autonumber + participant master_route as master-route + participant owner_target as owner-memory-target + participant owner_writer as owner-ssd-writer + participant owner_ring as owner-ssd-ring + participant owner_engine as owner-io-engine + participant owner_media as owner-local-ssd + + master_route->>owner_writer: SsdReplicaPersistReq(key, put_id, target_addr, len) + owner_writer->>owner_target: copy first len bytes + owner_target-->>owner_writer: real payload + owner_writer->>owner_writer: zero padding and build aligned_len buffer + owner_writer->>owner_ring: reserve contiguous aligned_len bytes + owner_ring-->>owner_writer: shard_id + file_offset + Writing entry + owner_writer->>owner_engine: O_DIRECT write(aligned buffer) + owner_engine->>owner_media: submit aligned_len bytes + owner_media-->>owner_engine: write completion(actual bytes) + owner_engine-->>owner_writer: completion + + alt actual bytes == aligned_len + owner_writer->>owner_ring: Writing → Committed and acquire route commit pin + owner_ring-->>owner_writer: KvSsdPersistGuard + owner_writer->>master_route: SsdReplicaCommitReq(key, put_id, len) + master_route->>master_route: validate, then publish or ignore commit + master_route-->>owner_writer: SsdReplicaCommitResp + owner_writer->>owner_ring: release route commit pin + owner_writer-->>master_route: SsdReplicaPersistResp(persisted = true) + else partial write or IO failure + owner_writer->>owner_ring: abort Writing entry + owner_writer-->>master_route: SsdReplicaPersistResp(error) + end +``` + +### 6.3 读取：direct read 和 scratch fallback + +**direct read 是有条件的 owner 本地 fast path；任一对齐或容量条件不满足时，scratch read 负责兜底。** + +每个 chunk 的文件位置由 `entry.file_offset + offset` 得到。本地回填读入 requester target；远端回填读入 SSD source staging。 + +| 路径 | 选择条件 | owner 本地动作 | 复制边界 | +| --- | --- | --- | --- | +| direct read | 读入地址和文件 offset 都按 512 bytes 对齐，target capacity 覆盖对齐后的 IO 范围。按 chunk 读取时，当前 read 长度也必须对齐。 | `io_uring` 直接把 shard 内容读入本地目标地址。 | 这段 SSD-to-target 读取没有额外的 owner 本地 copy；这不代表整条 `get` 链路都是零拷贝。 | +| scratch read | 任一 direct read 条件不满足。 | 先读入 512-byte 对齐的 `AlignedBuffer`，再把当前 chunk 的真实 payload 复制到目标地址。 | 多一次 owner 本地 CPU copy。 | + +两条路径都只向下游暴露真实 payload 长度。`O_DIRECT` padding 不进入 transfer 长度，也不进入用户 `MemHolder.len`。 + +下图只画 SSD source owner 内部的读盘。`owner-read-target` 始终位于 SSD source owner 本地：本地回填时，它同时是 requester target；远端回填时，它是 source staging。远端 requester 的最终 target 不在图中，后续 push 见第 5.3 节。 + +```mermaid +sequenceDiagram + autonumber + participant owner_reader as owner-ssd-reader + participant owner_ring as owner-ssd-ring + participant owner_engine as owner-io-engine + participant owner_media as owner-local-ssd + participant owner_scratch as owner-scratch-buffer + participant owner_target as owner-read-target + + owner_reader->>owner_ring: pin committed entry and request chunk plan + owner_ring-->>owner_reader: file_offset + len + aligned read len + owner_reader->>owner_reader: check address, offset, length, and capacity alignment + + alt direct read 条件全部成立 + owner_reader->>owner_engine: read shard directly into owner-read-target + owner_engine->>owner_media: submit O_DIRECT read + owner_media-->>owner_engine: IO complete, target memory filled + owner_engine-->>owner_reader: completion + owner_reader-->>owner_target: publish current chunk's real len + else 任一对齐条件不成立 + owner_reader->>owner_scratch: acquire aligned buffer + owner_reader->>owner_engine: read shard into owner-scratch-buffer + owner_engine->>owner_media: submit O_DIRECT read + owner_media-->>owner_engine: IO complete, scratch memory filled + owner_engine-->>owner_reader: completion + owner_reader->>owner_target: copy only current chunk's real payload bytes + end + owner_reader->>owner_ring: release read pin + Note over owner_reader,owner_target: transfer 和 MemHolder 只看到真实 len，padding 留在 owner 内部 +``` + +### 6.4 IO 调度：让回填读优先，同时保留后台写进展 + +**同一个 device 上，SSD 回填读位于 `get` 的同步等待路径，SSD persist 写则在内存副本发布后异步进行。** 因此当前调度给读更多的 SQE 提交机会，但仍持续推进写。写长期不进展会延迟 SSD 副本 commit，并让 `Writing` entry 和对齐 buffer 继续占用资源；持续读下的绝对读优先还可能拖住 ring 后续复用空间。`read_weight = 2` 表达的是这种有界的读偏置。 + +调度分成两层，各自处理不同问题： + +| 层次 | 当前机制 | 设计原因 | +| --- | --- | --- | +| device 任务层 | 每个去重后的 device 独立持有 writer queue、reader queue 和 `UringIoEngine`。writer task 管理 ring 分配以及 `Writing -> Committed/abort`；reader task 管理 chunk read 和 completion 复查。 | 在当前 owner 本地调度范围内，实际 device 是 IO 竞争边界。独立队列避免忙 device 在软件调度层阻住空闲 device，也让读写使用不同的上层并发上限和状态机。 | +| `io_uring` 提交层 | 每个 `io_uring` 线程各有 read/write channel，在 `io_depth` 内选择下一个 SQE。 | 分开 channel 后，引擎可以在真正占用 SQE 槽位时仲裁读写，避免后到的回填读只能排在单一 FIFO 中的后台写之后。 | + +`io_depth` 在设备并行度与在途资源占用之间设置上界：多个已提交 SQE 可以让 SSD 持续有工作可做，上限则避免单个 `io_uring` 线程无界占用 buffer 和 completion context。真正能到达这一层的并发量，还会受 writer/reader task 各自 inflight 上限的约束。 + +`read_weight = 2` 只在第二层生效。每个 `io_uring` 线程独立维护 `read_inflight` 和 `write_inflight`：前者小于等于后者的两倍时，先尝试 read channel；超过后，先尝试 write channel。首选 channel 为空时会立即取另一个 channel 的任务，不会为了维持比例而让 SQE 槽位空闲。 + +这个 `2:1` 是固定的软调度偏置，不是硬上限或队列积压阈值。等号分支会再提交一个 read，completion 速率也会持续改变 inflight 组成；因此它不保证单个 device 上的精确 `2:1` 比例。它也不考察队列长度、请求大小、等待时间或 deadline。权重只决定下一个尚未提交的 SQE；已提交的 IO 不会被抢占，`io_depth` 已满时，新的回填读仍需等待 CQE 释放槽位。所以这是一个低成本的提交偏置，不构成读延迟 SLA。 + +**`SingleBuffer` 是对当前数据形态的特化。** 写路径已经把 payload 放入连续的对齐 `AlignedBuffer`；direct read 和 scratch read 的目标也都是一段连续内存。默认 `KvSsdUringMode::SingleBuffer` 因而直接提交 `IORING_OP_READ` / `IORING_OP_WRITE`，避免构造并保活只含一个元素的 iovec。`Iovec` 保留为测试和对照路径。这项特化没有消除 `O_DIRECT` 对齐、padding、scratch copy 或网络传输，其收益范围仅限于 owner 本地 IO 提交层。 + +read completion 复查与 read pin 也承担不同职责： + +| 机制 | 作用时段 | 职责 | +| --- | --- | --- | +| read pin | 从获取 committed entry 到本次读取结束。 | 阻止 `tail` 越过正在使用的文件区间，预防 offset 在 IO 期间被复用。 | +| completion 复查 | CQE 返回后。 | 检查读取开始时捕获的 `entry.begin` 是否仍不早于当前 `tail`。已失效时删除 entry 并返回 `KeyNotFound`。 | + +read pin 负责预防覆盖，completion 复查负责拦住异常状态下的结果发布。复查无法撤销已经写入 target 的 IO，因而不能取代 pin；它是跨过异步提交边界后的最后一道不变量检查。在完整 `get` 路径中，返回错误后的 SSD stage 会进入第 7.2 节的 revoke 与有界重新选路，不会向用户交付这次 target。 + +下图聚焦单个 device 内部的调度，不涉及 `master-*` 或 `external-*` 角色。`owner-kernel-io` 表示该 owner 文件上的 `io_uring` SQE/CQE 交互，不表示 owner 独占内核。 + +```mermaid +sequenceDiagram + autonumber + participant owner_write as owner-write-path + participant owner_read as owner-read-path + participant owner_engine as owner-io-engine + participant owner_kernel as owner-kernel-io + participant owner_media as owner-local-ssd + participant owner_ring as owner-ssd-ring + + par write task arrives + owner_write->>owner_ring: reserve offset and create Writing entry + owner_ring-->>owner_write: shard_id + file_offset + owner_write->>owner_engine: enqueue write request + and read task arrives + owner_read->>owner_ring: pin committed entry and resolve offset + owner_ring-->>owner_read: read plan + pin + owner_read->>owner_engine: enqueue read request + end + + loop while this io_uring thread has a free SQE slot + owner_engine->>owner_engine: choose next unsent IO by the per-thread inflight mix, then fall back to the other channel + Note over owner_engine,owner_kernel: read_weight does not preempt submitted SQEs + alt write request selected + owner_engine->>owner_kernel: submit WRITE SQE + owner_kernel->>owner_media: O_DIRECT write + owner_media-->>owner_kernel: device completion + owner_kernel-->>owner_engine: write CQE + owner_engine-->>owner_write: actual bytes or error + owner_write->>owner_ring: commit or abort Writing entry + else read request selected + owner_engine->>owner_kernel: submit READ SQE + owner_kernel->>owner_media: O_DIRECT read + owner_media-->>owner_kernel: device completion + owner_kernel-->>owner_engine: read CQE + owner_engine-->>owner_read: actual bytes or error + owner_read->>owner_ring: check captured entry.begin is not before current tail + alt captured range 仍有效 + owner_ring-->>owner_read: range valid, return chunk bytes + else captured range 已失效 + owner_ring-->>owner_read: stale + owner_read->>owner_ring: remove stale entry + owner_read->>owner_read: return KeyNotFound + end + end + end +``` + +## 7. 一致性：沿用 key-version 和 holder 生命周期 + +**SSD 副本只有在 master route 和 owner ring 都有效时才可读。** `get_holding` 保存的是 requester target 的 master 侧 `Allocation` guard；远端 source staging guard 只保留在本次 `InflightGetInfo` 中。 + +一次 SSD 回填的完成语义包含以下四个方面： + +- **`GetDoneReq` 调用方** + - 本地 SSD source：requester 读盘并完成 target 写入后调用。 + - 远端 SSD source：SSD owner 读盘并 push 到 requester target 后调用。 +- **最终 holder 与 allocation 生命周期** + - `inflight_gets.req_node_id` 始终指向 requester owner，不受 `GetDoneReq` 调用方影响。 + - master 把 requester target guard 写入 `get_holding` 并生成 `holder_id`；target bytes 仍位于 requester owner 的 segment。 + - 远端 source staging guard 不进入 `get_holding`，只随 in-flight get 生命周期释放。 +- **远端响应转发** + - SSD owner 把 `GetDoneResp` 中的 `holder_id`、`allocation_mode`、错误码和 server time 写入 `SsdStageReadResp.done_*`。 + - requester 使用这些字段构造 owner 侧 `MemoryInfo` 和 `ExternalMemHolderInfo` 响应。external 再构造公共 `MemHolder`，用户看不到 SSD stage RPC。 +- **durable promotion** + - `GetStart` 尝试预留 durable slot。成功时 `allocation_mode = DurableReplica`，否则为 `Temporary`。 + - `GetDone` 只在 allocation mode 为 durable、route `put_id` 仍匹配、且 requester tomb tag 仍 live 时，才把 target guard 写入 entry 的 `memory`。 + - promotion 成功后，SSD 回填结果成为新的内存副本。 + +### 7.1 owner 驱逐后批量清理 route + +为了让 master 中因物理驱逐产生的陈旧 route 尽快收敛，同时尽量减少逐条通知产生的控制面通信，owner 先完成物理驱逐，再聚合对应的 key-version，批量通知 master 清理逻辑 route： + +- **驱逐输入**：ring 推进 `tail` 时，一次返回本轮覆盖掉的全部 committed key-version。 +- **batch 触发** + - 第一条记录到达后，最多等待 10 ms。 + - 累计到 256 条时立即发送。 +- **RPC 失败**：当前 batch 留在 task 中重试。 +- **队列或控制面不可用** + - 有界队列满、进程退出或 master 持续不可用，都不阻塞 SSD 写入。 + - 可能残留的陈旧 route 由第 7.2 节的读取兜底处理。 + +批量 RPC 只携带待清理的 key-version： + +```rust +pub struct SsdReplicaEviction { + pub key: String, // Identify the route entry to inspect. + pub put_id: PutIDForAKey, // Restrict cleanup to the evicted key version. +} + +pub struct BatchSsdReplicaEvictReq { + pub evictions: Vec, // Carry one owner-side eviction batch. +} +``` + +master 收到 batch 后按以下规则清理 route： + +- **owner 身份**：请求不重复传 `node_id`，master 直接使用 RPC 发送方身份。 +- **版本门禁**：只有当前 route `put_id` 与驱逐项匹配，master 才清除该 owner entry 的 `ssd`。 +- **清理粒度** + - entry 仍有内存副本：只清除 `ssd`。 + - `memory` 和 `ssd` 都为空：删除 owner entry。 + - 整个 key-version 没有 live 副本：删除 `kv_routes`。 +- **prefix index** + - 在一次写锁内批量清理。 + - 同名 key 已写入新 route：保留新版本 prefix entry。 +- **master 职责边界**：master 不参与 SSD offset 分配，也不决定本地驱逐顺序；只接收 owner 的驱逐结果并收敛逻辑 route。 + +### 7.2 陈旧 SSD route 的读取兜底 + +**SSD stage 失败后，requester 先撤销当前 source，再有界重新选路。** 重新选路可以继续使用其他 live source；没有副本时返回 miss。连续 3 次 SSD stage 失败后返回 error。任何失败分支都不会向用户暴露半成品 holder。 + +SSD stage 失败时，请求方会调用： + +```rust +GetRevokeReq { + get_id, // Revoke this in-flight get and release its allocations. + drop_ssd_source: true, // Clear the failed SSD option from the current route. +} +``` + +master 按下列顺序撤销失败 source： + +- **释放本次 in-flight**：先删除 `inflight_gets`，让 target 和可选 staging 进入释放路径。 +- **限定 route 删除条件** + - `drop_ssd_source == true`。 + - `source_kind == GetSourceKind::Ssd`；memory source 不会因该标志被删除。 + - route `put_id` 与本次 in-flight get 仍一致；不撤销新版本 route。 +- **清理失败 source** + - 把 `src_node_id` 对应 entry 的 `ssd` 设为 `None`。 + - entry 也没有内存副本：删除该 entry。 + - 整个 key-version 已无 live 副本：删除 `kv_routes`；启用 prefix index 时，再异步清理对应索引。 + - 失败 source 清理与主动 batch 驱逐共用同一个条件删除函数。 +- **重新选路** + - `GetRevokeResp` 返回后，请求方重新发送 `GetStartReq`。 + - 下一次可能命中另一份内存副本、SSD 副本、同 key 新版本，或返回 miss。 + - 陈旧 SSD source 最多触发 3 次重新选路，避免无限循环。 + - SSD stage RPC 失败和不合法的 `ssd_stage_len` 都进入同一套 revoke 与重选流程。 +- **allocation 与崩溃清理** + - 远端 staging 保存在 `source_allocation`；本地回填时该字段为 `None`。 + - `GetDone`、`GetRevoke` 或 in-flight TTL 驱逐都会释放远端 staging。 + - requester 在 in-flight 阶段崩溃：节点 tomb 阻止副本 promotion 和旧节点复用。`inflight_gets` 最迟由 60 秒 TTL 释放 target/staging guard。 + - requester 在 `GetDone` 后离开：master 的 member-left 清理移除该节点的 `get_holding`；正常路径则由 owner 的 holder delete-ack batch 释放。 + - 远端 SSD owner 崩溃：其进程内 staging bytes 消失。master 中对应的 staging guard 随 `GetRevoke` 或 in-flight TTL 释放，节点 tomb 使旧 SSD route 不再可选。 + +在 master 接受 `GetDoneReq` 之前发生的读取或传输失败会走 `GetRevoke`，target guard 不进入 `get_holding`，用户也不会拿到半成品 payload。 + +另一个边界发生在 master 已提交 `GetDone`、但响应随后丢失时。此时 target guard 已写入 `get_holding`，requester 却可能没有收到 `holder_id`。 + +当前通用 holder 协议没有按 `get_id` 幂等查询已提交结果的接口。这类未交付 holder 无法由随后的 `GetRevoke` 立即回收，只能等待能够识别该 holding 的后续清理，例如 requester member-left。这不影响 SSD 数据完整性，但属于当前完成协议的资源回收边界。 + +下图展示成功完成、失败 revoke 和超时清理三条路径。远端路径由 `owner-ssd-source` 代替 requester 调用 `GetDoneReq`，但最终 target 始终位于 requester segment。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant master_route as master-route + participant owner_ssd as owner-ssd-source + + owner_requester->>owner_ssd: SsdStageReadReq 或本地 stage 调用 + alt SSD stage 成功，target 已完整 + owner_ssd-->>owner_requester: 本地回填完成或最后一个 chunk 已 push + alt SSD source 就在 requester owner + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，释放其他 in-flight 引用 + master_route-->>owner_requester: GetDoneResp(holder_id) + else SSD source 在其他 owner + owner_ssd->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，释放 staging 和其他 in-flight 引用 + master_route-->>owner_ssd: GetDoneResp(holder_id, allocation_mode) + owner_ssd-->>owner_requester: SsdStageReadResp(done_*) + end + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else SSD stage 失败 + owner_ssd-->>owner_requester: stage error + owner_requester->>master_route: GetRevokeReq(drop_ssd_source = true) + master_route->>master_route: remove inflight_gets and release target / staging references + alt source_kind = Ssd 且 route put_id 仍匹配 + master_route->>master_route: 清空当前 owner entry.ssd，必要时删除 entry / route + else source 类型或版本已变化 + master_route->>master_route: 保留当前 route + end + master_route-->>owner_requester: GetRevokeResp + loop 陈旧 SSD source 最多重新选路 3 次 + owner_requester->>master_route: GetStartReq + master_route-->>owner_requester: 其他 live source / 新版本 / miss + end + owner_requester-->>external_user: 重试结果、miss 或达到上限后的 error + end + + opt 请求未进入 GetDone 或 GetRevoke + master_route->>master_route: 60 s in-flight TTL 释放 target / staging guard + end + opt requester 在 GetDone 后离开 + master_route->>master_route: member-left 清理该节点的 get_holding + end +``` + +### 7.3 删除、驱逐与重启 + +- **覆盖写与 delete**：覆盖写生成新 `put_id`，旧 SSD late commit 因版本不匹配被忽略。显式 delete 移除整个 `kv_routes` 条目；节点 tomb 清理则删除该节点在 `node_replicas` 中的整个 entry。 +- **两层独立驱逐**：内存驱逐只清空 `memory`；SSD ring 驱逐或读取失败只清空 `ssd`。两个字段都为空时才删除 owner entry，整个 route 没有副本时才删除 `kv_routes`。 +- **物理 bytes 不等于可读副本**：route 被删除或版本不匹配后，旧 bytes 即使还在 shard 中也不会被公共 `get` 命中。陈旧 route 被选中时，`pin_read` 会复查 committed 状态和 offset，失败后 revoke source 并重新 `GetStart`。 +- **owner 重启**：shard 以 `create + truncate` 打开，并重建空的本地 SSD cache 和 ring。当前没有 WAL、checkpoint 或 shard 扫描。控制面尚未完成 tomb 清理时，读取兜底会撤销失效 SSD source，再重新选路或返回 miss。 + +## 8. 性能测试 + +本节先给出核心结果，再说明计数边界与实验设置，最后分别展开 CPU、GPU 和 SSD 路径证据。性能数字只在对应小节定义的硬件、数据集、输出模式和统计窗口内成立。 + +### 8.1 核心结果 + +- **CPU / GPU 使用同一 pressure 口径**：CPU 主图包含 192 条记录，覆盖原生引用 / handle 与 `bytes`、4 个 payload、4 个并发度、Fluxon、两种 Mooncake 拓扑和 SSD 关/开；GPU 主图包含 72 条记录，覆盖 `cuda`、3 个 payload、4 个并发度、3 种拓扑和 SSD 关/开。264 条记录全部为 `SUCCESS`、0 error、来源完整且 unknown 为 0。 +- **miss 是结果，不是过滤条件**：两侧都使用约 2.5 GiB 数据集和 2 GiB DRAM。柱高只统计 hit payload，红点显示 `hit / (hit + miss)`。SSD-off 的 CPU 命中率为 Fluxon **89.112%–94.390%**、Mooncake 独立 owner **23.759%–66.879%**、Mooncake 每进程 **22.488%–68.655%**；GPU 对应范围为 **89.326%–93.267%**、**34.557%–62.223%** 和 **35.430%–64.782%**。 +- **SSD-on 也保留实际 miss**：CPU 的 Fluxon 和 Mooncake 独立 owner 在 SSD-on 下均为 100% 命中，Mooncake 每进程为 **72.830%–100%**；GPU 对应为 100%、100% 和 **59.855%–100%**。因此 SSD-on/off 柱差同时包含命中集合、介质来源和系统调度差异，不能解释成单一 SSD 开销。 +- **SSD-on 的 hit 同时来自两层**：CPU 中 Fluxon 的 SSD 来源占 hit 的 **17.550%–48.650%**，Mooncake 独立 owner 为 **62.065%–82.385%**，Mooncake 每进程为 **59.341%–83.011%**；GPU 对应为 **20.242%–46.508%**、**64.837%–77.135%** 和 **56.073%–84.198%**。 +- **GPU 多并发已进入同一张主图**：Fluxon 在 SSD-off 下从 c2 到 c16 提高到 5.83、5.83 和 4.68 倍，SSD-on 下提高到 6.50、7.22 和 6.80 倍，分别对应 4、8、16 MiB。Mooncake 的两种拓扑也按相同的六柱顺序展示，不再另画一张来源图。 +- **同路径 Foyer 消融下原生 SSD 仍然更快**：64 条测量记录对应 32 个原生/Foyer 配对，全部为 `SUCCESS`、0 error、0 unknown、0 persist failure，两侧命中率均为 100%。当前 Foyer adapter 的 hit payload 带宽是原生实现的 **14.8%–65.9%**，配对中位数为 **37.2%**，32 个配对中没有胜出；P95 是原生实现的 **2.30–8.46 倍**。这是当前接入路径的结果，不外推为 Foyer 库的通用性能结论。 + +这些数字只覆盖下文定义的单机、固定容量、Zipf 分布、30 秒窗口和各自 fast path。实验不包括跨机器 RDMA、多物理盘条带化、重复运行的统计置信区间或真实业务 trace。 + +### 8.2 计数边界：原生引用 / handle、`bytes` 与 `cuda` + +Mooncake 没有 Fluxon 意义上的 holder。`holder` 只是 benchmark 配置 `kv_get_output=holder` 的归一化模式名，Fluxon 和 Mooncake 实际返回 `MemHolder` 与 `BufferHandle`。 + +Mooncake 常规 `Store.get()` 直接返回 Python `bytes`，因此结果包含完整 payload copy。`Store.get_buffer()` 返回 native `BufferHandle`，无需创建 Python `bytes` 即可取得 `ptr()`、`size()` 或 `memoryview`。原生引用 / handle 组把计数边界放在 host 内存引用返回之后。 + +#### 两侧原生引用的调用与生命周期差异 + +benchmark 中的实际分支等价于： + +```python +if backend == "MOONCAKE": + # Internal adapter calls Mooncake Store.get_buffer(). + result = mooncake_store.get_buffer_blocking(key) # Fetch a native BufferHandle. +else: + result = fluxon_store.get_blocking(key) # Fetch a native MemHolder. + +native_ref = result.unwrap() # Keep the native host reference alive while counting success. +``` + +两者只在本次 benchmark 的**计数边界**上对齐，类型语义并不相同： + +| 对比维度 | Fluxon | Mooncake | 对齐方式 | +| --- | --- | --- | --- | +| 实际返回类型 | `MemHolder` | `mooncake.store.BufferHandle` | 本文直接写真实类型，不写“Mooncake holder”。 | +| 底层调用 | `KvClient.get_blocking()` | `Store.get_buffer()`，由 benchmark adapter 包成 `get_buffer_blocking()` | 都在 native host buffer ready 后返回。 | +| 数据访问 | `MemHolder.access()` | `ptr()`、`size()`、`memoryview(BufferHandle)` | 原生引用模式不 materialize；bytes / CUDA 模式再继续消费。 | +| 生命周期语义 | 参与 Fluxon 的 `holder_id`、`get_holding` 和释放流程。 | 由 Mooncake Python binding 的 `BufferHandle` 管理返回 buffer 的生命周期。 | 只要求对象存活到本次操作计数完成。 | +| 不对齐的部分 | Fluxon holder 协议。 | 没有 Fluxon 的 `holder_id/get_holding/revoke` 协议。 | 不声称 API 或生命周期协议等价。 | + +原生引用模式仍然包含后端选择 SSD route 时从 SSD 到 host 的搬运。它只省去随后将完整 payload 物化为 Python `bytes` 的额外 copy。`cuda` 从同一个 native holder 取得 pageable host pointer，再交给 Fluxon 与 Mooncake 共用的双槽 H2D 流水线。只有对应 CUDA event 完成后，这次操作才计为成功。 + +下图覆盖一次 benchmark `get` 从发起到达到对应计数边界的过程。`external-backend-adapter` 是 benchmark 侧的统一适配层，封装完整的 Fluxon 或 Mooncake get；它不表示两套系统共享 Fluxon 的 owner/master 架构。 + +- **`external-benchmark-worker`**：发起 KV get，按 `holder/bytes/cuda` 模式消费结果，并在该模式的完成点计数。 +- **`external-backend-adapter`**：在同一调用面上包住 Fluxon 或 Mooncake 的完整 get，并返回已驻留 host 的 `MemHolder` 或 `BufferHandle`。 +- **`external-host-output`**：表示 pageable host 上的 native reference 生命周期，以及 `bytes` 模式的完整 payload 物化。 +- **`external-pinned-slot`**：每个 benchmark worker 复用的两个 pinned host slot 之一。 +- **`external-cuda-runtime`**：对应 slot 的 CUDA stream 和 event，负责 `cudaMemcpyAsync` 提交与完成通知。 +- **`external-gpu`**：接收 payload 的 GPU device buffer；这次实验不计 D2H 或 kernel 执行。 + +```mermaid +sequenceDiagram + autonumber + participant external_worker as external-benchmark-worker + participant external_backend as external-backend-adapter + participant external_host as external-host-output + participant external_pinned as external-pinned-slot + participant external_cuda as external-cuda-runtime + participant external_gpu as external-gpu + + external_worker->>external_backend: get(key) + external_backend-->>external_worker: MemHolder or BufferHandle with host-resident payload + external_worker->>external_host: retain native reference until this operation is counted + + alt holder output + external_worker->>external_worker: validate native reference and count success + external_worker->>external_host: release native reference + else bytes output + external_worker->>external_host: materialize the full payload as Python bytes + external_host-->>external_worker: bytes ready + external_worker->>external_worker: validate bytes and count success + external_worker->>external_host: release native reference + else cuda output + external_worker->>external_pinned: host memcpy from pageable payload into the next free slot + external_worker->>external_cuda: cudaMemcpyAsync(pinned slot, device buffer) + external_cuda->>external_gpu: start H2D DMA + par current H2D + external_gpu-->>external_cuda: matching CUDA event completes + external_cuda-->>external_worker: H2D complete + and next KV request + external_worker->>external_backend: get(next key) + external_backend-->>external_worker: next native host reference + end + external_worker->>external_worker: count current success after its event + external_worker->>external_host: release current native reference + end +``` + +| benchmark 输出模式 | Fluxon 映射 | Mooncake 映射 | 共同计数边界 | +| --- | --- | --- | --- | +| `holder`（归一化模式名） | `get_blocking()` 返回 `MemHolder`。 | `get_buffer_blocking()` 调用 Mooncake `Store.get_buffer()`，返回 `BufferHandle`；该对象暴露 `ptr()`、`size()` 和 buffer protocol。 | 原生引用已返回；不创建完整 Python `bytes`。 | +| `bytes` | `MemHolder.access()` 解码出 payload `bytes`，校验长度并触碰首尾字节。 | 从 `memoryview(BufferHandle)` 定位 payload，再调用 `bytes(...)` 完整复制，校验长度并触碰首尾字节。 | 完整 Python `bytes` 已生成；不含 GPU copy。 | +| `cuda` | 从 `MemHolder` 中取得 CPU DLPack payload 指针。 | 使用 `BufferHandle.ptr() + payload_offset` 取得 host 指针。 | 对应 CUDA event 已完成；不含 D2H 和 kernel 执行。 | + +来源计数与输出计数共用同一个 warmup 后统计窗口，但两套系统的观测方式不同： + +| 后端 | DRAM / SSD 归因方式 | 完整性判定 | 边界 | +| --- | --- | --- | --- | +| Fluxon | Rust get 路径保留 `GetSourceKind::Memory` / `GetSourceKind::Ssd`，经 external holder 传到 benchmark；每个成功操作单独计数。 | 每个 hit 都必须有且只有一个 source kind。 | 表示本次 get 选中并完成的 route source；它不是 NVMe 设备层 IO trace。 | +| Mooncake | 在统计窗口起止点采样 `get_offload_rpc_read_count()`；`0.3.11.post1` 的单 key `get_buffer()` 在选中 `LOCAL_DISK` 后会进入唯一的 offload-RPC 计数入口。counter 增量记为 SSD hit，成功 hit 减该增量记为 DRAM hit。 | counter 必须单调，增量不能超过该节点成功 hit。 | 这是窗口 counter 推断。本次 CPU 和 GPU 有效记录的采样时刻与窗口边界最大偏差分别为 13.047 ms 和 1.912 ms。 | + +Fluxon / Mooncake coordinator 只在每个节点的 `memory + ssd + unknown == hit`、观测完整且 `unknown == 0` 时发布完整来源结果。如果任一节点的计数不自洽，该节点的全部 hit 都归入 unknown，不根据剩余量猜测来源。miss 不分配 DRAM / SSD source，单独进入命中率分母。CPU 的 192 条记录和 GPU 的 72 条记录均为完整来源、unknown 为 0。 + +分层逻辑读带宽是整个统计窗口的平均值：`tier logical read GiB/s = tier hit 数 × 原始 payload bytes ÷ 窗口秒数 ÷ 2^30`。CPU 和 GPU 的统计窗口均为 30 秒。本次实验使用固定 payload，因此可以从来源操作数精确换算逻辑读取量。柱高不包含 miss；红点使用 `hit ÷ (hit + miss)`，两者共同描述本次 pressure workload。 + +bootstrap 在统计窗口前完成，窗口内的 workload 只包含 get。因此，已纳入记录的 DRAM 和 SSD 逻辑写带宽均为 `0.000 GiB/s`。 + +SSD 回填写入 host memory 的流量不重复计入 DRAM 逻辑写带宽。物理 DRAM 流量、SSD 读放大和回填写放大需要通过硬件计数器或设备 trace 观测。 + +#### CUDA 双槽 pipeline 与阶段指标口径 + +`cuda` 模式为每个 benchmark 线程固定复用 2 组 pinned host buffer、stream、event 和 GPU buffer，不增加新的 YAML 参数。线程先把 pageable source 复制到空闲 pinned slot，再调用 `cudaMemcpyAsync`。提交后不等待 event，下一次 KV get 可以和 H2D 重叠；两个 slot 都在途时，才等待最早提交的 event。 + +统计窗口结束时会 drain 所有在途 slot，但完成时间超出窗口的操作仍会被过滤。Mooncake 的 benchmark get 路径直接调用 `get_buffer_blocking()`，不再额外调用一次 `get_size()`。 + +两边只在 host pointer 的提取方式上不同，之后都进入同一个 H2D 流水线。`cuda_host_stage_us` 记录 pageable source 到复用 pinned slot 的 copy，`cuda_submit_us` 记录 `cudaMemcpyAsync` 调用，`cuda_h2d_event_us` 记录覆盖 H2D 的 CUDA event 区间。 + +`cuda_pipeline_residence_us` 从开始占用 slot 到前台确认 event 完成，包含 host staging、H2D 和完成轮询。它与下一次 KV get 存在重叠，不能和其他阶段机械相加。benchmark 会把 `MemHolder` 或 `BufferHandle` 连同相关 view 保留到 event 完成，使两边释放 native reference 的时点与成功计数边界一致。 + +### 8.3 实验设置 + +主图 case 在同一台测试机上串行执行，硬件为 Intel Xeon Platinum 8468、Samsung NVMe（XFS）和 NVIDIA H100 80 GB，GPU 链路为 PCIe 5.0 x16。bootstrap 后只读，workload 使用固定 value size 和 Zipf 分布。 + +CPU 和 GPU 主图统一使用 SSD-pressure 数据集：数据集约 2.5 GiB，DRAM 贡献为 2 GiB；SSD-on 再贡献 16 GiB SSD，SSD-off 不贡献 SSD。数据集、请求分布、worker 布局和统计窗口保持一致。由于数据集大于 DRAM，miss 和成功 hit 集合的变化就是实验要观察的结果；图同时报告命中率与 hit payload 带宽，不把开关差解释成纯介质开销。 + +| 组别 | DRAM 存储贡献 / 原始数据集 | payload / keyspace | worker 布局 | 总时长 / warmup / 统计窗口 | +| --- | --- | --- | --- | --- | +| CPU Fluxon / Mooncake pressure | 2 GiB / 约 2.5 GiB | `1/4/8/16 MiB - 128 B` / `2560/640/320/160`；原生引用 / `bytes` | c2=`1×2`、c4=`1×4`、c8=`2×4`、c16=`4×4` | 35 / 5 / 30 s | +| CPU Fluxon SSD backend 消融 | 2 GiB / 约 2.5 GiB；两侧再配置 16 GiB SSD | 与 CPU pressure 相同；`MemHolder` / `bytes` | 与 CPU Fluxon / Mooncake pressure 相同 | 35 / 5 / 30 s | +| GPU pressure 开关矩阵 | 2 GiB / 约 2.5 GiB | `4/8/16 MiB - 128 B` / `640/320/160`；`cuda` | c2=`1×2`、c4=`1×4`、c8=`2×4`、c16=`4×4` | 35 / 5 / 30 s | + +主图使用的 Fluxon wheel SHA-256 前缀为 `11118d3`，Mooncake 为 `0.3.11.post1`。同路径 SSD backend 消融使用 Fluxon wheel SHA-256 前缀 `4896d0e`，其中集成 Foyer 0.22.3 作为 test-only owner backend。后文的早期 GPU SSD 写盘证据使用 Fluxon wheel SHA-256 前缀 `54de857`，不与主图或 backend 消融记录合并。 + +表中的 DRAM 存储贡献是 KV 可分配容量，不等同于进程 RSS。主图对齐的是 2 GiB 逻辑 value 容量和 16 GiB 逻辑 SSD 容量；Fluxon 与 Mooncake 的 envelope、allocator、索引和元数据不同，物理 DRAM 占用并不相等。逻辑 payload 带宽也不包含后端编码和元数据。 + +同路径 SSD backend 消融中的 2 GiB 仍是 Fluxon owner 上层 DRAM。集成 Foyer 另外保留 1 B 内部 memory capacity，并关闭该层的 memory admission；这 1 B 不替代或扣减 Fluxon 的 2 GiB DRAM。 + +CPU 和 GPU 的 cN 都表示所有 benchmark 进程合计 N 个 worker 线程。Fluxon / Mooncake 的 c8 和 c16 分别使用 2 个和 4 个 benchmark 进程，因此结果同时包含多进程绕开 Python GIL 后带来的并行度；GPU 每个线程固定复用两个 pinned host、stream、event 和 device buffer slot。 + +CPU 和 GPU 主图中的 Mooncake 都覆盖两种拓扑： + +- **`DEDICATED_OWNER`**：一个常驻 owner 贡献 2 GiB DRAM 和 16 GiB SSD，所有 benchmark 进程都是 zero-contribution。 +- **`PER_BENCHMARK_PROCESS`**：不启动独立 owner，每个 benchmark 进程同时作为请求方和存储 endpoint。2 GiB DRAM 和 16 GiB SSD 按进程数等分：c2/c4 每个实例为 2 GiB + 16 GiB，c8 为 1 GiB + 8 GiB，c16 为 512 MiB + 4 GiB。 + +在 SSD-on 侧，每个 storage instance 使用独立 SSD 目录，但都位于同一块 Samsung NVMe 上。两种拓扑从 c2 到 c16 都保持 2 GiB DRAM + 16 GiB SSD 的集群逻辑总容量；SSD-off 侧保持同样的 2 GiB DRAM，去掉 16 GiB SSD 贡献。 + +Fluxon / Mooncake 每个拓扑的两侧使用相同的 DRAM 贡献、pressure 数据集、requester buffer、worker 布局和 transport。排除用于隔离 case 的 cluster ID、目录和端口后，Fluxon 只增删 owner `owner_kv_ssd`；Mooncake 只增删 `--enable_offload=true` 和 16 GiB SSD 贡献。 + +Mooncake 的 CPU requester buffer 在 SSD 关/开两侧均为 256 MiB，GPU 每个 benchmark 进程两侧均为 2 GiB。两侧也都使用 `tcp_pool_nodelay_backport`。 + +#### bootstrap、transport 与容量排除项 + +每个 case 都按单 PUT 串行 bootstrap。1、4、8、16 MiB payload 在每次 PUT 后分别等待 6、25、50、100 ms，把 payload 提交速率上限控制在约 160 MiB/s。该速率不足以在 Mooncake 的一个 10 秒 offload heartbeat 周期内填满 2 GiB DRAM 层，从而避免 bootstrap 成功与否取决于首次 heartbeat 和 DRAM 写满的先后顺序。800 MiB/s 的启动探针只用于发现这个竞态，没有进入最终结果。 + +主图构建和早期 GPU SSD 写盘证据所用构建都对 Mooncake 设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`，并在 client connect 与 server accept 两端补 `TCP_NODELAY`。结果中的 transport profile 记为 `tcp_pool_nodelay_backport`。 + +在只开启 connection pool、未补 `TCP_NODELAY` 的预跑中，`batch_get_into_offload_object_internal` 均值从 `2.22 ms` 增至 `7.37 ms`，P95 出现约 `46 ms` 台阶。最终结果不混入默认短连接或只开 connection pool 的运行，也不把这个 backport 口径写成原版 `0.3.11.post1` 的默认表现。 + +Mooncake SSD offload 的单 slice 上限是 `0xFFFFF0` bytes。CPU 和 GPU 的 16 MiB 组都将 benchmark payload 减去 128 bytes，加入 91-byte FlatDict envelope 后仍低于该上限。 + +Fluxon / Mooncake 主图和 Fluxon SSD backend 消融均关闭 capacity guard，Mooncake 的 eviction high watermark 设为 0.80。SSD 容量不是这些实验的限制因素。主图只纳入状态为 `SUCCESS`、0 error、来源计数完整且 unknown 为 0 的运行；pressure 中的 miss 在 SSD 关/开两侧都保留并进入命中率。 + +### 8.4 CPU 性能对比：`MemHolder` / `BufferHandle` 与 `bytes` + +下图的两行分别是原生引用 / handle 和 `bytes`，四列分别是 1/4/8/16 MiB。每个 payload 分面覆盖 c2/c4/c8/c16；每个并发度固定按 Fluxon、Mooncake 独立 owner、Mooncake 每进程 × SSD 关/开排列六根柱。蓝、橙、青分别区分三种实现组，浅色表示 DRAM、深色表示 SSD，红点使用右轴表示总命中率。各 payload 分面使用独立左轴，范围写在右上角。 + +![CPU 性能对比：Fluxon 与 Mooncake 两种拓扑在不同并发度下的 SSD 开关效果、数据来源分布及命中率](../../pics/kv_ssd_cpu_sweep.png) + +CPU 的 192 条记录覆盖 `2 输出模式 × 4 payload × 4 并发度 × 3 实现组 × 2 SSD 状态`。`fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_cpu_sweep.py` 与 GPU 脚本共同调用同目录下的 `kv_ssd_pressure_chart.py`。输入文件按命令行顺序覆盖同 case key；最终矩阵缺一条、存在 error、来源不完整或 unknown 非零都会直接失败。 + +| 实现组 | SSD-off 命中率 | SSD-on 命中率 | SSD-on 中 SSD 来源占 hit | +| --- | --- | --- | --- | +| Fluxon | 89.112%–94.390% | 100% | 17.550%–48.650% | +| Mooncake 独立 owner | 23.759%–66.879% | 100% | 62.065%–82.385% | +| Mooncake 每进程 | 22.488%–68.655% | 72.830%–100% | 59.341%–83.011% | + +结果支持三个直接判断： + +- **命中率和 hit payload 带宽必须一起看**：SSD-off 会产生快速 miss；Mooncake 每进程的部分 SSD-on case 也有 miss。柱高不包含这些 miss，因此单看柱高会遗漏成功集合的变化。 +- **开关差不是单项 SSD 开销**：Fluxon 原生引用的 SSD-on hit 带宽相对 SSD-off 变化为 -88.8%–-21.1%，`bytes` 为 -45.8%–+2.2%；Mooncake 两种拓扑同时包含正负变化。差值混合了命中率、DRAM/SSD 来源、回填、输出物化和调度，不能拆成一个内部阶段。 +- **输出模式需要分开比较**：原生引用 / handle 在 host reference ready 时计数，`bytes` 还包含完整 payload 物化。两行都按每次 hit 的完整逻辑 payload 计带宽；handle 行并没有读取或复制同等字节，因此其数值可以超过物理 DRAM 带宽，不能跨行用柱高计算复制效率或模式倍率。 + +统计窗口内只有 get，因此已完成记录的 DRAM 和 SSD 逻辑写带宽均为 `0.000 GiB/s`。图中报告逻辑 payload 流量，不代表物理 DRAM 或 NVMe 带宽。 + +#### 8.4.1 同一 Fluxon 路径下的 SSD backend 消融 + +[Foyer](https://github.com/foyer-rs/foyer/tree/v0.22.3) 是一个面向高性能和高并发场景的 Rust 内存/磁盘混合缓存库。它的 `HybridCache` 在一套接口下组合内存 cache 与磁盘 storage，并提供块式磁盘引擎、可替换的缓存算法和 IO engine。 + +对 Fluxon 而言，Foyer 是一个现成的通用实现参照。这组消融为一个直接的工程取舍提供依据：Fluxon 原生 SSD backend 在目标 workload 和既有数据路径中的功能与性能表现，是否足以支持继续自研和维护。若通用库已经满足需求，直接复用 Foyer 可以减少存储引擎本身的开发与维护。 + +要比较 SSD 存储层，不能直接把 Fluxon 的整套 cache 替换掉。Fluxon 的内存路径同时包含 master `route`、owner DRAM 副本、allocation 与 holder 生命周期、跨 owner 传输，以及 SSD 回填后的 promotion。如果连上层内存 cache 也替换为 Foyer，公共语义和数据路径会一起改变，结果无法归因到 SSD backend。因此，实验保留 Fluxon 的内存层及其上层路径，将替换边界统一收敛到 owner 内部的 `KvSsdStorage`：一侧使用 Fluxon 原生的 SSD 持久化与读取实现，另一侧接入 Foyer `HybridCache` 的 storage tier，并关闭 Foyer 自身的 memory admission。 + +Foyer 分支通过 test-only 配置启用： + +```yaml +test_spec_config: + kv_ssd_storage_backend: foyer +``` + +默认值仍为 `native`。Foyer 分支当前只接受一个 SSD root，并在配置阶段拒绝 `kv_ssd_uring_mode: iovec`。两组都调用相同的公共 KV API，经过相同的 FlatDict、route、transport、PyO3/GIL 边界和 benchmark 进程布局；写入 SSD backend 的都是 Fluxon target 中包含 FlatDict envelope 的完整编码 value。owner 上层 DRAM 都是 2 GiB，SSD 都是 16 GiB，且两种 SSD backend 都使用 direct I/O。图中的逻辑带宽仍只按原始 payload bytes 计数。 + +Foyer 内部另设 1 B memory capacity、1 个 memory shard，并用 filter 关闭 memory admission。这里关闭的是 SSD backend 内部可能形成的额外内存层，不影响 Fluxon owner 的 2 GiB DRAM。专项测试在 persist 和重复 get 后观测到 Foyer memory usage 为 0，所有读取来源均为 Foyer storage tier。 + +| 对比项 | Fluxon 原生 SSD | Fluxon Foyer SSD | +| --- | --- | --- | +| SSD 组织 | Fluxon ring、定长 shard 和 `io_uring`。 | Foyer 0.22.3 `HybridCache`，`WriteOnInsertion`、`PsyncIoEngine` 和 64 MiB block。 | +| 读取落点 | 满足对齐条件时直接读入 Fluxon target；远端路径可让 SSD chunk read 与 push 重叠。 | Foyer 先返回包含完整 value 的 `Vec`，再复制到 Fluxon target；远端路径取得完整 value 后才逐 chunk 交给 push。 | +| persist 并发 | 使用原生 writer queue 和 device worker。 | forced storage admission；`enqueue → wait → may_contains` 串行化，避免 Foyer byte queue 满时静默丢弃写入。 | +| 驱逐收敛 | ring 覆盖后批量通知 master 删除对应 SSD route。 | Foyer 0.22.3 没有向当前集成暴露等价的 SSD eviction callback；陈旧 route 只能在读取失败时撤销。 | + +这个 workload 在 bootstrap 后只读，因此 Foyer persist 串行化不进入 30 秒统计窗口。约 2.5 GiB 数据集也小于 16 GiB SSD 容量；本次结果不评估 Foyer 容量驱逐后的 route 收敛。图中的浅色仍表示 Fluxon 上层 DRAM route，深色表示 Fluxon SSD route；它不把 Foyer 内部 memory cache 作为第三层来源。 + +图按 `MemHolder` / `bytes` 分成两行，每行覆盖 1 / 4 / 8 / 16 MiB；每个并发度放置原生与 Foyer 两根柱。蓝/紫区分 SSD 实现，浅/深区分 DRAM 与 SSD 来源，红点表示命中率。 + +![Fluxon 同路径 SSD backend 消融：原生实现与 Foyer](../../pics/kv_ssd_foyer_ablation.png) + +PNG 由 `fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_foyer_ablation.py` 读取原始汇总生成。脚本要求 64 条记录的 case key 完整，校验容量、并发布局、统计窗口、来源计数和配对控制项后才绘图。这里的 64 是测量记录数，即 32 个原生/Foyer 配对，不表示 64 种部署。 + +所有配对的命中率都是 100%，因此这组差距没有 miss 数量优势可以解释。P95 的同向变化表明差距同时出现在延迟侧。`bytes` 会在两侧共同追加完整 payload 物化，其中位比值高于 `MemHolder` 与公共物化成本稀释 backend 相对差异相容，但单次矩阵不足以分解这部分因果贡献。 + +当前两条读取路径的明确差异是：Foyer adapter 先取得完整 `Vec` 再复制到 Fluxon target；原生实现可以直接读入 target，远端时还保留 chunk read 与 push 重叠。同时还存在 `PsyncIoEngine` / `io_uring`、完整 value 交付和调度方式等差异。这些差异共同属于当前 adapter 对比，本次数据无法把柱高差逐项分配给某个内部阶段。 + +配对 case 对相同的 thread id 和 operation index 使用同一套确定性 Zipf key 选择。统计边界仍是固定 30 秒窗口，不是固定操作数；较慢的一侧会停在同一确定性序列的更短前缀，实测 DRAM/SSD 来源比例也可能随之变化。因此，这组数据回答指定 pressure workload 下的整体 backend 表现，不把柱高比值解释成单次纯 SSD get 的加速比。 + +### 8.5 GPU 性能对比：多并发 SSD 开关与 Mooncake 存储拓扑 + +上一节的 CPU 测试以 KV 数据到达 CPU 内存为计数终点，分别观察原生引用 / handle 和完整 `bytes` 物化。AI 推理或训练实际消费 KV 数据时，数据通常还需要继续搬到 GPU。因此，本节把计数终点延伸到对应 CUDA event 完成，比较多并发下 SSD 开关和 Mooncake 两种存储拓扑对 GPU 可用数据吞吐的影响。该边界包含 H2D 传输，不包含 GPU kernel 执行或 D2H。 + +GPU 与 CPU 使用同一种画法，只少了原生引用 / `bytes` 两行。4、8、16 MiB 各占一个分面；每个分面都覆盖 c2/c4/c8/c16，每个并发度再按 Fluxon、Mooncake 独立 owner、Mooncake 每进程 × SSD 关/开排列六根柱。蓝、橙、青区分三种实现，浅色表示 DRAM、深色表示 SSD，红点显示命中率；六组 case 都从 native host buffer 进入同一条 `cuda` H2D 路径。 + +![GPU 性能对比：三种拓扑在不同并发度下的 SSD 开关效果、数据来源分布及命中率](../../pics/kv_ssd_gpu_sweep.png) + +GPU PNG 由 `fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_gpu_sweep.py` 读取一个或多个 pressure 汇总，并与 CPU 脚本共同调用同目录下的 `kv_ssd_pressure_chart.py`。脚本校验最终 72 条记录完整后直接使用 Pillow 绘制；各 payload 分面使用独立左轴，命中率共用 0%–100% 右轴。 + +| 拓扑 | SSD-off 命中率 | SSD-on 命中率 | SSD-on 中 SSD 来源占 hit | +| --- | --- | --- | --- | +| Fluxon | 89.326%–93.267% | 100% | 20.242%–46.508% | +| Mooncake 独立 owner | 34.557%–62.223% | 100% | 64.837%–77.135% | +| Mooncake 每进程 | 35.430%–64.782% | 59.855%–100% | 56.073%–84.198% | + +下表把图中的并发扩展收敛为 c16/c2 hit payload 带宽倍率，三个 payload 的顺序都是 4 / 8 / 16 MiB： + +| 拓扑 | SSD-off c16/c2 | SSD-on c16/c2 | +| --- | --- | --- | +| Fluxon | 5.83 / 5.83 / 4.68 倍 | 6.50 / 7.22 / 6.80 倍 | +| Mooncake 独立 owner | 1.33 / 1.50 / 1.10 倍 | 1.73 / 1.37 / 1.35 倍 | +| Mooncake 每进程 | 2.67 / 2.01 / 1.23 倍 | 3.02 / 2.07 / 1.77 倍 | + +Fluxon 的 SSD-on hit 带宽相对 SSD-off 变化为 -54.6%–-27.9%，Mooncake 独立 owner 为 -32.0%–+37.3%，Mooncake 每进程为 -62.6%–-17.5%。这些变化必须与上表的命中率一起解释。c16 下，每进程相对独立 owner 的 hit 带宽在 SSD-off 下分别高 125.7%、107.3% 和 97.9%，SSD-on 下高 82.4%、64.5% 和 57.9%；这些是包含 endpoint 本地性、进程数和 transport 调度的整体拓扑结果，不能归因到单个内部环节。 + +主矩阵体现了三个结果： + +- **Fluxon 能随并发扩展**：4、8、16 MiB 从 c2 到 c16 分别提高到 7.18、7.68 和 5.33 倍。 +- **16 MiB 在 c8 后开始收益递减**：c8 到 c16 只再提高 14.0%，backend get 均值从 `3.082 ms` 增至 `7.367 ms`，P95 从 `24.007 ms` 增至 `51.150 ms`。H2D 之前的路径已开始制约吞吐。 +- **Mooncake 拓扑会影响结果**：每进程贡献模式在 c16 下相对独立 owner 分别提高 91.3%、103.7% 和 66.8%。两种模式的逻辑容量相同，差异主要出现在 GPU copy 之前。来源诊断显示 8/16 MiB 的每进程贡献模式具有更高 SSD 命中率，因此更高吞吐无法由更高 DRAM 命中率解释；独立 owner 串行点、endpoint 本地性和 transport 调度成本仍未拆分。 + +c16 下，Fluxon 相对两种 Mooncake 拓扑中更快的一种，在 4、8、16 MiB 组中分别为 3.83、5.61 和 6.73 倍。这些比值只覆盖第 8.2–8.3 节定义的计数边界、单机 fast path 和固定容量。 + +## 9. 总结 + +Fluxon KV 将本地 SSD 定位为运行期回填层，不改变用户侧 `put`、`get` 和 `delete` 契约。master 维护 key-version 级 `route`，并统一调度用于最终 value 副本和临时传输的内存 allocation。owner 承载真实 bytes，管理 SSD 文件位置、IO 和 ring 生命周期。 + +`put` 操作在内存副本发布后返回，SSD persist 在后台完成。`get` 操作优先读取内存。没有可读内存副本时，本地 SSD 回填复用 requester target，远端回填则按 chunk 从 source staging push 到 requester target。两条路径最后都进入同一个 `GetDone` 和 holder 交付链。 + +同路径消融进一步表明，在本文固定的 pressure workload 下，当前 Foyer adapter 的 hit payload 带宽中位数为原生 SSD 的 37.2%，P95 为 2.30–8.46 倍，32 个配对中没有胜出。两侧均为 100% 命中，但 Foyer adapter 多了完整 `Vec` 到 Fluxon target 的复制，也没有原生路径的 direct-target 与 chunk read/push 重叠。这一结果限定于当前 adapter，不代表 Foyer 库的一般性能。 + +当前 SSD 层不提供冷启动恢复，也不向用户暴露独立 API。专项测试验证了强制回填和陈旧 route 兜底，并报告限定条件下的端到端逻辑 payload 表现。这些结果不代表裸 SSD 带宽，也不外推到跨机器、多 owner 或多盘场景。 + +将本地 SSD 纳入同一套 KV 路由、生命周期和观测体系，是 Fluxon 构建统一 AI 数据面的一次关键扩展：内存继续承担低延迟热路径，SSD 在不改变 `put`、`get` 和 `delete` 契约的前提下承接更大的运行期工作集，命中数据还可以沿既有路径继续送往 GPU。我们希望这项能力减少 AI 推理和训练系统在分层缓存、回填传输与数据搬运上的重复建设，让开发者把更多精力放在模型和业务创新上。Fluxon 已基于 Apache License 2.0 开源，GitHub 仓库地址：https://github.com/Tele-AI/Fluxon。欢迎社区一起推进面向 AI 场景的数据流动基础设施。 diff --git "a/fluxon_doc_cn/design/fluxon_0_\351\205\215\347\275\256\346\200\273\350\247\210.md" "b/fluxon_doc_cn/design/fluxon_0_\351\205\215\347\275\256\346\200\273\350\247\210.md" index a4c5865..8c5e779 100644 --- "a/fluxon_doc_cn/design/fluxon_0_\351\205\215\347\275\256\346\200\273\350\247\210.md" +++ "b/fluxon_doc_cn/design/fluxon_0_\351\205\215\347\275\256\346\200\273\350\247\210.md" @@ -193,7 +193,7 @@ test_runner_ui: bootstrap_phases: - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: [etcd, greptime, tikv_pd, tikv] ``` @@ -208,7 +208,7 @@ retention: repos: - addr: git@github.com:Tele-AI/fluxon.git follow: - - branch: big_step2 + - branch: example-feature-branch run: name_prefix: fluxon_ci commands: @@ -260,8 +260,8 @@ instance_key: my-owner-1 # 只要 dram > 0，就进入 owner 分支 contribute_to_cluster_pool_size: - # 容量按 16 MiB 对齐 - dram: 1677721600 + # 容量解析后向下对齐到 16 MiB + dram: "1.5GB" vram: {} fluxonkv_spec: @@ -282,6 +282,11 @@ fluxonkv_spec: large_file_paths: - /var/lib/fluxon/large + # 可选；和 large_file_paths 长度一致，限制每个 KV SSD backing tier root 的 payload 用量 + # 缺省或 null 时不启用 KV SSD + large_limit_size: + - "4GB" + # 可选 p2p_listen_port: 31001 @@ -315,8 +320,10 @@ fluxonkv_spec: - `monitoring` 在 master 上必填。 - `master_ui` 依赖 `monitoring`，并作为嵌入式 monitor HTTP 服务启动。 -- `contribute_to_cluster_pool_size` 里的容量都按 16 MiB 对齐；`dram = 0` 但 `vram` 非 0 会被拒绝，避免半 owner 半 external 的模糊状态。 +- size 字段可以写整数 bytes，也可以写 `"512.9B"`、`"1.5GB"` 这类字符串；`MB/GB` 按二进制单位解析，等价于 `MiB/GiB`，小数结果按 bytes 向下截断。 +- `contribute_to_cluster_pool_size` 里的容量解析后都向下对齐到 16 MiB；`dram = 0` 但 `vram` 非 0 会被拒绝，避免半 owner 半 external 的模糊状态。 - owner 模式要求 `contribute_to_cluster_pool_size.dram > 0`，并且必须显式提供 `etcd_addresses`、`sub_cluster`、`large_file_paths`。 +- `fluxonkv_spec.large_limit_size` 只允许 owner 配置；出现时必须和 `large_file_paths` 长度一致，每项解析后必须大于或等于 512 bytes。 - zero-contribution `external` 模式禁止再写 owner 专属字段；运行时会从 owner `shared.json` 补齐这部分信息。 - `share_mem_path` 会拼成 `cluster_name` 作用域路径；`mmap.file`、`shared.json` 和 peer metadata 都位于这个 cluster-scoped 目录下。 - `test_spec_config.side_transfer_role = worker` 不是第三套 YAML，而是 zero-contribution client 的子分支；它强制 `TransferEngineType::P2p`，并关闭 transfer RPC fast path。 @@ -444,11 +451,11 @@ http_listen_addr: 0.0.0.0:18080 network: subnet_whitelist: - 127.0.0.0/8 - - 10.0.0.0/24 + - 192.0.2.0/24 primary_ip_to_extended_ips: - 10.0.0.10: - - 10.0.0.11 - - 10.0.0.12 + 192.0.2.10: + - 192.0.2.11 + - 192.0.2.12 ``` 以及协议/传输分支这两个输入： diff --git "a/fluxon_doc_cn/design/fluxonfs_videoreader\346\216\245\345\217\243\350\256\276\350\256\241.md" "b/fluxon_doc_cn/design/fluxonfs_videoreader\346\216\245\345\217\243\350\256\276\350\256\241.md" new file mode 100644 index 0000000..9594c97 --- /dev/null +++ "b/fluxon_doc_cn/design/fluxonfs_videoreader\346\216\245\345\217\243\350\256\276\350\256\241.md" @@ -0,0 +1,397 @@ +# FluxonFS VideoReader 接口设计 + +## 执行摘要 + +- `VideoReader(path)` 这类第三方解码入口在 C/C++ 层直接访问文件路径，绕不过 Python 侧 `open()` patcher。FluxonFS 需要提供一个原生 `VideoReader` 接口，把 FFmpeg 的随机读请求接到 FluxonFS 的 `stat / read_chunk / cache` 链路上。 +- 单 reader 公共入口为 `FluxonFsPatcher.open_video_reader(...)`。调用方传 `export_name + relpath + 输出尺寸`，返回 `FluxonFsVideoReader`，再调用 `read_frames_numpy(indices)` 得到 `uint8 NHWC` 数组。 +- 训练 dataloader 的推荐入口为 `FluxonFsPatcher.open_video_reader_pool(...)`。pool 在当前 Python 进程内按 `(export_name, relpath, height, width, num_threads, request_identity)` 管理 reader LRU，复用 reader 级 page cache，调用方不直接维护 reader 缓存。 +- 底层只保留 `RangeSource` 读取路径：FFmpeg 通过 `AVIOContext` 回调向 FluxonFS 请求字节范围。export 路由、身份校验、metadata cache、KV piece cache、异步 backfill、磁盘缓存等能力都收束在 FluxonFS 内部，VideoReader 不新增缓存协议或本地文件分支。 +- 解码在训练进程本地完成。FS agent 只提供字节范围读取和缓存填充，不承担视频解码和帧数组传输。 +- v1 先支持 CPU FFmpeg 解码。NVDEC / GPU decode 可以作为后续显式 backend 分支加入，不进入首版公共契约。 + +## 1. 背景与目标 + +训练数据集当前用 `decord.VideoReader(path).get_batch(...).asnumpy()` 解码视频。这个路径的问题在于： + +- FFmpeg 在 native 层按路径访问文件系统，Python 的 `open()` / `read()` patcher 无法接管它。 +- `VideoReader(...)` 初始化阶段可能扫描 packet / keyframe，远端文件会承受额外顺序读和 seek。 +- 同一视频多个窗口重复读取时，现有探针和训练路径会重复打开、重复索引、重复 seek/decode。 + +本文设计的目标是在 FluxonFS 内提供一个面向视频解码的读接口： + +- 让 FFmpeg 的 `read / seek` 请求走 FluxonFS 的远端读和缓存能力。 +- 把视频读取的公共接口固定在 FluxonFS 语义上，而不是暴露某个第三方解码库的文件路径约定。 +- 保持解码本地化，避免把未压缩帧通过 FS agent RPC 传输。 +- 对缓存、身份、文件版本和失败语义给出稳定边界。 + +## 2. 非目标 + +- 不在 FS agent 服务端解码视频。 +- 不复刻 `decord.VideoReader` 的全部 API。 +- 不新增独立于 FluxonFS export 配置的 per-video 配置文件。 +- 不为 VideoReader 增加环境变量参数传递。 +- 不为视频路径新增第二套 KV cache key 布局。 +- 不在首版支持写入视频、转码、音频流处理或 GPU decode。 + +## 3. 当前实现基础 + +FluxonFS 已有的读链路可以直接作为 VideoReader 的字节源： + +| 现有模块 | 当前职责 | VideoReader 复用方式 | +| --- | --- | --- | +| `FluxonFsPatcher` | 当前 Python 进程内的 FS 挂载、cache config 和请求身份入口 | 提供 `open_video_reader(...)` 和 `open_video_reader_pool(...)` 公共入口 | +| `FluxonFsExport` | 定义 export 根目录、路由、cache key prefix、metadata TTL 和 backfill 策略 | VideoReader 只接受 `export_name + relpath`，不绕过 export | +| `read_chunk` RPC | agent 端按 `offset + length` 读取远端文件 | FFmpeg `AVIOContext` 的 read 回调最终落到这里 | +| `remote_read_chunk_by_handle...` | client agent 侧带 metadata、KV cache 和 miss policy 的范围读 | 作为 `FluxonFsVideoByteSource` 的主读接口 | +| KV piece cache | 以 `(export, relpath, size, mtime_ns, piece_index)` 组织固定大小 piece | 随机读、重复读和跨 reader 复用 | +| FluxonFS 内部缓存层 | KV piece cache、异步 backfill、磁盘缓存等内部优化 | VideoReader 不直接选择或配置，只通过范围读接口受益 | + +当前 Python patcher 对 `open()` 的支持仍然保留，但它不是视频接口的主路径。FFmpeg 从 C/C++ 层读取文件时，不会自动进入 Python `open()` patcher。 + +## 4. 公共接口 + +单 reader 公共入口放在 `FluxonFsPatcher` 上： + +```python +with patcher.open_video_reader( + export_name="train-videos", + relpath="bucket/a/b/sample.mp4", + height=480, + width=832, + num_threads=8, +) as reader: + frames = reader.read_frames_numpy([0, 8, 16, 24]) +``` + +稳定契约： + +| 项目 | 契约 | +| --- | --- | +| 入口 | `FluxonFsPatcher.open_video_reader(...)` | +| 定位文件 | `export_name: str` + `relpath: str` | +| 输出尺寸 | `height: int` + `width: int`，必须为正数 | +| 解码线程 | `num_threads: int`，必须为正数 | +| 读取方法 | `read_frames_numpy(indices: Sequence[int]) -> numpy.ndarray` | +| 返回数组 | `dtype=uint8`，shape 为 `(len(indices), height, width, 3)` | +| 生命周期 | reader 支持 context manager；`close()` 后所有读取失败 | +| 身份 | 打开 reader 时捕获 `patcher` 当前 request identity；reader 生命周期内不跟随后续身份变更 | + +单 reader 只提供 `read_frames_numpy` 一个读取方法。不要同时暴露 `get_batch()`、`get_batch_numpy()`、`read_frames()` 等同义入口。 + +训练 dataloader 应优先使用进程内 pool： + +```python +pool = patcher.open_video_reader_pool(max_readers=32) +frames = pool.read_frames_numpy( + export_name="train-videos", + relpath="bucket/a/b/sample.mp4", + height=480, + width=832, + num_threads=8, + indices=[0, 8, 16, 24], +) +``` + +业务已经把多个样本整理成 batch 时，应把每个 clip 的 frame plan 作为 `FluxonFsVideoReadRequest` 传给 pool： + +```python +requests = [ + FluxonFsVideoReadRequest( + export_name="train-videos", + relpath="bucket/a/b/sample.mp4", + height=480, + width=832, + num_threads=8, + indices=(0, 8, 16, 24), + ), + FluxonFsVideoReadRequest( + export_name="train-videos", + relpath="bucket/a/b/sample.mp4", + height=480, + width=832, + num_threads=8, + indices=(12, 20, 28, 36), + ), +] +results = pool.read_many_numpy_with_stats(requests) +``` + +pool 的稳定契约： + +| 项目 | 契约 | +| --- | --- | +| 入口 | `FluxonFsPatcher.open_video_reader_pool(max_readers=32)` | +| LRU key | `(export_name, relpath, height, width, num_threads, request_identity)` | +| 容量 | `max_readers: int`，必须为正数；活跃 reader 不会被强制关闭 | +| 读取方法 | `read_frames_numpy(...) -> numpy.ndarray` | +| 指标方法 | `read_frames_numpy_with_stats(...) -> FluxonFsVideoReadResult` | +| 批量方法 | `read_many_numpy_with_stats(requests: Sequence[FluxonFsVideoReadRequest]) -> list[FluxonFsVideoReadResult]` | +| 生命周期 | pool 支持 context manager；`close()` 后所有新读取失败 | +| 线程语义 | 一个 cached reader 同一时刻只租给一个调用；并发请求可打开多个 reader | + +`read_many_numpy_with_stats(...)` 在 pool 内部按 LRU key 分组。同一视频、同一输出尺寸和同一 request identity 的多个 clip 会合并为一次 native `read_frames_numpy(combined_indices)` 调用，然后按请求顺序切片返回。这个 batch 只改变调度粒度；每个请求的输出 shape、dtype 和错误语义仍按单请求契约处理。 + +pool 只管理 reader 对象生命周期、reader 级 page cache 复用和同 reader key 的批量调度，不新增文件缓存协议。FluxonFS 的 KV piece cache、异步 backfill、磁盘缓存和权限检查仍然只在 RangeSource 读链路内生效。 + +构建要求： + +- native 解码实现需要用 `fluxon_pyo3` 的 `fluxon_fs_video_ffmpeg` feature 构建，并且构建环境提供 FFmpeg 开发库。 +- 未启用该 feature 时，Python 公共入口仍存在，但打开 reader 会返回明确错误。这是构建期能力开关，不是第二条运行时读取路径。 + +## 5. 架构 + +```mermaid +flowchart LR + A[Python 训练代码] --> P[FluxonFsVideoReaderPool 可选] + P --> B[FluxonFsPatcher.open_video_reader] + A --> B + B --> C[PyO3 FluxonFsVideoReader] + C --> D[FFmpeg Decoder] + D --> E[AVIOContext read / seek callback] + E --> F[FluxonFsVideoByteSource] + F --> G[reader page cache] + G --> H[remote_read_chunk_by_handle] + H --> I[FluxonFS 内部缓存层] + I --> J[read_chunk RPC] + J --> K[fs_agent export 根目录] + D --> L[numpy.ndarray uint8 NHWC] +``` + +这张图里的关键边界： + +- Python 只负责创建 reader、传入 frame index 和接收 numpy 数组。 +- FFmpeg 回调不调用 Python，避免 GIL 和 Python callback 成为高频 I/O 路径。 +- `FluxonFsVideoByteSource` 运行在 Rust/native 层，直接复用 FluxonFS agent 的读接口。 +- VideoReader 只看到范围读接口；KV、磁盘或其他缓存策略都属于 FluxonFS 内部实现。 + +## 6. 详细设计 + +### 6.1 打开 reader + +`open_video_reader(...)` 的打开流程： + +1. 校验 `export_name`、`relpath`、`height`、`width`、`num_threads`。 +2. 从 `FluxonFsPatcher` 对应的 Rust agent 捕获当前 request identity。 +3. 调用现有远端 `stat` 获取 `size`、`mtime_ns`、`is_file`。 +4. 构造文件签名 `sig = (size, mtime_ns)`。 +5. 创建 reader 级 `FluxonFsVideoByteSource` 和 page cache，绑定 `export_name`、`relpath`、`size`、`mtime_ns`、request identity。 +6. 首版在 `read_frames_numpy(...)` 内创建 FFmpeg `AVIOContext` / decoder，并把 `read_packet`、`seek` 回调绑定到 byte source。 + +打开阶段只产生一个 reader handle。后续每次 `read_frames_numpy(...)` 复用同一个 byte source 和 reader 级 page cache；当前实现暂不复用 FFmpeg decoder context。pool 的批量入口会把同一 reader key 下的多个 clip 合并为一次 `read_frames_numpy(...)`，从而减少重复从视频头顺序解码到目标帧的成本。 + +### 6.2 FFmpeg 读回调 + +FFmpeg 的读回调以当前位置为基础请求一段 bytes： + +```text +read_packet(buf, want) + -> source.read_at(current_offset, want) + -> current_offset += bytes_read +``` + +seek 回调只更新当前位置： + +```text +seek(offset, whence) + -> source.seek(offset, whence) +``` + +`AVSEEK_SIZE` 返回打开时 stat 得到的 `size`。如果 FFmpeg 请求超过 EOF，返回实际可读长度；如果请求已经在 EOF 之后，返回 EOF。 + +### 6.3 `FluxonFsVideoByteSource` + +`FluxonFsVideoByteSource` 是内部 Rust 对象，不作为 Python 公共类型暴露。它负责把 FFmpeg 的小范围读合并成 FluxonFS 适合的 page 读取。 + +| 字段 | 作用 | +| --- | --- | +| `export_name` | export 名称 | +| `relpath` | export 内相对路径 | +| `size` | 打开时文件大小 | +| `mtime_ns` | 打开时文件 mtime | +| `request_identity` | 打开 reader 时捕获的身份 | +| `page_bytes` | 内部页大小，默认跟随现有 piece / read chunk 约束 | +| `page_cache` | reader 级 LRU，减少 FFmpeg 小读和重复 seek 的 RPC 次数 | + +VideoReader 只保留 `RangeSource` 路径。字节源的读取顺序如下： + +1. 命中 reader 级 `page_cache`：直接切片返回。 +2. 命中 KV piece cache：读取 piece 后填入 `page_cache`。 +3. 未命中 KV：调用 `remote_read_chunk_by_handle...`，按当前 export 的 cache 策略决定是否 backfill。 + +### 6.4 缓存复用策略 + +VideoReader 不新增独立缓存配置。它复用已有 export 字段： + +| 配置 / 状态 | 复用方式 | +| --- | --- | +| `cache_kv_key_prefix` | KV piece cache key 前缀 | +| `cache_bytes_field_key` | KV 中保存 piece bytes 的字段 | +| `cache_max_bytes` | 判断对象是否允许进入 KV cache | +| `metadata_cache_ttl_ms` | 复用现有 stat / metadata cache 行为 | +| `async_backfill_enabled` | KV miss 后是否异步补齐 piece | +| FluxonFS 内部磁盘缓存 | 作为 `remote_read_chunk_by_handle...` 之后的内部优化，不向 VideoReader 暴露 FD / path | + +缓存 key 继续使用现有签名： + +```text +export_name + relpath + size + mtime_ns +``` + +这个签名的含义是“打开 reader 时看到的文件版本”。如果文件内容变化但 `size` 和 `mtime_ns` 没变，FluxonFS 无法稳定识别这是新版本；上游数据集需要保证训练视频文件在读取期间不可变。 + +### 6.5 单一读取路径 + +内部只保留 `RangeSource`： + +```text +FFmpeg read / seek + -> AVIOContext callback + -> FluxonFsVideoByteSource.read_at(offset, length) + -> remote_read_chunk_by_handle(...) + -> FluxonFS 内部缓存层 / read_chunk RPC +``` + +VideoReader 不接收本地 FD / path，不判断是否整文件 materialize，也不直接操作磁盘缓存。后续如果 FluxonFS 需要把某些 range read 映射到本地磁盘缓存，那也是 `remote_read_chunk_by_handle(...)` 内部的实现变化，不改变 VideoReader 契约。 + +### 6.6 帧读取 + +`read_frames_numpy(indices)` 的执行流程： + +1. 校验 `indices` 非负，且能转换为 frame index 列表。 +2. 首版按容器读取顺序从视频开头解码到最大目标帧，保证 frame index 语义稳定。 +3. 后续可以在明确 PTS / frame index 映射后加入 keyframe seek 优化。 +4. 对输出帧执行 resize / pixel format 转换。 +5. 写入一个 numpy-owned `uint8` 输出数组。 + +首版输出固定为 NHWC RGB： + +```text +(frame_count, height, width, 3) +``` + +这一路径仍然会 materialize 帧数组。设计目标是减少压缩视频字节读取和重复远端 I/O，不宣称帧输出零拷贝。 + +## 7. 生命周期与所有权 + +```mermaid +sequenceDiagram + participant Py as Python caller + participant Patcher as FluxonFsPatcher + participant Reader as FluxonFsVideoReader + participant Source as FluxonFsVideoByteSource + participant Agent as FluxonFS Agent + + Py->>Patcher: open_video_reader(export, relpath, size args) + Patcher->>Agent: stat(export, relpath, identity) + Agent-->>Patcher: size, mtime_ns + Patcher->>Reader: create native reader + Reader->>Source: bind export, relpath, size, mtime_ns, identity + Py->>Reader: read_frames_numpy(indices) + Reader->>Source: read_at(offset, length) + Source->>Agent: remote_read_chunk_by_handle(...) + Agent-->>Source: bytes + Source-->>Reader: bytes + Reader-->>Py: numpy.ndarray + Py->>Reader: close() +``` + +| 对象 | 所有者 | 生命周期 | +| --- | --- | --- | +| `FluxonFsPatcher` | Python 调用方 | 必须长于 reader | +| `FluxonFsVideoReaderPool` | Python 调用方 / DataLoader worker | 进程内 LRU；通常随 worker 生命周期创建和销毁 | +| `FluxonFsVideoReader` | Python 调用方 | context manager 或显式 `close()` | +| `FluxonFsVideoByteSource` | native reader | 随 reader 创建和销毁 | +| FFmpeg decoder context | native reader | 首版随每次 `read_frames_numpy(...)` 创建和销毁 | +| numpy 输出数组 | Python 调用方 | 每次 `read_frames_numpy` 独立返回 | + +reader 关闭后，byte source 和 reader 级 page cache 释放；FFmpeg context 由每次 `read_frames_numpy(...)` 创建并在该次调用结束时释放。FluxonFS 的共享 KV cache 和内部磁盘缓存按现有 agent 生命周期管理。 + +## 8. 失败语义与不变量 + +| 条件 | 行为 | +| --- | --- | +| `export_name` 不存在 | 打开 reader 失败 | +| `relpath` 不存在或不是文件 | 打开 reader 失败 | +| 权限不足 | 打开或读取失败，错误保留 FluxonFS 权限语义 | +| `height / width / num_threads <= 0` | 打开 reader 失败 | +| `indices` 包含负数 | `read_frames_numpy` 失败 | +| FFmpeg 无法识别容器或视频流 | `read_frames_numpy` 失败 | +| 读取中远端短读 | 当前读取失败 | +| reader 已关闭 | 后续读取失败 | + +关键不变量： + +- 所有远端数据访问都必须经过 FluxonFS export 和 access check。 +- FFmpeg 高频读回调不能调用 Python。 +- 缓存 key 和失效语义沿用 FluxonFS 当前 `(export, relpath, size, mtime_ns)` 签名。 +- 单 reader 公共接口只暴露一种读取方法：`read_frames_numpy`。 +- pool 公共接口提供单请求读取和批量读取；批量读取的输入必须是显式的 `FluxonFsVideoReadRequest`。 +- v1 返回 numpy-owned 数组，不把 native decoder 内部 buffer 暴露给 Python 长期持有。 + +## 9. 当前实现与专用 fast path 的边界 + +### 当前实现可直接复用 + +- `read_chunk` RPC 的 `offset + length` 文件读取。 +- agent 侧 `remote_read_chunk_by_handle...` 的 KV cache / remote read 分支。 +- `FluxonFsS3KvMissPolicy::{RemoteRead, StageToKvThenRead}` 的有限 miss 策略。 +- `async_backfill_enabled` 对 KV piece 的异步补齐。 +- FluxonFS 内部缓存能力，包括 KV piece cache、异步 backfill 和磁盘缓存。 + +### 需要新增 + +- `fluxon_py/fluxon_fs/video.py`：Python 公共 facade 和类型导出。 +- `FluxonFsPatcher.open_video_reader(...)`：单 reader 公共入口。 +- `FluxonFsPatcher.open_video_reader_pool(...)`：训练 dataloader 推荐入口，负责进程内 reader LRU。 +- `FluxonFsVideoReadRequest` 和 `read_many_numpy_with_stats(...)`：pool 级批量调度入口，负责同 reader key 多 clip 合并。 +- native `FluxonFsVideoReader`：持有 byte source、reader 级 page cache 和输出数组构造逻辑，并驱动 FFmpeg decoder。 +- native `FluxonFsVideoByteSource`：FFmpeg `AVIOContext` 回调背后的 FluxonFS 范围读适配层。 +- 针对视频 decode 的指标：open 耗时、range read 次数、KV hit / miss、FluxonFS 内部缓存 hit / miss、decode 耗时。 + +首版实现边界： + +- `open_video_reader(...)` 校验参数、捕获 request identity，并通过 FluxonFS agent 做 `stat`。 +- reader 生命周期内复用 `RangeSource` 和 reader 级 page cache。 +- `read_frames_numpy(...)` 每次创建 FFmpeg context，通过 `AVIOContext` 回调读取 FluxonFS range，并返回 numpy-owned `uint8 NHWC` 数组。 +- `FluxonFsVideoReaderPool` 在 Python 进程内复用 idle reader；活跃 reader 不参与 LRU 关闭，释放后再参与淘汰。 +- `read_many_numpy_with_stats(...)` 只在 Python pool 层合并同 reader key 的请求；native 层仍只有一个 frame-index batch decode 入口。 +- 帧定位先采用顺序解码。keyframe seek、持久 decoder、GPU decode 和指标上报是后续显式优化项。 + +### 专用 fast path + +首版不为 VideoReader 增加专用 fast path。唯一热路径是 `RangeSource`。缓存优化只能在 FluxonFS 内部发生，例如 KV piece 命中、异步 backfill 或内部磁盘缓存命中；这些优化不得改变 `open_video_reader(...)` 和 `read_frames_numpy(...)` 的公共契约。 + +未来如果加入 GPU decode，应作为新的 decode backend 分支，例如： + +```text +DecodeBackend = CpuFfmpeg | CudaNvdec +``` + +这个分支只改变解码后端，不改变 FluxonFS 字节读取和缓存契约。 + +## 10. 验证计划 + +首版验证按三层做： + +| 层级 | 验证内容 | +| --- | --- | +| 字节源单测 | `read_at`、seek、EOF、跨 page 读取、重复读命中 page cache | +| FluxonFS 集成测试 | export 权限、metadata cache、KV piece cache、内部磁盘缓存、文件变更签名 | +| 视频过程测试 | 同一视频多窗口读取、和 `decord.VideoReader(path)` 的帧 shape / dtype / 基本像素一致性、timeout 行为 | + +性能指标必须绑定测试条件。至少记录： + +- 视频文件大小、编码格式、分辨率、帧数。 +- export 位置和底层存储类型。 +- cache 冷启动 / 热启动。 +- reader open 时间。 +- FFmpeg range read 次数和总读取 bytes。 +- KV cache hit / miss。 +- FluxonFS 内部缓存 hit / miss。 +- `read_frames_numpy` 端到端耗时。 + +## 11. 关键结论 + +FluxonFS VideoReader 的核心价值是把 FFmpeg 的随机字节读取纳入 FluxonFS 的 export、权限和缓存链路。单 reader 接口提供 `FluxonFsPatcher.open_video_reader(...) -> FluxonFsVideoReader`；训练 dataloader 使用 `FluxonFsPatcher.open_video_reader_pool(...)` 复用 reader。底层统一走 `RangeSource`，缓存能力收束在 FluxonFS 内部。 + +这条设计保留训练进程本地解码，避免远端传输未压缩帧；同时让同一视频的重复窗口读取可以命中 FluxonFS cache，减少 FFmpeg 对远端文件系统的重复直接访问。 diff --git "a/fluxon_doc_cn/design/kv_1_\346\246\202\350\247\210\344\270\216\345\210\206\345\261\202.md" "b/fluxon_doc_cn/design/kv_1_\346\246\202\350\247\210\344\270\216\345\210\206\345\261\202.md" index 7975c9c..2faa190 100644 --- "a/fluxon_doc_cn/design/kv_1_\346\246\202\350\247\210\344\270\216\345\210\206\345\261\202.md" +++ "b/fluxon_doc_cn/design/kv_1_\346\246\202\350\247\210\344\270\216\345\210\206\345\261\202.md" @@ -63,12 +63,26 @@ pub struct OneKvNodesRoutes { pub put_id: PutIDForAKey, // 这个 key-version 绑定的 lease；None 表示非 lease key。 pub lease_id: Option, - // 这个已提交版本当前所有 live replica。 - pub nodes_replicas: RwLock>, + // 按 owner 索引这个版本的内存与 SSD 副本。 + pub node_replicas: RwLock>, // 限制 get 驱动的 durable replica 提升并发数。 pub get_durable_slots_used: AtomicU32, } +pub struct KvNodeReplicas { + // 标识当前 NodeID 对应的节点实例。 + pub tomb_tag: NodeTombTag, + // 当前内存副本；驱逐后独立设为 None。 + pub memory: Option>, + // 当前 SSD 副本；失败或驱逐后独立设为 None。 + pub ssd: Option, +} + +pub struct KvSsdReplicaInfo { + // SSD 中真实 payload 的长度。 + pub len: u64, +} + pub struct InflightPutInfo { // 放置策略最终选中的目标节点。 pub node_id: NodeID, @@ -106,7 +120,7 @@ pub struct InflightGetInfo { - `put_id`：当前实现里用于区分同一 key 不同版本的版本标识，形状是 `(put_time_ms, put_version)`，不是数学意义上的全局唯一 ID。 - `lease_id`：这个版本是否绑定 lease。`None` 表示非 lease key，`Some(id)` 表示受 lease 管理。 -- `nodes_replicas`：该版本当前有哪些副本，每个副本对应哪个 node、哪块 allocation、当前 tomb 状态如何。 +- `node_replicas`：按 node 保存一份公共 tomb tag，并用 `memory` / `ssd` 两个 `Option` 独立记录两层副本。内存或 SSD 单独驱逐时只清理自己的字段；两个字段都为空时删除该 node entry。 这意味着： diff --git "a/fluxon_doc_cn/design/kv_2_\350\260\203\347\224\250\346\227\266\345\272\217.md" "b/fluxon_doc_cn/design/kv_2_\350\260\203\347\224\250\346\227\266\345\272\217.md" index 45c0b4d..bded54c 100644 --- "a/fluxon_doc_cn/design/kv_2_\350\260\203\347\224\250\346\227\266\345\272\217.md" +++ "b/fluxon_doc_cn/design/kv_2_\350\260\203\347\224\250\346\227\266\345\272\217.md" @@ -35,8 +35,8 @@ pub struct OneKvNodesRoutes { pub put_id: PutIDForAKey, // 该版本是否绑定 lease。 pub lease_id: Option, - // 该版本当前所有 live replica。 - pub nodes_replicas: RwLock>, + // 按 owner 索引该版本的内存与 SSD 副本。 + pub node_replicas: RwLock>, ... } ``` @@ -126,8 +126,8 @@ pub struct OwnerHoldingGetInfo { pub struct OneKvNodesRoutes { pub put_id: PutIDForAKey, pub lease_id: Option, - // get_start 从这里挑选源副本。 - pub nodes_replicas: RwLock>, + // get_start 先选 memory，再从同一快照选择 SSD。 + pub node_replicas: RwLock>, // 限制 DurableReplica 提升并发数。 pub get_durable_slots_used: AtomicU32, ... @@ -218,8 +218,8 @@ pub enum DeleteKeyInfo { pub struct OneKvNodesRoutes { pub put_id: PutIDForAKey, - // delete 后需要按旧路由枚举哪些节点仍有 replica / cache。 - pub nodes_replicas: RwLock>, + // delete 后按旧路由枚举仍有内存 replica / cache 的节点。 + pub node_replicas: RwLock>, ... } ``` diff --git "a/fluxon_doc_cn/design/kv_3_\345\217\202\346\225\260\344\270\216\345\271\266\345\217\221.md" "b/fluxon_doc_cn/design/kv_3_\345\217\202\346\225\260\344\270\216\345\271\266\345\217\221.md" index 7cacd5b..b7ca5c9 100644 --- "a/fluxon_doc_cn/design/kv_3_\345\217\202\346\225\260\344\270\216\345\271\266\345\217\221.md" +++ "b/fluxon_doc_cn/design/kv_3_\345\217\202\346\225\260\344\270\216\345\271\266\345\217\221.md" @@ -108,13 +108,14 @@ client 侧 `get` 的热路径是： ### master 路由访问：短读锁 + 复制快照 -当前 `kv_routes` 是 `DashMap>`，而 `nodes_replicas` 是 `RwLock>`。 +当前 `kv_routes` 是 `DashMap>`，而 `node_replicas` 是 `RwLock>`。每个 node entry 用独立的 `memory` / `ssd` 选项表示两层副本。 典型做法是： - 先从 `kv_routes` 取出 `Arc`。 -- 用很短的读锁把 `nodes_replicas` clone 成局部 `HashMap` 快照。 -- 后续选源副本、处理 tomb、决定分配模式时都基于快照继续。 +- 用很短的读锁把 `node_replicas` clone 成局部 `HashMap` 快照。 +- 后续先选择 `memory` 副本；没有可用内存副本时，再从同一快照选择 `ssd` 副本。 +- tomb 清理删除整个 node entry；内存和 SSD 驱逐只更新各自的 `Option`，两个字段都为空时才删除 entry。 这样做的目的不是绝对无锁，而是： diff --git "a/fluxon_doc_cn/design/kv_5_SSD\345\255\230\345\202\250\350\256\276\350\256\241.md" "b/fluxon_doc_cn/design/kv_5_SSD\345\255\230\345\202\250\350\256\276\350\256\241.md" new file mode 100644 index 0000000..552edb2 --- /dev/null +++ "b/fluxon_doc_cn/design/kv_5_SSD\345\255\230\345\202\250\350\256\276\350\256\241.md" @@ -0,0 +1,1305 @@ +# KV 设计 5 - DRAM / SSD 多级缓存 + +## 稳定结论 + +Fluxon KV 将存储节点（owner）的本地 SSD 接入既有 key-version 副本体系，作为 DRAM 内存层的运行期回填层（backing store）。`put` 操作在内存副本可读后返回，系统随后在后台异步写入 SSD；`get` 操作优先查询内存，仅在没有可读内存副本时从 SSD 回填。 + +该设计不改变公共契约：用户继续调用 `put`、`get` 和 `delete`，`get` 仍返回 `MemHolder`。master 管理 key-version 级逻辑路由和内存 allocation 生命周期；owner 管理本地 SSD 文件位置、IO 和 ring 驱逐。SSD 不向用户暴露独立 API，也不承担冷启动恢复。 + +当工作集大于 DRAM 时，SSD 层可以保留已被驱逐的内存副本，扩大运行期可命中工作集。实际命中率和吞吐收益取决于访问分布、并发度、介质性能和回填开销，不由接入 SSD 这一事实单独保证。 + +本文档将控制节点（master）维护的 key-version 级副本索引称为路由表（`route`），将贡献内存或 SSD 容量的节点称为存储节点（owner）。设计要点如下： + +| 设计方面 | 核心要点说明 | +| --- | --- | +| 用户契约 | 继续使用 `put`、`get` 和 `delete` 操作；`get` 仍返回 `MemHolder`。 | +| 控制面与资源 | master 维护 `route`，并统一调度用于最终 value 副本和临时传输的内存分配（allocation）。owner 提供 memory segment，承载真实 bytes，并管理本地 SSD。 | +| 数据路径 | 写入的同步完成点是内存副本发布；读取优先使用内存，SSD 只在内存无可读副本时回填。 | +| 恢复边界 | SSD 仅作为运行期缓存。owner 重启后重建空 shard 和 ring，不扫描旧 shard，也不从 master 的 `route` 中恢复旧 SSD bytes。 | + +## 1. 架构概览：master 管理路由和内存分配，owner 管理本地 SSD + +内存 allocation 分为两类：一类承载最终 value 副本，另一类供跨 owner 传输临时使用，例如远端 SSD 回填的 source staging。master 统一调度这两类区间，因为它们的分配、保活和释放都与 `route` 及 in-flight 请求状态联动。allocation 对应的真实 bytes 仍位于 owner 注册的 memory segment。 + +SSD allocation 只负责最终 value 副本，因此只涉及 owner 本地 shard/ring 的文件区间。文件 offset、覆盖保护和 IO 调度都是本地状态，因此 owner 负责 SSD 分配和 ring 驱逐。 + +| 对象 | 资源或状态归属 | 分配与驱逐职责 | +| --- | --- | --- | +| `route` | master 持有 key-version 级副本索引。 | master 选择内存副本的放置位置并维护内存 allocation 生命周期。在内存或 SSD 副本失效后，master 收敛逻辑路由。 | +| 内存 | bytes 位于 owner 注册的 memory segment。 | master 从这些 segment 分配两类区间：内存副本区间和通信临时区间。master 同时管理对应 allocation 的保活与驱逐。 | +| SSD | bytes 位于 owner 本地的 shard 和 ring。 | owner 分配 SSD 文件区间并执行 ring 驱逐，再把 commit 或驱逐结果通知 master 更新 `route`。 | + +在 `route` 实现中，master 的 `OneKvNodesRoutes` 按 key-version 组织路由。`node_replicas` 是一张以 owner 为 key 的 map，每个 entry 用 `memory` 和 `ssd` 记录两层副本。二者共享当前 route 的 `put_id`，不能跨版本复用。SSD 文件位置、读写队列、文件 IO 对齐和覆盖保护都由 owner 管理。 + +### 1.1 角色与图示约定 + +全文时序图用角色前缀标明逻辑边界；前缀不表示独立进程。 + +| 前缀 | 表示的职责 | +| --- | --- | +| `external-*` | 集群外的公共 API 调用方或 benchmark 消费方，不贡献 Fluxon 存储容量。 | +| `master-*` | master 内的控制面，管理 `route`、in-flight 状态、master 侧 `Allocation` 保活句柄（guard）及 holding 生命周期，不中转 payload。 | +| `owner-*` | owner 侧的请求入口、内存或 SSD 存储组件，以及本地数据传输端点。多个 `owner-*` 角色可以位于同一 owner 进程，也可以分布在不同 owner。 | + +后续时序图同时涉及请求协调、bytes 位置、allocation 保活和数据搬运，四者不能混用： + +| 关系 | 文中的规范说法 | 当前实现落点 | +| --- | --- | --- | +| 请求由谁协调 | `owner-requester` 接收并推进本次公共请求。 | owner 侧 `ClientKvApi`；external 包装层在图中通常折叠。 | +| payload bytes 位于哪里 | bytes 位于 source/target owner 已注册的 memory segment，或 SSD owner 的本地 shard。 | RPC plan 中的 node、base address 和 absolute address。 | +| allocation 由谁保活 | master 持有表示 owner segment 区间的 RAII `Allocation` guard（保活句柄）。 | `InflightPutInfo`、`OneKvNodesRoutes.memory`、`InflightGetInfo` 和 `get_holding`。 | +| 数据由谁搬运 | memory get 由 requester pull；远端 SSD 回填由 SSD source push。 | owner 侧 `ClientTransferEngine`。 | + +`owner-requester` 只表示请求协调者，不代表固定 buffer，也不持有 master 侧的 `Allocation` 对象。它在不同操作中的数据面职责如下： + +| 操作 | owner 侧 range | master 侧 guard | 完成结果 | +| --- | --- | --- | --- | +| `put` | 常规路径的 source/staging bytes 位于 requester owner；side-transfer 路径位于单独的 `owner-put-source`。最终 target 由 master 选择，可能与 source 同 owner，也可能位于 `owner-put-target`。 | 本地 placement 用一个 `InflightPutAllocation::Local` guard 同时表示 source 和 target；远端 placement 用 `Remote { src, target }` 两个 guard。 | `PutDone` 把 target guard 转入 memory route；`put` 不产生 holder。 | +| `get` | route 路径的最终 target bytes 位于 requester owner。已有本地 replica 时直接复用；否则 memory source 或 SSD source 把 payload 写入该 target。 | `InflightGetInfo.allocation` 保活 target；`GetDone` 后转入 master `get_holding`。memory source 或远端 SSD source staging 的额外 guard 位于 `source_allocation`。 | owner 返回 holder metadata，external 在自己的进程中构造 `MemHolder`。 | +| `delete` | 没有 payload source、staging 或 target range。 | 不创建 payload `Allocation`。 | 只更新控制面 route，并执行既有 holder 失效流程。 | + +后续 `put` / `get` 时序从 owner 进入 master route 路径开始，折叠 external 包装 RPC。owner 返回共享内存 offset、长度和 `holder_id` 等 metadata，external 再据此构造 `MemHolder`，形成完整的公共结果。公共 `get` 操作还可能先命中 owner 本地的 `get_cached_info`，直接复用已有 `MemoryInfo`。第 4–7 节描述本地缓存未命中后进入 `GetStart` 的路径。 + +时序图默认采用常规 owner 请求入口。专用 side-transfer worker 会把请求处理者与 put source owner 分开：`PutStartReq.source_node_id` 指向与 worker 共享 mmap 的 owner，source/staging bytes 位于该 owner 的 segment。这个 fast path 不改变 master 选择 put target、`PutDone` 发布 target 的主契约。 + +### 1.2 四条主链路 + +第 3–7 节将沿这四条链路展开 allocation、传输和失败清理。 + +| 链路 | 核心流程 | 完成或收敛点 | +| --- | --- | --- | +| 写入与 SSD commit | master 分配 source/target；`PutDone` 发布内存副本；owner 随后异步写入 SSD。 | 公共 `put` 操作在内存副本发布后完成。SSD commit 只有在版本、memory 和 tomb 条件仍成立时才进入 `route`。 | +| 内存命中读取 | requester 复用本地副本，或从远端 memory owner pull 到最终 target。 | `GetDone` 把 target guard 转入 `get_holding`，external 根据 metadata 构造 `MemHolder`。 | +| SSD 回填读取 | 当前版本没有可读内存副本时，本地路径直接读入 requester target；远端路径读入 source staging 后按 chunk push。 | 两条路径复用同一个 `GetDone` 和 holder 交付链。 | +| SSD 驱逐与陈旧 route 兜底 | owner 覆盖 ring entry 后批量通知 master；读取端仍会复查本地 entry。 | 批量通知收敛正常驱逐；陈旧 source 读取失败时执行 revoke 并有界重新选路。 | + +payload 不经过 master。master 只维护逻辑路由和 allocation 生命周期；SSD shard、offset、IO 对齐与 ring 状态只在 owner 本地解析。 + +### 1.3 设计取舍与当前边界 + +引入 SSD 后，需要先明确三项职责边界： + +1. **同步语义**：SSD 不改变 `put` / `get` 操作的同步语义。 +2. **路由信息**：master 的 `route` 不记录 SSD 物理布局。 +3. **恢复能力**：SSD 不承担重启后的数据恢复。 + +在这些边界内，本设计保持现有公共 API 和内存副本语义，利用 owner 本地 SSD 扩大运行期缓存容量，并在没有可读内存副本时提供回填来源。文件布局和 IO 复杂度均封装在 owner 内部。 + +因此，当前 SSD 被定位为运行期回填层。它不向用户暴露独立的 SSD API，也不支持单 value 跨 device 条带化。SSD 文件布局、IO 调度和 ring 生命周期均由 owner 管理，不进入公共 API 或 master route。 + +| 设计选择 | 作用 | 代价 | 当前边界 | +| --- | --- | --- | --- | +| master 不保存 SSD offset | 跨节点控制面只处理 key-version、owner、payload 长度和 tomb tag 等逻辑信息。 | master 无法据此检查本地物理布局；陈旧 route 要靠节点 tomb、批量清理与读取兜底收敛。 | **route 是逻辑索引，不是恢复元数据**。 | +| `PutDone` 不等待 SSD | 同步写入只等待内存副本，SSD 背压留在 owner 本地 persist 路径。 | SSD 副本是 late commit，需要用 `put_id` 拒绝旧版本 commit，并在 master 处理 commit RPC 的窗口 pin 住本地 entry。 | **内存副本 ready 是同步提交点**。 | +| owner 本地驱逐 | ring 决定物理覆盖顺序，再批量通知 master 清理逻辑 route；发送失败的当前 batch 会重试。 | 有界队列满、进程退出或 master 持续不可用时，master 可能保留陈旧 route。 | **读取失败会 revoke 并有限重试**。 | +| `O_DIRECT + io_uring`，而非 `mmap` / PageCache | KV SSD 路径绕过 PageCache，直接提交对齐 IO。 | 实现必须协调 staging 地址、文件 offset、IO 长度、aligned buffer、ring 覆盖和 read pin。 | **文件 IO 实现封装在 owner**。 | +| 本地 SSD 回填 | SSD owner 与 requester owner 相同时复用最终 target，不申请独立 source staging。 | target capacity 必须覆盖对齐 IO；地址不满足 direct read 条件时仍会经过 scratch copy。 | **direct 条件成立时省掉本地二次 copy**。 | +| SSD owner 主动推送 | chunk ready 后由 SSD owner 直接推到 requester target，读盘和网络传输可重叠。 | requester 必须先通过 `SsdStageReadReq` 提供 target 地址和 `get_id`。 | **省去 stage-ready 后的一次 RTT**。 | +| value 级 device 选择 | 单 value 使用一个 device 上的连续文件区间，不需要跨 device 组合读结果。 | 单个大 value 不跨 device 条带化。 | **多 device 并行发生在 value 粒度**。 | +| 启动时重建空 shard 和 ring | 明确 SSD 只作为运行期缓存：不扫描 shard，也不通过 WAL 或 checkpoint 恢复 entry。 | 重启后需要由新写入重新形成 SSD 副本。 | **不支持冷启动恢复**。 | +| scratch 读路径 | 不满足 direct read 对齐时仍能正确读取真实 payload。 | 该 chunk 多一次本地 copy。 | **padding 不进入用户 holder**。 | + +## 2. 配置与观测：容量按介质分层可见 + +KV SSD 复用 owner 已有的 `large_file_paths`，不增加第二套路径参数，以保持配置项尽可能简洁。`large_limit_size` 与这些路径一一对应；运行时按实际 device 去重，并分别上报内存与 SSD 的容量和占用。 + +### 2.1 配置入口 + +配置仍位于 owner 的 `fluxonkv_spec`；`large_limit_size` 与 `large_file_paths` 等长： + +```yaml +fluxonkv_spec: + large_file_paths: [/data/fluxon_large0, /data/fluxon_large1] # Derive one KV SSD cache root from each path. + large_limit_size: ["4GB", "8GB"] # Limit the matching cache roots in the same order. +``` + +- **启用条件** + - `large_limit_size` 未配置或为 `null`：不启用 KV SSD。 + - `large_limit_size` 已配置：必须与 `large_file_paths` 长度一致，数组下标一一对应。 +- **容量值格式** + - 可以使用整数 bytes。 + - 也可以使用 `"512.9B"`、`"1.5GB"` 这类 size 字符串。 + - `MB/GB` 按二进制单位解析，等价于 `MiB/GiB`；小数结果按 bytes 向下截断。 + - 解析结果必须大于或等于 512 bytes，满足当前 `O_DIRECT` 对齐约束。 +- **external 约束** + - 不贡献存储容量的 external 不能声明 `large_limit_size`。 + - external 只继承 owner 暴露的共享 bundle 和运行时路径信息，不贡献 SSD 容量。 + +### 2.2 路径派生与设备去重 + +owner 不读取单独的 SSD 路径参数。每个 SSD cache root 都从对应的 `large_file_paths[i]` 派生： + +```text +/_cluster_kv_ssd_storage// +``` + +- **创建 root**：owner 创建派生目录，再通过 `metadata.dev()` 获取实际 device 身份。 +- **同设备去重** + - 多个 path 指向同一实际 device：只保留第一个 root。 + - 如果不去重，同一块盘上的两个目录会被误认为两条独立 IO 路径，导致调度层高估并行度。 + +### 2.3 容量、ring 与 per-device 对象 + +- **容量作用域** + - `large_limit_size[i]` 只约束 `large_file_paths[i]` 派生的 KV SSD cache root。 + - 它限制 key-value payload 在 KV SSD ring 中的用量，不代表整块 SSD 的物理容量。 + - 日志、profile、FS disk cache 等其他大文件资产不受该值约束。 +- **ring 空间不足** + + **ring 是构建在定长 shard 文件之上的循环空间分配器和本地 key-version 索引。** 每个 shard 用 `head` 为新 value 分配对齐后的连续区间；写到文件末尾时，物理 offset 回到文件开头继续复用空间。`tail` 随覆盖边界向前推进，回收最旧且已经可以覆盖的 entry。ring 同时记录 entry 的 `Writing/Committed` 状态、read pin 和 route commit pin，遇到仍在写入、读取或提交 route 的区间时会等待，不会直接覆盖。 + + ring 的作用是在固定容量内持续复用 SSD cache 空间，并把被覆盖的 key-version 交给 owner 后台通知 master 收敛逻辑 `route`。文中的 ring 与 `io_uring` 职责不同：前者管理文件区间、entry 状态和驱逐顺序，后者负责提交实际的文件读写。第 6.2 节再展开状态转换和覆盖保护。 + + 选择 ring 的直接工程动机，是让每个 shard 内的写入尽量保持追加式、对齐且连续。`head` 单向推进，空间用完后再回绕覆盖，省去了随机空洞分配、碎片合并和后台 compaction。现代 NVMe 也能处理随机 IO；这里更看重 `O_DIRECT` 连续写入、固定容量复用和回收路径的可预测性。 + + Moka 与 ring 解决的是不同层次的问题： + + | 机制 | 当前职责 | 无法替代的另一层能力 | + | --- | --- | --- | + | Moka | master 侧按 weight 管理内存副本的准入和驱逐。 | 不分配 SSD 文件 offset，也不管理 `Writing/Committed`、read pin 或 route commit pin。 | + | ring | owner 侧分配连续 SSD 区间，并按物理写入顺序回收和覆盖。 | 不提供 Moka 的访问热度与内存准入策略。 | + + 如果直接用 Moka 选择任意 SSD key 驱逐，被删除 value 的文件区间会形成不连续空洞，仍需额外实现 extent allocator、空洞合并或 compaction。Moka 可以叠加为逻辑淘汰策略，但不能替代 ring 对 SSD 物理空间的管理。 + + - 目标区域可以覆盖：ring 推进 tail 并驱逐旧 entry。 + - 目标区域仍有未完成写入、active read pin 或 route commit pin：当前写入等待空间释放。 +- **`KvSsdStorage`** + - 持有一组去重后的 device root。 + - 每个 device root 都独立持有： + - **shard 文件集合**：存放 ring 分配的连续 SSD offset。 + - **writer queue**：接收本 device 的写任务。 + - **reader queue**：接收本 device 的读任务。 + - **`UringIoEngine`**：提交本 device shard 文件上的 `read/write`；`Iovec` 只作为测试和对照路径保留。 +- **shard 到 device 的映射** + - 每个 device root 的 `large_limit_size` 会划分给该 device 的多个 shard。 + - `shard_to_device` 记录 shard 所属 device。 + - 读路径根据 committed entry 的 `shard_id` 找回对应 reader queue。 + +### 2.4 能力标记与容量观测 + +启用 KV SSD 的 owner 会在 `ClusterMember.metadata` 中发布 `kv_ssd_storage=true`。master 在 `PutStart` 时读取当前 live client members，并按下表确定 target 候选集和后续 SSD persist 决策： + +| live member 状态 | put target 候选集 | `PutDone` 后的处理 | +| --- | --- | --- | +| 没有 owner 声明 `kv_ssd_storage=true` | 所有具有已注册内存 allocator 的 live owner。 | `persist_to_ssd=false`，跳过 `SsdReplicaPersistReq`，按纯内存模式完成。 | +| 至少一个 owner 声明 `kv_ssd_storage=true` | 同时具有内存 allocator 和该标记的 live owner。 | `persist_to_ssd=true`，向选中的 target owner 异步发送 `SsdReplicaPersistReq`。 | + +master 把这次决策保存在 `InflightPutInfo` 中，`PutDone` 不重新推断。该标记只表达 owner 是否具备 KV SSD 能力，不暴露本地 shard、offset 或容量；容量和占用仍通过指标单独上报。target owner 收到 persist 请求后还会检查本地 `KvSsdStorage`，用于兜住 member capability 快照与 owner 实际状态之间的竞态。 + +owner 按资源类型分别上报 capacity/used，UI 也分别展示内存和 SSD： + +- **`memory_segment`** + - 指标：`kvcache_segment_capacity_bytes` / `kvcache_segment_used_bytes`。 + - 含义：owner 可直接承载 `MemHolder` 的内存 segment 容量和占用。 +- **`kv_ssd`** + - 指标：`kv_ssd_capacity_bytes` / `kv_ssd_used_bytes`。 + - 含义：owner 本地 SSD 介质层的容量和 ring 已用空间。 + +## 3. 写入：先内存可见，再异步进入 SSD + +**`put` 操作的同步完成点是内存副本发布。** `PutDoneResp` 返回后，用户已经可以读到该副本。SSD 对齐、写盘、ring 空间等待和 route 提交都在后台执行。 + +SSD persist 复用内存 target 所在 owner 的本地地址，因此 master 选择内存 target 时也同时决定是否生成 SSD 副本以及由哪个 owner 写盘。这里分为三层： + +| 层次 | 行为 | 分支定位 | +| --- | --- | --- | +| 能力上报 | owner 已通过 member metadata 上报 `kv_ssd_storage=true`。内存 segment 注册继续只描述 master 可分配的地址区间；SSD 文件区间仍由 owner 本地管理。 | master 已经能判断 owner 是否具备 KV SSD 能力，无需通过 persist RPC 探测。 | +| master 决策 | live SSD-capable owner 集合为空时使用纯内存 placement；集合非空时，target 必须同时具有内存 allocator 和 SSD capability。决策结果以 `persist_to_ssd` 写入 `InflightPutInfo`。 | 无 SSD owner 时不发送 persist RPC；启用 SSD placement 时不会主动选择无 SSD owner。 | +| owner 侧防御 | target owner 在处理 persist 请求时再次检查本地 `KvSsdStorage`。 | `persisted=false` 只覆盖 capability 快照过期、owner 重启或重新注册等竞态，不参与正常能力探测。 | + +一次写入分为三个阶段： + +| 阶段 | 动作 | 完成后的状态 | +| --- | --- | --- | +| 1. 填充 target | master 根据 live member capability 和已注册内存 allocator 生成 placement plan，并在 `InflightPutInfo` 中记录 `persist_to_ssd`。系统先将 payload 写入 `owner-put-source` 的 range。对于本地 placement，系统直接复用该 range；对于远端 placement，系统再将 payload 传输到 `owner-put-target`。 | payload 已到 target 地址，SSD persist 决策已经固定，尚未对其他读取发布。 | +| 2. 发布内存副本 | `owner-requester` 发送 `PutDoneReq`。master 将 target guard 转入 route 的 `memory`，并释放不再需要的 source staging。 | master 返回 `PutDoneResp`；公共 `put` 操作完成，内存副本可读。 | +| 3. 按 placement 决策处理 SSD | `persist_to_ssd=false` 时跳过 SSD RPC；为 `true` 时，master 保活 target 并发送 `SsdReplicaPersistReq`。target owner 再确认本地 `KvSsdStorage` 可用，然后对齐 payload、写入 shard，并通过 `SsdReplicaCommitReq` 提交 SSD route。 | 纯内存模式直接进入内存 cache；SSD commit 校验通过后 `ssd` 字段可用。竞态兜底、persist 失败或过期 commit 都不回滚第 2 阶段。 | + +下图按公共 `put` 操作的语义展开这三个阶段。图中省略 external 使用的 `ExternalPutStartReq`、共享内存写入和 `ExternalPutTransferEndReq` 包装层。常规路径将 `owner-put-source` 合并到 `owner-requester`；side-transfer 的分离边界见第 1.1 节。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant owner_put_target as owner-put-target + participant master_route as master-route + participant owner_ssd as owner-ssd-storage + participant owner_ring as owner-ssd-ring + + external_user->>owner_requester: put(key, value) + owner_requester->>master_route: PutStartReq + master_route->>master_route: 读取 live member 的 kv_ssd_storage capability + alt 存在 SSD-capable owner + master_route->>master_route: target = memory allocator ∩ SSD owner；persist_to_ssd = true + else 不存在 SSD-capable owner + master_route->>master_route: target = memory allocator；persist_to_ssd = false + end + master_route->>master_route: 分配并保活 InflightPutAllocation + master_route-->>owner_requester: source / target node-address plan + alt target 与 requester/source 同 owner + owner_requester->>owner_requester: payload 写入本地 target range + else target 在其他 owner + owner_requester->>owner_put_target: source/staging → target range + end + owner_requester->>master_route: PutDoneReq + master_route->>master_route: target guard 转入 memory route；远端 source guard 释放 + + par 返回公共 put + master_route-->>owner_requester: PutDoneResp + owner_requester-->>external_user: put 完成，内存副本可读 + and PutDone 后的介质处理 + alt persist_to_ssd = false + master_route->>master_route: 跳过 SSD RPC；内存副本进入驱逐 cache + else persist_to_ssd = true + master_route->>master_route: 保活 target allocation 并预留 cache capacity + master_route->>owner_ssd: SsdReplicaPersistReq(key, put_id, target_addr, len) + alt owner 本地无 KvSsdStorage（capability 状态竞态兜底） + owner_ssd-->>master_route: SsdReplicaPersistResp(persisted = false) + else owner 本地 KvSsdStorage 可用 + owner_ssd->>owner_ring: 对齐、分配 offset 并写 shard + alt 本地 persist 失败 + owner_ring-->>owner_ssd: write / ring error + owner_ssd-->>master_route: persist error + else 本地 persist 成功 + owner_ring-->>owner_ssd: Committed entry + KvSsdPersistGuard + owner_ssd->>master_route: SsdReplicaCommitReq + master_route->>master_route: 校验 key、put_id、memory 和 tomb tag + alt route 仍然有效 + master_route->>master_route: 更新对应 owner entry 的 ssd 字段 + master_route-->>owner_ssd: SsdReplicaCommitResp(OK) + else route 已删除或版本已变化 + master_route->>master_route: 忽略 late commit，不更新 route + master_route-->>owner_ssd: SsdReplicaCommitResp(OK) + end + owner_ssd->>owner_ring: 释放 route commit pin + owner_ssd-->>master_route: SsdReplicaPersistResp(persisted = true) + end + end + master_route->>master_route: 释放临时预留，内存副本加入驱逐 cache + end + end +``` + +除 SSD capability 兜底外，后台路径还需要处理三类并发风险： + +| 风险 | 机制 | 保证 | +| --- | --- | --- | +| persist 期间 target 被释放或容量超卖 | `Arc` 保活 target，capacity reservation 按 allocation capacity 扣减 cache 预算。RAII guard 覆盖成功、失败、取消和 master 关闭。 | persist 始终引用有效 allocation，容量账本与生命周期一致。 | +| 旧版本完成得更晚 | master 用 `put_id` 校验 late commit。 | 旧 SSD 副本不会进入新版本 route。 | +| 本地已写完，master 尚未处理完 commit RPC | `KvSsdPersistGuard` 在该窗口期 pin 住 ring entry。 | commit RPC 返回前，对应 SSD bytes 不会被 ring 覆盖；返回 `OK` 不代表 route 一定已发布。 | + +### 3.1 route 数据结构和 SSD commit 条件 + +`OneKvNodesRoutes` 中与写入可见性直接相关的结构如下，其他字段已省略： + +```rust +pub struct OneKvNodesRoutes { + pub put_id: PutIDForAKey, // Bind all replicas to one key version. + // ... + pub node_replicas: RwLock>, // Index both cache layers by owner. + // ... +} + +pub struct KvNodeReplicas { + pub tomb_tag: NodeTombTag, // Track this owner incarnation once. + pub memory: Option>, // Hold the readable memory replica when present. + pub ssd: Option, // Record the readable SSD replica when present. +} + +pub struct KvSsdReplicaInfo { + pub len: u64, // Preserve the real payload length. +} +``` + +`node_replicas` 中的每个 owner 只对应一个 entry： + +- **owner 身份**：来自 `HashMap` 的 `NodeID` key，不在 entry 内重复保存。 +- **节点实例**：`tomb_tag` 绑定本次 owner incarnation；节点离开或重启后，旧副本不能被新实例复用。 +- **介质状态**：`memory` 和 `ssd` 独立更新。内存驱逐只清空 `memory`，SSD ring 驱逐或读取失败只清空 `ssd`；两者都为空时才删除 entry。 + +owner 本地写盘成功后仍需提交 `SsdReplicaCommitReq`。master 只在以下条件全部满足时写入 `KvSsdReplicaInfo { len }`： + +- 当前 `kv_routes` 仍有该 key，且 route 的 `put_id` 与请求一致。 +- 对应 owner 的内存副本仍然存在。 +- owner 的 tomb tag 仍然 live。 + +`SsdReplicaCommitResp` 当前不区分“已发布”和“已忽略”。route 缺失、版本变化、memory 缺失或节点 tomb 时，master 会忽略 commit，但仍返回 `OK`。因此，owner 侧的 `persisted = true` 只表示本地 entry 已写入且 commit RPC 无协议错误，不能确认 master route 一定包含该 SSD 副本。 + +SSD persist 可能晚于同 key 的后续覆盖写。`put_id` 校验会丢弃旧版本 late commit，避免它进入新版本 route。`KvSsdPersistGuard` 从本地 commit 开始保活 entry，直到 master 处理完这次 commit RPC；无论 route 最终发布还是请求被忽略，guard 都会释放。 + +## 4. 读取：内存优先，SSD 只做回填 + +**只有当前 key-version 没有可读内存副本时，master 才选 SSD。** 以下流程从 owner 本地 `get_cached_info` 未命中后开始。进入 route 路径后，`GetStartReq` 按下表为 `get` 操作选择 source： + +| 分支 | 触发条件 | payload 路径 | 完成方式 | +| --- | --- | --- | --- | +| 内存 | route 中存在可读 `memory` 副本。 | source 在 requester owner 时，master 复用对应的 replica guard，requester 无需搬运 payload；否则由 requester 从 `owner-memory` pull 到最终 target。 | requester 发送 `GetDoneReq`。owner 返回 holder metadata，external 再构造公共 `MemHolder`。 | +| 本地 SSD | 没有可读内存副本，且 SSD source 就在 requester owner。 | shard 读到最终 target。满足 direct read 条件时直接写入；否则经过 owner 内部 scratch buffer 复制一次。 | requester 发送 `GetDoneReq`，后续 holder 交付与内存分支相同。 | +| 远端 SSD | 没有可读内存副本，且 SSD source 在其他 owner。 | `owner-ssd-source` 按 chunk 读入本地 staging，每个 chunk ready 后主动 push 到 requester target。 | SSD source owner 发送 `GetDoneReq`；requester 用 `done_*` 字段构造同一套 holder metadata，再由 external 构造公共 `MemHolder`。 | +| 未命中 | 当前 key-version 没有可读的 memory 或 SSD 副本。 | 不传输 payload。 | 返回 `KeyNotFound`。 | + +本文把“回填”限定为“把 SSD payload 写入 requester 的最终 target”。staging 是远端 SSD source owner 上的临时、对齐内存；它不会成为用户的 `MemHolder`。 + +下图展开四条 route 分支、数据面 range 的位置和 master 侧 allocation 控制关系。图中已折叠 owner 本地 `get_cached_info` 未命中之前的包装层。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant master_route as master-route + participant owner_memory as owner-memory + participant owner_ssd as owner-ssd-source + participant owner_media as owner-local-ssd + + external_user->>owner_requester: get(key) + owner_requester->>master_route: GetStartReq + master_route->>master_route: 读当前 key-version 的 node_replicas 快照 + + alt 存在可用 memory 副本 + master_route-->>owner_requester: memory source + requester target plan + alt source 就在 requester owner + owner_requester->>owner_requester: 复用本地 replica range + else source 在其他 owner + owner_requester->>owner_memory: transfer_data_no_copy (pull) + owner_memory-->>owner_requester: payload 写入最终 target + end + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_requester: GetDoneResp(holder_id) + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else 无 memory 副本，但存在可用 SSD 副本 + master_route->>master_route: 在 requester segment 分配并保活最终 target + alt SSD source 在其他 owner + master_route->>master_route: 在 SSD source segment 分配并保活对齐 staging + end + master_route-->>owner_requester: SSD source + target + stage plan + alt SSD source 就在 requester owner + owner_requester->>owner_ssd: 本地 load_and_push_kv_from_ssd 调用 + owner_ssd->>owner_media: pin entry 并读取 shard (direct or scratch) + owner_media-->>owner_ssd: IO 完成，最终 target 已写入 + owner_ssd-->>owner_requester: 本地读取完成 + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_requester: GetDoneResp(holder_id) + else SSD source 在其他 owner + owner_requester->>owner_ssd: SsdStageReadReq(get_id, stage, target, len) + loop 每个 payload chunk + owner_ssd->>owner_media: pin 并读入 source staging + owner_media-->>owner_ssd: IO 完成，staging chunk ready + owner_ssd->>owner_requester: push chunk 到 target offset + end + owner_ssd->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，必要时 promotion + master_route-->>owner_ssd: GetDoneResp(holder_id) + owner_ssd-->>owner_requester: SsdStageReadResp(done_*) + end + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else 当前 key-version 没有可读副本 + master_route-->>owner_requester: KeyNotFound + owner_requester-->>external_user: miss + end +``` + +三处实现边界需要单独说明： + +- **本地 SSD 回填**：`src_addr == target_addr`，控制面不申请 source staging。地址、offset 和容量满足 direct read 条件时，`KvSsdStorage` 直接写最终 target；否则内部 scratch fallback 会增加一次本地 copy。 +- **远端 SSD staging**：master 先计算 `ssd_stage_len = align_ssd_io_len(len)`，再从 SSD source owner 已注册的 segment 中额外分配 `SSD_ALIGNMENT - 1` bytes，从 allocation range 内取得 512-byte 对齐的 `src_addr`。requester 只发送一次 `SsdStageReadReq`，后续 chunk 均由 SSD owner 主动 push。 +- **陈旧 SSD route**：读取失败后先 revoke 当前 source，再重新选路；该失败不会向用户返回半成品 holder。 + +### 4.1 target 分配与有界背压 + +`get` target 遇到内存压力时按以下顺序处理： + +1. **遍历 allocator**：先尝试所有可用 allocator，不在单个 allocator 上连续空转。 +2. **推进回收**：运行该 owner 的 replica Moka cache、`inflight_puts` 和 `inflight_gets` maintenance tasks，然后立即重试。`Cache::remove()` 会让 entry 立即不可查询，但 RAII value 可能仍在维护队列中保活 allocation。maintenance 完成后，空间才真正释放。 +3. **等待外部释放**：仍无空间时，让出执行权并等待 route 删除、owner 缓存失效 RPC 或 holder ACK。重试间隔为 2 ms，总等待上限为 5 s。 +4. **有界失败**：5 s 后仍无空间才返回带 total/free capacity 的 `NoSpace`；等待不会扩大 owner 的配置容量。 + +新 target 在传输期间尚未进入驱逐 cache，但已经占用 requester owner 的 segment。master 为可提升为内存副本的 durable target 立即预留 allocation capacity，并通过 `InflightGetInfo` 或 allocation 的 on-drop callback 保活 reservation。 + +这份预留与异步 SSD persist、lease 共用 `cache_reserved_bytes`，避免把在途 `get` 占用误算为空闲预算。`GetRevoke`、60 秒 in-flight TTL 或 `GetStartResp` 发送失败都会释放它。 + +temporary target 不进入 Moka。`GetDone` 后，master 的 `get_holding` 继续持有 target guard，直到 owner 侧最后一个 `MemoryInfo` 释放并向 master ACK，或 member-left 清理该 requester。requester owner 的非 cache 空间因此要容纳并发 target，以及尚未走完释放链的 temporary holder。 + +### 4.2 `GetDone` 与 target 生命周期 + +- **非 lease 的 durable promotion 成功**：master 先更新 route，再释放独立的临时预留，并在同一次请求处理中按相同 weight 插入 Moka，避免 route 与容量索引不同步。 +- **带 lease 的 durable promotion 成功**：target 不进入 Moka；reservation 已绑定到 allocation 的 on-drop callback，随该 leased replica 的 allocation 一起存活。 +- **无法 promotion**：route 版本变化、节点 tomb 或 route 已删除时，target 降级为 temporary holder 并释放 durable slot。非 lease 路径的独立 reservation 随 `InflightGetInfo` 释放；lease 路径的 reservation 已绑定 allocation，会保留到 temporary holder 释放。 +- **所有成功路径**：`GetDone` 都把 requester target guard 写入 master 的 `get_holding`；远端 SSD source staging 和 memory source 的额外保活引用随 `InflightGetInfo` 释放。 + +与 allocation 生命周期直接相关的字段如下： + +```rust +pub struct InflightGetInfo { + pub put_id: PutIDForAKey, // Guard the selected key version. + pub src_node_id: NodeID, // Identify the selected source owner. + pub allocation: Arc, // Hold the requester target until GetDone. + pub source_allocation: Option>, // Keep a memory source or remote SSD staging alive. + pub route: Arc, // Keep the route and its durable slot alive. + pub allocation_mode: GetAllocationMode, // Decide whether promotion is allowed. + pub source_kind: GetSourceKind, // Select memory or SSD cleanup semantics. + pub cache_capacity_reservation: Option>, // Reserve durable target capacity. + // ... Other request identity and length fields. +} +``` + +`InflightGetInfo` 的核心作用是在一次 `get` 内同时保活 target、source 和 route： + +- **`allocation`** 始终表示 requester segment 中的 target；`GetDone` 后由 `OwnerHoldingGetInfo` 继续保活。用户侧 `MemHolder` 引用这段 bytes，但不持有这个 master 侧 Rust 对象。 +- **`source_allocation`** 是 master 侧的保活引用：memory 路径用它保活内存副本，远端 SSD 路径用它保活 staging，本地 SSD 回填时为 `None`。`GetDone`、`GetRevoke` 或 in-flight TTL 清理都会释放该引用。 +- **`route` 与容量预留** 共同保证 durable promotion 期间的 route 生命周期和容量预算。非 lease 成功路径在同一次请求中更新 route、释放独立 guard 并插入 Moka；lease 路径则由 allocation callback 延续 reservation。 +- **`put_id`、`src_node_id` 和 `source_kind`** 限定失败清理范围，避免撤销新版本 route 或错误的介质副本。 + +external 公共 `get` 的 holder 交付跨越三层对象。只有 master 的 `get_holding` 持有 allocation guard，不能把三层对象统称为“holder 持有 allocation”： + +| 层级 | 实际持有 | 正常释放链 | +| --- | --- | --- | +| master | `get_holding[(requester_owner_id, holder_id)]` 中的 `OwnerHoldingGetInfo`，内部持有 requester target 的 `Arc`。 | owner 的 delete-ack batch 通过 `BatchDeleteAckReq` 到达后移除；requester member-left 时也会批量清理。 | +| requester owner | `external_get_holding[(external_client_id, holder_id)]` 中的 `Arc`。它引用 target metadata，未持有 master 进程里的 Rust guard。 | external holder 释放后收到 `ExternalDeleteAckReq` 并移除该引用；owner 侧最后一个 `MemoryInfo` 释放时把记录放入 delete-ack batch，再由 `BatchDeleteAckReq` 通知 master。`get_cached_info` 若仍有引用，master holding 会继续保留。 | +| external process | 根据共享内存 base、`ExternalMemHolderInfo.offset/len` 和 `holder_id` 构造的公共 `MemHolder`。 | 对象释放时向 requester owner 发送 `ExternalDeleteAckReq`；它不直接访问或释放 master 的 `Allocation` 对象。 | + +holder 释放只移除 `get_holding` 这条引用。`ReuseReplica` 或 durable promotion 成功后，route 的 `memory` 仍持有独立的 `Arc`；该副本要经过缓存驱逐或 route 清理才会释放。temporary target 没有这条 route 引用。 + +## 5. 读取子步骤——回填传输：本地复用 target，远端 source 主动 push + +上一节已经确定了读取 source、requester target 及其生命周期；进入实际搬运阶段后，memory 与 SSD 链路在 source 形态上分开。memory route 指向 owner 已注册 memory segment 中的现成 bytes，requester 取得 source 和 target 地址后即可发起 pull。SSD route 指向 shard 文件中的 entry，磁盘内容还不是 transfer engine 可直接使用的内存 source，必须先由 SSD owner 完成本地读盘。 + +最终 target 始终位于 requester owner 的 memory segment，并由 master 侧 guard 保活。当 SSD owner 与 requester owner 为同一节点时，本地 IO 可直接写入 target。当两者不同时，本地 IO 无法跨 owner 写入 segment。此时，SSD owner 先用本地 source staging 接收 chunk，待 chunk ready 后主动 push 到远端 target。这个方向既由 SSD IO 与跨节点传输的边界决定，也让远端路径可以重叠读盘和网络传输。 + +### 5.1 本地回填复用最终 target + +SSD owner 与 requester owner 在同一节点时，master 令 `src_addr == target_addr`，并把 target allocation 的实际 capacity 放进 `ssd_stage_len`。owner 随后调用 `load_into_addr`：满足 direct read 条件时，读盘直接写 target；否则由 `KvSsdStorage` 通过 scratch buffer 完成。两条分支都不需要控制面申请独立 source allocation。 + +```rust +if peer_id.is_none() && stage_addr == target_addr { + return store + .load_into_addr( + key, // Select the logical KV entry. + put_id, // Select the exact committed key version. + target_addr, // Read into the final range in the requester segment. + len, // Expose only the real payload length. + stage_len, // Provide the target allocation capacity for aligned IO. + ) + .await; +} +``` + +### 5.2 远端传输方向 + +两条远端路径的数据都从 source 流向 requester target，差别在 transfer 发起方： + +- **内存：requester pull**：requester 已从 `GetStartResp` 取得远端 `src_addr` 和本地 `target_addr`，由它发起从 peer source 到本地 target 的 transfer。 +- **SSD：owner push**：requester 只发一次 `SsdStageReadReq`。SSD owner 每读好一个 chunk 就 push 到 requester target，省去 stage-ready 后的一次 RTT，并让读盘与网络传输重叠。 + +`transfer_data_no_copy` 的方向位决定 peer 端地址在本次传输中是 source 还是 target。该参数只描述数据方向，不表示 peer 持有 master 侧的 `Allocation` guard。下面是删减后的接口： + +```rust +pub async fn transfer_data_no_copy( + &self, // Use this node's transfer engine. + peer_node: Option, // Select the remote peer; None means local transfer. + peer_src_or_target: bool, // True: peer address is source; false: it is target. + src_addr: u64, // Provide the absolute source address. + target_addr: u64, // Provide the absolute target address. + len: u64, // Transfer this many real payload bytes. + seg_guard: Option, // Optionally keep the local segment alive. +) -> KvResult; +``` + +方向位同时决定 peer 地址的语义和本地 segment guard： + +- **`true`**：peer 地址是 source，当前节点 pull 到本地 target；owner 侧 segment guard 守住 `target_addr`。 +- **`false`**：peer 地址是 target，当前节点从本地 source push；owner 侧 segment guard 守住 `src_addr`。 + +远端 memory get 和远端 SSD 回填分别使用这两个方向： + +```rust +// Memory get: requester owner initiates transfer with peer as source. +client_transfer_engine + .transfer_data_no_copy( + Some(memory_owner_id), // Contact the owner that holds the memory replica. + true, // Interpret the peer address as the source. + remote_src_addr, // Read from the replica range in the peer's segment. + requester_target_addr, // Write into the requester's local target. + len, // Move only the real payload bytes. + None, // Let the engine acquire the local target guard. + ) + .await?; + +// SSD refill: SSD owner pushes a ready chunk to requester target. +let chunk_target_addr = requester_target_addr + .checked_add(chunk.offset) // Place this chunk at its payload offset. + .ok_or_else(|| KvError::Api(ApiError::InvalidArgument { // Convert overflow into a KV error. + detail: "chunk target addr overflow".to_string(), // Explain which address calculation failed. + }))?; + +client_transfer_engine + .transfer_data_no_copy( + Some(requester_owner_id), // Contact the requester whose segment contains the target. + false, // Interpret the peer address as the target. + chunk.stage_addr, // Read from this SSD owner's ready staging chunk. + chunk_target_addr, // Write to the matching offset in requester target. + chunk.len, // Move this chunk's real payload length. + None, // Let the engine acquire the local source guard. + ) + .await?; +``` + +只有远端 SSD 回填会发送 `SsdStageReadReq`。请求把 key-version、`get_id`、staging 地址与容量、requester target 位置和真实 payload 长度交给 SSD owner： + +```rust +pub struct SsdStageReadReq { + pub key: String, // Identify the logical KV entry. + pub put_id: PutIDForAKey, // Reject data from a different key version. + pub get_id: u64, // Bind the stage operation to its in-flight get. + pub stage_addr: u64, // Point to aligned staging on the SSD owner. + pub stage_len: u64, // State the available aligned staging capacity. + pub target_node_id: NodeIDString, // Identify the requester whose segment contains the target. + pub target_addr: u64, // Point to the final target range in that segment. + pub len: u64, // Preserve the real payload length. +} +``` + +### 5.3 chunk pipeline 与背压 + +远端 SSD owner 把读盘和传输拆成两个并发 future： + +- **producer**：按 chunk 从 SSD 读到 source staging。 +- **consumer**：收到 ready chunk 后立刻 push 到 requester target。 + +两者通过有界 `mpsc` ready queue 连接。默认 chunk 大小为 4 MiB；producer 最多并发 4 个 SSD read，consumer 最多保留 4 个 transfer inflight，ready queue 容纳 8 个已就绪 chunk 描述。第一个 chunk 读好后即可开始网络传输。 + +#### producer / consumer 实现片段 + +```rust +let ready_queue_capacity = DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT + .saturating_mul(2) // Buffer two transfer windows between reader and sender. + .max(1); // Keep the bounded channel valid for any configured window. +let (chunk_tx, chunk_rx) = + ::tokio::sync::mpsc::channel(ready_queue_capacity); // Connect producer and consumer. + +let producer = store.load_into_addr_chunks( + key, // Select the logical KV entry. + put_id, // Select the exact committed version. + stage_addr, // Read chunks into aligned source staging. + len, // Stop after the real payload length. + stage_len, // Enforce the staging allocation capacity. + DEFAULT_READ_TRANSFER_PIPELINE_CHUNK_BYTES, // Bound each SSD read to 4 MiB. + DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT, // Allow up to four concurrent SSD reads. + chunk_tx, // Publish each ready chunk to the sender. +); + +let consumer = self.transfer_loaded_ssd_chunks( + peer_id, // Select the remote requester; None means the target is local. + target_addr, // Use the final target range in the requester segment. + chunk_rx, // Consume chunks as soon as their SSD reads complete. +); +let (producer_res, consumer_res) = ::tokio::join!( + producer, // Run SSD reads concurrently with network transfers. + consumer, // Push ready chunks without waiting for the full value. +); +match (producer_res, consumer_res) { + (Ok(()), Ok(())) => Ok(()), + (_, Err(err)) => Err(err), + (Err(err), _) => Err(err), +} +``` + +consumer 只在 transfer inflight 少于 4 时从 ready queue 取新 chunk。达到上限后，有界 queue 会把背压传给 producer；已提交的 transfer 完成后，consumer 再继续取 chunk。 + +这条 pipeline 中有三种不同的保活机制： + +| 机制 | 保护对象 | 释放时点 | +| --- | --- | --- | +| ring read pin | SSD entry 与文件 offset，防止尚未完成的 SSD read 被 tail 覆盖。 | producer 完成全部 SSD read，并把最后一个 ready chunk 送入有界 queue 后。 | +| `InflightGetInfo.source_allocation` | master 侧 guard，保活 SSD owner segment 中的整块 source staging range。 | `GetDone`、`GetRevoke` 或 in-flight TTL。 | +| transfer segment guard | 某个在途 push 使用的本地 staging 地址范围。 | 对应 chunk transfer 完成。 | + +ring read pin 只覆盖 SSD 读取阶段，不覆盖整个网络 push 窗口。producer 释放 pin 后，尚未完成的 push 由 master 侧 staging guard 和每次 transfer 的 segment guard 保护。 + +下图聚焦远端 SSD 回填内部的 pipeline；`GetStart/GetDone` 的完整流程见第 4 节。`owner-requester-target` 只表示数据传输端点，其 allocation 生命周期仍由 master 的 in-flight 状态和 `get_holding` 管理。 + +```mermaid +sequenceDiagram + autonumber + participant owner_reader as owner-ssd-reader + participant owner_media as owner-local-ssd + participant owner_queue as owner-ready-queue + participant owner_sender as owner-ssd-sender + participant owner_target as owner-requester-target + + owner_reader->>owner_media: pin committed ring entry + par read producer + loop 按 4 MiB chunk 读取 + owner_reader->>owner_media: read chunk N into source staging + owner_media-->>owner_reader: IO complete, chunk N ready + owner_reader->>owner_queue: enqueue chunk N + end + owner_reader->>owner_media: 最后一个 ready chunk 入队后 release ring pin + and transfer consumer + loop 按 payload offset 推送 + owner_queue-->>owner_sender: deliver chunk N + owner_sender->>owner_target: push chunk N + owner_target-->>owner_sender: transfer complete + opt transfer inflight 已达 4 + owner_sender->>owner_queue: pause dequeue + owner_queue-->>owner_reader: bounded-channel backpressure + end + end + end + Note over owner_sender,owner_target: 未完成的 push 继续由 master staging guard 和 transfer guard 保活 +``` + +## 6. owner 本地机制：对齐、ring 和 IO 调度 + +**SSD 介质的控制引擎留在 owner。** SSD 文件 offset、`O_DIRECT` 文件 IO 对齐、ring 覆盖保护和 `io_uring` 调度都由 owner 处理。master 仍负责分配 requester target 和跨节点 staging，并传递相应地址、容量与真实 payload 长度。用户只看到真实 payload 语义。 + +master 使用已注册的 owner segment allocator 为每次回填分配 requester target，并创建对应 guard；跨节点回填还会分配 SSD source staging。bytes 位于相应 owner 的 segment，RAII guard 位于 master。master 不记录 SSD 文件布局，用户侧仍通过 `MemHolder` 访问真实 payload。 + +### 6.1 写入：区分真实长度和对齐长度 + +**`len` 是 payload 长度，`aligned_len` 是 owner 本地 IO 长度。** target owner 从请求携带的 target 地址复制 payload，再构造 512-byte 对齐的写入 buffer。route 中的 `Arc` 在此期间保活 target 地址范围。 + +| 字段 | 含义 | 使用范围 | +| --- | --- | --- | +| `len` | 真实 payload 长度。 | master route、`SsdReplicaCommitReq`、transfer 长度和 `MemHolder.len`。 | +| `aligned_len` | `align_up(len, 512)` 得到的对齐长度。 | owner 本地 ring 分配、shard offset 推进和 `O_DIRECT` read/write。 | + +例如，1000-byte payload 的 `aligned_len` 是 1024 bytes： + +1. owner 创建 1024-byte 的对齐 `AlignedBuffer`，先置零，再把前 1000 bytes 填入 payload。 +2. ring 分配、shard offset 推进和 SSD 写入都使用 1024 bytes。 +3. `SsdIndexEntry` 同时记录 `len = 1000` 和 `aligned_len = 1024`；owner 向 master 提交的仍是 `len = 1000`。 +4. 后续读取会校验 `entry.len`，transfer 和 `MemHolder` 只暴露前 1000 bytes。末尾 24 bytes padding 留在 owner 内部。 + +### 6.2 ring：从 `Writing` 到 `Committed` + +**只有完整写入的 `Committed` entry 才能参与读取；写入、route 提交或读取期间，ring 不会覆盖对应区间。** + +本地 SSD 引擎按 value 粒度选路： + +- 用 `next_write_device` round-robin 选择一个有效 device。 +- 把整个 value 发送到选中 device 的 writer queue。 +- 当前实现不把同一个 payload 拆到多块 device 做条带化。 + +因此，多 device 并行发生在 value 粒度：不同 value 可以落到不同 device；一个 value 内部仍是某个 shard 的连续 offset。 + +owner 本地 `SsdRingBuffer` 决定 SSD 文件位置。writer task 只在本 device 的 `shard_ids` 里选择 shard，并为对齐后的 payload 分配连续文件区间。ring 按以下规则分配： + +- `aligned_len = align_up(len, 512)`，ring 按对齐后长度分配空间。 +- 每个 shard 是环形空间，`head` 前进分配新写入，`tail` 表示可以覆盖到哪里。 +- 如果当前 shard 尾部剩余空间不足，会跳到文件开头继续找连续空间。 +- 如果推进 `tail` 会覆盖未完成写入、active read pin 或 route commit pin，返回 `BlockedByBusyIo`。 + +分配完成后，状态转换、覆盖保护和 route 清理遵循以下规则： + +- **`Writing`** + - offset 已分配，SSD 写入尚未成功。 + - 默认单 buffer 路径使用 `O_DIRECT` 和 `io_uring` write。 +- **`Committed`** + - 只有实际写入字节数等于 aligned buffer 长度，entry 才从 `Writing` 转为 `Committed`。 + - `pin_read` 只接受 committed 且 offset 仍有效的 entry，未完成写入不会被 `get` 选中。 +- **ring 覆盖保护** + - route commit pin：保护本地 commit 到 master 处理完 commit RPC 的窗口。 + - read pin：保护 SSD read 使用 entry 与 file offset 的窗口；远端 push 后半段由 staging allocation 和 transfer guard 保活。 +- **tail 驱逐通知** + - tail 推进时，ring 返回本次被覆盖的全部 committed `KvSsdKey`。 + - writer 把它们送入有界失效队列，不在 SSD 写入热路径逐条调用 master。 + - owner batch task 按 256 条或 10 ms 的阈值聚合 RPC；第 7 节展开 master 的条件清理和读取兜底。 + +下图展示 owner 本地 `Writing → Committed` 与 master 发布 SSD route 之间的时序。`owner-memory-target` 只有前 `len` bytes 是真实 payload，其地址范围由 master route 中的 `Arc` 保活。 + +```mermaid +sequenceDiagram + autonumber + participant master_route as master-route + participant owner_target as owner-memory-target + participant owner_writer as owner-ssd-writer + participant owner_ring as owner-ssd-ring + participant owner_engine as owner-io-engine + participant owner_media as owner-local-ssd + + master_route->>owner_writer: SsdReplicaPersistReq(key, put_id, target_addr, len) + owner_writer->>owner_target: copy first len bytes + owner_target-->>owner_writer: real payload + owner_writer->>owner_writer: zero padding and build aligned_len buffer + owner_writer->>owner_ring: reserve contiguous aligned_len bytes + owner_ring-->>owner_writer: shard_id + file_offset + Writing entry + owner_writer->>owner_engine: O_DIRECT write(aligned buffer) + owner_engine->>owner_media: submit aligned_len bytes + owner_media-->>owner_engine: write completion(actual bytes) + owner_engine-->>owner_writer: completion + + alt actual bytes == aligned_len + owner_writer->>owner_ring: Writing → Committed and acquire route commit pin + owner_ring-->>owner_writer: KvSsdPersistGuard + owner_writer->>master_route: SsdReplicaCommitReq(key, put_id, len) + master_route->>master_route: validate, then publish or ignore commit + master_route-->>owner_writer: SsdReplicaCommitResp + owner_writer->>owner_ring: release route commit pin + owner_writer-->>master_route: SsdReplicaPersistResp(persisted = true) + else partial write or IO failure + owner_writer->>owner_ring: abort Writing entry + owner_writer-->>master_route: SsdReplicaPersistResp(error) + end +``` + +### 6.3 读取：direct read 和 scratch fallback + +**direct read 是有条件的 owner 本地 fast path；任一对齐或容量条件不满足时，scratch read 负责兜底。** + +每个 chunk 的文件位置由 `entry.file_offset + offset` 得到。本地回填读入 requester target；远端回填读入 SSD source staging。 + +| 路径 | 选择条件 | owner 本地动作 | 复制边界 | +| --- | --- | --- | --- | +| direct read | 读入地址和文件 offset 都按 512 bytes 对齐，target capacity 覆盖对齐后的 IO 范围。按 chunk 读取时，当前 read 长度也必须对齐。 | `io_uring` 直接把 shard 内容读入本地目标地址。 | 这段 SSD-to-target 读取没有额外的 owner 本地 copy；这不代表整条 `get` 链路都是零拷贝。 | +| scratch read | 任一 direct read 条件不满足。 | 先读入 512-byte 对齐的 `AlignedBuffer`，再把当前 chunk 的真实 payload 复制到目标地址。 | 多一次 owner 本地 CPU copy。 | + +两条路径都只向下游暴露真实 payload 长度。`O_DIRECT` padding 不进入 transfer 长度，也不进入用户 `MemHolder.len`。 + +下图只画 SSD source owner 内部的读盘。`owner-read-target` 始终位于 SSD source owner 本地：本地回填时，它同时是 requester target；远端回填时，它是 source staging。远端 requester 的最终 target 不在图中，后续 push 见第 5.3 节。 + +```mermaid +sequenceDiagram + autonumber + participant owner_reader as owner-ssd-reader + participant owner_ring as owner-ssd-ring + participant owner_engine as owner-io-engine + participant owner_media as owner-local-ssd + participant owner_scratch as owner-scratch-buffer + participant owner_target as owner-read-target + + owner_reader->>owner_ring: pin committed entry and request chunk plan + owner_ring-->>owner_reader: file_offset + len + aligned read len + owner_reader->>owner_reader: check address, offset, length, and capacity alignment + + alt direct read 条件全部成立 + owner_reader->>owner_engine: read shard directly into owner-read-target + owner_engine->>owner_media: submit O_DIRECT read + owner_media-->>owner_engine: IO complete, target memory filled + owner_engine-->>owner_reader: completion + owner_reader-->>owner_target: publish current chunk's real len + else 任一对齐条件不成立 + owner_reader->>owner_scratch: acquire aligned buffer + owner_reader->>owner_engine: read shard into owner-scratch-buffer + owner_engine->>owner_media: submit O_DIRECT read + owner_media-->>owner_engine: IO complete, scratch memory filled + owner_engine-->>owner_reader: completion + owner_reader->>owner_target: copy only current chunk's real payload bytes + end + owner_reader->>owner_ring: release read pin + Note over owner_reader,owner_target: transfer 和 MemHolder 只看到真实 len，padding 留在 owner 内部 +``` + +### 6.4 IO 调度：让回填读优先，同时保留后台写进展 + +**同一个 device 上，SSD 回填读位于 `get` 的同步等待路径，SSD persist 写则在内存副本发布后异步进行。** 因此当前调度给读更多的 SQE 提交机会，但仍持续推进写。写长期不进展会延迟 SSD 副本 commit，并让 `Writing` entry 和对齐 buffer 继续占用资源；持续读下的绝对读优先还可能拖住 ring 后续复用空间。`read_weight = 2` 表达的是这种有界的读偏置。 + +调度分成两层，各自处理不同问题： + +| 层次 | 当前机制 | 设计原因 | +| --- | --- | --- | +| device 任务层 | 每个去重后的 device 独立持有 writer queue、reader queue 和 `UringIoEngine`。writer task 管理 ring 分配以及 `Writing -> Committed/abort`；reader task 管理 chunk read 和 completion 复查。 | 在当前 owner 本地调度范围内，实际 device 是 IO 竞争边界。独立队列避免忙 device 在软件调度层阻住空闲 device，也让读写使用不同的上层并发上限和状态机。 | +| `io_uring` 提交层 | 每个 `io_uring` 线程各有 read/write channel，在 `io_depth` 内选择下一个 SQE。 | 分开 channel 后，引擎可以在真正占用 SQE 槽位时仲裁读写，避免后到的回填读只能排在单一 FIFO 中的后台写之后。 | + +`io_depth` 在设备并行度与在途资源占用之间设置上界：多个已提交 SQE 可以让 SSD 持续有工作可做，上限则避免单个 `io_uring` 线程无界占用 buffer 和 completion context。真正能到达这一层的并发量，还会受 writer/reader task 各自 inflight 上限的约束。 + +`read_weight = 2` 只在第二层生效。每个 `io_uring` 线程独立维护 `read_inflight` 和 `write_inflight`：前者小于等于后者的两倍时，先尝试 read channel；超过后，先尝试 write channel。首选 channel 为空时会立即取另一个 channel 的任务，不会为了维持比例而让 SQE 槽位空闲。 + +这个 `2:1` 是固定的软调度偏置，不是硬上限或队列积压阈值。等号分支会再提交一个 read，completion 速率也会持续改变 inflight 组成；因此它不保证单个 device 上的精确 `2:1` 比例。它也不考察队列长度、请求大小、等待时间或 deadline。权重只决定下一个尚未提交的 SQE；已提交的 IO 不会被抢占，`io_depth` 已满时，新的回填读仍需等待 CQE 释放槽位。所以这是一个低成本的提交偏置，不构成读延迟 SLA。 + +**`SingleBuffer` 是对当前数据形态的特化。** 写路径已经把 payload 放入连续的对齐 `AlignedBuffer`；direct read 和 scratch read 的目标也都是一段连续内存。默认 `KvSsdUringMode::SingleBuffer` 因而直接提交 `IORING_OP_READ` / `IORING_OP_WRITE`，避免构造并保活只含一个元素的 iovec。`Iovec` 保留为测试和对照路径。这项特化没有消除 `O_DIRECT` 对齐、padding、scratch copy 或网络传输，其收益范围仅限于 owner 本地 IO 提交层。 + +read completion 复查与 read pin 也承担不同职责： + +| 机制 | 作用时段 | 职责 | +| --- | --- | --- | +| read pin | 从获取 committed entry 到本次读取结束。 | 阻止 `tail` 越过正在使用的文件区间，预防 offset 在 IO 期间被复用。 | +| completion 复查 | CQE 返回后。 | 检查读取开始时捕获的 `entry.begin` 是否仍不早于当前 `tail`。已失效时删除 entry 并返回 `KeyNotFound`。 | + +read pin 负责预防覆盖，completion 复查负责拦住异常状态下的结果发布。复查无法撤销已经写入 target 的 IO，因而不能取代 pin；它是跨过异步提交边界后的最后一道不变量检查。在完整 `get` 路径中，返回错误后的 SSD stage 会进入第 7.2 节的 revoke 与有界重新选路，不会向用户交付这次 target。 + +下图聚焦单个 device 内部的调度，不涉及 `master-*` 或 `external-*` 角色。`owner-kernel-io` 表示该 owner 文件上的 `io_uring` SQE/CQE 交互，不表示 owner 独占内核。 + +```mermaid +sequenceDiagram + autonumber + participant owner_write as owner-write-path + participant owner_read as owner-read-path + participant owner_engine as owner-io-engine + participant owner_kernel as owner-kernel-io + participant owner_media as owner-local-ssd + participant owner_ring as owner-ssd-ring + + par write task arrives + owner_write->>owner_ring: reserve offset and create Writing entry + owner_ring-->>owner_write: shard_id + file_offset + owner_write->>owner_engine: enqueue write request + and read task arrives + owner_read->>owner_ring: pin committed entry and resolve offset + owner_ring-->>owner_read: read plan + pin + owner_read->>owner_engine: enqueue read request + end + + loop while this io_uring thread has a free SQE slot + owner_engine->>owner_engine: choose next unsent IO by the per-thread inflight mix, then fall back to the other channel + Note over owner_engine,owner_kernel: read_weight does not preempt submitted SQEs + alt write request selected + owner_engine->>owner_kernel: submit WRITE SQE + owner_kernel->>owner_media: O_DIRECT write + owner_media-->>owner_kernel: device completion + owner_kernel-->>owner_engine: write CQE + owner_engine-->>owner_write: actual bytes or error + owner_write->>owner_ring: commit or abort Writing entry + else read request selected + owner_engine->>owner_kernel: submit READ SQE + owner_kernel->>owner_media: O_DIRECT read + owner_media-->>owner_kernel: device completion + owner_kernel-->>owner_engine: read CQE + owner_engine-->>owner_read: actual bytes or error + owner_read->>owner_ring: check captured entry.begin is not before current tail + alt captured range 仍有效 + owner_ring-->>owner_read: range valid, return chunk bytes + else captured range 已失效 + owner_ring-->>owner_read: stale + owner_read->>owner_ring: remove stale entry + owner_read->>owner_read: return KeyNotFound + end + end + end +``` + +## 7. 一致性：沿用 key-version 和 holder 生命周期 + +**SSD 副本只有在 master route 和 owner ring 都有效时才可读。** `get_holding` 保存的是 requester target 的 master 侧 `Allocation` guard；远端 source staging guard 只保留在本次 `InflightGetInfo` 中。 + +一次 SSD 回填的完成语义包含以下四个方面： + +- **`GetDoneReq` 调用方** + - 本地 SSD source：requester 读盘并完成 target 写入后调用。 + - 远端 SSD source：SSD owner 读盘并 push 到 requester target 后调用。 +- **最终 holder 与 allocation 生命周期** + - `inflight_gets.req_node_id` 始终指向 requester owner，不受 `GetDoneReq` 调用方影响。 + - master 把 requester target guard 写入 `get_holding` 并生成 `holder_id`；target bytes 仍位于 requester owner 的 segment。 + - 远端 source staging guard 不进入 `get_holding`，只随 in-flight get 生命周期释放。 +- **远端响应转发** + - SSD owner 把 `GetDoneResp` 中的 `holder_id`、`allocation_mode`、错误码和 server time 写入 `SsdStageReadResp.done_*`。 + - requester 使用这些字段构造 owner 侧 `MemoryInfo` 和 `ExternalMemHolderInfo` 响应。external 再构造公共 `MemHolder`，用户看不到 SSD stage RPC。 +- **durable promotion** + - `GetStart` 尝试预留 durable slot。成功时 `allocation_mode = DurableReplica`，否则为 `Temporary`。 + - `GetDone` 只在 allocation mode 为 durable、route `put_id` 仍匹配、且 requester tomb tag 仍 live 时，才把 target guard 写入 entry 的 `memory`。 + - promotion 成功后，SSD 回填结果成为新的内存副本。 + +### 7.1 owner 驱逐后批量清理 route + +为了让 master 中因物理驱逐产生的陈旧 route 尽快收敛，同时尽量减少逐条通知产生的控制面通信，owner 先完成物理驱逐，再聚合对应的 key-version，批量通知 master 清理逻辑 route： + +- **驱逐输入**：ring 推进 `tail` 时，一次返回本轮覆盖掉的全部 committed key-version。 +- **batch 触发** + - 第一条记录到达后，最多等待 10 ms。 + - 累计到 256 条时立即发送。 +- **RPC 失败**：当前 batch 留在 task 中重试。 +- **队列或控制面不可用** + - 有界队列满、进程退出或 master 持续不可用，都不阻塞 SSD 写入。 + - 可能残留的陈旧 route 由第 7.2 节的读取兜底处理。 + +批量 RPC 只携带待清理的 key-version： + +```rust +pub struct SsdReplicaEviction { + pub key: String, // Identify the route entry to inspect. + pub put_id: PutIDForAKey, // Restrict cleanup to the evicted key version. +} + +pub struct BatchSsdReplicaEvictReq { + pub evictions: Vec, // Carry one owner-side eviction batch. +} +``` + +master 收到 batch 后按以下规则清理 route： + +- **owner 身份**：请求不重复传 `node_id`，master 直接使用 RPC 发送方身份。 +- **版本门禁**：只有当前 route `put_id` 与驱逐项匹配，master 才清除该 owner entry 的 `ssd`。 +- **清理粒度** + - entry 仍有内存副本：只清除 `ssd`。 + - `memory` 和 `ssd` 都为空：删除 owner entry。 + - 整个 key-version 没有 live 副本：删除 `kv_routes`。 +- **prefix index** + - 在一次写锁内批量清理。 + - 同名 key 已写入新 route：保留新版本 prefix entry。 +- **master 职责边界**：master 不参与 SSD offset 分配，也不决定本地驱逐顺序；只接收 owner 的驱逐结果并收敛逻辑 route。 + +### 7.2 陈旧 SSD route 的读取兜底 + +**SSD stage 失败后，requester 先撤销当前 source，再有界重新选路。** 重新选路可以继续使用其他 live source；没有副本时返回 miss。连续 3 次 SSD stage 失败后返回 error。任何失败分支都不会向用户暴露半成品 holder。 + +SSD stage 失败时，请求方会调用： + +```rust +GetRevokeReq { + get_id, // Revoke this in-flight get and release its allocations. + drop_ssd_source: true, // Clear the failed SSD option from the current route. +} +``` + +master 按下列顺序撤销失败 source： + +- **释放本次 in-flight**：先删除 `inflight_gets`，让 target 和可选 staging 进入释放路径。 +- **限定 route 删除条件** + - `drop_ssd_source == true`。 + - `source_kind == GetSourceKind::Ssd`；memory source 不会因该标志被删除。 + - route `put_id` 与本次 in-flight get 仍一致；不撤销新版本 route。 +- **清理失败 source** + - 把 `src_node_id` 对应 entry 的 `ssd` 设为 `None`。 + - entry 也没有内存副本：删除该 entry。 + - 整个 key-version 已无 live 副本：删除 `kv_routes`；启用 prefix index 时，再异步清理对应索引。 + - 失败 source 清理与主动 batch 驱逐共用同一个条件删除函数。 +- **重新选路** + - `GetRevokeResp` 返回后，请求方重新发送 `GetStartReq`。 + - 下一次可能命中另一份内存副本、SSD 副本、同 key 新版本，或返回 miss。 + - 陈旧 SSD source 最多触发 3 次重新选路，避免无限循环。 + - SSD stage RPC 失败和不合法的 `ssd_stage_len` 都进入同一套 revoke 与重选流程。 +- **allocation 与崩溃清理** + - 远端 staging 保存在 `source_allocation`；本地回填时该字段为 `None`。 + - `GetDone`、`GetRevoke` 或 in-flight TTL 驱逐都会释放远端 staging。 + - requester 在 in-flight 阶段崩溃：节点 tomb 阻止副本 promotion 和旧节点复用。`inflight_gets` 最迟由 60 秒 TTL 释放 target/staging guard。 + - requester 在 `GetDone` 后离开：master 的 member-left 清理移除该节点的 `get_holding`；正常路径则由 owner 的 holder delete-ack batch 释放。 + - 远端 SSD owner 崩溃：其进程内 staging bytes 消失。master 中对应的 staging guard 随 `GetRevoke` 或 in-flight TTL 释放，节点 tomb 使旧 SSD route 不再可选。 + +在 master 接受 `GetDoneReq` 之前发生的读取或传输失败会走 `GetRevoke`，target guard 不进入 `get_holding`，用户也不会拿到半成品 payload。 + +另一个边界发生在 master 已提交 `GetDone`、但响应随后丢失时。此时 target guard 已写入 `get_holding`，requester 却可能没有收到 `holder_id`。 + +当前通用 holder 协议没有按 `get_id` 幂等查询已提交结果的接口。这类未交付 holder 无法由随后的 `GetRevoke` 立即回收，只能等待能够识别该 holding 的后续清理，例如 requester member-left。这不影响 SSD 数据完整性，但属于当前完成协议的资源回收边界。 + +下图展示成功完成、失败 revoke 和超时清理三条路径。远端路径由 `owner-ssd-source` 代替 requester 调用 `GetDoneReq`，但最终 target 始终位于 requester segment。 + +```mermaid +sequenceDiagram + autonumber + participant external_user as external-user + participant owner_requester as owner-requester + participant master_route as master-route + participant owner_ssd as owner-ssd-source + + owner_requester->>owner_ssd: SsdStageReadReq 或本地 stage 调用 + alt SSD stage 成功，target 已完整 + owner_ssd-->>owner_requester: 本地回填完成或最后一个 chunk 已 push + alt SSD source 就在 requester owner + owner_requester->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，释放其他 in-flight 引用 + master_route-->>owner_requester: GetDoneResp(holder_id) + else SSD source 在其他 owner + owner_ssd->>master_route: GetDoneReq + master_route->>master_route: target guard 转入 get_holding，释放 staging 和其他 in-flight 引用 + master_route-->>owner_ssd: GetDoneResp(holder_id, allocation_mode) + owner_ssd-->>owner_requester: SsdStageReadResp(done_*) + end + owner_requester-->>external_user: ExternalMemHolderInfo；external 本地构造 MemHolder + else SSD stage 失败 + owner_ssd-->>owner_requester: stage error + owner_requester->>master_route: GetRevokeReq(drop_ssd_source = true) + master_route->>master_route: remove inflight_gets and release target / staging references + alt source_kind = Ssd 且 route put_id 仍匹配 + master_route->>master_route: 清空当前 owner entry.ssd，必要时删除 entry / route + else source 类型或版本已变化 + master_route->>master_route: 保留当前 route + end + master_route-->>owner_requester: GetRevokeResp + loop 陈旧 SSD source 最多重新选路 3 次 + owner_requester->>master_route: GetStartReq + master_route-->>owner_requester: 其他 live source / 新版本 / miss + end + owner_requester-->>external_user: 重试结果、miss 或达到上限后的 error + end + + opt 请求未进入 GetDone 或 GetRevoke + master_route->>master_route: 60 s in-flight TTL 释放 target / staging guard + end + opt requester 在 GetDone 后离开 + master_route->>master_route: member-left 清理该节点的 get_holding + end +``` + +### 7.3 删除、驱逐与重启 + +- **覆盖写与 delete**：覆盖写生成新 `put_id`，旧 SSD late commit 因版本不匹配被忽略。显式 delete 移除整个 `kv_routes` 条目；节点 tomb 清理则删除该节点在 `node_replicas` 中的整个 entry。 +- **两层独立驱逐**：内存驱逐只清空 `memory`；SSD ring 驱逐或读取失败只清空 `ssd`。两个字段都为空时才删除 owner entry，整个 route 没有副本时才删除 `kv_routes`。 +- **物理 bytes 不等于可读副本**：route 被删除或版本不匹配后，旧 bytes 即使还在 shard 中也不会被公共 `get` 命中。陈旧 route 被选中时，`pin_read` 会复查 committed 状态和 offset，失败后 revoke source 并重新 `GetStart`。 +- **owner 重启**：shard 以 `create + truncate` 打开，并重建空的本地 SSD cache 和 ring。当前没有 WAL、checkpoint 或 shard 扫描。控制面尚未完成 tomb 清理时，读取兜底会撤销失效 SSD source，再重新选路或返回 miss。 + +## 8. 性能测试 + +本节先给出核心结果，再说明计数边界与实验设置，最后分别展开 CPU、GPU 和 SSD 路径证据。性能数字只在对应小节定义的硬件、数据集、输出模式和统计窗口内成立。 + +### 8.1 核心结果 + +- **CPU / GPU 使用同一 pressure 口径**：CPU 主图包含 192 条记录，覆盖原生引用 / handle 与 `bytes`、4 个 payload、4 个并发度、Fluxon、两种 Mooncake 拓扑和 SSD 关/开；GPU 主图包含 72 条记录，覆盖 `cuda`、3 个 payload、4 个并发度、3 种拓扑和 SSD 关/开。264 条记录全部为 `SUCCESS`、0 error、来源完整且 unknown 为 0。 +- **miss 是结果，不是过滤条件**：两侧都使用约 2.5 GiB 数据集和 2 GiB DRAM。柱高只统计 hit payload，红点显示 `hit / (hit + miss)`。SSD-off 的 CPU 命中率为 Fluxon **89.112%–94.390%**、Mooncake 独立 owner **23.759%–66.879%**、Mooncake 每进程 **22.488%–68.655%**；GPU 对应范围为 **89.326%–93.267%**、**34.557%–62.223%** 和 **35.430%–64.782%**。 +- **SSD-on 也保留实际 miss**：CPU 的 Fluxon 和 Mooncake 独立 owner 在 SSD-on 下均为 100% 命中，Mooncake 每进程为 **72.830%–100%**；GPU 对应为 100%、100% 和 **59.855%–100%**。因此 SSD-on/off 柱差同时包含命中集合、介质来源和系统调度差异，不能解释成单一 SSD 开销。 +- **SSD-on 的 hit 同时来自两层**：CPU 中 Fluxon 的 SSD 来源占 hit 的 **17.550%–48.650%**，Mooncake 独立 owner 为 **62.065%–82.385%**，Mooncake 每进程为 **59.341%–83.011%**；GPU 对应为 **20.242%–46.508%**、**64.837%–77.135%** 和 **56.073%–84.198%**。 +- **GPU 多并发已进入同一张主图**：Fluxon 在 SSD-off 下从 c2 到 c16 提高到 5.83、5.83 和 4.68 倍，SSD-on 下提高到 6.50、7.22 和 6.80 倍，分别对应 4、8、16 MiB。Mooncake 的两种拓扑也按相同的六柱顺序展示，不再另画一张来源图。 +- **同路径 Foyer 消融下原生 SSD 仍然更快**：64 条测量记录对应 32 个原生/Foyer 配对，全部为 `SUCCESS`、0 error、0 unknown、0 persist failure，两侧命中率均为 100%。当前 Foyer adapter 的 hit payload 带宽是原生实现的 **14.8%–65.9%**，配对中位数为 **37.2%**，32 个配对中没有胜出；P95 是原生实现的 **2.30–8.46 倍**。这是当前接入路径的结果，不外推为 Foyer 库的通用性能结论。 + +这些数字只覆盖下文定义的单机、固定容量、Zipf 分布、30 秒窗口和各自 fast path。实验不包括跨机器 RDMA、多物理盘条带化、重复运行的统计置信区间或真实业务 trace。 + +### 8.2 计数边界：原生引用 / handle、`bytes` 与 `cuda` + +Mooncake 没有 Fluxon 意义上的 holder。`holder` 只是 benchmark 配置 `kv_get_output=holder` 的归一化模式名，Fluxon 和 Mooncake 实际返回 `MemHolder` 与 `BufferHandle`。 + +Mooncake 常规 `Store.get()` 直接返回 Python `bytes`，因此结果包含完整 payload copy。`Store.get_buffer()` 返回 native `BufferHandle`，无需创建 Python `bytes` 即可取得 `ptr()`、`size()` 或 `memoryview`。原生引用 / handle 组把计数边界放在 host 内存引用返回之后。 + +#### 两侧原生引用的调用与生命周期差异 + +benchmark 中的实际分支等价于： + +```python +if backend == "MOONCAKE": + # Internal adapter calls Mooncake Store.get_buffer(). + result = mooncake_store.get_buffer_blocking(key) # Fetch a native BufferHandle. +else: + result = fluxon_store.get_blocking(key) # Fetch a native MemHolder. + +native_ref = result.unwrap() # Keep the native host reference alive while counting success. +``` + +两者只在本次 benchmark 的**计数边界**上对齐，类型语义并不相同： + +| 对比维度 | Fluxon | Mooncake | 对齐方式 | +| --- | --- | --- | --- | +| 实际返回类型 | `MemHolder` | `mooncake.store.BufferHandle` | 本文直接写真实类型，不写“Mooncake holder”。 | +| 底层调用 | `KvClient.get_blocking()` | `Store.get_buffer()`，由 benchmark adapter 包成 `get_buffer_blocking()` | 都在 native host buffer ready 后返回。 | +| 数据访问 | `MemHolder.access()` | `ptr()`、`size()`、`memoryview(BufferHandle)` | 原生引用模式不 materialize；bytes / CUDA 模式再继续消费。 | +| 生命周期语义 | 参与 Fluxon 的 `holder_id`、`get_holding` 和释放流程。 | 由 Mooncake Python binding 的 `BufferHandle` 管理返回 buffer 的生命周期。 | 只要求对象存活到本次操作计数完成。 | +| 不对齐的部分 | Fluxon holder 协议。 | 没有 Fluxon 的 `holder_id/get_holding/revoke` 协议。 | 不声称 API 或生命周期协议等价。 | + +原生引用模式仍然包含后端选择 SSD route 时从 SSD 到 host 的搬运。它只省去随后将完整 payload 物化为 Python `bytes` 的额外 copy。`cuda` 从同一个 native holder 取得 pageable host pointer，再交给 Fluxon 与 Mooncake 共用的双槽 H2D 流水线。只有对应 CUDA event 完成后，这次操作才计为成功。 + +下图覆盖一次 benchmark `get` 从发起到达到对应计数边界的过程。`external-backend-adapter` 是 benchmark 侧的统一适配层，封装完整的 Fluxon 或 Mooncake get；它不表示两套系统共享 Fluxon 的 owner/master 架构。 + +- **`external-benchmark-worker`**：发起 KV get，按 `holder/bytes/cuda` 模式消费结果，并在该模式的完成点计数。 +- **`external-backend-adapter`**：在同一调用面上包住 Fluxon 或 Mooncake 的完整 get，并返回已驻留 host 的 `MemHolder` 或 `BufferHandle`。 +- **`external-host-output`**：表示 pageable host 上的 native reference 生命周期，以及 `bytes` 模式的完整 payload 物化。 +- **`external-pinned-slot`**：每个 benchmark worker 复用的两个 pinned host slot 之一。 +- **`external-cuda-runtime`**：对应 slot 的 CUDA stream 和 event，负责 `cudaMemcpyAsync` 提交与完成通知。 +- **`external-gpu`**：接收 payload 的 GPU device buffer；这次实验不计 D2H 或 kernel 执行。 + +```mermaid +sequenceDiagram + autonumber + participant external_worker as external-benchmark-worker + participant external_backend as external-backend-adapter + participant external_host as external-host-output + participant external_pinned as external-pinned-slot + participant external_cuda as external-cuda-runtime + participant external_gpu as external-gpu + + external_worker->>external_backend: get(key) + external_backend-->>external_worker: MemHolder or BufferHandle with host-resident payload + external_worker->>external_host: retain native reference until this operation is counted + + alt holder output + external_worker->>external_worker: validate native reference and count success + external_worker->>external_host: release native reference + else bytes output + external_worker->>external_host: materialize the full payload as Python bytes + external_host-->>external_worker: bytes ready + external_worker->>external_worker: validate bytes and count success + external_worker->>external_host: release native reference + else cuda output + external_worker->>external_pinned: host memcpy from pageable payload into the next free slot + external_worker->>external_cuda: cudaMemcpyAsync(pinned slot, device buffer) + external_cuda->>external_gpu: start H2D DMA + par current H2D + external_gpu-->>external_cuda: matching CUDA event completes + external_cuda-->>external_worker: H2D complete + and next KV request + external_worker->>external_backend: get(next key) + external_backend-->>external_worker: next native host reference + end + external_worker->>external_worker: count current success after its event + external_worker->>external_host: release current native reference + end +``` + +| benchmark 输出模式 | Fluxon 映射 | Mooncake 映射 | 共同计数边界 | +| --- | --- | --- | --- | +| `holder`（归一化模式名） | `get_blocking()` 返回 `MemHolder`。 | `get_buffer_blocking()` 调用 Mooncake `Store.get_buffer()`，返回 `BufferHandle`；该对象暴露 `ptr()`、`size()` 和 buffer protocol。 | 原生引用已返回；不创建完整 Python `bytes`。 | +| `bytes` | `MemHolder.access()` 解码出 payload `bytes`，校验长度并触碰首尾字节。 | 从 `memoryview(BufferHandle)` 定位 payload，再调用 `bytes(...)` 完整复制，校验长度并触碰首尾字节。 | 完整 Python `bytes` 已生成；不含 GPU copy。 | +| `cuda` | 从 `MemHolder` 中取得 CPU DLPack payload 指针。 | 使用 `BufferHandle.ptr() + payload_offset` 取得 host 指针。 | 对应 CUDA event 已完成；不含 D2H 和 kernel 执行。 | + +来源计数与输出计数共用同一个 warmup 后统计窗口，但两套系统的观测方式不同： + +| 后端 | DRAM / SSD 归因方式 | 完整性判定 | 边界 | +| --- | --- | --- | --- | +| Fluxon | Rust get 路径保留 `GetSourceKind::Memory` / `GetSourceKind::Ssd`，经 external holder 传到 benchmark；每个成功操作单独计数。 | 每个 hit 都必须有且只有一个 source kind。 | 表示本次 get 选中并完成的 route source；它不是 NVMe 设备层 IO trace。 | +| Mooncake | 在统计窗口起止点采样 `get_offload_rpc_read_count()`；`0.3.11.post1` 的单 key `get_buffer()` 在选中 `LOCAL_DISK` 后会进入唯一的 offload-RPC 计数入口。counter 增量记为 SSD hit，成功 hit 减该增量记为 DRAM hit。 | counter 必须单调，增量不能超过该节点成功 hit。 | 这是窗口 counter 推断。本次 CPU 和 GPU 有效记录的采样时刻与窗口边界最大偏差分别为 13.047 ms 和 1.912 ms。 | + +Fluxon / Mooncake coordinator 只在每个节点的 `memory + ssd + unknown == hit`、观测完整且 `unknown == 0` 时发布完整来源结果。如果任一节点的计数不自洽，该节点的全部 hit 都归入 unknown，不根据剩余量猜测来源。miss 不分配 DRAM / SSD source，单独进入命中率分母。CPU 的 192 条记录和 GPU 的 72 条记录均为完整来源、unknown 为 0。 + +分层逻辑读带宽是整个统计窗口的平均值：`tier logical read GiB/s = tier hit 数 × 原始 payload bytes ÷ 窗口秒数 ÷ 2^30`。CPU 和 GPU 的统计窗口均为 30 秒。本次实验使用固定 payload，因此可以从来源操作数精确换算逻辑读取量。柱高不包含 miss；红点使用 `hit ÷ (hit + miss)`，两者共同描述本次 pressure workload。 + +bootstrap 在统计窗口前完成，窗口内的 workload 只包含 get。因此，已纳入记录的 DRAM 和 SSD 逻辑写带宽均为 `0.000 GiB/s`。 + +SSD 回填写入 host memory 的流量不重复计入 DRAM 逻辑写带宽。物理 DRAM 流量、SSD 读放大和回填写放大需要通过硬件计数器或设备 trace 观测。 + +#### CUDA 双槽 pipeline 与阶段指标口径 + +`cuda` 模式为每个 benchmark 线程固定复用 2 组 pinned host buffer、stream、event 和 GPU buffer，不增加新的 YAML 参数。线程先把 pageable source 复制到空闲 pinned slot，再调用 `cudaMemcpyAsync`。提交后不等待 event，下一次 KV get 可以和 H2D 重叠；两个 slot 都在途时，才等待最早提交的 event。 + +统计窗口结束时会 drain 所有在途 slot，但完成时间超出窗口的操作仍会被过滤。Mooncake 的 benchmark get 路径直接调用 `get_buffer_blocking()`，不再额外调用一次 `get_size()`。 + +两边只在 host pointer 的提取方式上不同，之后都进入同一个 H2D 流水线。`cuda_host_stage_us` 记录 pageable source 到复用 pinned slot 的 copy，`cuda_submit_us` 记录 `cudaMemcpyAsync` 调用，`cuda_h2d_event_us` 记录覆盖 H2D 的 CUDA event 区间。 + +`cuda_pipeline_residence_us` 从开始占用 slot 到前台确认 event 完成，包含 host staging、H2D 和完成轮询。它与下一次 KV get 存在重叠，不能和其他阶段机械相加。benchmark 会把 `MemHolder` 或 `BufferHandle` 连同相关 view 保留到 event 完成，使两边释放 native reference 的时点与成功计数边界一致。 + +### 8.3 实验设置 + +主图 case 在同一台测试机上串行执行，硬件为 Intel Xeon Platinum 8468、Samsung NVMe（XFS）和 NVIDIA H100 80 GB，GPU 链路为 PCIe 5.0 x16。bootstrap 后只读，workload 使用固定 value size 和 Zipf 分布。 + +CPU 和 GPU 主图统一使用 SSD-pressure 数据集：数据集约 2.5 GiB，DRAM 贡献为 2 GiB；SSD-on 再贡献 16 GiB SSD，SSD-off 不贡献 SSD。数据集、请求分布、worker 布局和统计窗口保持一致。由于数据集大于 DRAM，miss 和成功 hit 集合的变化就是实验要观察的结果；图同时报告命中率与 hit payload 带宽，不把开关差解释成纯介质开销。 + +| 组别 | DRAM 存储贡献 / 原始数据集 | payload / keyspace | worker 布局 | 总时长 / warmup / 统计窗口 | +| --- | --- | --- | --- | --- | +| CPU Fluxon / Mooncake pressure | 2 GiB / 约 2.5 GiB | `1/4/8/16 MiB - 128 B` / `2560/640/320/160`；原生引用 / `bytes` | c2=`1×2`、c4=`1×4`、c8=`2×4`、c16=`4×4` | 35 / 5 / 30 s | +| CPU Fluxon SSD backend 消融 | 2 GiB / 约 2.5 GiB；两侧再配置 16 GiB SSD | 与 CPU pressure 相同；`MemHolder` / `bytes` | 与 CPU Fluxon / Mooncake pressure 相同 | 35 / 5 / 30 s | +| GPU pressure 开关矩阵 | 2 GiB / 约 2.5 GiB | `4/8/16 MiB - 128 B` / `640/320/160`；`cuda` | c2=`1×2`、c4=`1×4`、c8=`2×4`、c16=`4×4` | 35 / 5 / 30 s | + +主图使用的 Fluxon wheel SHA-256 前缀为 `11118d3`，Mooncake 为 `0.3.11.post1`。同路径 SSD backend 消融使用 Fluxon wheel SHA-256 前缀 `4896d0e`，其中集成 Foyer 0.22.3 作为 test-only owner backend。后文的早期 GPU SSD 写盘证据使用 Fluxon wheel SHA-256 前缀 `54de857`，不与主图或 backend 消融记录合并。 + +表中的 DRAM 存储贡献是 KV 可分配容量，不等同于进程 RSS。主图对齐的是 2 GiB 逻辑 value 容量和 16 GiB 逻辑 SSD 容量；Fluxon 与 Mooncake 的 envelope、allocator、索引和元数据不同，物理 DRAM 占用并不相等。逻辑 payload 带宽也不包含后端编码和元数据。 + +同路径 SSD backend 消融中的 2 GiB 仍是 Fluxon owner 上层 DRAM。集成 Foyer 另外保留 1 B 内部 memory capacity，并关闭该层的 memory admission；这 1 B 不替代或扣减 Fluxon 的 2 GiB DRAM。 + +CPU 和 GPU 的 cN 都表示所有 benchmark 进程合计 N 个 worker 线程。Fluxon / Mooncake 的 c8 和 c16 分别使用 2 个和 4 个 benchmark 进程，因此结果同时包含多进程绕开 Python GIL 后带来的并行度；GPU 每个线程固定复用两个 pinned host、stream、event 和 device buffer slot。 + +CPU 和 GPU 主图中的 Mooncake 都覆盖两种拓扑： + +- **`DEDICATED_OWNER`**：一个常驻 owner 贡献 2 GiB DRAM 和 16 GiB SSD，所有 benchmark 进程都是 zero-contribution。 +- **`PER_BENCHMARK_PROCESS`**：不启动独立 owner，每个 benchmark 进程同时作为请求方和存储 endpoint。2 GiB DRAM 和 16 GiB SSD 按进程数等分：c2/c4 每个实例为 2 GiB + 16 GiB，c8 为 1 GiB + 8 GiB，c16 为 512 MiB + 4 GiB。 + +在 SSD-on 侧，每个 storage instance 使用独立 SSD 目录，但都位于同一块 Samsung NVMe 上。两种拓扑从 c2 到 c16 都保持 2 GiB DRAM + 16 GiB SSD 的集群逻辑总容量；SSD-off 侧保持同样的 2 GiB DRAM，去掉 16 GiB SSD 贡献。 + +Fluxon / Mooncake 每个拓扑的两侧使用相同的 DRAM 贡献、pressure 数据集、requester buffer、worker 布局和 transport。排除用于隔离 case 的 cluster ID、目录和端口后，Fluxon 只增删 owner `owner_kv_ssd`；Mooncake 只增删 `--enable_offload=true` 和 16 GiB SSD 贡献。 + +Mooncake 的 CPU requester buffer 在 SSD 关/开两侧均为 256 MiB，GPU 每个 benchmark 进程两侧均为 2 GiB。两侧也都使用 `tcp_pool_nodelay_backport`。 + +#### bootstrap、transport 与容量排除项 + +每个 case 都按单 PUT 串行 bootstrap。1、4、8、16 MiB payload 在每次 PUT 后分别等待 6、25、50、100 ms，把 payload 提交速率上限控制在约 160 MiB/s。该速率不足以在 Mooncake 的一个 10 秒 offload heartbeat 周期内填满 2 GiB DRAM 层，从而避免 bootstrap 成功与否取决于首次 heartbeat 和 DRAM 写满的先后顺序。800 MiB/s 的启动探针只用于发现这个竞态，没有进入最终结果。 + +主图构建和早期 GPU SSD 写盘证据所用构建都对 Mooncake 设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`，并在 client connect 与 server accept 两端补 `TCP_NODELAY`。结果中的 transport profile 记为 `tcp_pool_nodelay_backport`。 + +在只开启 connection pool、未补 `TCP_NODELAY` 的预跑中，`batch_get_into_offload_object_internal` 均值从 `2.22 ms` 增至 `7.37 ms`，P95 出现约 `46 ms` 台阶。最终结果不混入默认短连接或只开 connection pool 的运行，也不把这个 backport 口径写成原版 `0.3.11.post1` 的默认表现。 + +Mooncake SSD offload 的单 slice 上限是 `0xFFFFF0` bytes。CPU 和 GPU 的 16 MiB 组都将 benchmark payload 减去 128 bytes，加入 91-byte FlatDict envelope 后仍低于该上限。 + +Fluxon / Mooncake 主图和 Fluxon SSD backend 消融均关闭 capacity guard，Mooncake 的 eviction high watermark 设为 0.80。SSD 容量不是这些实验的限制因素。主图只纳入状态为 `SUCCESS`、0 error、来源计数完整且 unknown 为 0 的运行；pressure 中的 miss 在 SSD 关/开两侧都保留并进入命中率。 + +### 8.4 CPU 性能对比：`MemHolder` / `BufferHandle` 与 `bytes` + +下图的两行分别是原生引用 / handle 和 `bytes`，四列分别是 1/4/8/16 MiB。每个 payload 分面覆盖 c2/c4/c8/c16；每个并发度固定按 Fluxon、Mooncake 独立 owner、Mooncake 每进程 × SSD 关/开排列六根柱。蓝、橙、青分别区分三种实现组，浅色表示 DRAM、深色表示 SSD，红点使用右轴表示总命中率。各 payload 分面使用独立左轴，范围写在右上角。 + +![CPU 性能对比：Fluxon 与 Mooncake 两种拓扑在不同并发度下的 SSD 开关效果、数据来源分布及命中率](../../pics/kv_ssd_cpu_sweep.png) + +CPU 的 192 条记录覆盖 `2 输出模式 × 4 payload × 4 并发度 × 3 实现组 × 2 SSD 状态`。`fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_cpu_sweep.py` 与 GPU 脚本共同调用同目录下的 `kv_ssd_pressure_chart.py`。输入文件按命令行顺序覆盖同 case key；最终矩阵缺一条、存在 error、来源不完整或 unknown 非零都会直接失败。 + +| 实现组 | SSD-off 命中率 | SSD-on 命中率 | SSD-on 中 SSD 来源占 hit | +| --- | --- | --- | --- | +| Fluxon | 89.112%–94.390% | 100% | 17.550%–48.650% | +| Mooncake 独立 owner | 23.759%–66.879% | 100% | 62.065%–82.385% | +| Mooncake 每进程 | 22.488%–68.655% | 72.830%–100% | 59.341%–83.011% | + +结果支持三个直接判断： + +- **命中率和 hit payload 带宽必须一起看**：SSD-off 会产生快速 miss；Mooncake 每进程的部分 SSD-on case 也有 miss。柱高不包含这些 miss，因此单看柱高会遗漏成功集合的变化。 +- **开关差不是单项 SSD 开销**：Fluxon 原生引用的 SSD-on hit 带宽相对 SSD-off 变化为 -88.8%–-21.1%，`bytes` 为 -45.8%–+2.2%；Mooncake 两种拓扑同时包含正负变化。差值混合了命中率、DRAM/SSD 来源、回填、输出物化和调度，不能拆成一个内部阶段。 +- **输出模式需要分开比较**：原生引用 / handle 在 host reference ready 时计数，`bytes` 还包含完整 payload 物化。两行都按每次 hit 的完整逻辑 payload 计带宽；handle 行并没有读取或复制同等字节，因此其数值可以超过物理 DRAM 带宽，不能跨行用柱高计算复制效率或模式倍率。 + +统计窗口内只有 get，因此已完成记录的 DRAM 和 SSD 逻辑写带宽均为 `0.000 GiB/s`。图中报告逻辑 payload 流量，不代表物理 DRAM 或 NVMe 带宽。 + +#### 8.4.1 同一 Fluxon 路径下的 SSD backend 消融 + +[Foyer](https://github.com/foyer-rs/foyer/tree/v0.22.3) 是一个面向高性能和高并发场景的 Rust 内存/磁盘混合缓存库。它的 `HybridCache` 在一套接口下组合内存 cache 与磁盘 storage，并提供块式磁盘引擎、可替换的缓存算法和 IO engine。 + +对 Fluxon 而言，Foyer 是一个现成的通用实现参照。这组消融为一个直接的工程取舍提供依据：Fluxon 原生 SSD backend 在目标 workload 和既有数据路径中的功能与性能表现，是否足以支持继续自研和维护。若通用库已经满足需求，直接复用 Foyer 可以减少存储引擎本身的开发与维护。 + +要比较 SSD 存储层，不能直接把 Fluxon 的整套 cache 替换掉。Fluxon 的内存路径同时包含 master `route`、owner DRAM 副本、allocation 与 holder 生命周期、跨 owner 传输，以及 SSD 回填后的 promotion。如果连上层内存 cache 也替换为 Foyer，公共语义和数据路径会一起改变，结果无法归因到 SSD backend。因此，实验保留 Fluxon 的内存层及其上层路径，将替换边界统一收敛到 owner 内部的 `KvSsdStorage`：一侧使用 Fluxon 原生的 SSD 持久化与读取实现，另一侧接入 Foyer `HybridCache` 的 storage tier，并关闭 Foyer 自身的 memory admission。 + +Foyer 分支通过 test-only 配置启用： + +```yaml +test_spec_config: + kv_ssd_storage_backend: foyer +``` + +默认值仍为 `native`。Foyer 分支当前只接受一个 SSD root，并在配置阶段拒绝 `kv_ssd_uring_mode: iovec`。两组都调用相同的公共 KV API，经过相同的 FlatDict、route、transport、PyO3/GIL 边界和 benchmark 进程布局；写入 SSD backend 的都是 Fluxon target 中包含 FlatDict envelope 的完整编码 value。owner 上层 DRAM 都是 2 GiB，SSD 都是 16 GiB，且两种 SSD backend 都使用 direct I/O。图中的逻辑带宽仍只按原始 payload bytes 计数。 + +Foyer 内部另设 1 B memory capacity、1 个 memory shard，并用 filter 关闭 memory admission。这里关闭的是 SSD backend 内部可能形成的额外内存层，不影响 Fluxon owner 的 2 GiB DRAM。专项测试在 persist 和重复 get 后观测到 Foyer memory usage 为 0，所有读取来源均为 Foyer storage tier。 + +| 对比项 | Fluxon 原生 SSD | Fluxon Foyer SSD | +| --- | --- | --- | +| SSD 组织 | Fluxon ring、定长 shard 和 `io_uring`。 | Foyer 0.22.3 `HybridCache`，`WriteOnInsertion`、`PsyncIoEngine` 和 64 MiB block。 | +| 读取落点 | 满足对齐条件时直接读入 Fluxon target；远端路径可让 SSD chunk read 与 push 重叠。 | Foyer 先返回包含完整 value 的 `Vec`，再复制到 Fluxon target；远端路径取得完整 value 后才逐 chunk 交给 push。 | +| persist 并发 | 使用原生 writer queue 和 device worker。 | forced storage admission；`enqueue → wait → may_contains` 串行化，避免 Foyer byte queue 满时静默丢弃写入。 | +| 驱逐收敛 | ring 覆盖后批量通知 master 删除对应 SSD route。 | Foyer 0.22.3 没有向当前集成暴露等价的 SSD eviction callback；陈旧 route 只能在读取失败时撤销。 | + +这个 workload 在 bootstrap 后只读，因此 Foyer persist 串行化不进入 30 秒统计窗口。约 2.5 GiB 数据集也小于 16 GiB SSD 容量；本次结果不评估 Foyer 容量驱逐后的 route 收敛。图中的浅色仍表示 Fluxon 上层 DRAM route，深色表示 Fluxon SSD route；它不把 Foyer 内部 memory cache 作为第三层来源。 + +图按 `MemHolder` / `bytes` 分成两行，每行覆盖 1 / 4 / 8 / 16 MiB；每个并发度放置原生与 Foyer 两根柱。蓝/紫区分 SSD 实现，浅/深区分 DRAM 与 SSD 来源，红点表示命中率。 + +![Fluxon 同路径 SSD backend 消融：原生实现与 Foyer](../../pics/kv_ssd_foyer_ablation.png) + +PNG 由 `fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_foyer_ablation.py` 读取原始汇总生成。脚本要求 64 条记录的 case key 完整，校验容量、并发布局、统计窗口、来源计数和配对控制项后才绘图。这里的 64 是测量记录数，即 32 个原生/Foyer 配对，不表示 64 种部署。 + +所有配对的命中率都是 100%，因此这组差距没有 miss 数量优势可以解释。P95 的同向变化表明差距同时出现在延迟侧。`bytes` 会在两侧共同追加完整 payload 物化，其中位比值高于 `MemHolder` 与公共物化成本稀释 backend 相对差异相容，但单次矩阵不足以分解这部分因果贡献。 + +当前两条读取路径的明确差异是：Foyer adapter 先取得完整 `Vec` 再复制到 Fluxon target；原生实现可以直接读入 target，远端时还保留 chunk read 与 push 重叠。同时还存在 `PsyncIoEngine` / `io_uring`、完整 value 交付和调度方式等差异。这些差异共同属于当前 adapter 对比，本次数据无法把柱高差逐项分配给某个内部阶段。 + +配对 case 对相同的 thread id 和 operation index 使用同一套确定性 Zipf key 选择。统计边界仍是固定 30 秒窗口，不是固定操作数；较慢的一侧会停在同一确定性序列的更短前缀，实测 DRAM/SSD 来源比例也可能随之变化。因此，这组数据回答指定 pressure workload 下的整体 backend 表现，不把柱高比值解释成单次纯 SSD get 的加速比。 + +### 8.5 GPU 性能对比：多并发 SSD 开关与 Mooncake 存储拓扑 + +上一节的 CPU 测试以 KV 数据到达 CPU 内存为计数终点，分别观察原生引用 / handle 和完整 `bytes` 物化。AI 推理或训练实际消费 KV 数据时，数据通常还需要继续搬到 GPU。因此，本节把计数终点延伸到对应 CUDA event 完成，比较多并发下 SSD 开关和 Mooncake 两种存储拓扑对 GPU 可用数据吞吐的影响。该边界包含 H2D 传输，不包含 GPU kernel 执行或 D2H。 + +GPU 与 CPU 使用同一种画法，只少了原生引用 / `bytes` 两行。4、8、16 MiB 各占一个分面；每个分面都覆盖 c2/c4/c8/c16，每个并发度再按 Fluxon、Mooncake 独立 owner、Mooncake 每进程 × SSD 关/开排列六根柱。蓝、橙、青区分三种实现，浅色表示 DRAM、深色表示 SSD，红点显示命中率；六组 case 都从 native host buffer 进入同一条 `cuda` H2D 路径。 + +![GPU 性能对比：三种拓扑在不同并发度下的 SSD 开关效果、数据来源分布及命中率](../../pics/kv_ssd_gpu_sweep.png) + +GPU PNG 由 `fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_gpu_sweep.py` 读取一个或多个 pressure 汇总，并与 CPU 脚本共同调用同目录下的 `kv_ssd_pressure_chart.py`。脚本校验最终 72 条记录完整后直接使用 Pillow 绘制；各 payload 分面使用独立左轴，命中率共用 0%–100% 右轴。 + +| 拓扑 | SSD-off 命中率 | SSD-on 命中率 | SSD-on 中 SSD 来源占 hit | +| --- | --- | --- | --- | +| Fluxon | 89.326%–93.267% | 100% | 20.242%–46.508% | +| Mooncake 独立 owner | 34.557%–62.223% | 100% | 64.837%–77.135% | +| Mooncake 每进程 | 35.430%–64.782% | 59.855%–100% | 56.073%–84.198% | + +下表把图中的并发扩展收敛为 c16/c2 hit payload 带宽倍率，三个 payload 的顺序都是 4 / 8 / 16 MiB： + +| 拓扑 | SSD-off c16/c2 | SSD-on c16/c2 | +| --- | --- | --- | +| Fluxon | 5.83 / 5.83 / 4.68 倍 | 6.50 / 7.22 / 6.80 倍 | +| Mooncake 独立 owner | 1.33 / 1.50 / 1.10 倍 | 1.73 / 1.37 / 1.35 倍 | +| Mooncake 每进程 | 2.67 / 2.01 / 1.23 倍 | 3.02 / 2.07 / 1.77 倍 | + +Fluxon 的 SSD-on hit 带宽相对 SSD-off 变化为 -54.6%–-27.9%，Mooncake 独立 owner 为 -32.0%–+37.3%，Mooncake 每进程为 -62.6%–-17.5%。这些变化必须与上表的命中率一起解释。c16 下，每进程相对独立 owner 的 hit 带宽在 SSD-off 下分别高 125.7%、107.3% 和 97.9%，SSD-on 下高 82.4%、64.5% 和 57.9%；这些是包含 endpoint 本地性、进程数和 transport 调度的整体拓扑结果，不能归因到单个内部环节。 + +主矩阵体现了三个结果： + +- **Fluxon 能随并发扩展**：4、8、16 MiB 从 c2 到 c16 分别提高到 7.18、7.68 和 5.33 倍。 +- **16 MiB 在 c8 后开始收益递减**：c8 到 c16 只再提高 14.0%，backend get 均值从 `3.082 ms` 增至 `7.367 ms`，P95 从 `24.007 ms` 增至 `51.150 ms`。H2D 之前的路径已开始制约吞吐。 +- **Mooncake 拓扑会影响结果**：每进程贡献模式在 c16 下相对独立 owner 分别提高 91.3%、103.7% 和 66.8%。两种模式的逻辑容量相同，差异主要出现在 GPU copy 之前。来源诊断显示 8/16 MiB 的每进程贡献模式具有更高 SSD 命中率，因此更高吞吐无法由更高 DRAM 命中率解释；独立 owner 串行点、endpoint 本地性和 transport 调度成本仍未拆分。 + +c16 下，Fluxon 相对两种 Mooncake 拓扑中更快的一种，在 4、8、16 MiB 组中分别为 3.83、5.61 和 6.73 倍。这些比值只覆盖第 8.2–8.3 节定义的计数边界、单机 fast path 和固定容量。 + +## 9. 设计结论 + +Fluxon KV 将本地 SSD 定位为运行期回填层，不改变用户侧 `put`、`get` 和 `delete` 契约。master 维护 key-version 级 `route`，并统一调度用于最终 value 副本和临时传输的内存 allocation。owner 承载真实 bytes，管理 SSD 文件位置、IO 和 ring 生命周期。 + +`put` 操作在内存副本发布后返回，SSD persist 在后台完成。`get` 操作优先读取内存。没有可读内存副本时，本地 SSD 回填复用 requester target，远端回填则按 chunk 从 source staging push 到 requester target。两条路径最后都进入同一个 `GetDone` 和 holder 交付链。 + +同路径消融进一步表明，在本文固定的 pressure workload 下，当前 Foyer adapter 的 hit payload 带宽中位数为原生 SSD 的 37.2%，P95 为 2.30–8.46 倍，32 个配对中没有胜出。两侧均为 100% 命中，但 Foyer adapter 多了完整 `Vec` 到 Fluxon target 的复制，也没有原生路径的 direct-target 与 chunk read/push 重叠。这一结果限定于当前 adapter，不代表 Foyer 库的一般性能。 + +当前 SSD 层不提供冷启动恢复，也不向用户暴露独立 API。专项测试验证了强制回填和陈旧 route 兜底，并报告限定条件下的端到端逻辑 payload 表现。这些结果不代表裸 SSD 带宽，也不外推到跨机器、多 owner 或多盘场景。 diff --git "a/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 3 - \346\226\207\346\241\243\345\206\231\344\275\234\350\247\204\347\272\246.md" "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 3 - \346\226\207\346\241\243\345\206\231\344\275\234\350\247\204\347\272\246.md" index 75ed5b7..448abb0 100644 --- "a/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 3 - \346\226\207\346\241\243\345\206\231\344\275\234\350\247\204\347\272\246.md" +++ "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 3 - \346\226\207\346\241\243\345\206\231\344\275\234\350\247\204\347\272\246.md" @@ -5,9 +5,16 @@ - 让读者能在最短时间内抓住稳定结论。 - 让文档和代码、公开接口、实际行为保持一致。 +需要把这些规则落实为逐轮审校和 one-shot 提示词时，参见 [技术文档审校](./开发者%20-%205%20-%20技术文档审校.md)。 + ## 1. 通用规则 - 先写结论，再写展开。读者应在开头 30 秒内知道“这篇文档在回答什么问题”。 +- 开头保持在一个抽象层，只交代问题、核心要点和正文结构；局部实现、实验口径和分项证据留在对应章节。 +- 标题、导语和表头使用与内容匹配的标签。摘要写“核心要点”，能力盘点写“当前状态”，只有经过论证的判断才写“结论”。 +- 按主体归并职责和动作。介绍多个角色时，集中写完一个角色再写下一个角色，避免读者来回寻找同一主体的职责。 +- 重要的架构判断应在附近说明原因或机制，尤其是职责归属、生命周期和数据方向；不要只给结果，让理由散落到后文。 +- 主张强度必须与证据强度一致。已验证的行为和结果可以直接陈述；设计目标或预期收益使用“有望”“预计”等限定；证据不足的结果应删除或明确限定。 - 优先使用工程上自然的术语，少用自造概念。除非真的必要，不要写“根对象”“第一层分支”“authority object”这类词。 - 面向读者写，不要把作者的思维过程原样摊开。删除“这一节不讨论什么”“这个分支为什么属于上一层”这类模板套话。 - 如果一段文字可以改成表格、列表或图，就不要硬写成长段线性描述。 diff --git "a/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 5 - \346\212\200\346\234\257\346\226\207\346\241\243\345\256\241\346\240\241.md" "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 5 - \346\212\200\346\234\257\346\226\207\346\241\243\345\256\241\346\240\241.md" new file mode 100644 index 0000000..e0158a0 --- /dev/null +++ "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 5 - \346\212\200\346\234\257\346\226\207\346\241\243\345\256\241\346\240\241.md" @@ -0,0 +1,218 @@ +# 开发者 - 5 - 技术文档审校 + +本文把上游 `copy-editing` skill 中适合技术写作的检查方法整理为 Fluxon 文档审校流程。它用于技术事实和结构已经确认后的文字审校，不替代仓库文档规约，也不能改变公共接口、运行行为或性能结论。 + +## 1. 来源与适用边界 + +| 项目 | 内容 | +| --- | --- | +| 上游项目 | [`coreyhaines31/marketingskills`](https://github.com/coreyhaines31/marketingskills) | +| 上游 skill | [`copy-editing`](https://github.com/coreyhaines31/marketingskills/tree/0ba2a7fafc7b0827a261bd518e87cbda18e6675f/skills/copy-editing) | +| skill 版本 | `2.0.0` | +| 固定提交 | `0ba2a7fafc7b0827a261bd518e87cbda18e6675f` | +| 上游定位 | 审校已有的营销与转化文案，保留核心信息和作者语气。 | +| 本文用途 | 提取清晰度、语气、理由、证据和具体性检查，适配 Fluxon 技术文档。 | +| 许可 | MIT，完整声明见文末。 | + +使用时按以下优先级处理冲突： + +1. 当前任务要求、仓库 `AGENTS.md` 和[文档写作规约](./开发者%20-%203%20-%20文档写作规约.md)。 +2. 代码、公共接口、测试和已验证的运行行为。 +3. 本文的审校流程。 + +文字更顺不能成为修改技术含义的理由。遇到事实不确定、接口与文档不一致或证据不足时，先核对实现，再决定修改文档还是代码。 + +## 2. 上游七轮检查如何用于技术文档 + +上游 skill 把审校拆成七轮。Fluxon 只直接采用其中四轮，改写两轮，跳过一轮： + +| 上游检查 | Fluxon 用法 | 处理方式 | +| --- | --- | --- | +| Clarity | 检查长句、指代、未定义术语、上下文缺口和被限定条件埋住的主结论。 | 直接采用。 | +| Voice and Tone | 统一正式程度、角色名、组件名和中英文术语，保持自然的工程表达。 | 直接采用。 | +| So What | 对架构选择补充原因、影响和代价，回答“为什么这样设计”。 | 改写为工程动机检查，不写营销收益。 | +| Prove It | 用代码路径、类型签名、测试、指标或受限实验支撑行为与性能判断。 | 直接采用。 | +| Specificity | 写清范围、抽象层、前提、失败条件和不包含的路径。 | 直接采用。 | +| Heightened Emotion | 渲染痛点、焦虑、愿景或情绪。 | 跳过。技术文档优先准确和可验证。 | +| Zero Risk | 营销 CTA、保证和风险逆转。 | 只保留用户操作中的前置条件、失败结果和下一步，不引入营销 CTA。 | + +上游的英文词数、主动语态和短段落建议都是启发式规则。中文句长、代码标识和复杂所有权关系应按可读性判断，不机械套用固定阈值。 + +## 3. 审校流程 + +### 3.1 先固定技术事实 + +审校前先确认： + +- 文档类型和目标读者。 +- 公共契约、当前实现和专用 fast path 的边界。 +- 关键类型、配置键、返回值和失败语义是否与代码一致。 +- 行为、所有权和性能判断覆盖的完整路径。 +- 是否存在需要同步修改的中英文页面。 + +如果这些问题还没有答案，先做技术 review。copy editing 不能替代实现核对。 + +### 3.2 第一轮：结构与清晰度 + +先检查读者能否在开头快速回答三个问题： + +1. 这篇文档解决什么问题。 +2. 最重要的稳定信息是什么。 +3. 正文将按什么顺序展开。 + +随后逐节检查： + +- 标题是否直接表达本节判断或任务。 +- 每节是否先给当前抽象层的稳定信息，再进入字段、分支和数字。 +- 开头列出的轴是否在正文中逐项回应。 +- 同一事实是否被角色表、链路表、图例和总结重复解释。 +- 局部实现细节是否过早出现在导读或架构概览中。 + +### 3.3 第二轮：语气与术语 + +- 一个概念只保留一个规范名称、拼写和大小写。 +- 角色职责用明确主语表达，例如“master 维护 `route`”“owner 管理本地 SSD”。 +- 删除“稳定结论”“当前结论”等与实际内容不匹配的标签，直接写“核心要点”“当前状态”或具体判断。 +- 避免模板腔、宣传语和没有对象的“提升”“增强”“优化”。 +- 保留必要的英文代码术语，不为了中文化而改写公共类型或字段名。 + +### 3.4 第三轮：理由、证据与边界 + +每个重要判断都检查四件事： + +| 检查项 | 要回答的问题 | +| --- | --- | +| 理由 | 为什么选择这条职责边界、数据方向或生命周期。 | +| 机制 | 哪个对象、字段或调用链实现了该行为。 | +| 范围 | 判断覆盖到哪一步、哪个抽象层和哪些分支。 | +| 排除项 | 哪些成本、路径、恢复能力或部署条件没有包含在内。 | + +架构判断应把原因放在判断附近。例如，内存 allocation 归 master 调度时，应立即说明它同时覆盖最终 value 副本和跨 owner 传输临时内存；不能只列职责归属。 + +性能判断还必须绑定硬件、数据集、并发、输出边界、统计窗口和运行次数。端到端逻辑 payload 带宽不能写成裸 SSD 带宽。 + +### 3.5 第四轮：收缩与去重 + +按以下顺序处理重复内容： + +1. 保留首次出现的稳定判断。 +2. 保留能显著降低理解成本的一张表或一张图。 +3. 后续章节只展开新增机制、条件或失败路径。 +4. 总结回收公共契约、核心数据流和当前边界，不复述实验表格。 + +不要为了让每节“看起来完整”而重复角色定义。已经在总览中定义的角色，后续图前只补充该图中特有的含义。 + +### 3.6 第五轮：发布前回查 + +完成一轮修改后，从后向前回查，确认后续编辑没有破坏前面的判断： + +- 术语、章节号和交叉引用仍然一致。 +- 表格列数、代码围栏、折叠块和 Mermaid 图可以解析。 +- 中英文页面表达相同的契约和边界。 +- 所有精确标识仍使用正确名称。 +- 修改没有引入新的兼容路径、配置入口或未验证事实。 +- 文档站构建通过。 + +## 4. 协作审校与 one-shot 示例 + +### 4.1 从表面问题回溯根因 + +审校意见应同时给出位置、问题、原因和建议改法。只说“读起来不顺”不足以支持修改。 + +| 项目 | 示例 | +| --- | --- | +| 位置 | 开头的实验说明。 | +| 问题 | 在读者理解文章结构前提前展开统计口径。 | +| 原因 | 局部证据抢占导读层级。 | +| 改法 | 开头只保留全文结构，把统计条件留在实验节。 | + +对于明确、局部且不会改变技术语义的问题，可以直接修改并回查。会改变契约、结论范围或章节主线的修改，应先说明判断依据和影响。 + +一轮完整审校还应归纳多个表面问题背后的共同原因。这样才能修正全文，并把经验复用于下一篇文档。下面以 KV SSD 文章的一组实际修改为例： + +| 表面问题 | 根因 | 审校决策 | +| --- | --- | --- | +| 用“先记住四个稳定结论”引出内容摘要。 | 标签强度超过内容；摘要尚未构成经过论证的结论。 | 改为“四个核心要点”。 | +| 先写 owner 的物理资源，再写 master，最后又回到 owner 的 SSD 调度。 | 同一主体的职责被拆散，读者需要来回重组角色边界。 | 按主体归并，先集中说明 master，再集中说明 owner。 | +| 只写“内存 allocation 调度归 master”。 | 职责判断缺少就地理由，读者无法判断该边界是否覆盖完整。 | 紧跟说明内存 allocation 同时包含最终 value 副本和临时传输分配，两者都与 `route` 和 in-flight 状态联动。 | +| 开头提前解释第 9 节的实验分组、命中条件和统计口径。 | 局部证据与导读不在同一抽象层，细节遮住了全文结构。 | 开头只保留问题、核心要点和阅读路径，把实验条件留在第 9 节。 | +| 把“覆盖更大的运行期工作集”直接写成既成结果。 | 设计预期被写成已验证结果，主张强度超过现有证据，也没有直接落到读者关心的运行收益。 | 保留公共 API 不变这一范围，改为“有望提高缓存命中率，并更充分地利用存储带宽”。 | + +这些修改共同服务于两个目标：让读者按正确层级建立系统模型，并确保每个判断都不越过其证据边界。 + +### 4.2 可复用的 one-shot 示例 + +下面的提示词适用于技术事实已经核对、需要完成一轮结构与语言审校的文档： + +```text +目标文档：<文档路径> + +请对目标技术文档完成一轮审校并直接修改。保持公共 API、实现事实、实验数字和已有结论边界不变。 + +重点检查： +1. 开头只说明问题、核心要点和全文结构，不提前展开局部实现或实验方法。 +2. 标题、导语和表头的标签必须匹配内容，不把摘要或能力状态统称为“结论”。 +3. 按主体归并职责；介绍多个角色时，集中写完一个角色再写下一个角色。 +4. 每个重要架构判断附近都要说明原因或机制，尤其是职责归属、生命周期和数据方向。 +5. 主张强度必须匹配证据：已验证结果可以直接陈述；设计目标或预期收益使用“有望”“预计”等限定；证据不足的结果不写成事实。 +6. 删除重复的角色说明、错层细节和不增加新信息的总结，同时保留必要的范围、前提和排除项。 +7. 保持术语、类型名、接口名、章节号、交叉引用、表格和代码围栏一致。 + +完成后： +- 用“表面问题 → 根因 → 修改决策”概括主要改动，不逐句复述。 +- 如存在对应的中英文页面，同步更新。 +- 回查公共契约、技术事实和性能口径，并运行仓库规定的文档站构建。 +``` + +## 5. 不采用的上游做法 + +- 不用情绪渲染、FOMO、痛点放大或销售话术增强技术文章。 +- 不为每项实现强行补“用户收益”。只在理由影响理解或取舍时说明结果。 +- 不添加没有来源的数字、时间承诺、案例或比较级。 +- 不把英文“每句不超过 25 个词”等经验机械转换为中文字符限制。 +- 不要求所有句子改成主动语态；所有权和数据流清晰度优先。 +- 不使用虚构专家 persona 打分代替代码、测试和同层级链路核对。 + +## 6. 更新上游快照 + +本文固定到上游提交 `0ba2a7f`。更新时： + +1. 对比上游 `skills/copy-editing/SKILL.md` 和 `references/` 的变化。 +2. 只吸收适合技术文档、且不与 Fluxon 规约冲突的部分。 +3. 同步更新中英文文档、skill 版本和固定提交。 +4. 重新执行文档站构建。 + +不要自动用上游内容覆盖本文；上游目标是营销文案，本文目标是准确、可验证的技术文档。 + +## 7. 第三方许可 + +本文基于 Corey Haines 的 `copy-editing` skill 整理，并针对 Fluxon 技术文档进行了修改。上游采用 MIT License。 + +
+MIT License + +```text +MIT License + +Copyright (c) 2025 Corey Haines + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` + +
diff --git "a/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 6 - Code Review \350\247\204\347\272\246.md" "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 6 - Code Review \350\247\204\347\272\246.md" new file mode 100644 index 0000000..437593b --- /dev/null +++ "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 6 - Code Review \350\247\204\347\272\246.md" @@ -0,0 +1,114 @@ +# 开发者 - 6 - Code Review 规约 + +Code Review 规约必须把架构判断转换成可执行约束。单独写“职责要清晰”或“生命周期顺序要正确”只是结论；完整规约还要说明何时适用、必须做什么、如何验收、产生什么收益，以及如何限制修复本身的副作用。 + +## 1. 规约条目的固定结构 + +本文每条规约都包含以下字段： + +| 字段 | 要回答的问题 | +| --- | --- | +| 触发场景 | 出现什么代码形态、依赖或行为变化时应用本规约？ | +| 必须动作 | Author 和 Reviewer 必须建立或修改什么约束？ | +| 验收标准 | 哪些可观察后置条件成立才算完成？ | +| 预期收益 | 规约降低什么风险，或获得什么工程能力？ | +| 副作用边界 | 为实现规约不得额外引入什么复杂度、耦合或行为变化？ | +| 验证证据 | 用什么测试、类型约束或运行结果证明验收标准？ | + +文中的“必须”对应合并前要求；“应”允许在 PR 中给出明确理由后偏离；“可以”表示不影响合并的实现选择。 + +## 2. 按改动场景选择规约 + +Reviewer 先按改动信号匹配规约，不要求所有 PR 机械套用全部条目。 + +| 改动信号 | 主要适用规约 | +| --- | --- | +| 公共层开始访问内部字段，或资源在一个模块创建、另一个模块清理 | R1 职责与所有权、R4 可组合契约、R5 最小副作用修复 | +| 新增 `close`、后台任务、线程、进程、runtime 或嵌套 handle | R2 依赖生命周期、R3 状态 authority | +| 多个参与者共享 `closed`、event、lease 或 cleanup 状态 | R1 职责与所有权、R3 状态 authority | +| 新增 adapter、替换 backend、父模块包装子模块 | R1 职责与所有权、R4 可组合契约、R5 最小副作用修复 | +| close 与在途操作并发，或外层锁跨越内部阻塞调用 | R2 依赖生命周期、R3 状态 authority | +| 纯局部算法改动，不改变边界、资源或生命周期 | 不强制建立完整生命周期模型；按正确性和局部测试 Review | + +## 3. 可执行规约 + +### R1：跨模块状态和资源必须有唯一 owner + +| 字段 | 规约 | +| --- | --- | +| 触发场景 | 改动跨越公共层、组合层和内部模块；新增共享资源；出现重复 close；或外层直接修改内部状态。 | +| 必须动作 | 列出受影响状态和资源的创建者、借用者、共享者与最终释放者。每个资源指定唯一 final-release authority。公共层只表达稳定契约；组合层只安排模块顺序；内部模块维护自己的状态、不变量和 cleanup。 | +| 验收标准 | 每个资源都能回答“谁最后释放”；借用方不会释放 owner 资源；外层不通过内部字段完成 cleanup；替换内部实现时，公共层和组合层的生命周期逻辑不变。 | +| 预期收益 | 职责高内聚，避免重复释放和无人释放，缩小修改传播范围，使模块可独立测试和替换。 | +| 副作用边界 | 不为形式上的分层拆出无状态薄模块；不扩大公共 API；不把原本局部的资源提升为全局共享；不顺带迁移无关职责。 | +| 验证证据 | ownership 表或类型关系、公共契约测试，以及证明 borrower close 不破坏 owner 资源、独占 owner close 完成释放、共享 authority 遵守约定释放条件的测试。 | + +### R2：有依赖关系的生命周期必须按偏序执行 + +| 字段 | 规约 | +| --- | --- | +| 触发场景 | 对象依赖其他 runtime、transport、存储、worker、线程、进程或 handle 才能工作；新增启动、关闭、重启或热替换流程。 | +| 必须动作 | 写出依赖 DAG。若 `A` 依赖 `B`，初始化满足 `B → A`，销毁满足 `A → B`。关闭依次建立 admission barrier、传播 wake/cancel、收敛 in-flight、join 后台任务、逆依赖释放、发布完成。锁层级和 callback 方向必须与依赖方向兼容。 | +| 验收标准 | 关闭开始后不再接收新操作；被依赖资源释放前，所有依赖方已经静默；close 返回后，该生命周期边界内不再有任务、callback 或 handle 使用其拥有的资源；重复 close 能到达同一终态。 | +| 预期收益 | 获得确定性启停，避免 use-after-close、关闭死锁、悬挂后台任务和依赖提前释放。 | +| 副作用边界 | 只排序真实依赖；互不依赖的分支可以并行关闭。不要为了统一顺序引入全局锁或全局 coordinator；wake/cancel 只作用于本次关闭拥有的范围。 | +| 验证证据 | 真实生命周期测试、并发 close 与在途操作测试、超时测试，以及 close 返回后无后台活动的完成屏障测试。 | + +### R3：停止命令、进行状态和完成证明必须由明确 authority 管理 + +| 字段 | 规约 | +| --- | --- | +| 触发场景 | 一个布尔值或 event 同时被多个对象读写；停止请求与资源释放异步；close 可重试；或不同层都能发布 closed。 | +| 必须动作 | 为每个状态标明作用域、唯一 writer 和语义。区分停止命令、`Closing` 过程与 `Closed` 完成证明。同一作用域优先使用单调状态机；不同作用域的信号和完成状态分别由各自 owner 管理。只发信号的接口命名为 `request_shutdown()` 或等价语义。 | +| 验收标准 | 其他参与者发出 stop 不会让本模块跳过 cleanup；只有 owner 能发布本模块完成；状态只向终态推进；失败或重试不会重新开放入口，也不会永久跳过未完成步骤。 | +| 预期收益 | 消除共享布尔值的语义歧义，使关闭可重试、可等待、可观测，并避免提前返回和资源泄漏。 | +| 副作用边界 | 不为同一个事实增加多个 source of truth；简单同步对象不引入多余状态；能用有限 enum 表达时不扩散布尔组合；不保留旧状态别名形成兼容分支。 | +| 验证证据 | 状态转换测试、多参与者先后发 stop 的测试、部分构造和部分 cleanup 失败后的重试测试。 | + +### R4：子模块生命周期契约必须能够被父模块直接组合 + +| 字段 | 规约 | +| --- | --- | +| 触发场景 | 父模块包装子模块、新增 adapter 或 backend、一个 close 需要调用多个子模块，或错误需要跨层传播。 | +| 必须动作 | 为子模块定义生命周期前置条件、成功后置条件和失败后置条件。create / release、register / unregister、spawn / join、subscribe / cancel 必须成对。父模块只调用子模块契约并安排顺序，不复制其 cleanup 实现。独立 cleanup 即使部分失败也继续执行，并保留主错误和 teardown error。 | +| 验收标准 | 父模块不检查子模块内部字段就能决定下一步；子模块 close 返回后满足父模块逆序释放所需条件；替换一个符合契约的实现不要求修改父模块；失败结果能说明哪些后置条件尚未满足。 | +| 预期收益 | 局部正确性可以组合为系统正确性，backend 可替换，错误因果不丢失，模块可以独立演进。 | +| 副作用边界 | 不为假设中的未来实现提前设计通用框架；只抽象已经稳定的共同语义；特化 fast path 留在模块内部；不把所有错误压成无信息的统一返回值。 | +| 验证证据 | contract test、替代实现或 test double 的一致性测试、部分子模块关闭失败时其余 cleanup 仍执行的测试。 | + +### R5：修复必须落在不变量的 authority 处，并限制改动半径 + +| 字段 | 规约 | +| --- | --- | +| 触发场景 | 修复方案准备在外层增加判断、直接置空内部字段、增加全局 flag、兼容分支或新的配置入口，以绕过内部职责或生命周期问题。 | +| 必须动作 | 找到被破坏不变量的 owner，在 owner 内修复状态、cleanup 或契约；组合层只调整必要顺序。明确本次非目标，删除被新契约取代的 workaround，并只回归受影响边界。 | +| 验收标准 | 修复后跨边界知识减少或不增加；没有新增重复入口、并行配置通道或第二套生命周期路径；无关公共行为保持不变。 | +| 预期收益 | 修复根因而非症状，降低回归范围，避免兼容层和条件分支持续累积。 | +| 副作用边界 | 不借机重写无关模块；不以“统一架构”为由扩大迁移范围；若公共契约必须变化，单独说明迁移和影响，不能夹带在内部修复中。 | +| 验证证据 | 修复前反例、修复后不变量测试、受影响调用方回归，以及 diff 中没有新增旁路的检查。 | + +## 4. Review 输出与合并标准 + +Finding 应引用规约编号，并包含触发场景、违反的验收标准、影响和要求守住的副作用边界： + +```text +[R3][blocking] 共享 stop signal 被当成了当前模块 cleanup 的完成证明。 + +触发场景：另一个参与者先设置共享状态。 +违反标准：当前资源 owner 的 cleanup 被跳过，close 不再是 completion barrier。 +必须动作：分离共享停止命令与本模块完成状态，并由本模块发布完成。 +副作用边界：不要新增第二套公共 close API，也不要让外层直接清理内部资源。 +``` + +| 标记 | 合并标准 | +| --- | --- | +| `[blocking]` | 已有合理触发路径会破坏规约验收标准，修复后才能合并。 | +| `[major]` | 关键场景缺少约束或证据，补齐实现或验证后合并。 | +| `[minor]` | 不影响规约验收标准的局部维护性或诊断问题，可以明确接受后续处理。 | +| `[question]` | 触发场景、authority 或契约不清楚，回答后再定级。 | + +- **Request changes**：存在未解决的 `[blocking]` / `[major]`，或无法确定适用规约的 owner、依赖与验收标准。 +- **Approve**：所有匹配规约的必须动作已完成，验收标准有证据，预期收益成立且副作用保持在声明边界内。 +- **Comment**：仅剩澄清问题或不影响验收标准的建议。 + +文档类 PR 还应遵守[文档写作规约](./开发者%20-%203%20-%20文档写作规约.md)和[技术文档审校](./开发者%20-%205%20-%20技术文档审校.md)。 diff --git "a/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 7 - \344\272\213\344\273\266\350\256\242\351\230\205\344\270\216\345\205\250\351\207\217\345\277\253\347\205\247\350\247\204\347\272\246.md" "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 7 - \344\272\213\344\273\266\350\256\242\351\230\205\344\270\216\345\205\250\351\207\217\345\277\253\347\205\247\350\247\204\347\272\246.md" new file mode 100644 index 0000000..73d8531 --- /dev/null +++ "b/fluxon_doc_cn/dev_doc/\345\274\200\345\217\221\350\200\205 - 7 - \344\272\213\344\273\266\350\256\242\351\230\205\344\270\216\345\205\250\351\207\217\345\277\253\347\205\247\350\247\204\347\272\246.md" @@ -0,0 +1,84 @@ +# 开发者 - 7 - 事件订阅与全量快照规约 + +当组件通过“全量查询 + 增量事件”维护本地投影时，只允许使用下面两种顺序之一： + +1. 通用顺序：`subscribe/listen 就绪 → 拉取全量快照 → 安装快照 → 回放已缓冲的增量事件 → 发布 ready`。 +2. Revision 耦合顺序：数据源提供类似 etcd 的原子 revision 契约时，使用 `snapshot@R → watch@R+1`。 + +不满足 revision 耦合契约时，必须使用通用顺序。每次 `listen` 建立并确认就绪后，都必须进行一次全量同步；该要求同时覆盖首次启动和所有重新建立 `listen` 的恢复路径。 + +调用了 `listen()` 或启动了后台任务，不代表订阅已就绪。只有当数据源已确认订阅，或调用方已取得可用的 stream handle / cursor 时，才算 `listen` 建立完成。 + +## 1. 适用范围 + +本规约适用于任何由远端权威状态派生本地投影的场景，包括： + +- 集群成员与服务发现。 +- 路由表、配置镜像和资源索引。 +- controller / informer 类本地 cache。 +- 需要重建 watch 流的本地 cache。 + +本地投影的 owner 对快照安装、事件应用和 ready 状态负责。数据源仍是业务状态的唯一事实来源。 + +## 2. 通用顺序：listen 就绪后全量同步 + +```mermaid +sequenceDiagram + participant projection as local-projection-owner + participant source as authoritative-source + + projection->>source: subscribe / listen + source-->>projection: subscription ready (stream / cursor) + projection->>source: fetch full snapshot + source-->>projection: snapshot (optional revision R) + projection->>projection: install snapshot + loop buffered and live events + source-->>projection: incremental event + projection->>projection: apply by revision or idempotent version check + end + projection->>projection: publish ready +``` + +| 阶段 | 必须动作 | 验收条件 | +| --- | --- | --- | +| 订阅 | 等待数据源确认订阅已建立。fire-and-forget `spawn()` 不构成就绪证明。 | 自订阅就绪起的每次状态变化都能进入该订阅流。 | +| 快照 | 订阅就绪后拉取一次全量状态。 | 快照覆盖数据源定义的一致性时点上的完整对象集。 | +| 安装 | 由本地投影 owner 原子替换基线，或在唯一 writer 中完成替换。 | 增量处理不会与快照安装无序地并发写入同一投影。 | +| 回放 | 应用订阅就绪后缓冲的事件。有 revision 时只接受快照 revision 之后的事件；无 revision 时必须使用 generation / version 检查或幂等应用。 | 重复或陈旧事件不会回退快照中的新状态。 | +| Ready | 快照已安装，且订阅 backlog 已追至定义好的收敛点后，再发布 ready。 | ready 的读者不会观察到只有部分快照的中间态。 | + +## 3. 失败与重同步 + +- **订阅失败**：不开始全量拉取，不发布 ready。重试必须从建立新订阅开始。 +- **快照失败**：当次同步不成立。关闭或废弃当前订阅，再按完整顺序重试，避免无界缓冲。 +- **当前事件流失效**：lag、断线或 source epoch 变化都会使旧流失去完整性证明。撤销 ready 并废弃旧流；采用通用顺序时，新 `listen` 就绪后必须全量同步；采用 revision 耦合顺序时，重新执行 `snapshot@R → watch@R+1`。 +- **关闭**：先停止接收新事件，再 cancel / wake 并 join 事件任务，最后释放 stream 和投影依赖。close 成功返回后不得再有任务写入该投影。 + +固定 `sleep` 只能作为超时边界或调试手段，不能作为订阅就绪、快照安装或 backlog 收敛的证明。 + +## 4. 第二种顺序：snapshot@R → watch@R+1 + +如果数据源明确提供原子 revision 契约，可以使用以下顺序： + +`snapshot at revision R → watch from revision R + 1` + +该顺序必须同时满足： + +- 快照返回可用的 revision / cursor。 +- watch API 保证从指定 revision 后的第一个事件开始交付。 +- 数据源保证快照与事件日志使用同一个有序 revision 空间。 +- 测试覆盖 snapshot / watch 交界处的并发更新。 + +普通的 `fetch_all()` 加无起始 revision 的 `listen()` 不满足该顺序的前提。 + +## 5. Review 与验证清单 + +- [ ] 订阅就绪是可等待的 completion barrier，没有用 `spawn()` 成功替代订阅成功。 +- [ ] 通用顺序中的每次 `listen` 建立完成后都会拉取一次全量快照，包括首次启动和所有恢复路径。 +- [ ] 代码只实现两种合法顺序；不存在新 `listen` 直接沿用旧快照的第三条分支。 +- [ ] 使用 `snapshot@R → watch@R+1` 时，数据源和测试满足全部 revision 契约。 +- [ ] 快照安装和事件应用有唯一 writer，或有可证明的 revision 顺序。 +- [ ] 事件发生在订阅前、订阅后且快照前、快照后三个窗口时，最终投影都正确。 +- [ ] 重复事件和陈旧事件不会回退状态。 +- [ ] ready 只在快照与 backlog 都收敛后发布。 +- [ ] 测试用明确的状态条件等待，不用固定 `sleep` 代替同步。 diff --git a/fluxon_doc_en/blog/blog_1_Built for AI Data Movement - Fluxon Distributed Key-Value Cache, RPC, Message Queue, and File Object Cache Acceleration Layer.md b/fluxon_doc_en/blog/blog_1_Built for AI Data Movement - Fluxon Distributed Key-Value Cache, RPC, Message Queue, and File Object Cache Acceleration Layer.md new file mode 100644 index 0000000..c15f2c4 --- /dev/null +++ b/fluxon_doc_en/blog/blog_1_Built for AI Data Movement - Fluxon Distributed Key-Value Cache, RPC, Message Queue, and File Object Cache Acceleration Layer.md @@ -0,0 +1,196 @@ +# Fluxon: A Rust Data Plane for KV Cache, RPC, Message Queues, and S3-Compatible Object Caching + +![](../../pics/post_en.png) + +As GPU compute power continues to scale, the bottleneck in AI systems is expanding from individual operators into the data plane. Inference services need to reuse `KV Cache`, `latent cache`, and prefix cache across requests, processes, and nodes. Training pipelines need to pass intermediate state across heterogeneous resource pools. Model files, sample data, and `Checkpoint` data need to move reliably across remote access, local caching, and cross-cluster migration paths. + +Traditional approaches usually introduce separate caching systems, message queues, file systems, object storage gateways, and observability pipelines for these problems. Each system has its own connection, capacity, reclamation, routing, and observability logic. As model scale, request concurrency, training pipelines, and dataset size continue to grow, data-plane overhead keeps expanding and gradually consumes CPU, I/O, memory, and operational effort. + +Fluxon is a storage and transport integrated distributed system for the AI training and inference data plane. It brings distributed key-value caching, `RPC`, message queues, and an `S3`-compatible file object cache acceleration layer into one data plane acceleration foundation, so these high-frequency data objects can reuse one set of caching, transport, lease, capacity governance, and observability capabilities. + +## Origin: Pain from Elastic VAE Decoupling + +One early engineering motivation for Fluxon came from cross-resource-pool data handoff in `VAE` decoupled heterogeneous training. We wanted `Producer` and `Consumer` components to scale independently in different resource pools and hand off intermediate `Payload` objects asynchronously, instead of binding training components again through a fixed communication group. + +In this scenario, the boundary of `NCCL` is clear. It is a strong tool for synchronous collective communication inside a fixed member set, and it fits high-performance communication inside a training process group. However, when dynamic membership, asynchronous handoff, backpressure, message retention, and elastic scheduling across resource pools are required, the fixed-member communication model couples together the training components that we wanted to decouple. + +We also tried to use high-performance transport and caching systems such as `MooncakeStore`, which target `KV Cache`, for large `Payload` movement in training pipelines. They have clear value in specific cache scenarios. When directly generalized into a common AI training data plane, however, new engineering constraints appear. Under high load and long-running stability tests, if low-level memory lifetime and transport state are not handled completely, the business side may see process crashes or task interruptions. When the `RDMA` network jitters, training pipelines often struggle to recover if there is no automatic fallback to paths such as `TCP`. In addition, if low-level exceptions cannot be captured reliably and propagated upstream, `Producer` / `Consumer` components can enter inconsistent states, block for a long time, or require repeated manual diagnosis. + +These experiences made us realize that the AI data plane cannot only pursue peak performance on one segment of the transport path. It also needs disciplined error propagation, highly available network fallback, unified object lifecycle management, and a resource governance model that supports dynamic business membership. + +Because existing systems struggled to cover elastic decoupling, high-performance movement, and failure fallback at the same time, Fluxon chose to redesign this path from the data plane acceleration foundation. + +## Why a Unified AI Data Plane Is Needed + +Data objects in AI workloads are becoming larger, hotter, and more frequently moved across boundaries. + +On the inference side, high-frequency caches are no longer temporary objects inside a single process. In multi-replica inference services, world models, multimodal models, and long-context scenarios, these caches often need reuse across requests, processes, and nodes. If each instance independently manages cache residency, reclamation, and transport, GPU memory and host memory can be exhausted quickly by redundant data and duplicate residency. + +Training-side data flow is also becoming more complex. Cross-resource-pool `Producer` / `Consumer` handoff scenarios, represented by VAE decoupled heterogeneous training, can place `Producer` and `Consumer` components on different machines, different resource pools, or even different sub-clusters. If intermediate `Payload` objects can only move through traditional MQ, file landing, or object storage across multiple systems, the data path becomes longer and capacity governance becomes harder. + +The same applies on the storage side. High-resolution video, trajectory samples, model files, and `Checkpoint` data need to support remote access, local caching, `S3` forwarding, and cross-cluster migration at the same time. If file object caching and `KV` caching are split into two systems, large objects are repeatedly split, copied, landed, and re-indexed across systems. + +Fluxon focuses on the full AI data-plane lifecycle: how objects are allocated, where they are placed, how they are transferred across nodes, when they are evicted, how business processes connect to them, and how issues are located when something goes wrong. + +## Why Patchwork AI Data Planes Hit a Bottleneck + +A common traditional AI data-plane pattern is to use one KV or custom cache for high-frequency caching, send intermediate handoff through MQ or object storage, attach file data through another file system or object cache, and collect observability status from each component separately. + +This approach can assemble functionality quickly at an early stage, but several problems appear as the system grows. + +**First, local scenario experience is hard to transfer.** `MooncakeStore` provides a specialized design for `KV Cache` scenarios by binding cache semantics and `RDMA` transport to a specific path. In more general data-plane scenarios, low-level transport and capacity governance are hard to move over directly. + +**Second, unified resource control and scheduling are incomplete.** Framework-level caches such as `SGLang` `L2` keep in-framework indexes to provide lowest-latency access. External caches such as `MooncakeStore` as `L3` handle reuse across instances. These two layers often live in the same host `CPU` memory, while `L2` memory is hard to include in unified indexing, placement, and eviction governance. This also increases cache-crossing and object-handoff overhead. + +**Third, there is no shared-memory fast path for colocated processes.** When `MooncakeStore` is used as `L3`, for example, the data path is usually organized around `RDMA` / `TCP`, which fits cross-node pooling and remote prefetch. If colocated `Worker` or `Producer` / `Consumer` processes cannot directly enter the shared-memory data plane, object handoff still detours through the network transport stack. + +**Fourth, AI Infra lacks a dynamically elastic communication plane.** Collective communication libraries such as `NCCL` are strong tools for synchronous communication inside a fixed member set. In VAE decoupled heterogeneous training, cross-resource-pool `Producer` / `Consumer` handoff needs dynamic membership, asynchronous handoff, backpressure, message retention, and large `Payload` placement. A fixed-member communication model couples training components again and amplifies connection and recovery complexity. + +**Fifth, business processes are coupled with data-plane resource governance.** When business processes start, stop, scale, or fail dynamically, and also contribute capacity, manage object lifecycles, and maintain cross-node connections, data-plane capacity and connection topology fluctuate with business lifecycles. This affects cache reuse, failure recovery, and operational diagnosis. + +**Sixth, object lifecycle management is hard to converge.** When caches, messages, and file object access each maintain references, leases, eviction, routing, and reclamation state, it becomes difficult for a team to determine in one place whether an object is still being consumed, when it can be released, when it needs migration, or when indexes need rebuilding. The more components there are, the more state spreads across the business framework, cache layer, and transport layer. + +**Seventh, observability pipelines are fragmented.** If cache hits, `Owner` queues, transport paths, object materialization, and business-process latency are spread across separate systems, it is hard to locate the bottleneck quickly when a performance issue occurs. Teams often have to stitch clues together across multiple metric sets and logs, and diagnosis cost grows with the number of components. + +Fluxon is designed around these problems. It separates data-plane resources, object lifecycles, cross-node transport, and business integration into explicit abstractions, then governs them inside one data plane acceleration foundation. + +## Fluxon's Approach + +Fluxon is built on a storage and transport integrated data plane acceleration foundation and exposes three types of entrypoints upward: + +| Entrypoint | Target Scenario | Capability | +| --- | --- | --- | +| Distributed key-value cache / `RPC` | Inference cache, state sharing, service-to-service calls, tensor object reuse | Unified key-value read/write and inter-node `RPC`, serving high-frequency state caches and tensor object reuse | +| Message queue | Cross-resource-pool `Producer` / `Consumer` intermediate-state handoff, such as VAE decoupled heterogeneous training, and data-processing pipelines | Dynamically elastic `AI Infra` communication plane that reuses the key-value cache data plane for large `Payload` objects | +| File object cache acceleration layer | AI data, model files, `Checkpoint` data, remote object access, `S3` forwarding | `S3`-compatible file object cache acceleration layer that reuses the data plane acceleration foundation with `KV/MQ`, so tensors, message `Payload` objects, and file objects enter a unified caching and acceleration path | + +All three entrypoints reuse the same caching, transport, lease, capacity governance, object lifecycle, and observability capabilities. This means high-frequency data objects in these scenarios can all receive unified caching and acceleration inside the same data plane acceleration foundation. AI workloads do not need to build multiple separate data-plane paths. + +For colocated resource governance, Fluxon prefers to organize cache layers by real physical resources. Host `CPU` memory should first enter shared memory and unified object lifecycle governance, reducing the extra paths introduced by artificial `L2` / `L3` splits inside business frameworks. Cross-machine, cross-cluster, and remote object access should use distributed transport paths. + +This is the main difference between Fluxon and a single-point cache, single-point message queue, or single-point file interface. The core of Fluxon is the data plane acceleration foundation, and the APIs are the entrypoints that this foundation exposes for different scenarios. + +## Architecture Layers: Clear Role Boundaries + +Fluxon abstracts the control plane, data-plane resources, and business integration into three core roles. This layering decouples resource governance from business lifecycles and removes topology and connection bottlenecks for AI data-plane scale-out in large clusters. + +Master, the control plane, uniformly manages memory allocation, object placement, eviction, routing, and leases. It converges key decisions into the control plane, keeps object lifecycles consistent across multi-process and multi-node scenarios, and enables precise memory reclamation. + +Owner Client, the data-plane resource provider, is a resident node that contributes the shared memory pool and carries inter-Owner communication. Business processes connect to their local Owner, and the Owner performs cross-machine transport. This local proxy plus cross-machine backbone structure avoids connection storms caused by direct interconnection between business processes. Cross-machine connections are strictly converged between Owners, so network topology complexity, routing state, and data movement paths remain easier to control as the cluster grows. + +External Client, the business integration layer, carries dynamic access from inference services, MQ `Producer` / `Consumer`, FluxonFS, FluxonOps, and other components, and does not contribute cluster capacity. Because External Client does not provide physical capacity, frequent business-side start/stop, abnormal restart, or elastic scaling does not directly trigger data-plane capacity migration or `Rebalance` churn. Elasticity of compute services is isolated from stability of the data foundation, and the business layer can implement stateless scale-out more easily. + +Together, these three roles build the base stability and scalability of the AI data plane. Master, as the decision maker, converges global state and prevents object placement, routing, and reclamation logic from spreading into business processes. Owner Client, as the provider and carrier, anchors physical capacity and the cross-machine backbone, keeping cross-machine connections at the data-plane resource layer. External Client, as the accessor, carries elastic business workloads without disturbing the underlying topology. Once these boundaries are clear, Fluxon can let cache, message, and file object access truly reuse the same scalable data plane acceleration foundation without giving up the single-machine fast path. + +## Data Paths: Local Shared Memory, Cross-Node P2P, and Automatic Relay + +Fluxon covers local object handoff, cross-node object movement, and relay forwarding under complex network topologies. + +On the local path, business processes first connect to the Owner shared memory pool and reduce object-handoff cost through SHM / Busy Polling / Epoll UDS. High-frequency data objects can avoid unnecessary copying and reconstruction. + +On the cross-node path, Owners transfer data through P2P transport. The deployment can select `RDMA`, `TCP`, `QUIC`, and other paths, and Fluxon supports automatic relay forwarding across nodes and sub-clusters. When the source Owner and target Owner cannot communicate directly through the preferred link, the data plane can still complete transport through a relay path. Business processes do not need to understand complex network topology. They only connect to the local Owner, and cross-machine data movement converges to Owner-to-Owner paths. This layered architecture reduces cross-machine connection spread between business processes and further improves system scalability. + +This design allows Fluxon to serve multi-process reuse inside one machine, cross-node reuse inside a cluster, and cross-cluster data flow at the same time. The former focuses on shared memory and object handoff, while the latter focuses on connection convergence, routing adaptation, and cross-machine transport. + +## Core Engine: Rust for a Low-Overhead, Controllable Data Plane Acceleration Foundation + +Rising GPU compute and growing cluster scale make I/O and CPU explicit bottlenecks in AI systems. Connection handling, protocol encoding and decoding, high-concurrency transport, shared-memory management, and observability collection all sit on data-plane hot paths. If these hot paths are filled with interpreted execution, runtime scheduling, cross-language-boundary copies, and uncontrolled memory copying, the time saved on the GPU side can be consumed by data-plane overhead. + +Fluxon implements these key paths in Rust, with the goal of bringing concurrency safety, memory lifetimes, and system-call boundaries under stronger engineering constraints. + +- Concurrency: CPU-intensive paths are not constrained by the GIL, making it easier to use multi-core resources fully. +- Predictable latency: no GC pauses, reducing unpredictable jitter on hot paths. +- Long-running safety: ownership, lifetimes, and the type system constrain shared memory, connection state, and object references. +- Auditable code: strongly typed interfaces and explicit state machines make the low-level data plane easier to review by humans and tools. + +This system-level control of the low-level data plane is the foundation for efficient and safe movement of high-frequency data objects on the same data plane acceleration foundation. For the low-level data plane, this matters more than simply pursuing peak performance of one interface. + +## Unified Observability + +Fluxon's observability foundation uses GreptimeDB, which is also built in Rust. This matches Fluxon's low-level philosophy: GreptimeDB converges the three observability pillars, Metrics, Logs, and Traces, into one engine, just as Fluxon converges cache, message, and file object access into one data plane acceleration foundation. + +Fluxon collects metrics through the Prometheus protocol, combines them with tracing and structured logs, and provides a built-in GUI that clearly presents cluster topology, member status, key latencies, and queue depths. + +For a data-plane system, observability is part of governance, not a nice-to-have. Only when an issue can be located precisely to Owner, External Client, transport path, queue waiting, or object processing can the data-plane foundation be truly governable. Unified observability is a key part of closing the loop for Fluxon's data plane acceleration foundation. + +## FluxonKV: Distributed Key-Value Cache and RPC on a Shared Data Plane Acceleration Foundation + +Fluxon `KV/RPC` targets world model inference caches, state sharing, service-to-service calls, and tensor object reuse. In scenarios such as multi-view latent-space prediction, state extrapolation, and prefix-cache reuse, it covers a more general AI data plane than a single `KV Cache` scenario. + +KV and RPC share the same parameter organization, caching, and communication path. State storage, object reuse, and service-to-service calls do not need two separate paths. High-frequency calls and large-object reuse can be completed inside the same role model. + +![](../../pics/fluxon_kv.png) + +On the read path, Fluxon prioritizes local fast-path hits and advances metadata synchronization asynchronously in the background. The system reduces duplicate residency and memory waste across multiple cache tiers through hot-object reuse, and uses batched reclamation to converge fragmented control-plane interactions into batch operations. This reduces control-plane traffic and compute overhead, improving overall throughput and long-running behavior. + +For inference systems, this means high-frequency state caches, tensor objects, and service calls can enter the same data plane acceleration foundation without repeated conversion between a cache system and an RPC system. + +## FluxonMQ: Dynamically Elastic Communication Plane for AI Data Flow + +Fluxon `MQ` targets cross-resource-pool `Producer` / `Consumer` intermediate-state handoff and data-processing pipelines. VAE decoupled heterogeneous training is one representative scenario. When `Producer` and `Consumer` components are distributed across different machines, different resource pools, or even different sub-clusters, `MQ` converges message retention, capacity governance, and cross-cluster placement into a unified messaging layer. + +Traditional `MQ` systems are usually built on `TCP` and disk logs. Their original design is not oriented toward large `Tensor Payload` objects, so they lack a native high-speed `RDMA` data path. When intermediate state reaches dozens of `MB` or even `GB`, general-purpose `MQ` systems struggle to satisfy low-latency handoff, capacity governance, and high-bandwidth movement at the same time. + +Collective communication such as `NCCL` fits synchronous communication inside a fixed member set. Cross-resource-pool `Producer` / `Consumer` handoff needs dynamic membership, asynchronous handoff, backpressure, message retention, and large `Payload` placement. Specialized cache systems such as `MooncakeStore` perform strongly for `KV Cache` reuse, but their core cache-eviction semantics do not directly replace the `MQ` requirement that a message must remain before it is consumed. Their transport paths also usually depend on deployment-side choices between `RDMA` / `TCP`, making it difficult for them to take on automatic fallback, cross-cluster relay, and elastic handoff inside a data plane acceleration foundation. + +The core design philosophy here is: keep the control plane light, retaining only message shells, member topology, and Offset; keep the data plane heavy, moving large Payload objects directly through the KV data plane. This means Fluxon does not need to build a second large-object transport path for the message queue. + +`Producer` / `Consumer` components join dynamically as `External Client` instances without changing cluster capacity. They can scale with business load, while `Owner` remains stable as the resident data-plane resource provider. + +`Lease` binds message retention to the message channel, giving pre-consumption data retention an explicit time boundary. In cross-resource-pool and cross-sub-cluster scenarios, `Payload` placement can use the consumer-side location to shorten the prefetch path as much as possible. + +## FluxonFS: S3-Compatible File Object Cache Acceleration Layer + +Fluxon `FS` is positioned as an `S3`-compatible file object cache acceleration layer. It targets AI data, model files, `Checkpoint` data, high-resolution video, and trajectory samples, covering remote access, cache hits, `S3` forwarding, and cross-cluster migration. + +The key point of FS is reusing `KV/RPC` caching and communication capabilities. Files are split into `KeyValue` shards and enter Fluxon's caching, transport, and capacity-governance paths. In this way, files, objects, and KV caches converge from three fragmented systems into different entrypoints of the data plane acceleration foundation. + +For AI data platforms, this path can reduce the cost of switching file data across remote access, local caching, and cross-cluster migration. The upper layer still uses file object semantics, while the lower layer reuses unified data-plane capabilities. + +## Benchmark + +The current public benchmark charts cover three paths: RPC, KV, and FS, the file object cache acceleration layer. The following data is generated from specific test scenarios and focuses on architectural benefits and performance boundaries across different data paths. + +Note: Benchmark data is generated under specific topology and Payload sizes. Actual business gains depend on network environment, object size, and access pattern. + +### RPC Benchmark + +The RPC Benchmark compares ZeroRPC, Fluxon TCP Thread, and Fluxon RDMA under a 4 KB echo Payload, covering throughput and end-to-end latency. The charts show that Fluxon's RPC path significantly reduces latency and improves aggregate throughput in both P1 and P8 panels. + +![](../../pics/fluxon_rpc_bench.png) + +This result corresponds to the service-to-service call path. It shows that the parameter organization, routing, and communication foundation shared by RPC and KV can carry high-frequency small-Payload calls. End-to-end time for actual business handlers will still be affected by application logic. + +### KV Benchmark + +The KV Benchmark shows three scenarios: READ_AFFINITY, READ_ZIPF, and PUT_ONLY. Read-heavy scenarios are the current public result's strong area, especially for explaining the benefits from locality, hot-object reuse, and cross-node object location. + +![](../../pics/kv_benchmark_chart.png) + +In the pure-write PUT_ONLY scenario, the current performance constraint is mainly in the inflight metadata deduplication path rather than Payload transport itself. This is also one of the core directions for later optimization. + +### FS Benchmark: File Object Cache Acceleration + +`FS Benchmark` compares Fluxon `FS` with `Alluxio`, covering small-file read, large-file read, small-file write, and large-file write under cache warmup. The most prominent area in the chart is large-file write. Small-file read already shows an advantage, large-file read performance is roughly on par, and small-file write performance still has room for further optimization. + +![](../../pics/fs_benchmark_chart.png) + +This result corresponds to the file object cache acceleration layer. Fluxon FS benefits from reusing the KV/RPC data plane, while small-file writes are still affected by upper-layer open, write, close, and commit flows. + +## Why Open Source + +Fluxon open-sources a complete data plane acceleration foundation for the AI data plane: Rust core implementation, Python interfaces, distributed key-value cache, RPC, message queue, `S3`-compatible file object cache acceleration layer, deployment toolchain, test stack, and Benchmark are all organized in one project. + +We want developers to directly see how Fluxon organizes the control plane, data plane, business integration, observability, and tests, and to understand the data plane acceleration foundation behind these interfaces. AI infrastructure is moving from single-model, single-service, and single-cluster forms toward more complex data-flow patterns. Cache, message, and file object access all need to be considered again inside one data-plane path. + +Fluxon is open-sourced under Apache License 2.0. We hope to work with the community on AI inference caching, heterogeneous training, file object caching, cross-node transport, shared memory, Rust data planes, and observability systems. + +## Next Steps + +Fluxon is still evolving quickly. In the short term, we will continue to optimize `KV Cache` integration in `SGLang`. To meet `L2` requirements for extremely low-latency access, the interface layer will also gain new capabilities that make cooperation among in-framework indexing, local shared memory, and the data plane acceleration foundation more direct. + +In the longer term, Fluxon aims to become the data plane acceleration foundation inside AI systems: bringing AI workload data objects and cross-node transport into one governable and observable path. + +We expect algorithm engineers and model-serving developers to spend more energy on model innovation itself, instead of repeatedly paying for duplicated low-level data movement and reactive patches. Fluxon is designed to serve as this data plane acceleration foundation and support more complex, freer data movement in the AI era. + +Fluxon is developed by the AI Infra team at the Artificial Intelligence Research Institute of China Telecom (TeleAI), led by China Telecom Chief Scientist Professor Xuelong Li. Fluxon is open-sourced under Apache License 2.0. GitHub repository: https://github.com/Tele-AI/Fluxon. Developers interested in AI inference caching, heterogeneous training, Rust data planes, and distributed systems are welcome to participate. diff --git a/fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md b/fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md index 4577714..4616200 100644 --- a/fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md +++ b/fluxon_doc_en/dev_doc/Developer - 3 - Documentation Writing Rules.md @@ -5,9 +5,16 @@ This page defines how Fluxon user docs, developer docs, and design docs should b - Let readers reach the stable conclusion as quickly as possible. - Keep docs aligned with code, public contracts, and actual behavior. +For a sweep-by-sweep workflow and a reusable one-shot prompt, see [Technical Documentation Copy Editing](./Developer%20-%205%20-%20Technical%20Documentation%20Copy%20Editing.md). + ## 1. General Rules - Lead with the conclusion, then expand. A reader should know what the page answers within the first 30 seconds. +- Keep the opening at one abstraction level. State only the problem, core points, and document map; leave local implementation details, measurement methodology, and per-section evidence in their corresponding sections. +- Use labels that match the content in headings, introductions, and table headers. Call a summary “core points,” a capability inventory “current status,” and reserve “conclusion” for a claim supported by the preceding argument. +- Group responsibilities and actions by subject. When introducing multiple roles, finish one role before moving to the next so readers do not have to reconstruct a subject's responsibilities from scattered clauses. +- Place the rationale or mechanism close to an important architecture claim, especially for responsibility boundaries, lifetimes, and data direction. Do not state only the outcome and scatter its reason later in the document. +- Match claim strength to evidence strength. State verified behavior and results directly; qualify design goals or expected benefits with terms such as “expected,” “aims to,” or “may”; remove or explicitly bound outcomes that lack supporting evidence. - Prefer natural engineering terms. Avoid invented or template-heavy wording such as `root object`, `first-level branch`, or `authority object` unless the term is truly necessary. - Write for readers, not as a dump of the author's thinking process. Remove template filler such as "this section does not discuss" or "why this branch belongs to the previous layer." - If a paragraph can become a table, list, or diagram, do that instead of forcing a long linear explanation. diff --git a/fluxon_doc_en/dev_doc/Developer - 5 - Technical Documentation Copy Editing.md b/fluxon_doc_en/dev_doc/Developer - 5 - Technical Documentation Copy Editing.md new file mode 100644 index 0000000..74fd2b2 --- /dev/null +++ b/fluxon_doc_en/dev_doc/Developer - 5 - Technical Documentation Copy Editing.md @@ -0,0 +1,218 @@ +# Developer - 5 - Technical Documentation Copy Editing + +This document adapts the parts of the upstream `copy-editing` skill that are useful for technical writing into a Fluxon documentation review workflow. Use it after the technical facts and structure are established. It does not replace the repository documentation rules and cannot change public APIs, runtime behavior, or performance conclusions. + +## 1. Source and scope + +| Item | Details | +| --- | --- | +| Upstream project | [`coreyhaines31/marketingskills`](https://github.com/coreyhaines31/marketingskills) | +| Upstream skill | [`copy-editing`](https://github.com/coreyhaines31/marketingskills/tree/0ba2a7fafc7b0827a261bd518e87cbda18e6675f/skills/copy-editing) | +| Skill version | `2.0.0` | +| Pinned commit | `0ba2a7fafc7b0827a261bd518e87cbda18e6675f` | +| Upstream purpose | Edit existing marketing and conversion copy while preserving its core message and voice. | +| Purpose here | Adapt its clarity, voice, rationale, evidence, and specificity checks to Fluxon technical documentation. | +| License | MIT; the complete notice appears at the end of this document. | + +Resolve conflicts in this order: + +1. The current task, repository `AGENTS.md`, and the [documentation writing rules](./Developer%20-%203%20-%20Documentation%20Writing%20Rules.md). +2. Code, public contracts, tests, and verified runtime behavior. +3. This copy-editing workflow. + +Smoother prose is not a reason to change technical meaning. If a fact is uncertain, an interface disagrees with the documentation, or the available evidence is insufficient, verify the implementation before deciding whether to change the documentation or the code. + +## 2. Applying the seven upstream sweeps to technical documentation + +The upstream skill divides editing into seven sweeps. Fluxon directly adopts four, adapts two, and skips one: + +| Upstream sweep | Fluxon use | Decision | +| --- | --- | --- | +| Clarity | Check long sentences, references, undefined terms, missing context, and conclusions buried under qualifications. | Adopt. | +| Voice and Tone | Keep formality, role names, component names, and terminology consistent while using natural engineering language. | Adopt. | +| So What | Explain the reason, effect, and cost of an architecture choice; answer why the design works this way. | Adapt to engineering rationale; do not add marketing benefits. | +| Prove It | Support behavior and performance claims with code paths, type signatures, tests, metrics, or bounded experiments. | Adopt. | +| Specificity | State the scope, abstraction level, preconditions, failure conditions, and excluded paths. | Adopt. | +| Heightened Emotion | Amplify pain, anxiety, aspiration, or emotional impact. | Skip. Technical documentation prioritizes accuracy and verifiability. | +| Zero Risk | Add marketing CTAs, guarantees, and risk reversals. | Keep only operational preconditions, failure results, and next steps in user documentation. Do not add marketing CTAs. | + +The upstream limits for English word counts, active voice, and short paragraphs are heuristics. Judge Chinese sentence length, code identifiers, and complex ownership relationships by readability instead of applying fixed thresholds mechanically. + +## 3. Review workflow + +### 3.1 Establish the technical facts first + +Before editing, identify: + +- The document type and target reader. +- The boundaries between the public contract, current implementation, and specialized fast paths. +- Whether key types, configuration keys, return values, and failure semantics match the code. +- The complete path covered by behavior, ownership, and performance claims. +- The Chinese or English counterpart that must be updated at the same time. + +If these points are unresolved, perform a technical review first. Copy editing cannot replace implementation verification. + +### 3.2 Sweep one: structure and clarity + +First check whether a reader can answer three questions from the opening: + +1. What problem does this document address? +2. What is the most important stable information? +3. In what order will the document develop the topic? + +Then review each section: + +- Does the heading state the section's decision or task directly? +- Does the section lead with stable information at its current abstraction level before fields, branches, and measurements? +- Does the body answer the same set of concerns introduced at the beginning? +- Is the same fact repeated in a role table, flow table, diagram legend, and summary? +- Does local implementation detail appear too early in the introduction or architecture overview? + +### 3.3 Sweep two: voice and terminology + +- Use one canonical name, spelling, and capitalization for each concept. +- Give role boundaries explicit subjects, such as “master maintains `route`” and “owner manages local SSD.” +- Remove labels such as “stable conclusions” or “current conclusions” when the content is a summary or capability state. Use “core points,” “current status,” or the concrete statement instead. +- Avoid template language, promotional language, and unsupported claims to “improve,” “enhance,” or “optimize.” +- Keep required English code terms. Do not translate public types or fields merely to make the prose look more localized. + +### 3.4 Sweep three: rationale, evidence, and boundaries + +Check four properties for every important claim: + +| Property | Question to answer | +| --- | --- | +| Rationale | Why does this responsibility boundary, data direction, or lifetime exist? | +| Mechanism | Which object, field, or call path implements the behavior? | +| Scope | Which step, abstraction level, and branches does the claim cover? | +| Exclusions | Which costs, paths, recovery capabilities, or deployment conditions are outside the claim? | + +Place the reason close to the architecture claim. For example, when the master schedules memory allocations, immediately explain that these allocations cover both final value replicas and temporary cross-owner transfer memory. Listing ownership alone is insufficient. + +Bind performance claims to hardware, datasets, concurrency, output boundaries, measurement windows, and run counts. End-to-end logical payload bandwidth must not be described as raw SSD bandwidth. + +### 3.5 Sweep four: contraction and deduplication + +Resolve repetition in this order: + +1. Keep the first stable statement of the fact. +2. Keep one table or diagram when it materially lowers the cost of understanding. +3. In later sections, expand only new mechanisms, conditions, or failure paths. +4. In the summary, return to the public contract, core dataflow, and current boundaries without repeating experiment tables. + +Do not repeat role definitions merely to make every section appear self-contained. Once a role is defined in the overview, later diagram introductions should explain only meanings specific to that diagram. + +### 3.6 Sweep five: pre-publication recheck + +After a pass, work backward and confirm that later edits did not invalidate earlier decisions: + +- Terms, section numbers, and cross-references remain consistent. +- Tables, code fences, disclosure blocks, and Mermaid diagrams parse correctly. +- Chinese and English pages express the same contracts and boundaries. +- Exact identifiers still use their real names. +- The edit did not introduce a compatibility path, configuration entry, or unverified fact. +- The documentation site builds successfully. + +## 4. Collaborative review and one-shot example + +### 4.1 Trace surface issues to root causes + +An editing comment should identify the location, problem, reason, and proposed change. “This reads awkwardly” alone is not enough to justify an edit. + +| Item | Example | +| --- | --- | +| Location | The experiment paragraph in the opening. | +| Problem | It introduces measurement details before the reader understands the document structure. | +| Reason | Local evidence occupies the introduction's abstraction level. | +| Proposed change | Keep only the document map in the opening and move measurement conditions to the experiment section. | + +For clear, local issues that do not change technical meaning, edit directly and recheck the result. Before changing a contract, claim scope, or section narrative, explain the evidence and impact. + +A complete pass should also identify the common causes behind multiple surface issues. This lets the editor correct the full document and reuse the lesson on the next one. The following table uses a set of actual edits to the KV SSD article: + +| Surface issue | Root cause | Editing decision | +| --- | --- | --- | +| “First, remember four stable conclusions” introduces a summary. | The label is stronger than the content; a summary is not yet a supported conclusion. | Use “four core points.” | +| The prose states the owner's physical resources, switches to the master, and then returns to the owner's SSD scheduling. | One subject's responsibilities are fragmented, forcing readers to reconstruct the role boundary. | Group by subject: explain the master completely, then the owner. | +| The text only says that the master schedules memory allocations. | The responsibility claim lacks a nearby rationale, so readers cannot tell whether the boundary is complete. | Immediately state that memory allocations cover both final value replicas and temporary transfer memory, and that both interact with `route` and in-flight state. | +| The opening explains Section 9 experiment groups, hit conditions, and measurement methodology. | Local evidence and the introduction occupy different abstraction levels, hiding the document map. | Keep only the problem, core points, and reading path in the opening; leave experiment conditions in Section 9. | +| “Cover a larger runtime working set” is presented as an already realized result. | A design expectation is stated as a verified result, making the claim stronger than the evidence and leaving the reader-facing runtime benefit implicit. | Preserve the unchanged-public-API scope and say that the design is expected to raise cache hit rates and make fuller use of storage bandwidth. | + +These edits serve two shared goals: help readers build the system model at the right level and keep every claim within its evidence boundary. + +### 4.2 Reusable one-shot example + +Use the following prompt after the technical facts have been verified and the document needs a complete structural and language review: + +```text +Target document: + +Perform one complete copy-editing pass on the target technical document and edit it directly. Preserve public APIs, implementation facts, benchmark numbers, and existing claim boundaries. + +Check the following: +1. The opening states only the problem, core points, and document map. It does not preload local implementation detail or experiment methodology. +2. Labels in headings, introductions, and table headers match their content. Do not call summaries or capability states “conclusions.” +3. Group responsibilities by subject. When introducing multiple roles, finish one role before moving to the next. +4. Place the reason or mechanism next to every important architecture claim, especially responsibility boundaries, lifetimes, and data direction. +5. Match claim strength to evidence. State verified results directly; qualify design goals or expected benefits with “expected,” “aims to,” or “may”; do not present unsupported outcomes as facts. +6. Remove repeated role descriptions, details at the wrong abstraction level, and summaries that add no information. Preserve necessary scope, preconditions, and exclusions. +7. Keep terminology, type names, API names, section numbers, cross-references, tables, and code fences consistent. + +After editing: +- Summarize the main changes as “surface issue → root cause → editing decision” instead of listing every sentence edit. +- Update the corresponding Chinese or English page when one exists. +- Recheck public contracts, technical facts, and performance methodology, then run the repository's required documentation-site build. +``` + +## 5. Upstream practices not adopted here + +- Do not use emotional amplification, FOMO, exaggerated pain, or sales language in technical articles. +- Do not force a user benefit onto every implementation fact. Explain consequences only when they clarify rationale or trade-offs. +- Do not add unsourced numbers, time commitments, case studies, or comparative claims. +- Do not mechanically convert English guidance such as “no more than 25 words per sentence” into Chinese character limits. +- Do not force every sentence into active voice; ownership and dataflow clarity take priority. +- Do not use scores from invented expert personas as a substitute for code, tests, or same-level end-to-end tracing. + +## 6. Updating the upstream snapshot + +This document is pinned to upstream commit `0ba2a7f`. To update it: + +1. Compare upstream `skills/copy-editing/SKILL.md` and `references/` with the pinned revision. +2. Import only changes that apply to technical documentation and do not conflict with Fluxon rules. +3. Update the Chinese and English documents, skill version, and pinned commit together. +4. Rebuild the documentation site. + +Do not overwrite this document automatically with upstream content. The upstream target is marketing copy; this document targets accurate, verifiable technical writing. + +## 7. Third-party license + +This document is based on Corey Haines's `copy-editing` skill and has been modified for Fluxon technical documentation. The upstream work is licensed under the MIT License. + +
+MIT License + +```text +MIT License + +Copyright (c) 2025 Corey Haines + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` + +
diff --git a/fluxon_doc_en/dev_doc/Developer - 6 - Code Review Guidelines.md b/fluxon_doc_en/dev_doc/Developer - 6 - Code Review Guidelines.md new file mode 100644 index 0000000..0a3109f --- /dev/null +++ b/fluxon_doc_en/dev_doc/Developer - 6 - Code Review Guidelines.md @@ -0,0 +1,114 @@ +# Developer - 6 - Code Review Guidelines + +A code-review rule turns an architecture conclusion into an executable constraint. “Responsibilities should be clear” or “lifecycle order should be correct” is only a conclusion; a complete rule also states when it applies, what must be done, how completion is accepted, what benefit it creates, and how to limit the fix's side effects. + +## 1. Required structure of a rule + +Every rule in this document contains these fields: + +| Field | Question to answer | +| --- | --- | +| Trigger | Which code shape, dependency, or behavior change activates this rule? | +| Required action | Which constraint must the author and reviewer establish or change? | +| Acceptance criteria | Which observable postconditions mean the work is complete? | +| Expected benefit | Which risk is reduced or engineering capability is gained? | +| Side-effect boundary | Which extra complexity, coupling, or behavior change must the fix avoid? | +| Evidence | Which test, type constraint, or runtime result proves acceptance? | + +“Must” is a pre-merge requirement. “Should” permits deviation with an explicit rationale in the PR. “May” is an implementation choice that does not affect merge. + +## 2. Select rules by change scenario + +The reviewer routes change signals to relevant rules instead of mechanically applying every rule to every PR. + +| Change signal | Primary rules | +| --- | --- | +| A public layer starts accessing internal fields, or one module creates a resource and another cleans it up | R1 responsibility and ownership, R4 composable contracts, R5 minimal-side-effect fix | +| New `close`, background work, thread, process, runtime, or nested handle | R2 dependency lifecycle, R3 state authority | +| Multiple participants share `closed`, an event, a lease, or cleanup state | R1 responsibility and ownership, R3 state authority | +| New adapter, backend replacement, or parent wrapping a child | R1 responsibility and ownership, R4 composable contracts, R5 minimal-side-effect fix | +| Close races with an in-flight operation, or an outer lock spans an inner blocking call | R2 dependency lifecycle, R3 state authority | +| A purely local algorithm change does not alter boundaries, resources, or lifecycle | Do not force a complete lifecycle model; review local correctness and tests | + +## 3. Executable rules + +### R1: Cross-module state and resources have one owner + +| Field | Rule | +| --- | --- | +| Trigger | A change crosses public, composition, and internal layers; adds a shared resource; introduces duplicate close; or lets an outer layer mutate internal state. | +| Required action | List the creator, borrower, sharer, and final releaser of affected state and resources. Assign one final-release authority to each resource. The public layer expresses stable contracts, the composition layer orders modules, and the internal module maintains its state, invariants, and cleanup. | +| Acceptance criteria | Every resource answers “who releases last”; a borrower does not release owner resources; outer code does not manipulate internal fields to perform cleanup; replacing internals leaves public and composition lifecycle logic unchanged. | +| Expected benefit | Cohesive responsibility, no double release or leak, a smaller change-propagation radius, and independently testable and replaceable modules. | +| Side-effect boundary | Do not split stateless thin modules merely to create layers, expand public APIs, promote a local resource to global sharing, or move unrelated responsibilities. | +| Evidence | An ownership table or type relation, public contract tests, and tests proving borrower close preserves owner resources, exclusive-owner close completes release, and shared authority follows its declared release condition. | + +### R2: Dependent lifecycles execute in partial order + +| Field | Rule | +| --- | --- | +| Trigger | An object depends on a runtime, transport, store, worker, thread, process, or handle; or a change adds startup, shutdown, restart, or hot replacement. | +| Required action | Write the dependency DAG. If `A` depends on `B`, initialization satisfies `B → A` and destruction satisfies `A → B`. Shutdown establishes an admission barrier, propagates wake or cancel, converges in-flight work, joins background tasks, releases in reverse dependency order, and publishes completion. Lock hierarchy and callback direction remain compatible with dependency direction. | +| Acceptance criteria | Shutdown admits no new operations; all dependents are quiescent before a dependency releases; after close returns no task, callback, or handle inside that lifecycle boundary uses resources it owns; repeated close reaches the same terminal state. | +| Expected benefit | Deterministic lifecycle without use-after-close, shutdown deadlock, orphaned background work, or premature dependency release. | +| Side-effect boundary | Order only real dependencies; independent branches may close concurrently. Do not introduce a global lock or coordinator merely to unify order; wake or cancel remains scoped to resources owned by this shutdown. | +| Evidence | Real lifecycle tests, close racing with in-flight work, timeout tests, and completion-barrier tests showing no activity after close returns. | + +### R3: Stop commands, in-progress state, and completion proof have explicit authorities + +| Field | Rule | +| --- | --- | +| Trigger | Multiple objects read and write one boolean or event; stop request and release are asynchronous; close is retryable; or several layers can publish closed. | +| Required action | Assign every state a scope, sole writer, and meaning. Separate the stop command, `Closing` progress, and `Closed` completion proof. Prefer one monotonic state machine within one scope; signals and completion in different scopes remain owned by their respective modules. Name a signal-only API `request_shutdown()` or an equivalent explicit contract. | +| Acceptance criteria | Another participant's stop does not suppress this module's cleanup; only the owner publishes module completion; state advances only toward a terminal state; failure or retry neither reopens admission nor permanently skips unfinished work. | +| Expected benefit | No ambiguity in shared booleans, retryable and awaitable shutdown, better observability, and no premature return or resource leak. | +| Side-effect boundary | Do not add multiple sources of truth for one fact, introduce unnecessary states for a simple synchronous object, spread boolean combinations when a finite enum suffices, or keep old state aliases as compatibility branches. | +| Evidence | State-transition tests, tests where participants issue stop in different orders, and retry tests after partial construction or cleanup failure. | + +### R4: Child lifecycle contracts compose directly in parents + +| Field | Rule | +| --- | --- | +| Trigger | A parent wraps a child, an adapter or backend is added, one close invokes several children, or errors cross layers. | +| Required action | Define child lifecycle preconditions, success postconditions, and failure postconditions. Pair create / release, register / unregister, spawn / join, and subscribe / cancel. Parents invoke child contracts and order them without copying cleanup. Independent cleanup continues after partial failure, preserving the primary error and teardown errors. | +| Acceptance criteria | A parent decides the next step without inspecting child internals; child close provides the conditions required for reverse-order parent release; replacing a conforming implementation requires no parent change; failure reports which postconditions remain unsatisfied. | +| Expected benefit | Local correctness composes into system correctness, backends remain replaceable, failure causality survives, and modules evolve independently. | +| Side-effect boundary | Do not build a generic framework for hypothetical future implementations. Abstract only stable common semantics, keep specialized fast paths internal, and do not flatten all errors into an information-free result. | +| Evidence | Contract tests, consistency tests with an alternate implementation or test double, and tests proving remaining cleanup runs after one child close fails. | + +### R5: Fix at the invariant authority and limit the change radius + +| Field | Rule | +| --- | --- | +| Trigger | A proposed fix adds an outer conditional, clears internal fields directly, introduces a global flag, compatibility branch, or configuration entry to bypass an internal responsibility or lifecycle defect. | +| Required action | Locate the owner of the violated invariant and fix state, cleanup, or contract there; the composition layer changes only required ordering. State non-goals, remove workarounds superseded by the new contract, and limit regression testing to affected boundaries. | +| Acceptance criteria | Cross-boundary knowledge decreases or does not grow; no duplicate entrypoint, parallel configuration channel, or second lifecycle path is added; unrelated public behavior remains unchanged. | +| Expected benefit | Root-cause correction with a smaller regression radius and no accumulation of compatibility layers and conditional branches. | +| Side-effect boundary | Do not rewrite unrelated modules or expand migration under “architecture unification.” If a public contract must change, describe migration and impact separately instead of hiding it in an internal fix. | +| Evidence | A pre-fix counterexample, post-fix invariant tests, affected-caller regression, and review showing that the diff adds no bypass. | + +## 4. Review output and merge criteria + +A finding cites the rule and includes the trigger, violated acceptance criterion, impact, and side-effect boundary: + +```text +[R3][blocking] A shared stop signal is treated as proof that this module completed cleanup. + +Trigger: another participant sets the shared state first. +Violated criterion: cleanup owned by this module is skipped, so close is no longer a completion barrier. +Required action: separate shared stop intent from module-local completion, and let this module publish completion. +Side-effect boundary: do not add a second public close API or let the outer layer release internal resources directly. +``` + +| Label | Merge criterion | +| --- | --- | +| `[blocking]` | A plausible trigger violates acceptance criteria; fix before merge. | +| `[major]` | A critical scenario lacks a constraint or evidence; implement or verify before merge. | +| `[minor]` | A local maintainability or diagnostic issue does not affect acceptance criteria and may be explicitly deferred. | +| `[question]` | Trigger, authority, or contract is unclear; classify after clarification. | + +- **Request changes**: an unresolved `[blocking]` or `[major]` remains, or the applicable rule's owner, dependency, or acceptance criteria cannot be determined. +- **Approve**: required actions for all matched rules are complete, evidence satisfies acceptance criteria, expected benefits hold, and side effects stay inside the declared boundary. +- **Comment**: only clarification questions or suggestions that do not affect acceptance remain. + +Documentation PRs must also follow the [Documentation Writing Rules](./Developer%20-%203%20-%20Documentation%20Writing%20Rules.md) and [Technical Documentation Copy Editing](./Developer%20-%205%20-%20Technical%20Documentation%20Copy%20Editing.md). diff --git a/fluxon_doc_en/dev_doc/Developer - 7 - Event Subscription and Full Snapshot Guidelines.md b/fluxon_doc_en/dev_doc/Developer - 7 - Event Subscription and Full Snapshot Guidelines.md new file mode 100644 index 0000000..446d867 --- /dev/null +++ b/fluxon_doc_en/dev_doc/Developer - 7 - Event Subscription and Full Snapshot Guidelines.md @@ -0,0 +1,84 @@ +# Developer - 7 - Event Subscription and Full Snapshot Guidelines + +When a component maintains a local projection from a full query plus incremental events, it must use one of exactly two orderings: + +1. General ordering: `subscription/listener ready → fetch full snapshot → install snapshot → replay buffered incremental events → publish ready`. +2. Revision-coupled ordering: when the source provides an atomic revision contract such as etcd's, use `snapshot@R → watch@R+1`. + +The general ordering is mandatory unless the revision-coupled contract is available. Every `listen` that is established and confirmed ready must be followed by one full synchronization. This requirement applies both at initial startup and on every recovery path that establishes a new `listen`. + +Calling `listen()` or spawning a background task does not prove that the subscription is ready. A `listen` is established only after the source has acknowledged the subscription or the caller has obtained a usable stream handle or cursor. + +## 1. Scope + +This rule applies to local projections derived from remote authoritative state, including: + +- Cluster membership and service discovery. +- Route tables, configuration mirrors, and resource indexes. +- Controller- or informer-style local caches. +- Local caches whose watch streams may need to be rebuilt. + +The local projection owner owns snapshot installation, event application, and its ready state. The source remains the sole source of truth for business state. + +## 2. General ordering: full synchronization after listen readiness + +```mermaid +sequenceDiagram + participant projection as local-projection-owner + participant source as authoritative-source + + projection->>source: subscribe / listen + source-->>projection: subscription ready (stream / cursor) + projection->>source: fetch full snapshot + source-->>projection: snapshot (optional revision R) + projection->>projection: install snapshot + loop buffered and live events + source-->>projection: incremental event + projection->>projection: apply by revision or idempotent version check + end + projection->>projection: publish ready +``` + +| Phase | Required action | Acceptance criterion | +| --- | --- | --- | +| Subscribe | Wait until the source confirms that the subscription exists. A fire-and-forget `spawn()` is not readiness proof. | Every state change after subscription readiness can enter this subscription stream. | +| Snapshot | Fetch the complete state once after the subscription is ready. | The snapshot covers the complete object set at the source-defined consistency point. | +| Install | Let the local projection owner atomically replace the baseline, or perform replacement through its sole writer. | Incremental processing cannot race with snapshot installation as unordered writers to the same projection. | +| Replay | Apply events buffered since subscription readiness. With revisions, accept only events after the snapshot revision. Without revisions, use generation or version checks, or make application idempotent. | Duplicate or stale events cannot roll back newer state from the snapshot. | +| Ready | Publish ready only after the snapshot is installed and the subscription backlog reaches a defined convergence point. | A ready reader cannot observe a partially installed snapshot. | + +## 3. Failure and resynchronization + +- **Subscription failure**: do not start the full fetch or publish ready. A retry starts by establishing a new subscription. +- **Snapshot failure**: the synchronization attempt is invalid. Close or discard the current subscription and retry the complete sequence so buffering cannot grow without a bound. +- **Invalid current event stream**: lag, disconnect, or a source-epoch change invalidates the old stream's completeness proof. Revoke ready and discard the old stream. Under the general ordering, perform a full synchronization after the new `listen` becomes ready. Under the revision-coupled ordering, repeat `snapshot@R → watch@R+1`. +- **Shutdown**: stop event admission, cancel or wake and join the event task, then release the stream and projection dependencies. No task may write the projection after successful close returns. + +A fixed `sleep` may bound a timeout or aid debugging. It cannot prove subscription readiness, snapshot installation, or backlog convergence. + +## 4. Second ordering: snapshot@R → watch@R+1 + +If the source explicitly provides an atomic revision contract, the following sequence is valid: + +`snapshot at revision R → watch from revision R + 1` + +This ordering requires all of the following: + +- The snapshot returns a usable revision or cursor. +- The watch API guarantees delivery beginning with the first event after the requested revision. +- The source uses one ordered revision space for the snapshot and event log. +- Tests cover a concurrent update at the snapshot/watch boundary. + +An ordinary `fetch_all()` followed by `listen()` without a starting revision does not satisfy this ordering's prerequisites. + +## 5. Review and verification checklist + +- [ ] Subscription readiness is an awaitable completion barrier; successful `spawn()` is not treated as successful subscription. +- [ ] Under the general ordering, every established `listen` is followed by one full snapshot fetch, both at initial startup and on every recovery path. +- [ ] The code implements only the two valid orderings; there is no third branch that attaches a new `listen` directly to an old snapshot. +- [ ] When `snapshot@R → watch@R+1` is used, the source and tests satisfy the complete revision contract. +- [ ] Snapshot installation and event application have one writer or a provable revision order. +- [ ] The final projection is correct when an event occurs before subscription, after subscription but before snapshot, and after snapshot. +- [ ] Duplicate and stale events cannot roll state backward. +- [ ] Ready is published only after both the snapshot and backlog converge. +- [ ] Tests wait for explicit state conditions instead of using a fixed `sleep` as synchronization. diff --git a/fluxon_py/_api_ext_chan/mpmc.py b/fluxon_py/_api_ext_chan/mpmc.py index 4ddbc1e..f2e6c15 100644 --- a/fluxon_py/_api_ext_chan/mpmc.py +++ b/fluxon_py/_api_ext_chan/mpmc.py @@ -1,7 +1,7 @@ from concurrent.futures import thread from itertools import count from math import log -import os, time, struct, threading, json, random, typing, etcd3, copy, fcntl +import os, time, struct, threading, json, random, typing, etcd3, copy, fcntl, uuid from typing import Callable, Dict, Optional, Tuple, Any, List, Set, Union, cast from abc import ABC, abstractmethod @@ -95,7 +95,8 @@ MPMC_ATTACH_PAYLOAD_KEEPALIVE_RETRIES = 3 LOCAL_MEMBER_ID_RANGE_SIZE = 32 MPMC_CREATE_LOCK_TTL_SECONDS = 10 -MPMC_CREATE_LOCK_TIMEOUT_SECONDS = 10.0 +MPMC_CREATE_LOCK_TIMEOUT_SECONDS = 30.0 +_MPMC_CREATE_IN_PROGRESS_MESSAGE = "MPSC channel creation is already in progress" def new_etcd_client(api: KvClient) -> Result[etcd3.Etcd3Client, ApiError]: @@ -411,9 +412,9 @@ def _new_mpmc_role_key_prefix(mpmc_id: str, role: ChanRole) -> str: Get the key prefix for storing MPMC channel role. """ if role == ChanRole.PRODUCER: - return f"/mpmc_channels/producer/{mpmc_id}" + return f"/mpmc_channels/producer/{mpmc_id}/" elif role == ChanRole.CONSUMER: - return f"/mpmc_channels/consumer/{mpmc_id}" + return f"/mpmc_channels/consumer/{mpmc_id}/" else: raise ValueError(f"Invalid role: {role}") @@ -449,6 +450,21 @@ def _new_mpmc_ready_channels_prefix(mpmc_id: str) -> str: return f"/mpmc_channels/ready/{mpmc_id}/" # we need the / at the end for extracting mpsc_id from key +def _new_mpmc_create_reservations_prefix(mpmc_id: str) -> str: + """Get the prefix for in-flight sub-MPSC creation reservations.""" + return f"/mpmc_channels/create_reservations/{mpmc_id}/" + + +def _new_mpmc_create_reservation_key( + mpmc_id: str, role: ChanRole, member_id: int +) -> str: + """Get the member-owned reservation key for sub-MPSC creation.""" + return ( + f"{_new_mpmc_create_reservations_prefix(mpmc_id)}" + f"{role.value}/{member_id}" + ) + + def _extract_mpsc_id_from_ready_key(key: bytes, mpmc_id: str) -> str: """ Extract MPSC channel ID from a ready channel key. @@ -758,17 +774,12 @@ def get_meta(self) -> Result[Dict[str, Any], ApiError]: meta_object = json.loads(meta_data.decode()) return Result.new_ok(meta_object) - def get_mpsc_channels(self) -> Result[List[str], ApiError]: - """ - Get all MPSC channel IDs in this MPMC channel. - - Returns: - Result[List[str]]: List of MPSC channel IDs - """ + def _read_mpsc_channels_snapshot(self) -> Tuple[List[str], Optional[bytes]]: + """Read and validate the published sub-MPSC list and its raw value.""" channels_key = _new_mpmc_mpsc_channels_key(self.mpmc_id) channels_data, _ = self.etcd_client.get(channels_key) if channels_data is None: - return Result.new_ok([]) + return [], None raw = json.loads(channels_data.decode()) if not isinstance(raw, list): raise ValueError(f"invalid mpsc_channels value for mpmc_id={self.mpmc_id}: {raw!r}") @@ -784,6 +795,11 @@ def get_mpsc_channels(self) -> Result[List[str], ApiError]: raise ValueError(f"invalid mpsc_id element for mpmc_id={self.mpmc_id}: {item!r}") channels.append(item) + return channels, channels_data + + def get_mpsc_channels(self) -> Result[List[str], ApiError]: + """Get all published MPSC channel IDs in this MPMC channel.""" + channels, _ = self._read_mpsc_channels_snapshot() return Result.new_ok(channels) def get_remote_ready_channels(self) -> List[str]: @@ -840,6 +856,7 @@ def try_claim_ready_channel(self, mpsc_id: str) -> Result[bool, ApiError]: raise ValueError(f"mpmc_member_id is None for mpmc_id={self.mpmc_id}") ready_key = _new_mpmc_ready_channel_key(self.mpmc_id, mpsc_id) + expected_owner = str(self.mpmc_member_id).encode() try: success, _ = self.etcd_client.transaction( compare=[ @@ -848,7 +865,7 @@ def try_claim_ready_channel(self, mpsc_id: str) -> Result[bool, ApiError]: success=[ self.etcd_client.transactions.put( ready_key, - str(self.mpmc_member_id).encode(), + expected_owner, self.mpmc_member_lease, ) ], @@ -856,10 +873,38 @@ def try_claim_ready_channel(self, mpsc_id: str) -> Result[bool, ApiError]: ) return Result.new_ok(bool(success)) except Exception as e: + reconciliation_errors: List[str] = [] + for attempt in range(1, 4): + try: + ready_owner, _ = self.etcd_client.get(ready_key) + if ready_owner == expected_owner: + logging.warning( + "Ready claim response was lost but authority confirms ownership: " + "mpmc_id=%s mpsc_id=%s member_id=%s", + self.mpmc_id, + mpsc_id, + self.mpmc_member_id, + ) + return Result.new_ok(True) + if ready_owner is not None: + reconciliation_errors.append( + f"attempt={attempt}: observed_other_owner={ready_owner!r}" + ) + except Exception as read_error: + reconciliation_errors.append( + f"attempt={attempt}: {read_error}" + ) + if attempt < 3: + time.sleep(0.05 * attempt) + + # An absent read cannot fence the original create-if-absent txn. + # Revoke this member lease so a delayed put cannot become live. + self._fail_close_member(reason="ambiguous_ready_claim") return Result.new_error( ChanBindError( f"Failed to claim ready key for mpmc_id={self.mpmc_id}, " - f"mpsc_id={mpsc_id}: {e}" + f"mpsc_id={mpsc_id}: {e}; reconciliation_errors=" + f"{reconciliation_errors}" ) ) @@ -869,7 +914,24 @@ def _best_effort_delete_ready_channel(self, mpsc_id: str, *, reason: str) -> Non ready_key = _new_mpmc_ready_channel_key(self.mpmc_id, mpsc_id) try: - self.etcd_client.delete(ready_key) + expected_owner = str(self.mpmc_member_id).encode() + success, _ = self.etcd_client.transaction( + compare=[ + self.etcd_client.transactions.value(ready_key) + == expected_owner + ], + success=[self.etcd_client.transactions.delete(ready_key)], + failure=[], + ) + if not success: + logging.debug( + "Ready key cleanup skipped after %s because ownership changed: " + "mpmc_id=%s mpsc_id=%s member_id=%s", + reason, + self.mpmc_id, + mpsc_id, + self.mpmc_member_id, + ) except Exception as e: logging.warning( "Failed to delete ready key after %s for mpmc_id=%s mpsc_id=%s: %s", @@ -879,6 +941,45 @@ def _best_effort_delete_ready_channel(self, mpsc_id: str, *, reason: str) -> Non e, ) + def _bind_claimed_existing_unready_consumer( + self, + api: KvClient, + chan_config: Dict[str, int], + mpsc_id: str, + ) -> Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError]: + """Bind a consumer after this member has atomically claimed its ready key.""" + try: + mpsc_consumer = MPSCChanConsumer( + api, + mpsc_id, + chan_config, + self.etcd_client, + self.mpmc_member_lease, + self.mpmc_global_lease, + override_payload_lease_id=self.payload_lease_id, + parent_mpmc_id_opt=self.mpmc_id, + parent_mpmc_member_id_opt=self.mpmc_member_id, + ) + except Exception as e: + self._best_effort_delete_ready_channel( + mpsc_id, + reason="existing_unready_bind_failure", + ) + return Result.new_error( + ChanBindError( + f"Failed to bind claimed existing MPSC consumer for " + f"mpmc_id={self.mpmc_id}, mpsc_id={mpsc_id}: {e}" + ) + ) + + mpsc_consumer._mpmc_ready_claimed = True + logging.debug( + "Bound claimed existing unready MPSC consumer for mpmc_id=%s, mpsc_id=%s", + self.mpmc_id, + mpsc_id, + ) + return Result.new_ok(mpsc_consumer) + def _try_bind_existing_unready_consumer( self, api: KvClient, @@ -906,40 +1007,219 @@ def _try_bind_existing_unready_consumer( ) continue + return self._bind_claimed_existing_unready_consumer( + api, + chan_config, + mpsc_id, + ) + + return None + + def _get_create_reservation_count(self) -> int: + """Count live member-leased sub-MPSC creation reservations.""" + prefix = _new_mpmc_create_reservations_prefix(self.mpmc_id) + return sum(1 for _ in self.etcd_client.get_prefix(prefix)) + + def _best_effort_delete_create_reservation( + self, + reservation_key: str, + reservation_token: bytes, + *, + reason: str, + ) -> Tuple[bool, bool]: + """Fence this create attempt with a linearized reservation CAS. + + Returns ``(linearized, deleted_owned_token)``. A linearized compare is + enough to fence a pending publish, while an ambiguous reservation-create + RPC additionally requires ``deleted_owned_token=True`` or member-lease + revocation because a late create could otherwise appear after an absent + compare. + """ + errors: List[str] = [] + for attempt in range(1, 4): try: - mpsc_consumer = MPSCChanConsumer( - api, - mpsc_id, - chan_config, - self.etcd_client, - self.mpmc_member_lease, - self.mpmc_global_lease, - override_payload_lease_id=self.payload_lease_id, - parent_mpmc_id_opt=self.mpmc_id, - parent_mpmc_member_id_opt=self.mpmc_member_id, - ) - except Exception as e: - self._best_effort_delete_ready_channel( - mpsc_id, - reason="existing_unready_bind_failure", + success, _ = self.etcd_client.transaction( + compare=[ + self.etcd_client.transactions.value(reservation_key) + == reservation_token + ], + success=[self.etcd_client.transactions.delete(reservation_key)], + failure=[], ) - return Result.new_error( - ChanBindError( - f"Failed to bind claimed existing MPSC consumer for " - f"mpmc_id={self.mpmc_id}, mpsc_id={mpsc_id}: {e}" + if not success: + logging.debug( + "Create reservation cleanup skipped after %s because ownership changed: " + "mpmc_id=%s key=%s", + reason, + self.mpmc_id, + reservation_key, ) + return True, bool(success) + except Exception as e: + errors.append(f"attempt={attempt}: {e}") + if attempt < 3: + time.sleep(0.05 * attempt) + + # A live keepalive would otherwise preserve the stale reservation + # indefinitely. Stop this member so its lease can be revoked or expire. + logging.error( + "Failed to delete owned create reservation after %s; closing MPMC member: " + "mpmc_id=%s key=%s errors=%s", + reason, + self.mpmc_id, + reservation_key, + errors, + ) + self._fail_close_member( + reason=f"reservation_cleanup_failure:{reason}", + ) + return False, False + + def _fail_close_member(self, *, reason: str) -> None: + """Stop keepalive and revoke this member lease after an ambiguous create.""" + self.shutdown_ctl.closed = True + self._lm_mpmc_member = None + lease_id = int(self.mpmc_member_lease.id) + try: + self.etcd_client.revoke_lease(lease_id) + except Exception as e: + if "requested lease not found" not in str(e).lower(): + logging.error( + "Failed to revoke MPMC member lease after %s: " + "mpmc_id=%s member_id=%s lease_id=%s error=%s", + reason, + self.mpmc_id, + self.mpmc_member_id, + lease_id, + e, ) - mpsc_consumer._mpmc_ready_claimed = True - logging.debug( - "Bound claimed existing unready MPSC consumer for mpmc_id=%s, mpsc_id=%s", - self.mpmc_id, - mpsc_id, + @staticmethod + def _release_failed_new_mpsc( + mpsc_object: Union[MPSCChanConsumer, MPSCChanProducer], + *, + reason: str, + ) -> None: + """Release a process-local MPSC handle after publication failed.""" + try: + mpsc_object.release_local_handle().unwrap() + except Exception as e: + logging.warning( + "Failed to release new MPSC local handle after %s: %s", + reason, + e, ) - return Result.new_ok(mpsc_consumer) + def _reconcile_new_mpsc_publish( + self, + mpsc_id: str, + chan_role: ChanRole, + ) -> Optional[bool]: + """Resolve an ambiguous publish RPC using linearizable authority reads.""" + successful_absent_reads = 0 + for attempt in range(1, 6): + try: + current_mpscs, _ = self._read_mpsc_channels_snapshot() + if mpsc_id not in current_mpscs: + successful_absent_reads += 1 + elif chan_role == ChanRole.PRODUCER: + return True + else: + ready_key = _new_mpmc_ready_channel_key(self.mpmc_id, mpsc_id) + ready_owner, _ = self.etcd_client.get(ready_key) + expected_owner = str(self.mpmc_member_id).encode() + if ready_owner == expected_owner: + return True + logging.error( + "Published MPSC list contains %s but ready ownership is inconsistent: " + "mpmc_id=%s expected=%r actual=%r", + mpsc_id, + self.mpmc_id, + expected_owner, + ready_owner, + ) + return None + except Exception as e: + logging.warning( + "Publish reconciliation attempt %s/5 failed for mpmc_id=%s " + "mpsc_id=%s: %s", + attempt, + self.mpmc_id, + mpsc_id, + e, + ) + if attempt < 5: + time.sleep(0.1) + + if successful_absent_reads > 0: + return False return None + def _abort_unpublished_new_mpsc( + self, + mpsc_object: Union[MPSCChanConsumer, MPSCChanProducer], + *, + reason: str, + ) -> None: + """Release a failed create and delete its unpublished distributed metadata.""" + mpsc_id = mpsc_object.chan_id + self._release_failed_new_mpsc(mpsc_object, reason=reason) + + meta_key = _new_etcd_meta_key(mpsc_id) + abort_key = f"/channels/aborted/{mpsc_id}" + channel_prefix = f"/channels/{mpsc_id}/" + allocator_prefix = f"dist_id_allocator/channels/{mpsc_id}/" + cluster_lease_key = f"cluster_lease/id_allocator/channels/{mpsc_id}" + cleanup_errors: List[str] = [] + for attempt in range(1, 4): + try: + meta_data, _ = self.etcd_client.get(meta_key) + global_long_lease_id: Optional[int] = None + if meta_data is not None: + try: + meta_obj = json.loads(meta_data.decode()) + except Exception as meta_error: + logging.warning( + "Failed to decode unpublished MPSC meta before cleanup: " + "mpmc_id=%s mpsc_id=%s error=%s", + self.mpmc_id, + mpsc_id, + meta_error, + ) + else: + lease_id_value = meta_obj.get("global_long_lease_id") + if isinstance(lease_id_value, int) and lease_id_value > 0: + global_long_lease_id = lease_id_value + + # Fence every Rust-side create/bind write before deleting state. + # Channel ids are monotonic and never reused, so this marker is + # intentionally permanent until a full `/channels` reset. + self.etcd_client.put(abort_key, b"1") + self.etcd_client.delete_prefix(channel_prefix) + self.etcd_client.delete_prefix(allocator_prefix) + self.etcd_client.delete(meta_key) + self.etcd_client.delete(cluster_lease_key) + if global_long_lease_id is not None: + try: + self.etcd_client.revoke_lease(global_long_lease_id) + except Exception as revoke_error: + if "requested lease not found" not in str(revoke_error).lower(): + raise + return + except Exception as e: + cleanup_errors.append(f"attempt={attempt}: {e}") + if attempt < 3: + time.sleep(0.05 * attempt) + + logging.error( + "Failed to clean unpublished MPSC distributed metadata: mpmc_id=%s " + "mpsc_id=%s reason=%s errors=%s", + self.mpmc_id, + mpsc_id, + reason, + cleanup_errors, + ) + def _ensure_member_lease_alive(self) -> Result[OkNone, ApiError]: lease_id = int(self.mpmc_member_lease.id) endpoint = self._etcd_endpoints[0] if self._etcd_endpoints else None @@ -1127,6 +1407,89 @@ def try_existing_channels( return Result.new_ok(mpsc_object) create_error = create_result.unwrap_error() + if ( + isinstance(create_error, ChanCreateError) + and create_error.message == _MPMC_CREATE_IN_PROGRESS_MESSAGE + ): + # A reservation replaces the old behavior where contenders waited + # behind a long-held create lock. Wait outside that lock for the + # publisher, or retry reservation if the owner failed and cleaned up. + wait_deadline = time.monotonic() + MPMC_CREATE_LOCK_TIMEOUT_SECONDS + while time.monotonic() < wait_deadline: + if self.shutdown_ctl.closed: + return Result.new_error( + ChannelClosedError( + message="MPMC channel is closed.", + channel_id=self.mpmc_id, + ) + ) + + self._refresh_local_ready_state() + ready_channels = self.get_ready_channels() + unready_channels = self.unready_channels + existing_result = try_existing_channels( + ready_channels, + unready_channels, + ) + if existing_result is not None: + return existing_result + + try: + in_flight_creates = self._get_create_reservation_count() + except Exception as e: + return Result.new_error( + ChanCreateError( + f"Failed to inspect in-flight MPSC creation for " + f"MPMC channel {self.mpmc_id}: {e}" + ) + ) + + should_retry_create = in_flight_creates == 0 + if producer is None: + try: + active_consumers = self._get_active_consumer_count() + except Exception as e: + return Result.new_error( + ChanCreateError( + f"Failed to count active consumers while waiting for " + f"MPMC channel {self.mpmc_id}: {e}" + ) + ) + published_channel_count = len( + set(ready_channels).union(unready_channels) + ) + should_retry_create = active_consumers > ( + published_channel_count + in_flight_creates + ) + + if should_retry_create: + create_result = self.try_create_mpsc_channel( + api, + chan_config, + chan_role, + ) + if create_result.is_ok(): + mpsc_object = create_result.unwrap() + if producer is not None: + assert isinstance(mpsc_object, MPSCChanProducer) + producer._record_mpsc_producer(mpsc_object) + return Result.new_ok(mpsc_object) + + create_error = create_result.unwrap_error() + if not ( + isinstance(create_error, ChanCreateError) + and create_error.message == _MPMC_CREATE_IN_PROGRESS_MESSAGE + ): + break + + time.sleep(0.05) + else: + create_error = ChanCreateError( + f"{_MPMC_CREATE_IN_PROGRESS_MESSAGE} after waiting " + f"{MPMC_CREATE_LOCK_TIMEOUT_SECONDS:.0f}s for " + f"MPMC channel {self.mpmc_id}" + ) + if ( producer is not None and isinstance(create_error, ChanCreateError) @@ -1161,23 +1524,30 @@ def try_create_mpsc_channel( chan_config: Dict[str, int], chan_role: ChanRole, ) -> Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError]: + """Reserve, construct, and publish a new sub-MPSC. + + The create lock protects only authority reads, reservation accounting, + and the final list update. Rust-backed MPSC construction runs outside + the lock because it may perform lease probes and other network I/O. """ - Try to create a new MPSC channel and add it to this MPMC channel. - Uses etcd lock to ensure atomic creation. - Producer: only create if this is the first channel - Consumer: only create if active consumer count > existing MPSC channels - - Args: - api(KvClient): KV store API (required for creating MPSC objects) - api(KvClient): KV store API (required for creating MPSC objects) - chan_config(Dict[str, int]): Channel configuration (required for creating MPSC objects) - chan_role(ChanRole): Channel role (PRODUCER or CONSUMER) - - Returns: - Result[Union[MPSCChanConsumer, MPSCChanProducer]]: New MPSC object - """ + if chan_role not in (ChanRole.PRODUCER, ChanRole.CONSUMER): + return Result.new_error( + ChanCreateError(f"Invalid channel role: {chan_role}") + ) + if self.mpmc_member_id is None: + return Result.new_error( + ChanCreateError( + f"Cannot create MPSC for mpmc_id={self.mpmc_id}: member id is missing" + ) + ) + lock_key = f"/mpmc_channels/{self.mpmc_id}/create_lock" - + claimed_existing_mpsc_id: Optional[str] = None + reservation_key: Optional[str] = None + reservation_token: Optional[bytes] = None + reservation_create_confirmed = False + + # Phase 1: re-read authority state and reserve one creation slot. try: with EtcdLock( self._etcd_endpoints, @@ -1185,153 +1555,280 @@ def try_create_mpsc_channel( MPMC_CREATE_LOCK_TTL_SECONDS, MPMC_CREATE_LOCK_TIMEOUT_SECONDS, ): - # Get current channels mpsc_result = self.get_mpsc_channels() if not mpsc_result.is_ok(): error = mpsc_result.unwrap_error() - if error is not None: - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(error) - else: - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(ChanCreateError("Failed to get channels")) - + return Result.new_error( + error + if error is not None + else ChanCreateError("Failed to get channels") + ) + current_mpscs = mpsc_result.unwrap() - - # Role-specific constraints + in_flight_creates = self._get_create_reservation_count() + if chan_role == ChanRole.PRODUCER: - # Producer: only create if this is the first channel if current_mpscs: - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(ChanCreateError("Producer can only create the first channel")) - elif chan_role == ChanRole.CONSUMER: - # Consumer: lock-free snapshots are not authoritative. Re-read - # under the create lock and claim any existing unready MPSC - # before deciding to allocate a new one. + return Result.new_error( + ChanCreateError("Producer can only create the first channel") + ) + if in_flight_creates > 0: + return Result.new_error( + ChanCreateError(_MPMC_CREATE_IN_PROGRESS_MESSAGE) + ) + else: + # The lock-free caller already tried this path. Re-check under + # the lock so a concurrently published unready channel wins + # over allocating another one, but bind it after releasing the lock. ready_channel_set = set(self.get_remote_ready_channels()) current_unready_mpscs = [ mpsc_id for mpsc_id in current_mpscs if mpsc_id not in ready_channel_set ] - if current_unready_mpscs: - logging.debug( - "Consumer create-lock recheck found existing unready MPSCs for " - "mpmc_id=%s: %s", - self.mpmc_id, - current_unready_mpscs, - ) - existing_consumer_res = self._try_bind_existing_unready_consumer( - api, - chan_config, - current_unready_mpscs, - ) - if existing_consumer_res is not None: - return existing_consumer_res - - # Only create if the lock-protected recheck still found no - # claimable unready channel and active consumers outnumber - # the existing sub-MPSC count. - active_consumers = self._get_active_consumer_count() - if active_consumers <= len(current_mpscs): - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(ChanCreateError("Not enough active consumers to create new channel")) - else: - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(ChanCreateError(f"Invalid channel role: {chan_role}")) - - # Create MPSC object (let it handle its own ID allocation) - if chan_role == ChanRole.PRODUCER: - logging.debug(f"Creating new MPSC producer for MPMC channel {self.mpmc_id}") - mpsc_producer = MPSCChanProducer( - api, - None, - chan_config, - self.etcd_client, - # Tie producer membership keys to per-member lease so restarts do not collide. - self.mpmc_member_lease, - # Keep channel meta stable under the shared/global lease. - self.mpmc_global_lease, - # IMPORTANT: Always reuse the MPMC-shared kv payload lease for - # newly created sub MPSC channels so all producers bind payload - # keys under the SAME lease. Not doing this would cause subtle - # payload lease "split" across sub-channels even though they - # belong to the same MPMC. This is not a fallback; it's a - # required invariant for MPMC semantics. - override_payload_lease_id=self.payload_lease_id, - parent_mpmc_id_opt=self.mpmc_id, - parent_mpmc_member_id_opt=self.mpmc_member_id, - ) - # Record the created channel - mpsc_id = mpsc_producer.chan_id - current_mpscs.append(mpsc_id) - channels_key = _new_mpmc_mpsc_channels_key(self.mpmc_id) - self.etcd_client.put(channels_key, json.dumps(current_mpscs).encode(), self.mpmc_global_lease) - return Result.new_ok(mpsc_producer) - elif chan_role == ChanRole.CONSUMER: - logging.debug(f"Creating new MPSC consumer for MPMC channel {self.mpmc_id}") - mpsc_consumer: Optional[MPSCChanConsumer] = None - try: - mpsc_consumer = MPSCChanConsumer( - api, - None, - chan_config, - self.etcd_client, - # Tie consumer membership keys to per-member lease so restarts do not collide. - self.mpmc_member_lease, - # Keep channel meta stable under the shared/global lease. - self.mpmc_global_lease, - # Match producer-side semantics: new sub MPSC must reuse - # the shared kv payload lease of the parent MPMC. - override_payload_lease_id=self.payload_lease_id, - parent_mpmc_id_opt=self.mpmc_id, - parent_mpmc_member_id_opt=self.mpmc_member_id, - ) - except Exception as e: - logging.error(f"Fatal error creating MPSC consumer for MPMC channel {self.mpmc_id}: {e}") - return Result[ - Union[MPSCChanConsumer, MPSCChanProducer], ApiError - ].new_error( - ChanCreateError( - f"Failed to create MPSC consumer when try_create_mpsc_channel: {e}" + for mpsc_id in current_unready_mpscs: + claim_res = self.try_claim_ready_channel(mpsc_id) + if not claim_res.is_ok(): + return Result.new_error(claim_res.unwrap_error()) + if claim_res.unwrap(): + claimed_existing_mpsc_id = mpsc_id + break + + if claimed_existing_mpsc_id is None: + active_consumers = self._get_active_consumer_count() + if active_consumers <= len(current_mpscs) + in_flight_creates: + if in_flight_creates > 0: + return Result.new_error( + ChanCreateError(_MPMC_CREATE_IN_PROGRESS_MESSAGE) + ) + return Result.new_error( + ChanCreateError( + "Not enough active consumers to create new channel" + ) ) - ) - assert mpsc_consumer is not None - mpsc_id = mpsc_consumer.chan_id - channels_key = _new_mpmc_mpsc_channels_key(self.mpmc_id) - ready_key = _new_mpmc_ready_channel_key(self.mpmc_id, mpsc_id) - published_mpscs = current_mpscs.copy() - published_mpscs.append(mpsc_id) - success, _ = self.etcd_client.transaction( + + if claimed_existing_mpsc_id is None: + reservation_key = _new_mpmc_create_reservation_key( + self.mpmc_id, + chan_role, + self.mpmc_member_id, + ) + reservation_token = uuid.uuid4().hex.encode() + reserved, _ = self.etcd_client.transaction( compare=[ - self.etcd_client.transactions.create(ready_key) == 0 + self.etcd_client.transactions.create(reservation_key) == 0 ], success=[ self.etcd_client.transactions.put( - channels_key, - json.dumps(published_mpscs).encode(), - self.mpmc_global_lease, - ), - self.etcd_client.transactions.put( - ready_key, - str(self.mpmc_member_id).encode(), + reservation_key, + reservation_token, self.mpmc_member_lease, - ), + ) ], failure=[], ) - if not success: - try: - mpsc_consumer.release_local_handle().unwrap() - except Exception as e: - logging.debug(f"close leaked MPSC consumer error: {e}") - return Result[ - Union[MPSCChanConsumer, MPSCChanProducer], ApiError - ].new_error( - ChanCreateError( - f"Failed to publish claimed MPSC consumer {mpsc_id} for " - f"MPMC channel {self.mpmc_id}" - ) + reservation_create_confirmed = True + if not reserved: + return Result.new_error( + ChanCreateError(_MPMC_CREATE_IN_PROGRESS_MESSAGE) ) - mpsc_consumer._mpmc_ready_claimed = True - logging.debug(f"Created new MPSC consumer {mpsc_id} for MPMC channel {self.mpmc_id}") - return Result.new_ok(mpsc_consumer) - except Exception as e: - return Result[Union[MPSCChanConsumer, MPSCChanProducer], ApiError].new_error(ChanCreateError(f"Failed to create MPSC channel: {e}")) + if reservation_key is not None and reservation_token is not None: + _, deleted_owned_token = self._best_effort_delete_create_reservation( + reservation_key, + reservation_token, + reason="reserve_lock_exit_failure", + ) + if not reservation_create_confirmed and not deleted_owned_token: + # An unresolved create RPC may still arrive after a cleanup + # compare that observed the key as absent. Revoking the + # member lease prevents that late put from becoming live. + self._fail_close_member( + reason="ambiguous_reservation_create", + ) + if claimed_existing_mpsc_id is not None: + self._best_effort_delete_ready_channel( + claimed_existing_mpsc_id, + reason="claim_lock_exit_failure", + ) + return Result.new_error( + ChanCreateError(f"Failed to reserve MPSC channel creation: {e}") + ) + + if claimed_existing_mpsc_id is not None: + return self._bind_claimed_existing_unready_consumer( + api, + chan_config, + claimed_existing_mpsc_id, + ) + + assert reservation_key is not None + assert reservation_token is not None + + # Phase 2: run the expensive Rust-backed constructor without the create lock. + try: + if chan_role == ChanRole.PRODUCER: + logging.debug( + "Creating new MPSC producer outside create lock for MPMC channel %s", + self.mpmc_id, + ) + mpsc_object: Union[MPSCChanConsumer, MPSCChanProducer] = MPSCChanProducer( + api, + None, + chan_config, + self.etcd_client, + self.mpmc_member_lease, + self.mpmc_global_lease, + override_payload_lease_id=self.payload_lease_id, + parent_mpmc_id_opt=self.mpmc_id, + parent_mpmc_member_id_opt=self.mpmc_member_id, + ) + else: + logging.debug( + "Creating new MPSC consumer outside create lock for MPMC channel %s", + self.mpmc_id, + ) + mpsc_object = MPSCChanConsumer( + api, + None, + chan_config, + self.etcd_client, + self.mpmc_member_lease, + self.mpmc_global_lease, + override_payload_lease_id=self.payload_lease_id, + parent_mpmc_id_opt=self.mpmc_id, + parent_mpmc_member_id_opt=self.mpmc_member_id, + ) + except Exception as e: + self._best_effort_delete_create_reservation( + reservation_key, + reservation_token, + reason="constructor_failure", + ) + return Result.new_error( + ChanCreateError( + f"Failed to construct new MPSC {chan_role.value} for " + f"MPMC channel {self.mpmc_id}: {e}" + ) + ) + + # Phase 3: publish against the latest list, then delete the reservation + # in the same transaction. The raw-list compare also protects against a + # stale lock holder overwriting a newer list generation. + mpsc_id = mpsc_object.chan_id + published = False + publish_error: Optional[Exception] = None + try: + with EtcdLock( + self._etcd_endpoints, + lock_key, + MPMC_CREATE_LOCK_TTL_SECONDS, + MPMC_CREATE_LOCK_TIMEOUT_SECONDS, + ): + current_mpscs, current_channels_value = ( + self._read_mpsc_channels_snapshot() + ) + channels_key = _new_mpmc_mpsc_channels_key(self.mpmc_id) + compares = [ + self.etcd_client.transactions.value(reservation_key) + == reservation_token + ] + if current_channels_value is None: + compares.append( + self.etcd_client.transactions.create(channels_key) == 0 + ) + else: + compares.append( + self.etcd_client.transactions.value(channels_key) + == current_channels_value + ) + + success_ops = [ + self.etcd_client.transactions.put( + channels_key, + json.dumps(current_mpscs + [mpsc_id]).encode(), + self.mpmc_global_lease, + ) + ] + if chan_role == ChanRole.CONSUMER: + ready_key = _new_mpmc_ready_channel_key(self.mpmc_id, mpsc_id) + compares.append( + self.etcd_client.transactions.create(ready_key) == 0 + ) + success_ops.append( + self.etcd_client.transactions.put( + ready_key, + str(self.mpmc_member_id).encode(), + self.mpmc_member_lease, + ) + ) + success_ops.append( + self.etcd_client.transactions.delete(reservation_key) + ) + + published, _ = self.etcd_client.transaction( + compare=compares, + success=success_ops, + failure=[], + ) + except Exception as e: + publish_error = e + + if not published and publish_error is not None: + reconciled = self._reconcile_new_mpsc_publish(mpsc_id, chan_role) + if reconciled is True: + published = True + + if not published: + fenced, _ = self._best_effort_delete_create_reservation( + reservation_key, + reservation_token, + reason="publish_failure", + ) + detail = "transaction compare failed" if publish_error is None else str(publish_error) + reconciled_after_fence = ( + self._reconcile_new_mpsc_publish(mpsc_id, chan_role) + if fenced + else None + ) + if reconciled_after_fence is True: + published = True + elif fenced and reconciled_after_fence is False: + self._abort_unpublished_new_mpsc( + mpsc_object, + reason="publish_failure", + ) + else: + # Without a completed reservation CAS, a delayed publish may + # still commit. Preserve distributed metadata and stop this + # member instead of risking deletion of a published channel. + self._release_failed_new_mpsc( + mpsc_object, + reason="ambiguous_publish", + ) + self.shutdown_ctl.closed = True + self._lm_mpmc_member = None + return Result.new_error( + ChanCreateError( + f"Could not fence and reconcile MPSC {mpsc_id} publication " + f"for MPMC channel {self.mpmc_id}: {detail}" + ) + ) + + if not published: + return Result.new_error( + ChanCreateError( + f"Failed to publish new MPSC {mpsc_id} for MPMC channel " + f"{self.mpmc_id}: {detail}" + ) + ) + + if isinstance(mpsc_object, MPSCChanConsumer): + mpsc_object._mpmc_ready_claimed = True + logging.debug( + "Created and published new MPSC %s for MPMC channel %s", + mpsc_id, + self.mpmc_id, + ) + return Result.new_ok(mpsc_object) def _get_active_consumer_count(self) -> int: """ diff --git a/fluxon_py/config.py b/fluxon_py/config.py index 5861f64..7f13bf8 100644 --- a/fluxon_py/config.py +++ b/fluxon_py/config.py @@ -7,6 +7,7 @@ from typing import Dict, Any, List, Union, Tuple from abc import abstractmethod from collections.abc import Mapping +from decimal import Decimal, InvalidOperation from enum import Enum import sys import yaml @@ -24,6 +25,86 @@ def debug_print(*args): _YAML_KEY_TYPES = (str, int, float, bool) _YAML_SCALAR_TYPES = (str, int, float, bool, type(None)) +_U64_MAX = 2**64 - 1 +_SIZE_BYTE_UNITS = { + "": 1, + "b": 1, + "byte": 1, + "bytes": 1, + "kb": 1024, + "kib": 1024, + "mb": 1024**2, + "mib": 1024**2, + "gb": 1024**3, + "gib": 1024**3, + "tb": 1024**4, + "tib": 1024**4, +} +_CAPACITY_ALIGNMENT_BYTES = 16 * 1024 * 1024 + + +def _parse_size_bytes(value: Any, path: str) -> int: + if isinstance(value, bool): + raise ValueError(f"{path} must be bytes or a size string, got bool") + if isinstance(value, int): + if value < 0: + raise ValueError(f"{path} must be >= 0") + return value + if not isinstance(value, str): + raise ValueError(f"{path} must be bytes or a size string, got {type(value).__name__}") + + raw = value + match = re.fullmatch(r"\s*([0-9]+(?:\.[0-9]+)?|\.[0-9]+)\s*([A-Za-z]*)\s*", raw) + if match is None: + raise ValueError(f"{path} size must look like 16777216, 512MB, or 1.5GB") + number_text, unit_text = match.groups() + unit = unit_text.lower() + multiplier = _SIZE_BYTE_UNITS.get(unit) + if multiplier is None: + raise ValueError( + f"{path} size unit must be one of B, KB, MB, GB, TB, KiB, MiB, GiB, TiB" + ) + + try: + byte_value = Decimal(number_text) * Decimal(multiplier) + except InvalidOperation as exc: + raise ValueError(f"{path} size has invalid number: {raw}") from exc + bytes_int = int(byte_value) + if bytes_int < 0: + raise ValueError(f"{path} must be >= 0") + if bytes_int > _U64_MAX: + raise ValueError(f"{path} exceeds u64 bytes") + return bytes_int + + +def _align_capacity_down(value: int) -> int: + return value // _CAPACITY_ALIGNMENT_BYTES * _CAPACITY_ALIGNMENT_BYTES + + +def _normalize_size_bytes_fields(cfg: Dict[str, Any]) -> None: + contrib = cfg.get("contribute_to_cluster_pool_size") + if isinstance(contrib, dict): + if "dram" in contrib and contrib["dram"] is not None: + contrib["dram"] = _align_capacity_down( + _parse_size_bytes(contrib["dram"], "contribute_to_cluster_pool_size.dram") + ) + vram = contrib.get("vram") + if isinstance(vram, dict): + for gpu_id, size in list(vram.items()): + vram[gpu_id] = _align_capacity_down( + _parse_size_bytes(size, f"contribute_to_cluster_pool_size.vram.{gpu_id}") + ) + + spec = cfg.get("fluxonkv_spec") + if isinstance(spec, dict): + large_limit_size = spec.get("large_limit_size") + if large_limit_size is not None: + if not isinstance(large_limit_size, list): + raise ValueError("fluxonkv_spec.large_limit_size must be a list in owner mode") + spec["large_limit_size"] = [ + _parse_size_bytes(size, f"fluxonkv_spec.large_limit_size[{idx}]") + for idx, size in enumerate(large_limit_size) + ] def _to_plain_yaml_obj(value: Any, path: str) -> Any: @@ -67,9 +148,9 @@ def _yaml_template(): rdma_device_names: # Explicit RDMA devices for protocol config (['{str}'](optional)) pprof_duration_seconds: # Dump pprof flamegraph after N seconds (int(optional)) contribute_to_cluster_pool_size: # Capacity contributed to cluster pool (dict(optional)) - dram: 1677721600 # - DRAM contribution (int(multiple of 16777216)) + dram: 1677721600 # - DRAM contribution (size_bytes(multiple of 16777216)) vram: # - VRAM contribution per GPU (dict(dynamic_key)) - '{gpu_id}': 1677721600 # - Capacity for a given GPU id (int(multiple of 16777216)) + '{gpu_id}': 1677721600 # - Capacity for a given GPU id (size_bytes(multiple of 16777216)) test_spec_config: # Test-only config overrides (dict(optional)) disable_observability: false # Disable observe / OTLP background tasks (bool(optional)) disable_master_replica_cache: false # Disable master replica cache maintenance (bool(optional)) @@ -93,6 +174,8 @@ def _yaml_template(): side_transfer_worker_count: 0 # Owner-side worker count for side-transfer fanout (int(optional)) side_transfer_worker_p2p_port_base: # Optional owner-side worker port base (int(optional)) side_transfer_role: # worker (str(optional)) + kv_ssd_storage_backend: # native|foyer (str(optional)) + kv_ssd_uring_mode: # single_buffer|iovec; native backend only (str(optional)) # Notes: # - Zero-contribution mode is selected when contribute_to_cluster_pool_size is missing, # or when dram is 0 and all VRAM entries are 0. @@ -100,16 +183,21 @@ def _yaml_template(): # specific part, only one of [mooncake_spec,fluxonkv_spec] is required mooncake_spec: # mooncake 特定配置 (dict(optional)) + local_hostname: # Mooncake transfer/offload endpoint host (str(optional)) local_buffer_size: # 本地缓冲区大小 (int(multiple of 16777216)) metadata_server: # 元数据服务器地址 ('{str}://{str}:{int}/metadata') master_server_address: # 主服务器地址 ('{str}:{int}') etcd_addresses: # etcd地址列表, 注意是列表!!! (['{str}:{int}']) + enable_ssd_offload: false # Enable Mooncake SSD offload for this store (bool(optional)) + ssd_offload_path: # Local SSD offload directory path (str(optional)) + skip_get_size_on_get: false # Skip Mooncake get_size() before get(); intended for benchmark reads (bool(optional)) fluxonkv_spec: # fluxon kv specific config (dict(optional)) etcd_addresses: # Etcd address list ((None|['{str}:{int}'])) cluster_name: # Cluster name (str) share_mem_path: # Shared bundle path for mmap.file/shared.json/peer metadata (str) large_file_paths: # Owner-mode ordered large-file roots (['{str}'](optional)) + large_limit_size: # Owner-mode KV SSD byte limits aligned with large_file_paths (['{size_bytes}'](optional)) p2p_listen_port: # P2P QUIC listen port override (int(optional)) redis_compat: # Enable Redis protocol shim (dict(optional)) listen_addr: # TCP listen addr, e.g. "127.0.0.1:16379" (str) @@ -146,6 +234,8 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: "side_transfer_worker_count", "side_transfer_worker_p2p_port_base", "side_transfer_role", + "kv_ssd_storage_backend", + "kv_ssd_uring_mode", } unknown = sorted(set(raw.keys()) - allowed_keys) if unknown: @@ -172,6 +262,35 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: raise ValueError(f"{ctx}.{key} must be a bool") out[key] = value + kv_ssd_storage_backend = raw.get("kv_ssd_storage_backend") + if kv_ssd_storage_backend is not None: + if not isinstance(kv_ssd_storage_backend, str): + raise ValueError(f"{ctx}.kv_ssd_storage_backend must be a string") + allowed_kv_ssd_storage_backends = {"native", "foyer"} + if kv_ssd_storage_backend not in allowed_kv_ssd_storage_backends: + raise ValueError( + f"{ctx}.kv_ssd_storage_backend must be one of {sorted(allowed_kv_ssd_storage_backends)}, got {kv_ssd_storage_backend!r}" + ) + out["kv_ssd_storage_backend"] = kv_ssd_storage_backend + + kv_ssd_uring_mode = raw.get("kv_ssd_uring_mode") + if kv_ssd_uring_mode is not None: + if not isinstance(kv_ssd_uring_mode, str): + raise ValueError(f"{ctx}.kv_ssd_uring_mode must be a string") + allowed_kv_ssd_uring_modes = {"single_buffer", "iovec"} + if kv_ssd_uring_mode not in allowed_kv_ssd_uring_modes: + raise ValueError( + f"{ctx}.kv_ssd_uring_mode must be one of {sorted(allowed_kv_ssd_uring_modes)}, got {kv_ssd_uring_mode!r}" + ) + out["kv_ssd_uring_mode"] = kv_ssd_uring_mode + + if out.get("kv_ssd_storage_backend") == "foyer" and out.get( + "kv_ssd_uring_mode", "single_buffer" + ) != "single_buffer": + raise ValueError( + f"{ctx}.kv_ssd_uring_mode is only valid with {ctx}.kv_ssd_storage_backend=native" + ) + transport_mode = raw.get("transport_mode") transport_mode_was_explicit = transport_mode is not None side_transfer_role_raw = raw.get("side_transfer_role") @@ -311,7 +430,7 @@ def _is_zero_contribution_fluxonkv_config(cfg: Dict[str, Any]) -> bool: if not isinstance(contrib, dict): raise ValueError("contribute_to_cluster_pool_size must be a mapping when provided") - dram = int(contrib["dram"]) + dram = _parse_size_bytes(contrib["dram"], "contribute_to_cluster_pool_size.dram") vram_raw = contrib.get("vram") # Missing vram is normalized to "no GPU contribution". if vram_raw is None: @@ -322,8 +441,8 @@ def _is_zero_contribution_fluxonkv_config(cfg: Dict[str, Any]) -> bool: vram = vram_raw vram_is_zero = True - for _, value in vram.items(): - if int(value) != 0: + for gpu_id, value in vram.items(): + if _parse_size_bytes(value, f"contribute_to_cluster_pool_size.vram.{gpu_id}") != 0: vram_is_zero = False break if dram == 0 and not vram_is_zero: @@ -365,6 +484,7 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: "redis_compat", "sub_cluster", "large_file_paths", + "large_limit_size", ] for key in forbidden_spec_keys: if key in spec: @@ -376,7 +496,7 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: raise ValueError( "contribute_to_cluster_pool_size is required for owner mode (non-zero contribution)" ) - if int(contrib["dram"]) == 0: + if _parse_size_bytes(contrib["dram"], "contribute_to_cluster_pool_size.dram") == 0: raise ValueError("owner mode requires non-zero contribute_to_cluster_pool_size.dram") if "etcd_addresses" not in spec: @@ -404,6 +524,19 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: f"fluxonkv_spec.large_file_paths[{idx}] must be a non-empty string in owner mode" ) + if "large_limit_size" in spec and spec["large_limit_size"] is not None: + large_limit_size = spec["large_limit_size"] + if not isinstance(large_limit_size, list): + raise ValueError("fluxonkv_spec.large_limit_size must be a list in owner mode") + if len(large_limit_size) != len(large_file_paths): + raise ValueError( + "fluxonkv_spec.large_limit_size length must match fluxonkv_spec.large_file_paths length" + ) + for idx, limit in enumerate(large_limit_size): + limit_bytes = _parse_size_bytes(limit, f"fluxonkv_spec.large_limit_size[{idx}]") + if limit_bytes < 512: + raise ValueError(f"fluxonkv_spec.large_limit_size[{idx}] must be >= 512") + class FluxonKvClientConfig(): """Configuration class for KV Cache stores that reads from YAML config files.""" @@ -424,6 +557,7 @@ def __init__(self, config_dict: Dict[str, Any]): plain["test_spec_config"] = _normalize_test_spec_config( plain.get("test_spec_config"), "test_spec_config" ) + _normalize_size_bytes_fields(plain) # Backend selection contract: # - Exactly one backend spec must be provided (no fallback inference). @@ -491,6 +625,15 @@ def mooncake_spec_local_buffer_size(self): if "mooncake_spec" not in self.config_dict: return None return self.config_dict["mooncake_spec"]["local_buffer_size"] + + @property + def mooncake_spec_local_hostname(self) -> str: + if "mooncake_spec" not in self.config_dict: + return self.instance_key + raw = self.config_dict["mooncake_spec"].get("local_hostname") + if raw is None: + return self.instance_key + return str(raw) @property def mooncake_spec_metadata_server(self): @@ -509,6 +652,25 @@ def mooncake_spec_etcd_addresses(self): if "mooncake_spec" not in self.config_dict: return None return self.config_dict["mooncake_spec"]["etcd_addresses"] + + @property + def mooncake_spec_enable_ssd_offload(self) -> bool: + if "mooncake_spec" not in self.config_dict: + return False + return bool(self.config_dict["mooncake_spec"].get("enable_ssd_offload", False)) + + @property + def mooncake_spec_ssd_offload_path(self) -> str: + if "mooncake_spec" not in self.config_dict: + return "" + raw = self.config_dict["mooncake_spec"].get("ssd_offload_path", "") + return "" if raw is None else str(raw) + + @property + def mooncake_spec_skip_get_size_on_get(self) -> bool: + if "mooncake_spec" not in self.config_dict: + return False + return bool(self.config_dict["mooncake_spec"].get("skip_get_size_on_get", False)) @property def fluxonkv_spec_etcd_addresses(self): @@ -656,7 +818,7 @@ def ensure_zero_contribution_for_channel(self) -> None: f"got {type(cfg).__name__}" ) - dram = int(cfg["dram"]) + dram = _parse_size_bytes(cfg["dram"], "contribute_to_cluster_pool_size.dram") vram_raw = cfg.get("vram") if vram_raw is None: vram = {} @@ -671,7 +833,7 @@ def ensure_zero_contribution_for_channel(self) -> None: "Message-queue semantics require dynamic join/leave; non-zero contribution causes instability." ) for gpu_id, size in vram.items(): - if int(size) != 0: + if _parse_size_bytes(size, f"contribute_to_cluster_pool_size.vram.{gpu_id}") != 0: raise ValueError( f"For channel storage, contribute_to_cluster_pool_size must be zero-contribution. Current value: {cfg}. " f"Non-zero vram entry detected: gpu_id={gpu_id}. " @@ -750,8 +912,8 @@ def parse_type_recursive(type_str: str) -> Optional[Tuple[str, Dict[str, Any]]]: return parsed_type_name, merged_params return type_name, {"constraint": constraint} - # 6) Primitive types: str, int, bool, None - if type_str in ["str", "int", "bool", "None"]: + # 6) Primitive types: str, int, bool, size_bytes, None + if type_str in ["str", "int", "bool", "size_bytes", "None"]: return type_str, {} debug_print("type_str ", type_str, "not matched to any type") @@ -871,6 +1033,15 @@ def raise_validation_error(msg: str): raise_validation_error(f"Value must be multiple of {multiple}, got {value}") else: return None + + elif type_name == "size_bytes": + try: + bytes_value = _parse_size_bytes(value, path) + except ValueError as exc: + if raise_err: + raise_validation_error(str(exc)) + return None + return bytes_value elif type_name == "bool": if not isinstance(value, bool): @@ -944,6 +1115,25 @@ def raise_validation_error(msg: str): return None format_str = params["format"] + if format_str == "{int}": + for i, item in enumerate(value): + if not isinstance(item, int) or isinstance(item, bool): + if raise_err: + raise_validation_error(f"Expected int in list at index {i}, got {type(item).__name__}") + else: + return None + return value + if format_str == "{size_bytes}": + for i, item in enumerate(value): + try: + _parse_size_bytes(item, f"{path}[{i}]") + except ValueError as exc: + if raise_err: + raise_validation_error(str(exc)) + else: + return None + return value + regex_pattern = _convert_format_to_regex_pattern(format_str) for i, item in enumerate(value): if not isinstance(item, str): diff --git a/fluxon_py/fluxon_fs/__init__.py b/fluxon_py/fluxon_fs/__init__.py index 8dea96c..cde459a 100644 --- a/fluxon_py/fluxon_fs/__init__.py +++ b/fluxon_py/fluxon_fs/__init__.py @@ -20,6 +20,12 @@ ) from .patcher import FluxonFsPatcher from .bootstrap import install_patcher_from_master +from .video import ( + FluxonFsVideoReadRequest, + FluxonFsVideoReadResult, + FluxonFsVideoReader, + FluxonFsVideoReaderPool, +) def publish_export( @@ -91,6 +97,10 @@ def unpublish_export( "parse_master_config_from_file", "parse_master_panel_config_from_file", "FluxonFsPatcher", + "FluxonFsVideoReadRequest", + "FluxonFsVideoReadResult", + "FluxonFsVideoReader", + "FluxonFsVideoReaderPool", "install_patcher_from_master", "publish_export", "unpublish_export", diff --git a/fluxon_py/fluxon_fs/patcher.py b/fluxon_py/fluxon_fs/patcher.py index 1575b34..8ec3813 100644 --- a/fluxon_py/fluxon_fs/patcher.py +++ b/fluxon_py/fluxon_fs/patcher.py @@ -807,6 +807,36 @@ def mount_remote_dir( ) -> None: self._agent.mount_remote_dir(local_mount_dir_abs, export_name) + def open_video_reader( + self, + *, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int = 8, + ) -> Any: + from .video import FluxonFsVideoReader + + return FluxonFsVideoReader._open( + agent=self._agent, + export_name=export_name, + relpath=relpath, + height=height, + width=width, + num_threads=num_threads, + request_identity=self._request_identity, + ) + + def open_video_reader_pool(self, *, max_readers: int = 32) -> Any: + from .video import FluxonFsVideoReaderPool + + return FluxonFsVideoReaderPool( + agent=self._agent, + request_identity=self._request_identity, + max_readers=max_readers, + ) + def publish_export( self, *, diff --git a/fluxon_py/fluxon_fs/video.py b/fluxon_py/fluxon_fs/video.py new file mode 100644 index 0000000..1bbb23d --- /dev/null +++ b/fluxon_py/fluxon_fs/video.py @@ -0,0 +1,532 @@ +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Sequence +from dataclasses import dataclass +from threading import RLock +from typing import Any, Optional, Tuple + + +class FluxonFsVideoReader: + """Video reader backed by FluxonFS range reads.""" + + def __init__(self, inner: Any) -> None: + self._inner = inner + self._closed = False + + @classmethod + def _open( + cls, + *, + agent: Any, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int, + request_identity: Optional[Tuple[str, str]], + ) -> "FluxonFsVideoReader": + export_name = _require_non_empty_str(export_name, "export_name") + relpath = _require_non_empty_str(relpath, "relpath") + height = _require_positive_int(height, "height") + width = _require_positive_int(width, "width") + num_threads = _require_positive_int(num_threads, "num_threads") + inner = agent.open_video_reader( + export_name, + relpath, + height, + width, + num_threads, + request_identity, + ) + return cls(inner) + + def read_frames_numpy(self, indices: Sequence[int]) -> Any: + if self._closed: + raise RuntimeError("FluxonFsVideoReader is closed") + if not isinstance(indices, Sequence): + raise TypeError(f"indices must be a sequence of int, got {type(indices)}") + frame_indices = [] + for idx in indices: + if type(idx) is not int: + raise TypeError(f"indices must contain int values, got {type(idx)}") + if idx < 0: + raise ValueError("indices must be non-negative") + frame_indices.append(int(idx)) + return self._inner.read_frames_numpy(frame_indices) + + def stats(self) -> dict[str, int]: + if self._closed: + raise RuntimeError("FluxonFsVideoReader is closed") + stats = self._inner.stats() + if not isinstance(stats, dict): + raise TypeError(f"native video stats must be dict, got {type(stats)}") + out: dict[str, int] = {} + for key, value in stats.items(): + if not isinstance(key, str): + raise TypeError(f"native video stats keys must be str, got {type(key)}") + if type(value) is not int: + raise TypeError(f"native video stats values must be int, got {type(value)}") + out[key] = int(value) + return out + + def close(self) -> None: + if self._closed: + return + self._inner.close() + self._closed = True + + def __enter__(self) -> "FluxonFsVideoReader": + if self._closed: + raise RuntimeError("FluxonFsVideoReader is closed") + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.close() + + +@dataclass(frozen=True) +class FluxonFsVideoReadResult: + """Result for one pooled FluxonFS video decode.""" + + array: Any + reader_stats: dict[str, int] + pool_stats: dict[str, int] + reader_cache_hit: bool + batch_stats: dict[str, int] + + +@dataclass(frozen=True) +class FluxonFsVideoReadRequest: + """Request for one FluxonFS video clip decode.""" + + export_name: str + relpath: str + height: int + width: int + num_threads: int + indices: Sequence[int] + + def __post_init__(self) -> None: + object.__setattr__(self, "export_name", _require_non_empty_str(self.export_name, "export_name")) + object.__setattr__(self, "relpath", _require_non_empty_str(self.relpath, "relpath")) + object.__setattr__(self, "height", _require_positive_int(self.height, "height")) + object.__setattr__(self, "width", _require_positive_int(self.width, "width")) + object.__setattr__(self, "num_threads", _require_positive_int(self.num_threads, "num_threads")) + object.__setattr__(self, "indices", tuple(_validate_indices(self.indices))) + + +@dataclass(frozen=True) +class _VideoReaderKey: + export_name: str + relpath: str + height: int + width: int + num_threads: int + request_identity: Optional[Tuple[str, str]] + + +@dataclass +class _PooledReader: + key: _VideoReaderKey + reader: FluxonFsVideoReader + active: bool = False + retired: bool = False + + +class _ReaderLease: + def __init__( + self, + *, + pool: "FluxonFsVideoReaderPool", + reader_id: int, + entry: _PooledReader, + cache_hit: bool, + ) -> None: + self.pool = pool + self.reader_id = reader_id + self.entry = entry + self.cache_hit = cache_hit + + def __enter__(self) -> "_ReaderLease": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.pool._release(self.reader_id, self.entry) + + +class FluxonFsVideoReaderPool: + """Process-local LRU pool for FluxonFS video readers.""" + + def __init__( + self, + *, + agent: Any, + request_identity: Optional[Tuple[str, str]], + max_readers: int = 32, + ) -> None: + self._agent = agent + self._request_identity = _validate_request_identity(request_identity) + self._max_readers = _require_positive_int(max_readers, "max_readers") + self._lock = RLock() + self._entries: OrderedDict[int, _PooledReader] = OrderedDict() + self._next_reader_id = 1 + self._closed = False + self._open_count = 0 + self._close_count = 0 + self._evict_count = 0 + self._cache_hits = 0 + self._cache_misses = 0 + + def read_frames_numpy( + self, + *, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int, + indices: Sequence[int], + ) -> Any: + result = self.read_frames_numpy_with_stats( + export_name=export_name, + relpath=relpath, + height=height, + width=width, + num_threads=num_threads, + indices=indices, + ) + return result.array + + def read_frames_numpy_with_stats( + self, + *, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int, + indices: Sequence[int], + ) -> FluxonFsVideoReadResult: + key = _VideoReaderKey( + export_name=_require_non_empty_str(export_name, "export_name"), + relpath=_require_non_empty_str(relpath, "relpath"), + height=_require_positive_int(height, "height"), + width=_require_positive_int(width, "width"), + num_threads=_require_positive_int(num_threads, "num_threads"), + request_identity=self._request_identity, + ) + frame_indices = _validate_indices(indices) + + with self._acquire(key) as lease: + reader = lease.entry.reader + before = reader.stats() + array = reader.read_frames_numpy(frame_indices) + after = reader.stats() + + return FluxonFsVideoReadResult( + array=array, + reader_stats=_stats_delta(before, after), + pool_stats=self.stats(), + reader_cache_hit=lease.cache_hit, + batch_stats={ + "read_many_group_size": 1, + "read_many_group_frames": len(frame_indices), + "read_many_group_index": 0, + }, + ) + + def read_many_numpy_with_stats( + self, + requests: Sequence[FluxonFsVideoReadRequest], + ) -> list[FluxonFsVideoReadResult]: + read_requests = _validate_read_requests(requests) + if not read_requests: + return [] + + grouped: OrderedDict[_VideoReaderKey, list[tuple[int, FluxonFsVideoReadRequest]]] = OrderedDict() + for pos, request in enumerate(read_requests): + key = _VideoReaderKey( + export_name=request.export_name, + relpath=request.relpath, + height=request.height, + width=request.width, + num_threads=request.num_threads, + request_identity=self._request_identity, + ) + grouped.setdefault(key, []).append((pos, request)) + + output: list[Optional[FluxonFsVideoReadResult]] = [None] * len(read_requests) + for key, group in grouped.items(): + with self._acquire(key) as lease: + reader = lease.entry.reader + before = reader.stats() + counts = [len(request.indices) for _, request in group] + combined_indices: list[int] = [] + for _, request in group: + combined_indices.extend(request.indices) + combined_array = reader.read_frames_numpy(combined_indices) + after = reader.stats() + + _validate_combined_array(combined_array, len(combined_indices)) + reader_stats_parts = _split_stats_delta(_stats_delta(before, after), counts) + pool_stats = self.stats() + cursor = 0 + for group_index, ((pos, _request), reader_stats) in enumerate(zip(group, reader_stats_parts)): + count = counts[group_index] + array = combined_array[cursor : cursor + count] + cursor += count + output[pos] = FluxonFsVideoReadResult( + array=array, + reader_stats=reader_stats, + pool_stats=pool_stats, + reader_cache_hit=lease.cache_hit, + batch_stats={ + "read_many_group_size": len(group), + "read_many_group_frames": len(combined_indices), + "read_many_group_index": group_index, + }, + ) + + return _unwrap_results(output) + + def stats(self) -> dict[str, int]: + with self._lock: + active = sum(1 for entry in self._entries.values() if entry.active) + return { + "max_readers": self._max_readers, + "current_readers": len(self._entries), + "active_readers": active, + "open_count": self._open_count, + "close_count": self._close_count, + "evict_count": self._evict_count, + "reader_cache_hits": self._cache_hits, + "reader_cache_misses": self._cache_misses, + } + + def close(self) -> None: + readers_to_close: list[FluxonFsVideoReader] = [] + with self._lock: + if self._closed: + return + self._closed = True + for reader_id, entry in list(self._entries.items()): + entry.retired = True + if entry.active: + continue + readers_to_close.append(entry.reader) + self._entries.pop(reader_id, None) + self._close_count += 1 + for reader in readers_to_close: + reader.close() + + def __enter__(self) -> "FluxonFsVideoReaderPool": + with self._lock: + if self._closed: + raise RuntimeError("FluxonFsVideoReaderPool is closed") + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.close() + + def _acquire(self, key: _VideoReaderKey) -> _ReaderLease: + with self._lock: + if self._closed: + raise RuntimeError("FluxonFsVideoReaderPool is closed") + for reader_id, entry in self._entries.items(): + if entry.key == key and not entry.active and not entry.retired: + entry.active = True + self._entries.move_to_end(reader_id) + self._cache_hits += 1 + return _ReaderLease( + pool=self, + reader_id=reader_id, + entry=entry, + cache_hit=True, + ) + self._cache_misses += 1 + + reader = FluxonFsVideoReader._open( + agent=self._agent, + export_name=key.export_name, + relpath=key.relpath, + height=key.height, + width=key.width, + num_threads=key.num_threads, + request_identity=key.request_identity, + ) + + with self._lock: + if self._closed: + reader.close() + raise RuntimeError("FluxonFsVideoReaderPool is closed") + reader_id = self._next_reader_id + self._next_reader_id += 1 + entry = _PooledReader(key=key, reader=reader, active=True) + self._entries[reader_id] = entry + self._open_count += 1 + readers_to_close = self._evict_idle_locked() + + for old_reader in readers_to_close: + old_reader.close() + return _ReaderLease(pool=self, reader_id=reader_id, entry=entry, cache_hit=False) + + def _release(self, reader_id: int, entry: _PooledReader) -> None: + readers_to_close: list[FluxonFsVideoReader] = [] + with self._lock: + entry.active = False + if entry.retired or self._closed: + removed = self._entries.pop(reader_id, None) + if removed is not None: + readers_to_close.append(removed.reader) + self._close_count += 1 + elif reader_id in self._entries: + self._entries.move_to_end(reader_id) + readers_to_close.extend(self._evict_idle_locked()) + for reader in readers_to_close: + reader.close() + + def _evict_idle_locked(self) -> list[FluxonFsVideoReader]: + readers_to_close: list[FluxonFsVideoReader] = [] + while len(self._entries) > self._max_readers: + evicted = False + for reader_id, entry in list(self._entries.items()): + if entry.active: + continue + entry.retired = True + self._entries.pop(reader_id, None) + readers_to_close.append(entry.reader) + self._evict_count += 1 + self._close_count += 1 + evicted = True + break + if not evicted: + break + return readers_to_close + + +def _validate_indices(indices: Sequence[int]) -> list[int]: + if not isinstance(indices, Sequence): + raise TypeError(f"indices must be a sequence of int, got {type(indices)}") + frame_indices = [] + for idx in indices: + if type(idx) is not int: + raise TypeError(f"indices must contain int values, got {type(idx)}") + if idx < 0: + raise ValueError("indices must be non-negative") + frame_indices.append(int(idx)) + return frame_indices + + +def _validate_read_requests( + requests: Sequence[FluxonFsVideoReadRequest], +) -> list[FluxonFsVideoReadRequest]: + if not isinstance(requests, Sequence): + raise TypeError(f"requests must be a sequence of FluxonFsVideoReadRequest, got {type(requests)}") + out: list[FluxonFsVideoReadRequest] = [] + for request in requests: + if not isinstance(request, FluxonFsVideoReadRequest): + raise TypeError( + "requests must contain FluxonFsVideoReadRequest values, " + f"got {type(request)}" + ) + out.append(request) + return out + + +def _validate_request_identity(value: Optional[Tuple[str, str]]) -> Optional[Tuple[str, str]]: + if value is None: + return None + if not isinstance(value, tuple) or len(value) != 2: + raise TypeError("request_identity must be a (username, password) tuple or None") + username, password = value + if not isinstance(username, str) or not isinstance(password, str): + raise TypeError("request_identity values must be str") + return (username, password) + + +def _stats_delta(before: dict[str, int], after: dict[str, int]) -> dict[str, int]: + out: dict[str, int] = {} + for key, value in after.items(): + out[key] = int(value) - int(before.get(key, 0)) + return out + + +def _split_stats_delta(delta: dict[str, int], weights: list[int]) -> list[dict[str, int]]: + parts = [dict[str, int]() for _ in weights] + if not weights: + return parts + for key, value in delta.items(): + for index, part in enumerate(_split_int_by_weights(int(value), weights)): + parts[index][key] = part + return parts + + +def _split_int_by_weights(value: int, weights: list[int]) -> list[int]: + if not weights: + return [] + if value == 0: + return [0 for _ in weights] + + sign = -1 if value < 0 else 1 + abs_value = abs(value) + positive_weights = [max(0, int(weight)) for weight in weights] + total_weight = sum(positive_weights) + if total_weight <= 0: + positive_weights = [1 for _ in weights] + total_weight = len(weights) + + shares: list[int] = [] + remainders: list[tuple[int, int]] = [] + assigned = 0 + for index, weight in enumerate(positive_weights): + raw = abs_value * weight + share, remainder = divmod(raw, total_weight) + shares.append(share) + remainders.append((remainder, index)) + assigned += share + + remaining = abs_value - assigned + for _, index in sorted(remainders, reverse=True)[:remaining]: + shares[index] += 1 + return [sign * share for share in shares] + + +def _validate_combined_array(array: Any, expected_frames: int) -> None: + shape = getattr(array, "shape", None) + if not isinstance(shape, tuple) or len(shape) < 1: + raise TypeError(f"native video batch must expose tuple shape, got {type(shape)}") + if int(shape[0]) != expected_frames: + raise RuntimeError( + "native video batch returned unexpected frame count: " + f"expected={expected_frames} actual={shape[0]}" + ) + + +def _unwrap_results( + results: list[Optional[FluxonFsVideoReadResult]], +) -> list[FluxonFsVideoReadResult]: + out: list[FluxonFsVideoReadResult] = [] + for result in results: + if result is None: + raise RuntimeError("missing FluxonFS video batch result") + out.append(result) + return out + + +def _require_non_empty_str(value: Any, name: str) -> str: + if not isinstance(value, str): + raise TypeError(f"{name} must be str, got {type(value)}") + out = value.strip() + if not out: + raise ValueError(f"{name} must be non-empty") + return out + + +def _require_positive_int(value: Any, name: str) -> int: + if type(value) is not int: + raise TypeError(f"{name} must be int, got {type(value)}") + out = int(value) + if out <= 0: + raise ValueError(f"{name} must be positive") + return out diff --git a/fluxon_py/kvclient/fluxon.py b/fluxon_py/kvclient/fluxon.py index 1325e3d..4545ddc 100644 --- a/fluxon_py/kvclient/fluxon.py +++ b/fluxon_py/kvclient/fluxon.py @@ -283,6 +283,10 @@ def access(self) -> Result[Dict[str, Union[int, float, bool, str, bytes, DLPacke return Result.new_error(res.unwrap_error()) return Result.new_ok(res.unwrap()) + def _benchmark_source_kind(self) -> str: + """Return the storage source selected for this benchmark GET.""" + return str(self._inner_holder._benchmark_source_kind()) + # release() intentionally omitted for now. @@ -368,11 +372,13 @@ def put( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exists = bool(opts.reject_if_exists) if opts is not None else False inner_res = self._client.put( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exists=reject_if_exists, ) if not inner_res.is_ok(): err = inner_res.unwrap_error() @@ -416,11 +422,13 @@ def put_blocking( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exists = bool(opts.reject_if_exists) if opts is not None else False inner_res = self._client.put_blocking( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exists=reject_if_exists, ) if not inner_res.is_ok(): return Result.new_error(inner_res.unwrap_error()) @@ -469,11 +477,13 @@ def put_payload_from_ptr_blocking( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exists = bool(opts.reject_if_exists) if opts is not None else False inner_res = self._client.put_blocking( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exists=reject_if_exists, ) if not inner_res.is_ok(): return Result.new_error(inner_res.unwrap_error()) diff --git a/fluxon_py/kvclient/kvclient_interface.py b/fluxon_py/kvclient/kvclient_interface.py index f50db0f..8714519 100644 --- a/fluxon_py/kvclient/kvclient_interface.py +++ b/fluxon_py/kvclient/kvclient_interface.py @@ -29,9 +29,11 @@ class PutOptionalArgs: - lease_id: attach the written key to a lease on commit. - reject_if_inflight_same_key: ask Fluxon to fail-fast when the same key is already being written by another inflight put. + - reject_if_exists: create the key only when it does not already exist. """ lease_id: Optional[int] = None reject_if_inflight_same_key: bool = False + reject_if_exists: bool = False def support_mooncake(self) -> Tuple[bool, List[str]]: """ @@ -41,7 +43,7 @@ def support_mooncake(self) -> Tuple[bool, List[str]]: (supported: bool, unsupported_fields: list[str]) Notes: - - Mooncake is write-once; currently does not support lease binding. + - Mooncake natively supports reject_if_exists and does not support lease binding. """ unsupported: List[str] = [] if self.lease_id is not None: diff --git a/fluxon_py/kvclient/mooncake.py b/fluxon_py/kvclient/mooncake.py index 4457cc3..00ce21a 100644 --- a/fluxon_py/kvclient/mooncake.py +++ b/fluxon_py/kvclient/mooncake.py @@ -135,20 +135,25 @@ def __init__(self, config: "FluxonKvClientConfig"): self._rwlock = ReadWriteLock() self._renew_lock = threading.Lock() self._closed = False + self._skip_get_size_on_get = config.mooncake_spec_skip_get_size_on_get # config = self._config device_name = "" if config.protocol_type == "rdma": device_name = config.protocol_rdma_device_names - server_name = config.instance_key + server_name = config.mooncake_spec_local_hostname logging.info( "=============== Mooncake store setup args ===============\n" + f"instance_key: {config.instance_key}\n" f"server_name: {server_name}\n" f"metadata_server: {config.mooncake_spec_metadata_server}\n" f"master_server_address: {config.mooncake_spec_master_server_address}\n" f"contribute_to_cluster_pool_size: {config.contribute_to_cluster_pool_size}\n" f"local_buffer_size: {config.mooncake_spec_local_buffer_size}\n" + f"enable_ssd_offload: {config.mooncake_spec_enable_ssd_offload}\n" + f"ssd_offload_path: {config.mooncake_spec_ssd_offload_path}\n" + f"skip_get_size_on_get: {config.mooncake_spec_skip_get_size_on_get}\n" f"protocol_type: {config.protocol_type}\n" f"device_name: {device_name}\n" f"==============================================================\n" @@ -165,7 +170,7 @@ def __init__(self, config: "FluxonKvClientConfig"): def setup_store() -> None: try: - retcode = self._store.setup( + setup_args = [ server_name, config.mooncake_spec_metadata_server, config.contribute_to_cluster_pool_size["dram"], # Use DRAM only @@ -173,7 +178,19 @@ def setup_store() -> None: config.protocol_type, device_name, config.mooncake_spec_master_server_address, - ) + ] + if ( + config.mooncake_spec_enable_ssd_offload + or config.mooncake_spec_ssd_offload_path + ): + setup_args.extend( + [ + None, + config.mooncake_spec_enable_ssd_offload, + config.mooncake_spec_ssd_offload_path, + ] + ) + retcode = self._store.setup(*setup_args) except Exception as e: # pragma: no cover - defensive retcode = -1 errmsg = f"init_mooncake setup raised exception: {e}" @@ -416,6 +433,7 @@ def put( max_calls=1, period=5, ) + reject_if_exists = bool(opts.reject_if_exists) if opts is not None else False if not self._initialized: return Result.new_error( @@ -429,7 +447,7 @@ def put( def put_operation(key: str, values: Tuple[bytes, ...]): """ - Put operations explicitly remove old data before writing new data. + Put a value once, removing old data first only when overwrite is enabled. Args: key(str): The input key. @@ -450,29 +468,30 @@ def debug_values(): return values_debug def try_put(key: str, values: Tuple[bytes, ...]) -> Result[OkNone, ApiError]: - """Try one force-delete-then-put Mooncake write attempt.""" + """Try one Mooncake write attempt.""" try: - with self._rwlock.read_lock(): - remove_retcode = self._store.remove(key, True) + if not reject_if_exists: + with self._rwlock.read_lock(): + remove_retcode = self._store.remove(key, True) - logging.debug( - f"[put_operation] Force remove retcode: {remove_retcode} for key: {key}" - ) - if remove_retcode != 0: - remove_error = try_new_error_from_mooncake( - remove_retcode, - f"Remove operation failed for key '{key}'", - key=key, + logging.debug( + f"[put_operation] Force remove retcode: {remove_retcode} for key: {key}" ) - if not isinstance(remove_error, KeyNotFoundError): - logging.warning( - "=============== Mooncake store delete-before-put failed ===============\n" - f"key: {key}\n" - f"values: {debug_values()}\n" - f"error: {remove_error}\n" - "==============================================================\n" + if remove_retcode != 0: + remove_error = try_new_error_from_mooncake( + remove_retcode, + f"Remove operation failed for key '{key}'", + key=key, ) - return Result.new_error(remove_error) + if not isinstance(remove_error, KeyNotFoundError): + logging.warning( + "=============== Mooncake store delete-before-put failed ===============\n" + f"key: {key}\n" + f"values: {debug_values()}\n" + f"error: {remove_error}\n" + "==============================================================\n" + ) + return Result.new_error(remove_error) with self._rwlock.read_lock(): if len(values) == 1: @@ -573,15 +592,28 @@ def try_get(key: str) -> Result[MemHolder, ApiError]: try: # https://github.com/kvcache-ai/Mooncake/blob/e475a369fe45d528135b2b318d7e9464e1846222/docs/source/mooncake-store-api/python-binding.md?plain=1#L682 with self._rwlock.read_lock(): - datasize = self._store.get_size(key) - if datasize < 0: - logging.warning(f"[get_operation] Get failed with retcode:{datasize}") + if self._skip_get_size_on_get: + data: Optional[bytes] = self._store.get(key) + if data is None: + return Result.new_error( + KeyNotFoundError( + f"Get failed for key '{key}': Mooncake returned no bytes", + details={"key": key}, + key=key, + ) + ) + mem_holder = SimpleMemHolder(data) + return Result.new_ok(mem_holder) + else: + datasize = self._store.get_size(key) + if datasize < 0: + logging.warning(f"[get_operation] Get failed with retcode:{datasize}") - return Result.new_error( - try_new_error_from_mooncake( - datasize, f"Get failed for key '{key}'")) - logging.debug("[get_operation] The key exists.") - data: Optional[bytes] = self._store.get(key) + return Result.new_error( + try_new_error_from_mooncake( + datasize, f"Get failed for key '{key}'")) + logging.debug("[get_operation] The key exists.") + data = self._store.get(key) # mooncake store always return bytes assert data is not None, "Data should have bytes after exists checking!" @@ -618,6 +650,47 @@ def get_blocking(self, key: str) -> Result[MemHolder, ApiError]: return Result.new_error(get_result.unwrap_error()) return get_result.unwrap().wait() + def get_buffer_blocking(self, key: str) -> Result[Any, ApiError]: + """Return Mooncake's lifetime-bound native host buffer for a key.""" + if not self._initialized: + return Result.new_error( + GeneralError( + message="Store not initialized when get_buffer_blocking(). Call setup() first." + ) + ) + + def try_get_buffer(buffer_key: str) -> Result[Any, ApiError]: + try: + with self._rwlock.read_lock(): + buffer_handle = self._store.get_buffer(buffer_key) + if buffer_handle is None: + return Result.new_error( + KeyNotFoundError( + f"Get failed for key '{buffer_key}': Mooncake returned no buffer", + details={"key": buffer_key}, + key=buffer_key, + ) + ) + return Result.new_ok(buffer_handle) + except Exception as exc: + return Result.new_error(exception_to_error(exc)) + + return self.retry_operation(try_get_buffer, key) + + def _benchmark_offload_rpc_read_count(self) -> int: + """Return Mooncake's cumulative peer offload-RPC read count.""" + if not self._initialized: + raise RuntimeError( + "Store not initialized when reading the benchmark offload counter" + ) + count = self._store.get_offload_rpc_read_count() + if isinstance(count, bool) or not isinstance(count, int) or count < 0: + raise TypeError( + "Mooncake get_offload_rpc_read_count() must return a non-negative int; " + f"got {count!r}" + ) + return int(count) + def get_size(self, key: str) -> Result[int, ApiError]: """ diff --git a/fluxon_py/tests/fluxon_fs_transfer_tikv_support.py b/fluxon_py/tests/fluxon_fs_transfer_tikv_support.py index 7e38d04..19f0139 100644 --- a/fluxon_py/tests/fluxon_fs_transfer_tikv_support.py +++ b/fluxon_py/tests/fluxon_fs_transfer_tikv_support.py @@ -52,9 +52,9 @@ SHM_TEST_DIR_BASE = Path("/dev/shm/fxfs_t") SHM_TEST_DIR_CLEANUP_LOCK_PATH = SHM_TEST_DIR_BASE / ".cleanup.lock" TRANSFER_REPORT_BASE = Path("/tmp/fluxon_fs_transfer_reports") -ETCD_HARNESS_WORK_ROOT = Path("/mnt/nvme0/fluxon_fs_transfer_etcd") +ETCD_HARNESS_WORK_ROOT = Path("/tmp/fluxon-example/fluxon_fs_transfer_etcd") ETCD_HARNESS_LOCK_NAME = ".lock" -TIKV_HARNESS_WORK_ROOT = Path("/mnt/nvme0/fluxon_fs_transfer_tikv") +TIKV_HARNESS_WORK_ROOT = Path("/tmp/fluxon-example/fluxon_fs_transfer_tikv") TIKV_HARNESS_LOCK_NAME = ".lock" TIKV_HARNESS_PD_LEASE_SECS = 60 TIKV_HARNESS_UNIFIED_READPOOL_MAX_THREADS = 4 diff --git a/fluxon_py/tests/test_api_chan_mpmc/test_create_reservation_contract.py b/fluxon_py/tests/test_api_chan_mpmc/test_create_reservation_contract.py new file mode 100644 index 0000000..84e1909 --- /dev/null +++ b/fluxon_py/tests/test_api_chan_mpmc/test_create_reservation_contract.py @@ -0,0 +1,620 @@ +#!/usr/bin/env python3 +"""Contract tests for MPMC sub-channel creation reservations.""" + +from __future__ import annotations + +import importlib +import json +import sys +import threading +import types +import unittest +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional +from unittest import mock + + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +# This contract exercises only Python coordination logic. Keep importing that +# logic independent from the native extension and replace all native entrypoints +# used by the module before loading it. +_pyo3_loader = importlib.import_module("fluxon_py.tool.pyo3") +_pyo3_loader._FLUXON_PYO3_MODULE_LAZY = types.SimpleNamespace( + MpscContext=object, + LeaseManagerHandle=object, + EtcdLock=object, +) + +from fluxon_py._api_ext_chan import mpmc # noqa: E402 + + +@dataclass(frozen=True) +class _Compare: + kind: str + key: str + expected: object + + +@dataclass(frozen=True) +class _Operation: + kind: str + key: str + value: Optional[bytes] = None + + +class _CompareOperand: + def __init__(self, kind: str, key: str) -> None: + self.kind = kind + self.key = key + + def __eq__(self, expected: object) -> _Compare: # type: ignore[override] + return _Compare(self.kind, self.key, expected) + + +class _FakeTransactions: + def create(self, key: str) -> _CompareOperand: + return _CompareOperand("create", key) + + def value(self, key: str) -> _CompareOperand: + return _CompareOperand("value", key) + + def put(self, key: str, value: bytes, _lease: object) -> _Operation: + return _Operation("put", key, value) + + def delete(self, key: str) -> _Operation: + return _Operation("delete", key) + + +class _FakeMetadata: + def __init__(self, key: str) -> None: + self.key = key.encode() + + +class _FakeLeaseInfo: + TTL = 60 + + +class _FakeEtcd: + def __init__(self) -> None: + self.values: Dict[str, bytes] = {} + self.transactions = _FakeTransactions() + self.fail_publish = False + self.raise_after_publish_commit = False + self.defer_publish_until_cleanup = False + self.pending_publish_ops: Optional[List[_Operation]] = None + self.defer_reservation_create_until_revoke = False + self.pending_reservation_ops: Optional[List[_Operation]] = None + self.raise_after_ready_claim_commit = False + self.publish_attempts = 0 + self.revoked_lease_ids: List[int] = [] + + def get(self, key: str): + value = self.values.get(key) + return value, _FakeMetadata(key) if value is not None else None + + def get_prefix(self, prefix: str): + return [ + (value, _FakeMetadata(key)) + for key, value in sorted(self.values.items()) + if key.startswith(prefix) + ] + + def get_lease_info(self, _lease_id: int) -> _FakeLeaseInfo: + return _FakeLeaseInfo() + + def delete(self, key: str) -> bool: + return self.values.pop(key, None) is not None + + def put(self, key: str, value: bytes) -> None: + self.values[key] = value + + def delete_prefix(self, prefix: str) -> int: + keys = [key for key in self.values if key.startswith(prefix)] + for key in keys: + del self.values[key] + return len(keys) + + def revoke_lease(self, lease_id: int) -> None: + self.revoked_lease_ids.append(lease_id) + self.pending_reservation_ops = None + + def transaction(self, *, compare, success, failure): + del failure + is_publish = any( + item.kind == "put" and item.key.endswith("/mpsc_channels") + for item in success + ) + if is_publish: + self.publish_attempts += 1 + if self.fail_publish: + return False, [] + if self.defer_publish_until_cleanup: + self.pending_publish_ops = list(success) + raise RuntimeError("injected delayed publish response loss") + + is_reservation_create = any( + item.kind == "put" and "/create_reservations/" in item.key + for item in success + ) + if is_reservation_create and self.defer_reservation_create_until_revoke: + self.pending_reservation_ops = list(success) + raise RuntimeError("injected delayed reservation create response loss") + + is_reservation_cleanup = any( + item.kind == "delete" and "/create_reservations/" in item.key + for item in success + ) + if is_reservation_cleanup and self.pending_publish_ops is not None: + self._apply_operations(self.pending_publish_ops) + self.pending_publish_ops = None + + matches = all(self._compare_matches(item) for item in compare) + if not matches: + return False, [] + self._apply_operations(success) + if is_publish and self.raise_after_publish_commit: + raise RuntimeError("injected response loss after publish commit") + is_ready_claim = ( + not is_publish + and any( + item.kind == "put" and "/mpmc_channels/ready/" in item.key + for item in success + ) + ) + if is_ready_claim and self.raise_after_ready_claim_commit: + raise RuntimeError("injected response loss after ready claim commit") + return True, [] + + def _apply_operations(self, operations: List[_Operation]) -> None: + for operation in operations: + if operation.kind == "put": + assert operation.value is not None + self.values[operation.key] = operation.value + elif operation.kind == "delete": + self.values.pop(operation.key, None) + else: # pragma: no cover - catches incomplete fake coverage + raise AssertionError(f"unsupported operation: {operation!r}") + + def _compare_matches(self, item: _Compare) -> bool: + if item.kind == "create": + return (0 if item.key not in self.values else 1) == item.expected + if item.kind == "value": + return self.values.get(item.key) == item.expected + raise AssertionError(f"unsupported comparison: {item!r}") + + +@dataclass(frozen=True) +class _LockCall: + endpoints: List[str] + name: str + ttl_seconds: int + timeout_seconds: float + + +class _FakeLock: + def __init__(self, tracker: "_LockTracker") -> None: + self.tracker = tracker + + def __enter__(self) -> "_FakeLock": + self.tracker.depth += 1 + self.tracker.max_depth = max(self.tracker.max_depth, self.tracker.depth) + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + del exc_type, exc_value, traceback + self.tracker.depth -= 1 + if self.tracker.depth < 0: # pragma: no cover - fake invariant + raise AssertionError("negative fake lock depth") + + +class _LockTracker: + def __init__(self) -> None: + self.calls: List[_LockCall] = [] + self.depth = 0 + self.max_depth = 0 + + def new_lock( + self, + endpoints: List[str], + name: str, + ttl_seconds: int, + timeout_seconds: float, + ) -> _FakeLock: + self.calls.append( + _LockCall( + list(endpoints), + name, + ttl_seconds, + timeout_seconds, + ) + ) + return _FakeLock(self) + + +class _ReleaseResult: + def unwrap(self) -> None: + return None + + +class _FakeMpscState: + def __init__(self, lock_tracker: _LockTracker) -> None: + self.lock_tracker = lock_tracker + self.constructor_lock_depths: List[int] = [] + self.instances: List[object] = [] + self.next_channel_id = 1000 + self.fail_constructor = False + self.on_construct: Optional[Callable[[], None]] = None + + +def _new_fake_mpsc_consumer_type(state: _FakeMpscState): + class _FakeMpscConsumer: + def __init__(self, *args: object, **kwargs: object) -> None: + del kwargs + state.constructor_lock_depths.append(state.lock_tracker.depth) + if state.fail_constructor: + raise RuntimeError("injected MPSC constructor failure") + + requested_id = args[1] + if requested_id is None: + state.next_channel_id += 1 + requested_id = str(state.next_channel_id) + self.chan_id = str(requested_id) + self._mpmc_ready_claimed = False + self.release_count = 0 + state.instances.append(self) + + callback = state.on_construct + state.on_construct = None + if callback is not None: + callback() + + def get_chan_id(self) -> str: + return self.chan_id + + def release_local_handle(self) -> _ReleaseResult: + self.release_count += 1 + return _ReleaseResult() + + return _FakeMpscConsumer + + +class _FakeLease: + def __init__(self, lease_id: int) -> None: + self.id = lease_id + + +def _new_channel(etcd: _FakeEtcd, member_id: int) -> mpmc.MPMCChannel: + channel = object.__new__(mpmc.MPMCChannel) + channel.mpmc_id = "7" + channel.etcd_client = etcd + channel._etcd_endpoints = ["127.0.0.1:2379"] + channel.mpmc_member_id = member_id + channel.mpmc_member_lease = _FakeLease(100 + member_id) + channel._lm_mpmc_member = None + channel.mpmc_global_lease = _FakeLease(1) + channel.payload_lease_id = 10 + channel.shutdown_ctl = types.SimpleNamespace(closed=False) + channel.ready_channels = [] + channel.unready_channels = [] + channel._ready_channels_lock = threading.Lock() + channel.new_ready_channels_callback = None + channel.remove_ready_channels_callback = None + return channel + + +def _add_active_consumers(etcd: _FakeEtcd, count: int) -> None: + prefix = mpmc._new_mpmc_role_key_prefix("7", mpmc.ChanRole.CONSUMER) + for member_id in range(1, count + 1): + etcd.values[f"{prefix}{member_id}"] = b"active" + + +def _reservation_keys(etcd: _FakeEtcd) -> List[str]: + prefix = mpmc._new_mpmc_create_reservations_prefix("7") + return sorted(key for key in etcd.values if key.startswith(prefix)) + + +class TestCreateReservationContract(unittest.TestCase): + def test_ready_claim_response_loss_reconciles_owned_key(self) -> None: + etcd = _FakeEtcd() + etcd.raise_after_ready_claim_commit = True + channel = _new_channel(etcd, member_id=11) + + result = channel.try_claim_ready_channel("91") + + self.assertTrue(result.is_ok()) + self.assertTrue(result.unwrap()) + ready_key = mpmc._new_mpmc_ready_channel_key("7", "91") + self.assertEqual(etcd.values[ready_key], b"11") + self.assertFalse(channel.shutdown_ctl.closed) + + def test_ambiguous_ready_claim_with_other_owner_revokes_member(self) -> None: + etcd = _FakeEtcd() + channel = _new_channel(etcd, member_id=11) + ready_key = mpmc._new_mpmc_ready_channel_key("7", "91") + etcd.values[ready_key] = b"22" + + with mock.patch.object( + etcd, + "transaction", + side_effect=RuntimeError("injected ambiguous ready claim"), + ): + result = channel.try_claim_ready_channel("91") + + self.assertFalse(result.is_ok()) + _ = result.unwrap_error() + self.assertTrue(channel.shutdown_ctl.closed) + self.assertEqual(etcd.revoked_lease_ids, [111]) + self.assertEqual(etcd.values[ready_key], b"22") + + def test_existing_unready_consumer_does_not_construct_create_lock(self) -> None: + etcd = _FakeEtcd() + channels_key = mpmc._new_mpmc_mpsc_channels_key("7") + etcd.values[channels_key] = json.dumps(["91"]).encode() + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.get_next_available_channel(object(), {}) + + self.assertTrue(result.is_ok()) + bound = result.unwrap() + self.assertEqual(bound.chan_id, "91") + self.assertTrue(bound._mpmc_ready_claimed) + self.assertEqual(locks.calls, []) + self.assertEqual(mpsc_state.constructor_lock_depths, [0]) + + def test_reservation_caps_concurrent_create_and_success_cleans_it(self) -> None: + etcd = _FakeEtcd() + _add_active_consumers(etcd, 1) + primary = _new_channel(etcd, member_id=11) + contender = _new_channel(etcd, member_id=12) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + contender_results: List[Any] = [] + + def try_contending_create() -> None: + contender_results.append( + contender.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + ) + + mpsc_state.on_construct = try_contending_create + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = primary.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertTrue(result.is_ok()) + _ = result.unwrap() + self.assertEqual(len(contender_results), 1) + self.assertFalse(contender_results[0].is_ok()) + self.assertEqual( + contender_results[0].unwrap_error().message, + mpmc._MPMC_CREATE_IN_PROGRESS_MESSAGE, + ) + self.assertEqual(mpsc_state.constructor_lock_depths, [0]) + self.assertEqual(len(mpsc_state.instances), 1) + self.assertEqual(_reservation_keys(etcd), []) + channels_value = etcd.values[mpmc._new_mpmc_mpsc_channels_key("7")] + self.assertEqual(json.loads(channels_value.decode()), ["1001"]) + self.assertEqual(locks.depth, 0) + + def test_waiter_retries_when_capacity_opens_with_other_reservations_alive(self) -> None: + etcd = _FakeEtcd() + channel = _new_channel(etcd, member_id=11) + expected_consumer = types.SimpleNamespace(chan_id="1001") + first_result = mpmc.Result.new_error( + mpmc.ChanCreateError(mpmc._MPMC_CREATE_IN_PROGRESS_MESSAGE) + ) + second_result = mpmc.Result.new_ok(expected_consumer) + + with mock.patch.object( + channel, + "_ensure_member_lease_alive", + return_value=mpmc.Result.new_ok(mpmc.OK_NONE), + ), mock.patch.object( + channel, + "_refresh_local_ready_state", + return_value=None, + ), mock.patch.object( + channel, + "_get_create_reservation_count", + return_value=1, + ), mock.patch.object( + channel, + "_get_active_consumer_count", + return_value=2, + ), mock.patch.object( + channel, + "try_create_mpsc_channel", + side_effect=[first_result, second_result], + ) as try_create: + result = channel.get_next_available_channel(object(), {}) + + self.assertTrue(result.is_ok()) + self.assertIs(result.unwrap(), expected_consumer) + self.assertEqual(try_create.call_count, 2) + + def test_constructor_failure_cleans_reservation(self) -> None: + etcd = _FakeEtcd() + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + mpsc_state.fail_constructor = True + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertFalse(result.is_ok()) + error = result.unwrap_error() + self.assertIn("injected MPSC constructor failure", error.message) + self.assertEqual(mpsc_state.constructor_lock_depths, [0]) + self.assertEqual(_reservation_keys(etcd), []) + self.assertNotIn(mpmc._new_mpmc_mpsc_channels_key("7"), etcd.values) + + def test_ambiguous_reservation_create_revokes_member_lease(self) -> None: + etcd = _FakeEtcd() + etcd.defer_reservation_create_until_revoke = True + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertFalse(result.is_ok()) + error = result.unwrap_error() + self.assertIn("Failed to reserve MPSC channel creation", error.message) + self.assertTrue(channel.shutdown_ctl.closed) + self.assertEqual(etcd.revoked_lease_ids, [111]) + self.assertIsNone(etcd.pending_reservation_ops) + self.assertEqual(_reservation_keys(etcd), []) + self.assertEqual(mpsc_state.instances, []) + + def test_publish_failure_releases_handle_and_cleans_reservation(self) -> None: + etcd = _FakeEtcd() + etcd.fail_publish = True + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + def seed_new_mpsc_metadata() -> None: + mpsc_id = mpsc_state.instances[-1].chan_id + etcd.values[mpmc._new_etcd_meta_key(mpsc_id)] = json.dumps( + {"global_long_lease_id": 501} + ).encode() + etcd.values[f"/channels/{mpsc_id}/consumer/consumer_1"] = b"member" + etcd.values[f"dist_id_allocator/channels/{mpsc_id}/consumers"] = b"1" + etcd.values[f"cluster_lease/id_allocator/channels/{mpsc_id}"] = b"501" + + mpsc_state.on_construct = seed_new_mpsc_metadata + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertFalse(result.is_ok()) + _ = result.unwrap_error() + self.assertEqual(etcd.publish_attempts, 1) + self.assertEqual(len(mpsc_state.instances), 1) + self.assertEqual(mpsc_state.instances[0].release_count, 1) + self.assertEqual(_reservation_keys(etcd), []) + self.assertNotIn(mpmc._new_mpmc_mpsc_channels_key("7"), etcd.values) + self.assertNotIn(mpmc._new_etcd_meta_key("1001"), etcd.values) + self.assertFalse(any(key.startswith("/channels/1001/") for key in etcd.values)) + self.assertFalse( + any( + key.startswith("dist_id_allocator/channels/1001/") + for key in etcd.values + ) + ) + self.assertEqual(etcd.values["/channels/aborted/1001"], b"1") + self.assertEqual(etcd.revoked_lease_ids, [501]) + + def test_committed_publish_response_loss_is_reconciled_as_success(self) -> None: + etcd = _FakeEtcd() + etcd.raise_after_publish_commit = True + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertTrue(result.is_ok()) + consumer = result.unwrap() + self.assertEqual(consumer.release_count, 0) + self.assertTrue(consumer._mpmc_ready_claimed) + self.assertEqual(_reservation_keys(etcd), []) + channels_value = etcd.values[mpmc._new_mpmc_mpsc_channels_key("7")] + self.assertEqual(json.loads(channels_value.decode()), ["1001"]) + ready_key = mpmc._new_mpmc_ready_channel_key("7", "1001") + self.assertEqual(etcd.values[ready_key], b"11") + + def test_late_publish_commit_before_cleanup_fence_is_not_aborted(self) -> None: + etcd = _FakeEtcd() + etcd.defer_publish_until_cleanup = True + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertTrue(result.is_ok()) + consumer = result.unwrap() + self.assertEqual(consumer.release_count, 0) + self.assertTrue(consumer._mpmc_ready_claimed) + self.assertIsNone(etcd.pending_publish_ops) + channels_value = etcd.values[mpmc._new_mpmc_mpsc_channels_key("7")] + self.assertEqual(json.loads(channels_value.decode()), ["1001"]) + ready_key = mpmc._new_mpmc_ready_channel_key("7", "1001") + self.assertEqual(etcd.values[ready_key], b"11") + + def test_successful_create_uses_30_second_timeout_for_both_lock_phases(self) -> None: + etcd = _FakeEtcd() + _add_active_consumers(etcd, 1) + channel = _new_channel(etcd, member_id=11) + locks = _LockTracker() + mpsc_state = _FakeMpscState(locks) + fake_consumer = _new_fake_mpsc_consumer_type(mpsc_state) + + with mock.patch.object(mpmc, "EtcdLock", locks.new_lock), mock.patch.object( + mpmc, "MPSCChanConsumer", fake_consumer + ): + result = channel.try_create_mpsc_channel( + object(), {}, mpmc.ChanRole.CONSUMER + ) + + self.assertTrue(result.is_ok()) + _ = result.unwrap() + self.assertEqual(mpmc.MPMC_CREATE_LOCK_TIMEOUT_SECONDS, 30.0) + self.assertEqual(len(locks.calls), 2) + self.assertEqual([call.timeout_seconds for call in locks.calls], [30.0, 30.0]) + self.assertEqual(mpsc_state.constructor_lock_depths, [0]) + self.assertEqual(_reservation_keys(etcd), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/fluxon_py/tests/test_config.py b/fluxon_py/tests/test_config.py index 6de5180..e243f3d 100644 --- a/fluxon_py/tests/test_config.py +++ b/fluxon_py/tests/test_config.py @@ -48,6 +48,8 @@ def _build_checks(selected_test_id: Optional[str]) -> List[Tuple[str, Callable[[ ("fluxonkv_sub_cluster_config", test_fluxonkv_sub_cluster_config), ("fluxonkv_owner_requires_sub_cluster", test_fluxonkv_owner_requires_sub_cluster), ("fluxonkv_owner_requires_large_file_paths", test_fluxonkv_owner_requires_large_file_paths), + ("fluxonkv_large_limit_size_contract", test_fluxonkv_large_limit_size_contract), + ("fluxonkv_size_string_contract", test_fluxonkv_size_string_contract), ("fluxonkv_external_forbids_large_file_paths", test_fluxonkv_external_forbids_large_file_paths), ("fluxonkv_p2p_relay_removed", test_fluxonkv_p2p_relay_removed), ("fluxon_client_config_yaml_shape", test_fluxon_client_config_yaml_shape), @@ -55,6 +57,7 @@ def _build_checks(selected_test_id: Optional[str]) -> List[Tuple[str, Callable[[ ("fluxonkv_runtime_defaults_are_internal", test_fluxonkv_runtime_defaults_are_internal), ("fluxonkv_removed_rdma_config_keys", test_fluxonkv_removed_rdma_config_keys), ("fluxonkv_test_spec_config", test_fluxonkv_test_spec_config), + ("mooncake_ssd_offload_config", test_mooncake_ssd_offload_config), ("fluxon_pyo3_import_authority", test_fluxon_pyo3_import_authority), ] if selected_test_id is None: @@ -310,6 +313,76 @@ def test_fluxonkv_owner_requires_large_file_paths(): print(f"❌ FAIL: test_fluxonkv_owner_requires_large_file_paths - {e}") +def test_fluxonkv_large_limit_size_contract(): + """Ensure KV SSD limits align one-to-one with large_file_paths.""" + try: + valid = _owner_fluxonkv_base_config(tag="large_limit_size_valid") + valid["fluxonkv_spec"]["large_limit_size"] = [1048576] + rendered = FluxonKvClientConfig(valid).to_fluxon_kv_client_config_yaml_str() + loaded = yaml.safe_load(rendered) + assert loaded["fluxonkv_spec"]["large_limit_size"] == [1048576] + + mismatch = _owner_fluxonkv_base_config(tag="large_limit_size_mismatch") + mismatch["fluxonkv_spec"]["large_file_paths"] = [ + "/tmp/kvcache_large/large_limit_size_mismatch_a", + "/tmp/kvcache_large/large_limit_size_mismatch_b", + ] + mismatch["fluxonkv_spec"]["large_limit_size"] = [1048576] + try: + FluxonKvClientConfig(mismatch) + print("❌ FAIL: test_fluxonkv_large_limit_size_contract - length mismatch should be rejected") + return + except ValueError: + pass + + too_small = _owner_fluxonkv_base_config(tag="large_limit_size_too_small") + too_small["fluxonkv_spec"]["large_limit_size"] = [1] + try: + FluxonKvClientConfig(too_small) + print("❌ FAIL: test_fluxonkv_large_limit_size_contract - value below 512 should be rejected") + return + except ValueError: + pass + + external = { + "instance_key": "test_external", + "contribute_to_cluster_pool_size": {"dram": 0, "vram": {}}, + "fluxonkv_spec": { + "cluster_name": "test_cluster", + "share_mem_path": "/tmp/kvcache_shared_memory/test", + "large_limit_size": [1048576], + }, + } + try: + FluxonKvClientConfig(external) + print("❌ FAIL: test_fluxonkv_large_limit_size_contract - external large_limit_size should be rejected") + return + except ValueError: + pass + + print("✅ PASS: test_fluxonkv_large_limit_size_contract") + except Exception as e: + print(f"❌ FAIL: test_fluxonkv_large_limit_size_contract - {e}") + + +def test_fluxonkv_size_string_contract(): + """Ensure memory contribution and KV SSD limits accept size strings.""" + try: + cfg = _owner_fluxonkv_base_config(tag="size_string") + cfg["contribute_to_cluster_pool_size"]["dram"] = "1.1GB" + cfg["contribute_to_cluster_pool_size"]["vram"] = {"0": "0.5GB"} + cfg["fluxonkv_spec"]["large_limit_size"] = ["512.9B"] + config = FluxonKvClientConfig(copy.deepcopy(cfg)) + loaded = yaml.safe_load(config.to_fluxon_kv_client_config_yaml_str()) + assert loaded["contribute_to_cluster_pool_size"]["dram"] == 70 * 16 * 1024 * 1024 + assert loaded["contribute_to_cluster_pool_size"]["vram"]["0"] == 512 * 1024 * 1024 + assert loaded["fluxonkv_spec"]["large_limit_size"] == [512] + + print("✅ PASS: test_fluxonkv_size_string_contract") + except Exception as e: + print(f"❌ FAIL: test_fluxonkv_size_string_contract - {e}") + + def test_fluxonkv_external_forbids_large_file_paths(): """Ensure zero-contribution external config cannot declare owner-only large_file_paths.""" try: @@ -504,6 +577,23 @@ def test_fluxonkv_test_spec_config(): assert loaded["test_spec_config"]["enable_iceoryx_logs"] is True assert "transport_mode" not in loaded["test_spec_config"] + foyer_backend = copy.deepcopy(base) + foyer_backend["test_spec_config"] = {"kv_ssd_storage_backend": "foyer"} + config = FluxonKvClientConfig(foyer_backend) + loaded = yaml.safe_load(config.to_fluxon_kv_client_config_yaml_str()) + assert loaded["test_spec_config"]["kv_ssd_storage_backend"] == "foyer" + + invalid_foyer_uring = copy.deepcopy(foyer_backend) + invalid_foyer_uring["test_spec_config"]["kv_ssd_uring_mode"] = "iovec" + try: + FluxonKvClientConfig(invalid_foyer_uring) + print( + "❌ FAIL: test_fluxonkv_test_spec_config - foyer with iovec should be rejected" + ) + return + except ValueError: + pass + closed_backend = copy.deepcopy(base) closed_backend["test_spec_config"]["transport_mode"] = "transfer_only" closed_backend["test_spec_config"]["rdma_device_names"] = ["mlx5_0"] @@ -636,6 +726,36 @@ def test_fluxonkv_test_spec_config(): print(f"❌ FAIL: test_fluxonkv_test_spec_config - {e}") +def test_mooncake_ssd_offload_config(): + """Test Mooncake SSD offload fields are accepted and preserved.""" + try: + cfg = config_dict() + cfg["mooncake_spec"]["enable_ssd_offload"] = True + cfg["mooncake_spec"]["ssd_offload_path"] = "/tmp/mooncake-offload" + cfg["mooncake_spec"]["skip_get_size_on_get"] = True + cfg["mooncake_spec"]["local_hostname"] = "192.0.2.10" + config = FluxonKvClientConfig(cfg) + assert config.mooncake_spec_enable_ssd_offload is True + assert config.mooncake_spec_ssd_offload_path == "/tmp/mooncake-offload" + assert config.mooncake_spec_skip_get_size_on_get is True + assert config.mooncake_spec_local_hostname == "192.0.2.10" + loaded = yaml.safe_load(config.to_yaml_str()) + assert loaded["mooncake_spec"]["enable_ssd_offload"] is True + assert loaded["mooncake_spec"]["ssd_offload_path"] == "/tmp/mooncake-offload" + assert loaded["mooncake_spec"]["skip_get_size_on_get"] is True + assert loaded["mooncake_spec"]["local_hostname"] == "192.0.2.10" + + default_cfg = config_dict() + default_config = FluxonKvClientConfig(default_cfg) + assert default_config.mooncake_spec_enable_ssd_offload is False + assert default_config.mooncake_spec_ssd_offload_path == "" + assert default_config.mooncake_spec_skip_get_size_on_get is False + assert default_config.mooncake_spec_local_hostname == "test_instance" + print("✅ PASS: test_mooncake_ssd_offload_config") + except Exception as e: + print(f"❌ FAIL: test_mooncake_ssd_offload_config - {e}") + + def config_dict(): return { "instance_key": "test_instance", @@ -718,12 +838,13 @@ def test_to_yaml_str_roundtrip(cfg: dict): def test_verification(cfg): """Test configuration verification for invalid values.""" try: - invalid_config = copy.deepcopy(cfg) - invalid_config["contribute_to_cluster_pool_size"]["dram"] = 16777217 - FluxonKvClientConfig(invalid_config) - print("❌ FAIL: Invalid dram size should be rejected") - except ValueError: - print("✅ PASS: Invalid dram size correctly rejected") + aligned_config = copy.deepcopy(cfg) + aligned_config["contribute_to_cluster_pool_size"]["dram"] = 16777217 + config = FluxonKvClientConfig(aligned_config) + assert config.contribute_to_cluster_pool_size["dram"] == 16777216 + print("✅ PASS: Unaligned dram size correctly aligned") + except Exception as e: + print(f"❌ FAIL: Unaligned dram size should be aligned - {e}") try: invalid_config = copy.deepcopy(cfg) @@ -814,12 +935,13 @@ def test_verification(cfg): print(f"❌ FAIL: Valid GPU configuration should be accepted - {e}") try: - invalid_gpu_config = copy.deepcopy(cfg) - invalid_gpu_config["contribute_to_cluster_pool_size"]["vram"] = {"0": 16777217} - FluxonKvClientConfig(invalid_gpu_config) - print("❌ FAIL: Invalid GPU configuration should be rejected") - except ValueError: - print("✅ PASS: Invalid GPU configuration correctly rejected") + aligned_gpu_config = copy.deepcopy(cfg) + aligned_gpu_config["contribute_to_cluster_pool_size"]["vram"] = {"0": 16777217} + config = FluxonKvClientConfig(aligned_gpu_config) + assert config.contribute_to_cluster_pool_size["vram"]["0"] == 16777216 + print("✅ PASS: Unaligned GPU configuration correctly aligned") + except Exception as e: + print(f"❌ FAIL: Unaligned GPU configuration should be aligned - {e}") if __name__ == "__main__": diff --git a/fluxon_py/tests/test_fluxon_fs_video.py b/fluxon_py/tests/test_fluxon_fs_video.py new file mode 100644 index 0000000..6e76b46 --- /dev/null +++ b/fluxon_py/tests/test_fluxon_fs_video.py @@ -0,0 +1,313 @@ +import sys +import unittest +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + + +def main() -> None: + unittest.main() + + +from fluxon_py.fluxon_fs.video import ( # noqa: E402 + FluxonFsVideoReadRequest, + FluxonFsVideoReader, + FluxonFsVideoReaderPool, +) + + +class _FakeNativeVideoReader: + def __init__(self) -> None: + self.calls = [] + self.closed = False + self.read_at_calls = 0 + self.remote_read_bytes = 0 + + def read_frames_numpy(self, indices): + self.calls.append(list(indices)) + self.read_at_calls += 1 + self.remote_read_bytes += 1024 * len(indices) + return ("frames", tuple(indices)) + + def stats(self): + return { + "read_at_calls": self.read_at_calls, + "remote_read_bytes": self.remote_read_bytes, + } + + def close(self) -> None: + self.closed = True + + +class _FakeArray: + def __init__(self, values) -> None: + self.values = tuple(values) + self.shape = (len(self.values),) + self.nbytes = len(self.values) + + def __getitem__(self, item): + selected = self.values[item] + if isinstance(selected, tuple): + return _FakeArray(selected) + return _FakeArray((selected,)) + + def __eq__(self, other) -> bool: + if not isinstance(other, _FakeArray): + return False + return self.values == other.values + + +class _FakeArrayNativeVideoReader(_FakeNativeVideoReader): + def read_frames_numpy(self, indices): + self.calls.append(list(indices)) + self.read_at_calls += 1 + self.remote_read_bytes += 1024 * len(indices) + return _FakeArray(indices) + + +class _FakeAgent: + def __init__(self) -> None: + self.calls = [] + self.natives: list[_FakeNativeVideoReader] = [] + + def open_video_reader( + self, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int, + request_identity, + ): + self.calls.append( + ( + export_name, + relpath, + height, + width, + num_threads, + request_identity, + ) + ) + native = _FakeNativeVideoReader() + self.natives.append(native) + return native + + +class _FakeArrayAgent(_FakeAgent): + def open_video_reader( + self, + export_name: str, + relpath: str, + height: int, + width: int, + num_threads: int, + request_identity, + ): + self.calls.append( + ( + export_name, + relpath, + height, + width, + num_threads, + request_identity, + ) + ) + native = _FakeArrayNativeVideoReader() + self.natives.append(native) + return native + + +class TestFluxonFsVideoReaderFacade(unittest.TestCase): + def test_open_passes_strong_arguments_to_native_agent(self) -> None: + agent = _FakeAgent() + + reader = FluxonFsVideoReader._open( + agent=agent, + export_name="videos", + relpath="a/b.mp4", + height=480, + width=832, + num_threads=8, + request_identity=("user", "pass"), + ) + + self.assertIsInstance(reader, FluxonFsVideoReader) + self.assertEqual( + agent.calls, + [("videos", "a/b.mp4", 480, 832, 8, ("user", "pass"))], + ) + + def test_read_frames_numpy_validates_indices_and_closes_native(self) -> None: + native = _FakeNativeVideoReader() + reader = FluxonFsVideoReader(native) + + out = reader.read_frames_numpy([2, 0, 2]) + self.assertEqual(out, ("frames", (2, 0, 2))) + self.assertEqual(native.calls, [[2, 0, 2]]) + self.assertEqual( + reader.stats(), + { + "read_at_calls": 1, + "remote_read_bytes": 3072, + }, + ) + + with self.assertRaisesRegex(ValueError, "non-negative"): + reader.read_frames_numpy([-1]) + with self.assertRaisesRegex(TypeError, "int values"): + reader.read_frames_numpy([1, "2"]) # type: ignore[list-item] + with self.assertRaisesRegex(TypeError, "int values"): + reader.read_frames_numpy([True]) # type: ignore[list-item] + + reader.close() + self.assertTrue(native.closed) + with self.assertRaisesRegex(RuntimeError, "closed"): + reader.read_frames_numpy([0]) + with self.assertRaisesRegex(RuntimeError, "closed"): + reader.stats() + + def test_open_rejects_invalid_arguments_before_native_call(self) -> None: + agent = _FakeAgent() + + with self.assertRaisesRegex(ValueError, "export_name must be non-empty"): + FluxonFsVideoReader._open( + agent=agent, + export_name="", + relpath="a.mp4", + height=480, + width=832, + num_threads=8, + request_identity=None, + ) + with self.assertRaisesRegex(ValueError, "height must be positive"): + FluxonFsVideoReader._open( + agent=agent, + export_name="videos", + relpath="a.mp4", + height=0, + width=832, + num_threads=8, + request_identity=None, + ) + with self.assertRaisesRegex(TypeError, "num_threads must be int"): + FluxonFsVideoReader._open( + agent=agent, + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=True, # type: ignore[arg-type] + request_identity=None, + ) + + self.assertEqual(agent.calls, []) + + +class TestFluxonFsVideoReaderPool(unittest.TestCase): + def test_pool_reuses_idle_reader_and_returns_per_read_stats(self) -> None: + agent = _FakeAgent() + pool = FluxonFsVideoReaderPool( + agent=agent, + request_identity=("user", "pass"), + max_readers=2, + ) + + first = pool.read_frames_numpy_with_stats( + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=2, + indices=[0, 1], + ) + second = pool.read_frames_numpy_with_stats( + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=2, + indices=[2, 3, 4], + ) + + self.assertFalse(first.reader_cache_hit) + self.assertTrue(second.reader_cache_hit) + self.assertEqual(len(agent.natives), 1) + self.assertEqual(first.reader_stats["read_at_calls"], 1) + self.assertEqual(first.reader_stats["remote_read_bytes"], 2048) + self.assertEqual(second.reader_stats["read_at_calls"], 1) + self.assertEqual(second.reader_stats["remote_read_bytes"], 3072) + self.assertEqual(pool.stats()["reader_cache_hits"], 1) + self.assertEqual(pool.stats()["reader_cache_misses"], 1) + + pool.close() + self.assertTrue(agent.natives[0].closed) + + def test_pool_evicts_lru_idle_reader(self) -> None: + agent = _FakeAgent() + pool = FluxonFsVideoReaderPool(agent=agent, request_identity=None, max_readers=1) + + pool.read_frames_numpy( + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=2, + indices=[0], + ) + first_native = agent.natives[0] + pool.read_frames_numpy( + export_name="videos", + relpath="b.mp4", + height=480, + width=832, + num_threads=2, + indices=[0], + ) + + self.assertTrue(first_native.closed) + self.assertEqual(pool.stats()["current_readers"], 1) + self.assertEqual(pool.stats()["evict_count"], 1) + pool.close() + + def test_pool_read_many_coalesces_same_reader_key(self) -> None: + agent = _FakeArrayAgent() + pool = FluxonFsVideoReaderPool(agent=agent, request_identity=None, max_readers=2) + + results = pool.read_many_numpy_with_stats( + [ + FluxonFsVideoReadRequest( + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=2, + indices=(0, 4), + ), + FluxonFsVideoReadRequest( + export_name="videos", + relpath="a.mp4", + height=480, + width=832, + num_threads=2, + indices=(2, 6, 8), + ), + ] + ) + + self.assertEqual(len(agent.natives), 1) + self.assertEqual(agent.natives[0].calls, [[0, 4, 2, 6, 8]]) + self.assertEqual(results[0].array, _FakeArray((0, 4))) + self.assertEqual(results[1].array, _FakeArray((2, 6, 8))) + self.assertEqual(sum(result.reader_stats["read_at_calls"] for result in results), 1) + self.assertEqual(sum(result.reader_stats["remote_read_bytes"] for result in results), 5120) + self.assertEqual(results[0].batch_stats["read_many_group_size"], 2) + self.assertEqual(results[1].batch_stats["read_many_group_frames"], 5) + pool.close() + + +if __name__ == "__main__": + main() diff --git a/fluxon_py/tests/test_mooncake_put_contract.py b/fluxon_py/tests/test_mooncake_put_contract.py new file mode 100644 index 0000000..0ad25e3 --- /dev/null +++ b/fluxon_py/tests/test_mooncake_put_contract.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import sys +import threading +import types +import unittest +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + + +if "mooncake.store" not in sys.modules: + mooncake_module = types.ModuleType("mooncake") + mooncake_store_module = types.ModuleType("mooncake.store") + + class _MooncakeDistributedStore: + pass + + mooncake_store_module.MooncakeDistributedStore = _MooncakeDistributedStore + mooncake_module.store = mooncake_store_module + sys.modules["mooncake"] = mooncake_module + sys.modules["mooncake.store"] = mooncake_store_module + + +from fluxon_py.api_error import KeyAlreadyExistsError +from fluxon_py.kvclient.kvclient_interface import PutOptionalArgs +from fluxon_py.kvclient.mooncake import MooncakeStore + + +class _ReadWriteLock: + def read_lock(self): + return nullcontext() + + def write_lock(self): + return nullcontext() + + +class _NativeStore: + def __init__(self, *, put_retcode: int = 0) -> None: + self.put_retcode = put_retcode + self.calls: list[tuple[object, ...]] = [] + + def remove(self, key: str, force: bool) -> int: + self.calls.append(("remove", key, force)) + return 0 + + def put(self, key: str, value: bytes) -> int: + self.calls.append(("put", key, value)) + return self.put_retcode + + +class TestMooncakePutContract(unittest.TestCase): + def _new_store(self, native_store: _NativeStore) -> tuple[MooncakeStore, list[int]]: + store = object.__new__(MooncakeStore) + store._initialized = True + store._store = native_store + store._rwlock = _ReadWriteLock() + store._renew_lock = threading.Lock() + store._thread_pool = ThreadPoolExecutor(max_workers=1) + renew_calls: list[int] = [] + + def _unexpected_renew(): + renew_calls.append(1) + raise AssertionError("Mooncake write-once rejection must not renew the store") + + store.renew_store = _unexpected_renew + self.addCleanup(store._thread_pool.shutdown, wait=True) + return store, renew_calls + + def test_reject_if_exists_put_skips_remove(self) -> None: + native_store = _NativeStore() + store, renew_calls = self._new_store(native_store) + + result = store.put_blocking( + "key", + {"payload": b"value"}, + opts=PutOptionalArgs(reject_if_exists=True), + ) + + self.assertTrue(result.is_ok()) + _ = result.unwrap() + self.assertEqual([call[0] for call in native_store.calls], ["put"]) + self.assertEqual(renew_calls, []) + + def test_overwrite_put_removes_before_put(self) -> None: + for opts in (None, PutOptionalArgs(reject_if_exists=False)): + with self.subTest(opts=opts): + native_store = _NativeStore() + store, renew_calls = self._new_store(native_store) + + result = store.put_blocking("key", {"payload": b"value"}, opts=opts) + + self.assertTrue(result.is_ok()) + _ = result.unwrap() + self.assertEqual( + [call[0] for call in native_store.calls], + ["remove", "put"], + ) + self.assertEqual(renew_calls, []) + + def test_reject_if_exists_native_already_exists_does_not_renew(self) -> None: + native_store = _NativeStore(put_retcode=-705) + store, renew_calls = self._new_store(native_store) + + result = store.put_blocking( + "key", + {"payload": b"value"}, + opts=PutOptionalArgs(reject_if_exists=True), + ) + + self.assertFalse(result.is_ok()) + self.assertIsInstance(result.unwrap_error(), KeyAlreadyExistsError) + self.assertEqual([call[0] for call in native_store.calls], ["put"]) + self.assertEqual(renew_calls, []) + + +if __name__ == "__main__": + raise SystemExit(unittest.main()) diff --git a/fluxon_py/tests/test_pyo3_etcd.py b/fluxon_py/tests/test_pyo3_etcd.py index 4c79451..afad0a7 100644 --- a/fluxon_py/tests/test_pyo3_etcd.py +++ b/fluxon_py/tests/test_pyo3_etcd.py @@ -124,6 +124,36 @@ def test_lock_exclusive_and_context_manager_release(self) -> None: context_lease_id = int(held_lock.lease_id) self._assert_lease_not_live(context_lease_id) + def test_lock_keeps_lease_alive_past_ttl(self) -> None: + lock_name = self.prefix + "long_critical_section_lock" + ttl_seconds = 2 + lock_a = self.fluxon_pyo3.EtcdLock( + [self.endpoint], lock_name, ttl_seconds, 1.0 + ) + lock_b = self.fluxon_pyo3.EtcdLock( + [self.endpoint], lock_name, ttl_seconds, 0.5 + ) + + try: + self.assertTrue(lock_a.acquire()) + lease_id = int(lock_a.lease_id) + + time.sleep(ttl_seconds + 1.5) + + self.assertGreater(self.client.lease_ttl(lease_id), 0) + self.assertFalse(lock_b.acquire()) + + self.assertTrue(lock_a.release()) + self.assertTrue(lock_b.acquire(1.0)) + finally: + for lock in (lock_b, lock_a): + if not lock.held: + continue + try: + lock.release() + except RuntimeError: + pass + if __name__ == "__main__": raise SystemExit(unittest.main()) diff --git a/fluxon_rs/Cargo.lock b/fluxon_rs/Cargo.lock index a4b0ecd..d5bb591 100644 --- a/fluxon_rs/Cargo.lock +++ b/fluxon_rs/Cargo.lock @@ -485,6 +485,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -575,6 +577,12 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" +[[package]] +name = "cmsketch" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7ee2cfacbd29706479902b06d75ad8f1362900836aa32799eabc7e004bfd854" + [[package]] name = "cobs" version = "0.3.0" @@ -615,6 +623,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "cpp_demangle" version = "0.4.5" @@ -960,6 +979,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastant" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" +dependencies = [ + "small_ctor", + "web-time", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -1034,7 +1063,6 @@ dependencies = [ "bitcode", "bytes", "crossbeam", - "etcd-client", "fluxon_commu_closed_sdk_consumer", "fluxon_commu_contract", "fluxon_framework", @@ -1232,11 +1260,13 @@ dependencies = [ "fluxon_framework_compiled", "fluxon_observability", "fluxon_util", + "foyer", "futures", "hex", "hyper 0.14.32", "iceoryx2", "iceoryx2-cal", + "io-uring", "kanal", "lazy_static", "libc", @@ -1322,7 +1352,6 @@ dependencies = [ "base64 0.21.7", "chrono", "clap", - "etcd-client", "fluxon_cli", "fluxon_framework", "fluxon_kv", @@ -1372,6 +1401,7 @@ dependencies = [ "anyhow", "async-trait", "bytes", + "cc", "crossbeam-channel", "etcd-client", "fluxon_cli", @@ -1389,6 +1419,7 @@ dependencies = [ "libc", "limit_thirdparty", "parking_lot", + "pkg-config", "pyo3", "serde", "serde_json", @@ -1447,6 +1478,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "foreign-types" version = "0.3.2" @@ -1471,6 +1508,130 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "foyer" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0abc0b87814989efa711f9becd9f26969820e2d3905db27d10969c4bd45890" +dependencies = [ + "anyhow", + "equivalent", + "foyer-common", + "foyer-memory", + "foyer-storage", + "foyer-tokio", + "futures-util", + "mea", + "mixtrics", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-common" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3db80d5dece93adb7ad709c84578794724a9cba342a7e566c3551c7ec626789" +dependencies = [ + "anyhow", + "bincode", + "bytes", + "cfg-if", + "foyer-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "serde", + "twox-hash", +] + +[[package]] +name = "foyer-intrusive-collections" +version = "0.10.0-dev" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4fee46bea69e0596130e3210e65d3424e0ac1e6df3bde6636304bdf1ca4a3b" +dependencies = [ + "memoffset", +] + +[[package]] +name = "foyer-memory" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db907f40a527ca2aa2f40a5f68b32ea58aa70f050cd233518e9ffd402cfba6ce" +dependencies = [ + "anyhow", + "bitflags 2.9.1", + "cmsketch", + "equivalent", + "foyer-common", + "foyer-intrusive-collections", + "foyer-tokio", + "futures-util", + "hashbrown 0.16.1", + "itertools 0.14.0", + "mea", + "mixtrics", + "parking_lot", + "paste", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-storage" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1983f1db3d0710e9c9d5fc116d9202dccd41a2d1e032572224f1aff5520aa958" +dependencies = [ + "allocator-api2", + "anyhow", + "bytes", + "core_affinity", + "equivalent", + "fastant", + "foyer-common", + "foyer-memory", + "foyer-tokio", + "fs4", + "futures-core", + "futures-util", + "hashbrown 0.16.1", + "io-uring", + "itertools 0.14.0", + "libc", + "lz4", + "mea", + "parking_lot", + "pin-project", + "rand 0.9.2", + "serde", + "tracing", + "twox-hash", + "zstd", +] + +[[package]] +name = "foyer-tokio" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6577b05a7ffad0db555aedf00bfe52af818220fc4c1c3a7a12520896fc38627" +dependencies = [ + "tokio", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix 1.0.7", + "windows-sys 0.59.0", +] + [[package]] name = "futures" version = "0.3.31" @@ -1618,11 +1779,22 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 5.3.0", "wasi 0.14.2+wasi-0.2.4", "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "gimli" version = "0.31.1" @@ -1711,7 +1883,18 @@ checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" dependencies = [ "allocator-api2", "equivalent", - "foldhash", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", ] [[package]] @@ -2395,6 +2578,17 @@ dependencies = [ "str_stack", ] +[[package]] +name = "io-uring" +version = "0.7.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" +dependencies = [ + "bitflags 2.9.1", + "cfg-if", + "libc", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -2446,12 +2640,40 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.77" @@ -2586,6 +2808,25 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "matchers" version = "0.1.0" @@ -2601,6 +2842,15 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +[[package]] +name = "mea" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2640d335e7273dacdcf51044026139b2e269c3bb0dfc3f8cb3496b85e3f6a42c" +dependencies = [ + "slab", +] + [[package]] name = "memchr" version = "2.7.5" @@ -2668,6 +2918,16 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "mixtrics" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c46b5adfb7a3ae4996d327a5bdc90e78fec025806dd312bdbe6f07a755e0ec9" +dependencies = [ + "itertools 0.15.0", + "parking_lot", +] + [[package]] name = "moka" version = "0.12.11" @@ -2831,6 +3091,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "object" version = "0.36.7" @@ -3498,6 +3768,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.7.3" @@ -4206,9 +4482,15 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "slab" -version = "0.4.10" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04dc19736151f35336d325007ac991178d504a119863a2fcb3758cdb5e52c50d" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "small_ctor" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" [[package]] name = "smallvec" @@ -5080,6 +5362,15 @@ dependencies = [ "utf-8", ] +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" +dependencies = [ + "rand 0.9.2", +] + [[package]] name = "typenum" version = "1.18.0" @@ -5902,3 +6193,31 @@ dependencies = [ "quote", "syn 2.0.104", ] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/fluxon_rs/fluxon_cli/src/cli_renderer.rs b/fluxon_rs/fluxon_cli/src/cli_renderer.rs index 71ae73d..c6cffbe 100644 --- a/fluxon_rs/fluxon_cli/src/cli_renderer.rs +++ b/fluxon_rs/fluxon_cli/src/cli_renderer.rs @@ -617,13 +617,13 @@ pub fn render_cluster(snapshot: &ClusterSnapshot) -> String { )); for d in &o.devices { lines.push(format!( - "|-- device: {} used={} cap={} util={}", - d.device, d.used, d.cap, d.util + "|-- resource={} device={} used={} cap={} util={}", + d.resource_kind, d.device, d.used, d.cap, d.util )); } } } - push_box(&mut out, "Owner segment usage", &lines); + push_box(&mut out, "Owner storage usage", &lines); out.push('\n'); } diff --git a/fluxon_rs/fluxon_cli/src/lib.rs b/fluxon_rs/fluxon_cli/src/lib.rs index 11c9fcd..7fa5a38 100644 --- a/fluxon_rs/fluxon_cli/src/lib.rs +++ b/fluxon_rs/fluxon_cli/src/lib.rs @@ -47,12 +47,19 @@ use fluxon_commu::{ RdmaLinkLayer, RdmaPhysState, RdmaPortSnapshot, RdmaPortState, cluster_member_base_prefix, cluster_member_ext_prefix, cluster_owner_rdma_control_prefix, scan_etcd_prefix_paginated, }; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use serde::Deserialize; use std::collections::{BTreeMap, HashMap}; use std::io::{IsTerminal, Read, Write}; const MQ_META_KEY_PREFIX: &str = "/channels/meta/"; +fn managed_etcd_backend(endpoints: &[String]) -> anyhow::Result { + Ok(ManagedEtcdClient::acquire(EtcdEndpointSet::new( + endpoints.to_vec(), + )?)) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] #[serde(rename_all = "lowercase")] enum MqChanRoleWire { @@ -231,7 +238,9 @@ pub async fn load_transfer_engine_edges_for_cluster( etcd_endpoints: &[String], cluster_name: &str, ) -> anyhow::Result> { - let mut etcd = EtcdClient::connect(etcd_endpoints.to_vec(), None) + let etcd_backend = managed_etcd_backend(etcd_endpoints)?; + let mut etcd = etcd_backend + .client() .await .with_context(|| "connect etcd (transfer_link scan)".to_string())?; let p2p_prefix = transfer_link_p2p_prefix(cluster_name); @@ -1185,7 +1194,9 @@ async fn build_fs_cluster_snapshot( cfg: &MonitorConfig, warnings: &mut Vec, ) -> anyhow::Result { - let mut etcd = EtcdClient::connect(cfg.etcd_endpoints.clone(), None) + let etcd_backend = managed_etcd_backend(&cfg.etcd_endpoints)?; + let mut etcd = etcd_backend + .client() .await .with_context(|| "connect etcd".to_string())?; @@ -1402,6 +1413,7 @@ async fn build_fs_cluster_snapshot( container_memory_limit_bytes: None, members: Vec::new(), segment_devices: Vec::new(), + storage_devices: Vec::new(), }); if is_p2p_relay { @@ -1469,7 +1481,9 @@ pub async fn build_cluster_snapshot_with_prom_query_time( ) -> anyhow::Result { let mut warnings: Vec = Vec::new(); if cfg.member_kind == MemberKind::Mq { - let mut etcd = EtcdClient::connect(cfg.etcd_endpoints.clone(), None) + let etcd_backend = managed_etcd_backend(&cfg.etcd_endpoints)?; + let mut etcd = etcd_backend + .client() .await .with_context(|| "connect etcd".to_string())?; let mut mq = build_mq_snapshot(cfg, &mut warnings, &mut etcd).await?; @@ -1587,7 +1601,9 @@ pub async fn build_cluster_snapshot_with_prom_query_time( if cfg.member_kind == MemberKind::Fs { return build_fs_cluster_snapshot(cfg, &mut warnings).await; } - let mut etcd = EtcdClient::connect(cfg.etcd_endpoints.clone(), None) + let etcd_backend = managed_etcd_backend(&cfg.etcd_endpoints)?; + let mut etcd = etcd_backend + .client() .await .with_context(|| "connect etcd".to_string())?; @@ -2128,6 +2144,7 @@ pub async fn build_cluster_snapshot_with_prom_query_time( container_memory_limit_bytes: None, members: Vec::new(), segment_devices: Vec::new(), + storage_devices: Vec::new(), }); if is_p2p_relay { node_entry.is_p2p_relay = true; @@ -2214,72 +2231,160 @@ pub async fn build_cluster_snapshot_with_prom_query_time( "missing segment metrics for owner_client: owner_id={} (expected Prom series: kvcache_segment_*_bytes{{node=\"{}\",device}})", owner_id, owner_id )); - continue; + } else { + let mut missing_cap: Vec = Vec::new(); + let mut missing_used: Vec = Vec::new(); + let mut cap_zero: Vec = Vec::new(); + let mut used_gt_cap: Vec = Vec::new(); + for device in devices { + let cap = prom_maps + .seg_capacity_bytes_by_node_device + .get(&(owner_id.clone(), device.clone())) + .copied(); + let used = prom_maps + .seg_used_bytes_by_node_device + .get(&(owner_id.clone(), device.clone())) + .copied(); + if cap.is_none() { + missing_cap.push(device.clone()); + } + if used.is_none() { + missing_used.push(device.clone()); + } + if let Some(cap) = cap { + if cap == 0.0 { + cap_zero.push(device.clone()); + } + } + if let (Some(used), Some(cap)) = (used, cap) { + if used > cap { + used_gt_cap.push(device.clone()); + } + } + n.segment_devices.push(crate::model::SegmentDeviceSnapshot { + device: device.clone(), + seg_capacity_bytes: cap, + seg_used_bytes: used, + }); + n.storage_devices.push(crate::model::StorageDeviceSnapshot { + resource_kind: "memory_segment".to_string(), + device, + capacity_bytes: cap, + used_bytes: used, + }); + } + n.segment_devices.sort_by(|a, b| a.device.cmp(&b.device)); + + if !missing_cap.is_empty() { + warnings.push(format!( + "segment metrics missing capacity series for owner_id={} devices={}", + owner_id, + missing_cap.join(",") + )); + } + if !missing_used.is_empty() { + warnings.push(format!( + "segment metrics missing used series for owner_id={} devices={}", + owner_id, + missing_used.join(",") + )); + } + if !cap_zero.is_empty() { + warnings.push(format!( + "segment metrics has cap=0 for owner_id={} devices={}", + owner_id, + cap_zero.join(",") + )); + } + if !used_gt_cap.is_empty() { + warnings.push(format!( + "segment metrics has used>cap for owner_id={} devices={}", + owner_id, + used_gt_cap.join(",") + )); + } } - let mut missing_cap: Vec = Vec::new(); - let mut missing_used: Vec = Vec::new(); - let mut cap_zero: Vec = Vec::new(); - let mut used_gt_cap: Vec = Vec::new(); - for device in devices { + let mut ssd_devices: std::collections::BTreeSet = std::collections::BTreeSet::new(); + for ((node, device), _v) in &prom_maps.kv_ssd_capacity_bytes_by_node_device { + if node == &owner_id { + ssd_devices.insert(device.clone()); + } + } + for ((node, device), _v) in &prom_maps.kv_ssd_used_bytes_by_node_device { + if node == &owner_id { + ssd_devices.insert(device.clone()); + } + } + + let mut ssd_missing_cap: Vec = Vec::new(); + let mut ssd_missing_used: Vec = Vec::new(); + let mut ssd_cap_zero: Vec = Vec::new(); + let mut ssd_used_gt_cap: Vec = Vec::new(); + for device in ssd_devices { let cap = prom_maps - .seg_capacity_bytes_by_node_device + .kv_ssd_capacity_bytes_by_node_device .get(&(owner_id.clone(), device.clone())) .copied(); let used = prom_maps - .seg_used_bytes_by_node_device + .kv_ssd_used_bytes_by_node_device .get(&(owner_id.clone(), device.clone())) .copied(); if cap.is_none() { - missing_cap.push(device.clone()); + ssd_missing_cap.push(device.clone()); } if used.is_none() { - missing_used.push(device.clone()); + ssd_missing_used.push(device.clone()); } if let Some(cap) = cap { if cap == 0.0 { - cap_zero.push(device.clone()); + ssd_cap_zero.push(device.clone()); } } if let (Some(used), Some(cap)) = (used, cap) { if used > cap { - used_gt_cap.push(device.clone()); + ssd_used_gt_cap.push(device.clone()); } } - n.segment_devices.push(crate::model::SegmentDeviceSnapshot { + n.storage_devices.push(crate::model::StorageDeviceSnapshot { + resource_kind: "kv_ssd".to_string(), device, - seg_capacity_bytes: cap, - seg_used_bytes: used, + capacity_bytes: cap, + used_bytes: used, }); } - n.segment_devices.sort_by(|a, b| a.device.cmp(&b.device)); + n.storage_devices + .sort_by(|a, b| match a.resource_kind.cmp(&b.resource_kind) { + std::cmp::Ordering::Equal => a.device.cmp(&b.device), + o => o, + }); - if !missing_cap.is_empty() { + if !ssd_missing_cap.is_empty() { warnings.push(format!( - "segment metrics missing capacity series for owner_id={} devices={}", + "kv ssd metrics missing capacity series for owner_id={} devices={}", owner_id, - missing_cap.join(",") + ssd_missing_cap.join(",") )); } - if !missing_used.is_empty() { + if !ssd_missing_used.is_empty() { warnings.push(format!( - "segment metrics missing used series for owner_id={} devices={}", + "kv ssd metrics missing used series for owner_id={} devices={}", owner_id, - missing_used.join(",") + ssd_missing_used.join(",") )); } - if !cap_zero.is_empty() { + if !ssd_cap_zero.is_empty() { warnings.push(format!( - "segment metrics has cap=0 for owner_id={} devices={}", + "kv ssd metrics has cap=0 for owner_id={} devices={}", owner_id, - cap_zero.join(",") + ssd_cap_zero.join(",") )); } - if !used_gt_cap.is_empty() { + if !ssd_used_gt_cap.is_empty() { warnings.push(format!( - "segment metrics has used>cap for owner_id={} devices={}", + "kv ssd metrics has used>cap for owner_id={} devices={}", owner_id, - used_gt_cap.join(",") + ssd_used_gt_cap.join(",") )); } } diff --git a/fluxon_rs/fluxon_cli/src/model.rs b/fluxon_rs/fluxon_cli/src/model.rs index 4b9e14a..7673cd4 100644 --- a/fluxon_rs/fluxon_cli/src/model.rs +++ b/fluxon_rs/fluxon_cli/src/model.rs @@ -1020,6 +1020,14 @@ pub struct SegmentDeviceSnapshot { pub seg_used_bytes: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StorageDeviceSnapshot { + pub resource_kind: String, + pub device: String, + pub capacity_bytes: Option, + pub used_bytes: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct NodeSnapshot { pub node_key: String, @@ -1035,6 +1043,8 @@ pub struct NodeSnapshot { pub container_memory_limit_bytes: Option, pub members: Vec, pub segment_devices: Vec, + #[serde(default)] + pub storage_devices: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -1226,6 +1236,7 @@ pub struct OwnerSegmentUsageViewModel { #[derive(Debug, Clone)] pub struct OwnerSegmentDeviceUsageViewModel { + pub resource_kind: String, pub device: String, pub used: String, pub cap: String, @@ -1554,6 +1565,10 @@ pub struct MemberTableRowView { pub seg_used_sort: String, pub seg_cap_text: String, pub seg_cap_sort: String, + pub ssd_used_text: String, + pub ssd_used_sort: String, + pub ssd_cap_text: String, + pub ssd_cap_sort: String, } #[derive(Debug, Clone)] @@ -1567,6 +1582,7 @@ pub struct OwnerSegmentOwnerRowView { #[derive(Debug, Clone)] pub struct OwnerSegmentDeviceRowView { pub owner_id: String, + pub resource_kind: String, pub device: String, pub used: String, pub cap: String, @@ -1732,26 +1748,30 @@ pub fn build_cluster_view_model(snapshot: &ClusterSnapshot) -> ClusterViewModel }; let total_used_bytes = - sum_opt_f64_complete(n.segment_devices.iter().map(|d| d.seg_used_bytes)); + sum_opt_f64_complete(n.storage_devices.iter().map(|d| d.used_bytes)); let total_cap_bytes = - sum_opt_f64_complete(n.segment_devices.iter().map(|d| d.seg_capacity_bytes)); + sum_opt_f64_complete(n.storage_devices.iter().map(|d| d.capacity_bytes)); let (total_used, _total_used_status) = fmt_bytes_auto(total_used_bytes, false); let (total_cap, _total_cap_status) = fmt_bytes_auto(total_cap_bytes, false); let total_util = fmt_util_percent(total_used_bytes, total_cap_bytes); let mut devices: Vec = - Vec::with_capacity(n.segment_devices.len()); - for d in &n.segment_devices { - let (used, _used_status) = fmt_bytes_auto(d.seg_used_bytes, false); - let (cap, _cap_status) = fmt_bytes_auto(d.seg_capacity_bytes, false); + Vec::with_capacity(n.storage_devices.len()); + for d in &n.storage_devices { + let (used, _used_status) = fmt_bytes_auto(d.used_bytes, false); + let (cap, _cap_status) = fmt_bytes_auto(d.capacity_bytes, false); devices.push(OwnerSegmentDeviceUsageViewModel { + resource_kind: d.resource_kind.clone(), device: d.device.clone(), used, cap, - util: fmt_util_percent(d.seg_used_bytes, d.seg_capacity_bytes), + util: fmt_util_percent(d.used_bytes, d.capacity_bytes), }); } - devices.sort_by(|a, b| a.device.cmp(&b.device)); + devices.sort_by(|a, b| match a.resource_kind.cmp(&b.resource_kind) { + std::cmp::Ordering::Equal => a.device.cmp(&b.device), + o => o, + }); owners.push(OwnerSegmentUsageViewModel { owner_id, @@ -1847,6 +1867,28 @@ pub fn build_member_table_rows(snapshot: &ClusterSnapshot) -> Vec Option { + let mut sum = 0f64; + let mut seen = false; + for d in devices.iter().filter(|d| d.resource_kind == resource_kind) { + let value = if use_used_bytes { + d.used_bytes + } else { + d.capacity_bytes + }; + let Some(value) = value else { + return None; + }; + sum += value; + seen = true; + } + if seen { Some(sum) } else { None } + } + let visible_roles = snapshot.visible_member_roles.as_ref(); let mut rows: Vec = Vec::new(); for node in &snapshot.nodes { @@ -1882,6 +1924,18 @@ pub fn build_member_table_rows(snapshot: &ClusterSnapshot) -> Vec Vec OwnerSegmentTab continue; }; - let total_used_bytes = - sum_opt_f64_complete(n.segment_devices.iter().map(|d| d.seg_used_bytes)); + let total_used_bytes = sum_opt_f64_complete(n.storage_devices.iter().map(|d| d.used_bytes)); let total_cap_bytes = - sum_opt_f64_complete(n.segment_devices.iter().map(|d| d.seg_capacity_bytes)); + sum_opt_f64_complete(n.storage_devices.iter().map(|d| d.capacity_bytes)); let (total_used, _total_used_status) = fmt_bytes_auto(total_used_bytes, false); let (total_cap, _total_cap_status) = fmt_bytes_auto(total_cap_bytes, false); let total_util = fmt_util_percent(total_used_bytes, total_cap_bytes); @@ -2068,22 +2125,26 @@ pub fn build_owner_segment_tables(snapshot: &ClusterSnapshot) -> OwnerSegmentTab total_util, }); - for d in &n.segment_devices { - let (used, _used_status) = fmt_bytes_auto(d.seg_used_bytes, false); - let (cap, _cap_status) = fmt_bytes_auto(d.seg_capacity_bytes, false); + for d in &n.storage_devices { + let (used, _used_status) = fmt_bytes_auto(d.used_bytes, false); + let (cap, _cap_status) = fmt_bytes_auto(d.capacity_bytes, false); device_rows.push(OwnerSegmentDeviceRowView { owner_id: owner_id.clone(), + resource_kind: d.resource_kind.clone(), device: d.device.clone(), used, cap, - util: fmt_util_percent(d.seg_used_bytes, d.seg_capacity_bytes), + util: fmt_util_percent(d.used_bytes, d.capacity_bytes), }); } } owner_rows.sort_by(|a, b| a.owner_id.cmp(&b.owner_id)); device_rows.sort_by(|a, b| match a.owner_id.cmp(&b.owner_id) { - std::cmp::Ordering::Equal => a.device.cmp(&b.device), + std::cmp::Ordering::Equal => match a.resource_kind.cmp(&b.resource_kind) { + std::cmp::Ordering::Equal => a.device.cmp(&b.device), + o => o, + }, o => o, }); @@ -2092,3 +2153,152 @@ pub fn build_owner_segment_tables(snapshot: &ClusterSnapshot) -> OwnerSegmentTab device_rows, } } + +#[cfg(test)] +mod tests { + use super::*; + + fn owner_member(member_id: &str) -> MemberSnapshot { + MemberSnapshot { + member_id: member_id.to_string(), + role: MemberRole::OwnerClient, + node_start_time: 0, + is_side_transfer_worker: false, + hostname: None, + accessible_ip: None, + shared_mem_dir: None, + p2p_listen_port: None, + is_p2p_relay: false, + rdma_runtime_reported: false, + rdma_probe_error: None, + rdma_devices: Vec::new(), + rdma_ports: Vec::new(), + rdma_transfer_engine: None, + process_pid: None, + process_cmd: None, + sub_cluster: None, + product_uuid: None, + node_cpu_usage_percent: None, + node_cpu_logical_cores: None, + node_memory_usage_bytes: None, + node_memory_total_bytes: None, + container_memory_usage_bytes: None, + container_memory_limit_bytes: None, + process_resident_memory_bytes: None, + process_cpu_usage_percent: None, + tokio_num_workers: None, + tokio_alive_tasks: None, + tokio_global_queue_depth: None, + tokio_busy_percent: None, + tokio_max_worker_busy_percent: None, + tokio_park_unpark_rate_hz: None, + process_net_tx_mbps: None, + process_net_rx_mbps: None, + kv_put_rps: None, + kv_get_rps: None, + kv_put_bps: None, + kv_get_bps: None, + kv_put_latency_mean_us: None, + kv_put_latency_p95_us: None, + kv_put_latency_p99_us: None, + kv_get_latency_mean_us: None, + kv_get_latency_p95_us: None, + kv_get_latency_p99_us: None, + seg_capacity_bytes: None, + seg_used_bytes: None, + fs_read_rps: None, + fs_write_rps: None, + } + } + + fn minimal_kv_snapshot(node: NodeSnapshot) -> ClusterSnapshot { + ClusterSnapshot { + cluster_name: "test".to_string(), + member_kind: MemberKind::Kv, + etcd_endpoints: Vec::new(), + prometheus_base_url: "http://127.0.0.1:9090".to_string(), + warnings: Vec::new(), + visible_member_roles: None, + master_id: None, + master_network: None, + transfer_engine_edges: Vec::new(), + kv_peer_network: Vec::new(), + rdma_netdev_network: Vec::new(), + fs_mount_fs: Vec::new(), + shm_files: Vec::new(), + fs_export_registry: Vec::new(), + fs_mount_registry: Vec::new(), + kv_topology_owner_external_max: Vec::new(), + kv_topology_machine_external_max: Vec::new(), + kv_topology_sub_cluster_owner_owner_max: Vec::new(), + nodes: vec![node], + mq: None, + total_put_rps: None, + total_get_rps: None, + total_put_bps: None, + total_get_bps: None, + total_put_latency_mean_us: None, + total_put_latency_p95_us: None, + total_put_latency_p99_us: None, + total_get_latency_mean_us: None, + total_get_latency_p95_us: None, + total_get_latency_p99_us: None, + } + } + + #[test] + fn owner_storage_tables_include_memory_segment_and_kv_ssd() { + let snapshot = minimal_kv_snapshot(NodeSnapshot { + node_key: "node-a".to_string(), + hostname: None, + accessible_ip: None, + shared_mem_dir: None, + is_p2p_relay: false, + node_cpu_usage_percent: None, + node_cpu_logical_cores: None, + node_memory_usage_bytes: None, + node_memory_total_bytes: None, + container_memory_usage_bytes: None, + container_memory_limit_bytes: None, + members: vec![owner_member("owner-a")], + segment_devices: Vec::new(), + storage_devices: vec![ + StorageDeviceSnapshot { + resource_kind: "memory_segment".to_string(), + device: "seg0".to_string(), + capacity_bytes: Some(100.0), + used_bytes: Some(40.0), + }, + StorageDeviceSnapshot { + resource_kind: "kv_ssd".to_string(), + device: "dev:1:/ssd".to_string(), + capacity_bytes: Some(900.0), + used_bytes: Some(60.0), + }, + ], + }); + + let tables = build_owner_segment_tables(&snapshot); + + assert_eq!(tables.owner_rows.len(), 1); + assert_eq!(tables.owner_rows[0].owner_id, "owner-a"); + assert_eq!(tables.owner_rows[0].total_used, "100B"); + assert_eq!(tables.owner_rows[0].total_cap, "1.0KB"); + assert_eq!(tables.owner_rows[0].total_util, "10.0%"); + + let kinds = tables + .device_rows + .iter() + .map(|row| (row.resource_kind.as_str(), row.device.as_str())) + .collect::>(); + assert_eq!( + kinds, + vec![("kv_ssd", "dev:1:/ssd"), ("memory_segment", "seg0")] + ); + + let rows = build_member_table_rows(&snapshot); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].ssd_used_text, "60B"); + assert_eq!(rows[0].ssd_cap_text, "900B"); + } +} diff --git a/fluxon_rs/fluxon_cli/src/prom.rs b/fluxon_rs/fluxon_cli/src/prom.rs index 4406c50..c4d5d1c 100644 --- a/fluxon_rs/fluxon_cli/src/prom.rs +++ b/fluxon_rs/fluxon_cli/src/prom.rs @@ -280,6 +280,8 @@ pub struct PromSnapshotMaps { pub seg_capacity_bytes_by_node_device: HashMap<(String, String), f64>, pub seg_used_bytes_by_node_device: HashMap<(String, String), f64>, + pub kv_ssd_capacity_bytes_by_node_device: HashMap<(String, String), f64>, + pub kv_ssd_used_bytes_by_node_device: HashMap<(String, String), f64>, } impl PromSnapshotMaps { @@ -323,6 +325,8 @@ impl PromSnapshotMaps { get_latency_p99_us: HashMap::new(), seg_capacity_bytes_by_node_device: HashMap::new(), seg_used_bytes_by_node_device: HashMap::new(), + kv_ssd_capacity_bytes_by_node_device: HashMap::new(), + kv_ssd_used_bytes_by_node_device: HashMap::new(), } } } @@ -1405,5 +1409,24 @@ pub async fn collect_prom_snapshot( .await, ); + out.kv_ssd_capacity_bytes_by_node_device = take_node_device_metric( + &q( + prom, + warnings, + "kv_ssd_capacity_bytes_by_node_device", + fluxon_observability::keys::PROM_METRIC_KV_SSD_CAPACITY_BYTES, + ) + .await, + ); + out.kv_ssd_used_bytes_by_node_device = take_node_device_metric( + &q( + prom, + warnings, + "kv_ssd_used_bytes_by_node_device", + fluxon_observability::keys::PROM_METRIC_KV_SSD_USED_BYTES, + ) + .await, + ); + out } diff --git a/fluxon_rs/fluxon_cli/src/server.rs b/fluxon_rs/fluxon_cli/src/server.rs index ed6baec..026b186 100644 --- a/fluxon_rs/fluxon_cli/src/server.rs +++ b/fluxon_rs/fluxon_cli/src/server.rs @@ -24,6 +24,7 @@ use std::pin::Pin; use std::sync::Arc; use tokio::sync::{RwLock, watch}; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use fluxon_util::{ FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2, fluxon_cli_proxy_desc_etcd_key_v2, fluxon_cli_proxy_desc_etcd_service_prefix_v2, @@ -358,6 +359,7 @@ fn parse_member_kind(s: &str) -> Option { #[derive(Clone)] struct AppState { cfg: Arc, + etcd_backend: ManagedEtcdClient, log_schema_cache: Arc, proxy_client: hyper::Client, Body>, registered_panel_proxy_backend: Option, @@ -1161,7 +1163,7 @@ async fn index(State(st): State>, req: Request) -> Response } async fn landing(State(st): State>) -> Response { - let mut etcd = match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + let mut etcd = match st.etcd_backend.client().await { Ok(c) => c, Err(e) => { return text_response( @@ -1184,7 +1186,7 @@ async fn landing(State(st): State>) -> Response { } async fn api_clusters(State(st): State>) -> Response { - let mut etcd = match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + let mut etcd = match st.etcd_backend.client().await { Ok(c) => c, Err(e) => { return text_response( @@ -1329,7 +1331,7 @@ async fn topology_page( } let prefix = format!("/fluxon_fs_mount_registry/{}/mounts/", cluster_name); - match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + match st.etcd_backend.client().await { Ok(mut etcd) => { let mut mounts: Vec = Vec::new(); @@ -1416,7 +1418,7 @@ async fn topology_page( } let key = format!("/fluxon_fs_export_registry/{}/snapshot", cluster_name); - match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + match st.etcd_backend.client().await { Ok(mut etcd) => match etcd.get(key.clone(), None).await { Ok(resp) => { let mut exports: Vec = @@ -1540,7 +1542,7 @@ async fn ops_fluxon_kv_update_rdma_devices( } }; - let mut etcd = match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + let mut etcd = match st.etcd_backend.client().await { Ok(client) => client, Err(err) => { return text_response( @@ -1767,7 +1769,7 @@ async fn cli( Query(q): Query, ) -> Response { let Some(cluster_name) = q.cluster_name.as_ref() else { - let mut etcd = match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + let mut etcd = match st.etcd_backend.client().await { Ok(c) => c, Err(e) => { return text_response( @@ -2917,7 +2919,7 @@ async fn proxy_registered_service_impl( .map(|s| s.to_string()) .unwrap_or_else(|| "".to_string()); - let mut etcd = match EtcdClient::connect(st.cfg.etcd_endpoints.clone(), None).await { + let mut etcd = match st.etcd_backend.client().await { Ok(c) => c, Err(e) => { return text_response( @@ -3353,11 +3355,14 @@ pub async fn serve_http_from_tcp( listener: std::net::TcpListener, registered_panel_proxy_backend: Option, ) -> anyhow::Result<()> { + let etcd_backend = + ManagedEtcdClient::acquire(EtcdEndpointSet::new(cfg.etcd_endpoints.clone())?); let cfg = Arc::new(cfg); let log_schema_cache = Arc::new(LogSchemaCache::new()); let proxy_client = new_proxy_client(); let app = build_router(Arc::new(AppState { cfg, + etcd_backend, log_schema_cache, proxy_client, registered_panel_proxy_backend, @@ -3379,6 +3384,8 @@ pub async fn serve_http_with_shutdown_from_tcp( where F: std::future::Future + Send + 'static, { + let etcd_backend = + ManagedEtcdClient::acquire(EtcdEndpointSet::new(cfg.etcd_endpoints.clone())?); let cfg = Arc::new(cfg); let log_schema_cache = Arc::new(LogSchemaCache::new()); let proxy_client = new_proxy_client(); @@ -3390,6 +3397,7 @@ where let app = build_router(Arc::new(AppState { cfg, + etcd_backend, log_schema_cache, proxy_client, registered_panel_proxy_backend, diff --git a/fluxon_rs/fluxon_cli/src/web_renderer.rs b/fluxon_rs/fluxon_cli/src/web_renderer.rs index ec16ad7..c98b253 100644 --- a/fluxon_rs/fluxon_cli/src/web_renderer.rs +++ b/fluxon_rs/fluxon_cli/src/web_renderer.rs @@ -3167,6 +3167,7 @@ mod tests { test_member("ext-a", MemberRole::ExternalClient, false), ], segment_devices: Vec::new(), + storage_devices: Vec::new(), }], mq: None, total_put_rps: None, diff --git a/fluxon_rs/fluxon_cli/templates/monitor_table.html b/fluxon_rs/fluxon_cli/templates/monitor_table.html index f804278..e4d2024 100644 --- a/fluxon_rs/fluxon_cli/templates/monitor_table.html +++ b/fluxon_rs/fluxon_cli/templates/monitor_table.html @@ -350,6 +350,8 @@

Cluster {{ header.cluster_name }} get_avg + + {% endif %} @@ -373,6 +375,8 @@

Cluster {{ header.cluster_name }} get_avg + + @@ -530,6 +542,8 @@

Cluster {{ header.cluster_name }} + + @@ -557,6 +571,8 @@

Cluster {{ header.cluster_name }} Cluster {{ header.cluster_name }} + data-sort-seg-cap="{{ row.seg_cap_sort }}" + data-sort-ssd-used="{{ row.ssd_used_sort }}" + data-sort-ssd-cap="{{ row.ssd_cap_sort }}"> {{ row.node_key }} {{ row.member_id }} log @@ -606,6 +624,8 @@

Cluster {{ header.cluster_name }} {{ row.get_avg_text }} {{ row.seg_used_text }} {{ row.seg_cap_text }} + {{ row.ssd_used_text }} + {{ row.ssd_cap_text }} {% endfor %} @@ -614,7 +634,7 @@

Cluster {{ header.cluster_name }} - Owner segment usage + Owner storage usage
@@ -637,12 +657,13 @@

Cluster {{ header.cluster_name }} Device breakdown +
Resource breakdown

+ @@ -653,6 +674,7 @@

Cluster {{ header.cluster_name }} {{ d.owner_id }} +

@@ -1011,12 +1033,12 @@

Cluster {{ header.cluster_name }} {{ header.cluster_name }} {{ header.cluster_name }} {{ header.cluster_name }} {{ header.cluster_name }} {{ header.cluster_name }} , + _etcd_backend: ManagedEtcdClient, observe_handle: std::sync::OnceLock, ipc_bandwidth_attributor_handle: std::sync::OnceLock, self_info: Arc>, @@ -277,8 +278,21 @@ impl ClosedClusterManagerRuntime { .map_err(cluster_manager_closed_sdk_error)?; let (event_tx, _) = abroadcast::channel(100); let (self_rdma_resolved_tx, _) = watch::channel(current_rdma); + let endpoint_set = EtcdEndpointSet::new(arg.etcd_endpoints.clone()).map_err(|error| { + crate::ClusterError::EtcdConnection { + endpoints: arg.etcd_endpoints.clone(), + error: error.to_string(), + } + })?; + let etcd_backend = ManagedEtcdClient::acquire(endpoint_set); let transfer_link_p2p_snapshot_source = CommuTransferLinkP2pSnapshotSource::new( - connect_transfer_link_client(&arg.etcd_endpoints).await?, + etcd_backend + .client() + .await + .map_err(|error| crate::ClusterError::EtcdConnection { + endpoints: arg.etcd_endpoints.clone(), + error: error.to_string(), + })?, format!("/{}/transfer_link/p2p", cluster_name), ); let transfer_link_writer = make_runtime_transfer_link_writer(handle); @@ -287,6 +301,7 @@ impl ClosedClusterManagerRuntime { self_member_id, cluster_name, etcd_endpoints: arg.etcd_endpoints, + _etcd_backend: etcd_backend, observe_handle: std::sync::OnceLock::new(), ipc_bandwidth_attributor_handle: std::sync::OnceLock::new(), self_info: Arc::new(RwLock::new(self_info)), @@ -303,12 +318,11 @@ impl ClosedClusterManagerRuntime { transfer_link_writer, transfer_link_p2p_snapshot_source, }; - runtime.spawn_event_mirror(); runtime.spawn_self_rdma_mirror(); Ok(runtime) } - fn spawn_event_mirror(&self) { + fn spawn_event_mirror(&self, stream_handle: ClosedRuntimeHandle) { let handle = self.handle; let event_tx = self.event_tx.clone(); let self_member_id = self.self_member_id.clone(); @@ -316,9 +330,6 @@ impl ClosedClusterManagerRuntime { let prev_members = Arc::clone(&self.prev_members); let self_info = Arc::clone(&self.self_info); tokio::spawn(async move { - let Ok(stream_handle) = subscribe_cluster_manager_events(handle).await else { - return; - }; loop { match recv_cluster_event_stream(stream_handle).await { Ok(ClosedRuntimeClusterEventStreamItem::Event(event)) => { @@ -389,17 +400,6 @@ fn cluster_manager_closed_sdk_error( crate::ClusterError::Unknown(format!("closed sdk cluster-manager call failed: {error}")) } -async fn connect_transfer_link_client( - etcd_endpoints: &[String], -) -> crate::ClusterResult { - EtcdClient::connect(etcd_endpoints.to_vec(), None) - .await - .map_err(|error| crate::ClusterError::EtcdConnection { - endpoints: etcd_endpoints.to_vec(), - error: error.to_string(), - }) -} - fn make_runtime_transfer_link_writer( handle: ClosedRuntimeHandle, ) -> CommuTransferLinkEtcdWriterHandle { @@ -726,22 +726,44 @@ impl ClusterManager { ClosedRuntimeClusterManagerCall::Init2ForInitDag, ) .await?; - self.closed.is_watching.store(true, Ordering::Relaxed); - let self_info = closed_cluster_manager_member_call( + + // Subscribe before taking the initialization snapshot. Events that race with the + // snapshot stay buffered in the stream, so the local cache cannot permanently miss + // a member between the snapshot and subscription. + let stream_handle = subscribe_cluster_manager_events(self.closed.handle) + .await + .map_err(cluster_manager_closed_sdk_error)?; + let self_info = match closed_cluster_manager_member_call( self.closed.handle, ClosedRuntimeClusterManagerCall::GetSelfInfo, ) - .await?; + .await + { + Ok(self_info) => self_info, + Err(error) => { + let _ = drop_runtime_handle(stream_handle).await; + return Err(error); + } + }; refresh_closed_self_info_cache(&self.closed, self_info); - let members = closed_cluster_manager_members_call( + let members = match closed_cluster_manager_members_call( self.closed.handle, ClosedRuntimeClusterManagerCall::GetMembers, ) - .await?; + .await + { + Ok(members) => members, + Err(error) => { + let _ = drop_runtime_handle(stream_handle).await; + return Err(error); + } + }; *self.closed.members.write().expect("members cache poisoned") = members .into_iter() .map(|member| (member.id.to_string(), member)) .collect(); + self.closed.spawn_event_mirror(stream_handle); + self.closed.is_watching.store(true, Ordering::Relaxed); Ok(()) } } diff --git a/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs b/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs index 878e5c6..e5353a5 100644 --- a/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs +++ b/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs @@ -74,12 +74,10 @@ impl ClosedLocalSegmentLeaseRegistry { where G: Send + Sync + 'static, { - let boxed = self - .guards - .lock() - .await - .remove(&handle) - .ok_or_else(|| format!("closed sdk local segment lease handle {handle} not found"))?; + let boxed = + self.guards.lock().await.remove(&handle).ok_or_else(|| { + format!("closed sdk local segment lease handle {handle} not found") + })?; boxed.downcast::().map(|guard| *guard).map_err(|_| { format!( "closed sdk local segment lease handle {handle} has unexpected runtime guard type" @@ -461,7 +459,7 @@ impl ClientTransferEngineCore { len, seg_guard, ) - .await + .await } } @@ -482,7 +480,11 @@ impl ClientTransferEngineCore { let initial_local_segment_guard = match seg_guard { Some(guard) => Some(guard), None if runtime.supports_local_segment_transfer() => { - let local_addr = if peer_src_or_target { target_addr } else { src_addr }; + let local_addr = if peer_src_or_target { + target_addr + } else { + src_addr + }; match runtime.ensure_local_segment_guard(local_addr, None).await { Ok(guard) => Some(guard), Err(_) => None, diff --git a/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs b/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs index 6fab54e..caad34b 100644 --- a/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs +++ b/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs @@ -11,9 +11,9 @@ use fluxon_commu_contract::{ ClosedRuntimeCallRawObservedOutputView, ClosedRuntimeClusterEventStreamItem, ClosedRuntimeClusterManagerCall, ClosedRuntimeClusterManagerResponse, ClosedRuntimeClusterRdmaResolvedConfigStreamItem, ClosedRuntimeDesiredTransferPeer, - ClosedRuntimeDispatchRequestView, - ClosedRuntimeDispatchResponse, ClosedRuntimeDispatchTransportPolicy, ClosedRuntimeError, - ClosedRuntimeHandle, ClosedRuntimeHostCallbackHandle, ClosedRuntimeP2pCall, + ClosedRuntimeDispatchRequestView, ClosedRuntimeDispatchResponse, + ClosedRuntimeDispatchTransportPolicy, ClosedRuntimeError, ClosedRuntimeHandle, + ClosedRuntimeHostCallbackHandle, ClosedRuntimeP2pCall, ClosedRuntimeP2pCallRawObservedRequestView, ClosedRuntimeP2pResponse, ClosedRuntimeP2pSendResponseRawRequestView, ClosedRuntimePeerGen, ClosedRuntimeRawSlice, ClosedRuntimeRequest, ClosedRuntimeResponse, ClosedRuntimeTransferEngineCall, @@ -491,11 +491,15 @@ impl WireBodyPartsOwner { let (raw_lengths, raw_payload) = match raw_bytes.len() { 0 => (WireBodyRawLengths::Empty, WireBodyRawPayload::Empty), 1 => { - let part = raw_bytes.into_iter().next().expect("single raw part missing"); - let len = - u32::try_from(part.len()).map_err(|_| ClosedSdkConsumerError::RuntimeDecode { + let part = raw_bytes + .into_iter() + .next() + .expect("single raw part missing"); + let len = u32::try_from(part.len()).map_err(|_| { + ClosedSdkConsumerError::RuntimeDecode { detail: format!("wire raw part too large for u32 length: {}", part.len()), - })?; + } + })?; ( WireBodyRawLengths::Single([len]), WireBodyRawPayload::Single(part), @@ -849,8 +853,7 @@ fn decode_call_raw_observed_output_view( return Err(ClosedSdkConsumerError::RuntimeDecode { detail: format!( "closed SDK call_raw_observed serialize_part overflow: serialize_len={} full_len={}", - message_view.body.serialize_part.len, - message_view.body.full_body.len, + message_view.body.serialize_part.len, message_view.body.full_body.len, ), }); } @@ -860,21 +863,19 @@ fn decode_call_raw_observed_output_view( .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { detail: "closed SDK call_raw_observed raw_bytes length overflow".to_string(), })?; - let expected_full_len = - message_view - .body - .serialize_part - .len - .checked_add(raw_total) - .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { - detail: "closed SDK call_raw_observed body length overflow".to_string(), - })?; + let expected_full_len = message_view + .body + .serialize_part + .len + .checked_add(raw_total) + .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { + detail: "closed SDK call_raw_observed body length overflow".to_string(), + })?; if expected_full_len != message_view.body.full_body.len { return Err(ClosedSdkConsumerError::RuntimeDecode { detail: format!( "closed SDK call_raw_observed body length mismatch: expected={} full_len={}", - expected_full_len, - message_view.body.full_body.len, + expected_full_len, message_view.body.full_body.len, ), }); } @@ -923,9 +924,7 @@ fn decode_call_raw_observed_output_view( frame_recv_done_ts_us: message_view.local_observe.frame_recv_done_ts_us, dispatch_enqueued_ts_us: message_view.local_observe.dispatch_enqueued_ts_us, dispatch_started_ts_us: message_view.local_observe.dispatch_started_ts_us, - complete_pending_call_ts_us: message_view - .local_observe - .complete_pending_call_ts_us, + complete_pending_call_ts_us: message_view.local_observe.complete_pending_call_ts_us, }, }, observe: fluxon_commu_contract::ClosedRuntimeRpcCallTransportObserveTrace { @@ -1550,8 +1549,8 @@ async fn invoke_completion_async_with_keepalive( ) -> i32, ) -> Result<(i32, Bytes), ClosedSdkConsumerError> { let (sender, receiver) = tokio::sync::oneshot::channel::<(i32, Bytes)>(); - let user_data = Box::into_raw(Box::new(RuntimeCompletionState { sender, keepalive })) - .cast::(); + let user_data = + Box::into_raw(Box::new(RuntimeCompletionState { sender, keepalive })).cast::(); let submit_status = submit(user_data, Some(runtime_completion_callback)); if submit_status != 0 { unsafe { @@ -2082,7 +2081,9 @@ pub async fn p2p_call_raw_observed( ) .await?; match status_code { - FLUXON_COMMU_CLOSED_RUNTIME_RESULT_OK => decode_call_raw_observed_output_view(payload.as_ref()), + FLUXON_COMMU_CLOSED_RUNTIME_RESULT_OK => { + decode_call_raw_observed_output_view(payload.as_ref()) + } FLUXON_COMMU_CLOSED_RUNTIME_RESULT_ERR => { let error = bitcode::decode::(payload.as_ref()).map_err( |decode_error| ClosedSdkConsumerError::RuntimeDecode { diff --git a/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs b/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs index 1738ade..ca54a71 100644 --- a/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs +++ b/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs @@ -9,28 +9,23 @@ use std::time::{Duration, Instant}; use fluxon_fs_core::config::{ FS_AGENT_TRANSFER_STREAM_CLOSE_RPC_PATH, FS_AGENT_TRANSFER_STREAM_NEXT_RPC_PATH, - FS_AGENT_TRANSFER_STREAM_OPEN_RPC_PATH, - FS_MASTER_TRANSFER_SCHEDULER_HEARTBEAT_RPC_PATH, FS_MASTER_TRANSFER_SCHEDULER_RESULT_RPC_PATH, - FluxonFsTransferBatchCollectInfoWire, FluxonFsTransferBatchKind, - FluxonFsTransferCollectInfoKind, FluxonFsTransferDispositionWire, - FluxonFsTransferFailedFileReasonKindWire, - FluxonFsTransferReadStreamCloseWire, FluxonFsTransferReadStreamNextResultWire, - FluxonFsTransferReadStreamNextWire, FluxonFsTransferReadStreamOpenResultWire, - FluxonFsTransferReadStreamOpenWire, - FluxonFsTransferSkipEntryKind, FluxonFsTransferSkipEntryWire, - FluxonFsTransferManifestEntryWire, FluxonFsTransferManifestWire, - FluxonFsTransferScanMode, - FluxonFsTransferScanEventAckWire, FluxonFsTransferScanEventKindWire, - FluxonFsTransferScanEventWire, FluxonFsTransferScanLaunchResultWire, + FS_AGENT_TRANSFER_STREAM_OPEN_RPC_PATH, FS_MASTER_TRANSFER_SCHEDULER_HEARTBEAT_RPC_PATH, + FS_MASTER_TRANSFER_SCHEDULER_RESULT_RPC_PATH, FluxonFsTransferBatchCollectInfoWire, + FluxonFsTransferBatchKind, FluxonFsTransferCollectInfoKind, FluxonFsTransferDispositionWire, + FluxonFsTransferFailedFileReasonKindWire, FluxonFsTransferManifestEntryWire, + FluxonFsTransferManifestWire, FluxonFsTransferReadStreamCloseWire, + FluxonFsTransferReadStreamNextResultWire, FluxonFsTransferReadStreamNextWire, + FluxonFsTransferReadStreamOpenResultWire, FluxonFsTransferReadStreamOpenWire, FluxonFsTransferScanAssignmentWire, FluxonFsTransferScanBatchWire, - FluxonFsTransferScanChildUnitWire, FluxonFsTransferScanFrontier, + FluxonFsTransferScanChildUnitWire, FluxonFsTransferScanEventAckWire, + FluxonFsTransferScanEventKindWire, FluxonFsTransferScanEventWire, FluxonFsTransferScanFrontier, FluxonFsTransferScanFrontierDirEntry, FluxonFsTransferScanFrontierEntry, - FluxonFsTransferScanResultWire, - FluxonFsTransferSymlinkNoticeEntryWire, FluxonFsTransferWorkerCollectInfoResultWire, - FluxonFsTransferWorkerAssignmentWire, FluxonFsTransferWorkerFileResultWire, - FluxonFsTransferWorkerFailedFileResultWire, - FluxonFsTransferWorkerHeartbeatResultWire, FluxonFsTransferWorkerHeartbeatTelemetryWire, - FluxonFsTransferWorkerHeartbeatWire, + FluxonFsTransferScanLaunchResultWire, FluxonFsTransferScanMode, FluxonFsTransferScanResultWire, + FluxonFsTransferSkipEntryKind, FluxonFsTransferSkipEntryWire, + FluxonFsTransferSymlinkNoticeEntryWire, FluxonFsTransferWorkerAssignmentWire, + FluxonFsTransferWorkerCollectInfoResultWire, FluxonFsTransferWorkerFailedFileResultWire, + FluxonFsTransferWorkerFileResultWire, FluxonFsTransferWorkerHeartbeatResultWire, + FluxonFsTransferWorkerHeartbeatTelemetryWire, FluxonFsTransferWorkerHeartbeatWire, FluxonFsTransferWorkerLaunchResultWire, FluxonFsTransferWorkerResultAckWire, FluxonFsTransferWorkerResultWire, FluxonFsTransferWorkerStopReasonWire, transfer_collect_info_output_relpath, @@ -39,8 +34,8 @@ use fluxon_fs_core::retry::{ BackoffConfig, DEFAULT_WARN_INTERVAL_SECS, WarnConfig, next_backoff, should_warn, }; use fluxon_kv::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError}; -use fluxon_kv::user_api::flat_dict::{FlatDict, FlatValue}; use fluxon_kv::user_api::FluxonUserApi; +use fluxon_kv::user_api::flat_dict::{FlatDict, FlatValue}; use parking_lot::{Condvar, Mutex}; use super::{ @@ -202,16 +197,13 @@ fn transfer_scan_session_state() -> &'static Mutex { TRANSFER_SCAN_SESSION_STATE.get_or_init(|| Mutex::new(TransferScanSessionState::default())) } -fn cleanup_expired_transfer_scan_sessions( - state: &mut TransferScanSessionState, - now_unix_ms: i64, -) { - state - .root_dir_listing_sessions - .retain(|_, session| session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms); - state - .subtree_streaming_sessions - .retain(|_, session| session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms); +fn cleanup_expired_transfer_scan_sessions(state: &mut TransferScanSessionState, now_unix_ms: i64) { + state.root_dir_listing_sessions.retain(|_, session| { + session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms + }); + state.subtree_streaming_sessions.retain(|_, session| { + session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms + }); } fn same_root_continuation_scan_unit( @@ -301,10 +293,7 @@ fn flush_pending_root_direct_files_batch( return Ok(None); } let batch = build_direct_files_only_batch_from_entries_with_batch_id( - direct_files_only_batch_id_for_partition( - assignment, - session.next_direct_files_batch_index, - ), + direct_files_only_batch_id_for_partition(assignment, session.next_direct_files_batch_index), assignment, assignment.root_relpath.clone(), std::mem::take(&mut session.pending_direct_files), @@ -313,7 +302,8 @@ fn flush_pending_root_direct_files_batch( )?; session.pending_direct_bytes = 0; session.next_direct_files_batch_index = session.next_direct_files_batch_index.saturating_add(1); - session.emitted_direct_files_batch_count = session.emitted_direct_files_batch_count.saturating_add(1); + session.emitted_direct_files_batch_count = + session.emitted_direct_files_batch_count.saturating_add(1); Ok(Some(batch)) } @@ -414,7 +404,8 @@ fn open_transfer_root_dir_listing_session( root_dir_abs: &str, assignment: &FluxonFsTransferScanAssignmentWire, ) -> Result, FlatDict> { - let dir_abs = safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; + let dir_abs = + safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; let read_dir = match retry_after_target_path_chmod( dir_abs.as_path(), "root_read_dir", @@ -458,7 +449,10 @@ fn take_transfer_root_dir_listing_session( let now_unix_ms = chrono::Utc::now().timestamp_millis(); let mut state = transfer_scan_session_state().lock(); cleanup_expired_transfer_scan_sessions(&mut state, now_unix_ms); - if let Some(mut session) = state.root_dir_listing_sessions.remove(assignment.scan_unit_id.as_str()) { + if let Some(mut session) = state + .root_dir_listing_sessions + .remove(assignment.scan_unit_id.as_str()) + { if session.job_id == assignment.job_id && session.scan_epoch == assignment.scan_epoch && session.root_relpath == assignment.root_relpath @@ -507,7 +501,8 @@ fn open_transfer_subtree_streaming_session( if is_relpath_skipped(&assignment.skip_entries, assignment.root_relpath.as_str()) { return Ok(None); } - let dir_abs = safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; + let dir_abs = + safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; let root_md = retry_after_target_path_chmod( Path::new(root_dir_abs), "subtree_stream_root_symlink_metadata", @@ -790,7 +785,8 @@ fn collect_transfer_root_dir_listing_slice( assignment: &FluxonFsTransferScanAssignmentWire, deadline: Option, ) -> Result { - let Some(mut session) = take_transfer_root_dir_listing_session(root_dir_abs, assignment)? else { + let Some(mut session) = take_transfer_root_dir_listing_session(root_dir_abs, assignment)? + else { return Ok(TransferRootDirListingOutcome::Finished( build_finished_empty_transfer_scan_result(assignment), )); @@ -848,7 +844,8 @@ fn collect_transfer_root_dir_listing_slice( }; scanned_entries = scanned_entries.saturating_add(1); let name = ent.file_name().to_string_lossy().to_string(); - let child_relpath = normalize_child_relpath(assignment.root_relpath.as_str(), name.as_str()); + let child_relpath = + normalize_child_relpath(assignment.root_relpath.as_str(), name.as_str()); if is_relpath_skipped(&assignment.skip_entries, child_relpath.as_str()) { continue; } @@ -899,10 +896,12 @@ fn collect_transfer_root_dir_listing_slice( let size = md.len().min(i64::MAX as u64) as i64; session.root_visible_entries = true; session.root_total_bytes = session.root_total_bytes.saturating_add(size); - session.pending_direct_files.push(FluxonFsTransferScanFrontierEntry { - relpath: child_relpath, - size, - }); + session + .pending_direct_files + .push(FluxonFsTransferScanFrontierEntry { + relpath: child_relpath, + size, + }); session.pending_direct_bytes = session.pending_direct_bytes.saturating_add(size); if should_flush_direct_batch( assignment.batch_ready_bytes, @@ -910,7 +909,9 @@ fn collect_transfer_root_dir_listing_slice( session.pending_direct_files.len(), session.pending_direct_empty_dirs.len(), ) { - if let Some(batch) = flush_pending_root_direct_files_batch(assignment, &mut session)? { + if let Some(batch) = + flush_pending_root_direct_files_batch(assignment, &mut session)? + { direct_files_only_batches.push(batch); } } @@ -933,14 +934,18 @@ fn collect_transfer_root_dir_listing_slice( session.pending_direct_files.len(), session.pending_direct_empty_dirs.len(), ) { - if let Some(batch) = flush_pending_root_direct_files_batch(assignment, &mut session)? { + if let Some(batch) = + flush_pending_root_direct_files_batch(assignment, &mut session)? + { direct_files_only_batches.push(batch); } } } else { - session.direct_dirs.push(FluxonFsTransferScanFrontierDirEntry { - relpath: child_relpath, - }); + session + .direct_dirs + .push(FluxonFsTransferScanFrontierDirEntry { + relpath: child_relpath, + }); } } } @@ -1244,12 +1249,14 @@ impl TransferWorkerProgressWindow { fn record_written_bytes_and_maybe_ramp(&self, bytes: i64, now_unix_ms: i64) { let normalized = bytes.max(0); self.window_bytes.fetch_add(normalized, Ordering::SeqCst); - self.total_written_bytes.fetch_add(normalized, Ordering::SeqCst); + self.total_written_bytes + .fetch_add(normalized, Ordering::SeqCst); self.maybe_ramp(now_unix_ms); } fn record_materialized_empty_dir(&self) { - self.total_materialized_empty_dirs.fetch_add(1, Ordering::SeqCst); + self.total_materialized_empty_dirs + .fetch_add(1, Ordering::SeqCst); } fn total_materialized_empty_dirs(&self) -> i64 { @@ -1298,8 +1305,9 @@ impl TransferWorkerProgressWindow { } if previous_goodput > 0 { let delta = current_goodput.saturating_sub(previous_goodput); - let improvement_percent = - delta.saturating_mul(100).saturating_div(previous_goodput.max(1)); + let improvement_percent = delta + .saturating_mul(100) + .saturating_div(previous_goodput.max(1)); if improvement_percent < self.policy.min_improvement_percent { return; } @@ -1335,10 +1343,8 @@ impl TransferWorkerProgressWindow { .saturating_mul(1000) .saturating_div(window_elapsed_ms.max(1)) }; - self.peak_sample_goodput_bytes_per_sec.fetch_max( - window_goodput_bytes_per_sec.max(0), - Ordering::SeqCst, - ); + self.peak_sample_goodput_bytes_per_sec + .fetch_max(window_goodput_bytes_per_sec.max(0), Ordering::SeqCst); Some(TransferWorkerThroughputSample { window_started_unix_ms, window_elapsed_ms, @@ -1448,7 +1454,8 @@ impl TransferReadStreamActorOwned { data: Vec::new(), }); } - self.fill_prefetch_queue().map_err(|err| self.cache_terminal_error(err))?; + self.fill_prefetch_queue() + .map_err(|err| self.cache_terminal_error(err))?; let to_take = std::cmp::min(length as usize, (self.file_size - next_offset) as usize); let buf = self .take_prefetched_bytes(to_take) @@ -1457,7 +1464,8 @@ impl TransferReadStreamActorOwned { self.replay_offset = next_offset; self.replay_data = buf.clone(); self.next_offset = next_offset.saturating_add(buf.len() as i64); - self.fill_prefetch_queue().map_err(|err| self.cache_terminal_error(err))?; + self.fill_prefetch_queue() + .map_err(|err| self.cache_terminal_error(err))?; Ok(FluxonFsTransferReadStreamNextResultWire { stream_missing: false, data: buf, @@ -1598,7 +1606,11 @@ impl TransferReadStreamActorHandle { struct TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { log_context: TransferWorkerLogContext, @@ -1611,7 +1623,11 @@ where impl TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { fn new( @@ -1693,7 +1709,8 @@ where } fn progress_snapshot(&self) -> TransferWorkerProgressSnapshot { - self.progress.snapshot(chrono::Utc::now().timestamp_millis()) + self.progress + .snapshot(chrono::Utc::now().timestamp_millis()) } fn stop(&self) { @@ -1737,7 +1754,8 @@ impl TransferReadStreamRegistryHandle { }); } } - let full_path = safe_join_root(root_dir_abs, open.relpath.as_str()).map_err(resp_err_kverr)?; + let full_path = + safe_join_root(root_dir_abs, open.relpath.as_str()).map_err(resp_err_kverr)?; let file = open_file_with_target_path_chmod_retry(&full_path, "open_stream")?; let md = file.metadata().map_err(resp_err_io)?; let file_size = md.len().min(i64::MAX as u64) as i64; @@ -1764,12 +1782,15 @@ impl TransferReadStreamRegistryHandle { }); } state.streams.insert(stream_id.clone(), entry); - state.dedup_by_worker_file.insert(dedup_key, stream_id.clone()); + state + .dedup_by_worker_file + .insert(dedup_key, stream_id.clone()); drop(state); if let Err(resp) = TransferReadStreamActorHandle::start(stream_id.as_str(), actor) { let mut state = self.state.lock(); state.streams.remove(stream_id.as_str()); - state.dedup_by_worker_file + state + .dedup_by_worker_file .retain(|_, existing_stream_id| existing_stream_id != &stream_id); return Err(resp); } @@ -1811,7 +1832,9 @@ impl TransferReadStreamRegistryHandle { let Some(entry) = state.streams.remove(stream_id) else { return; }; - state.dedup_by_worker_file.retain(|_, existing_stream_id| existing_stream_id != stream_id); + state + .dedup_by_worker_file + .retain(|_, existing_stream_id| existing_stream_id != stream_id); entry.close(); } } @@ -1894,20 +1917,22 @@ fn decode_transfer_stream_open_result_payload( return Err(TransferWorkerRpcFailure::Fatal(resp.clone())); } Ok(FluxonFsTransferReadStreamOpenResultWire { - stream_id: require_str(resp, "stream_id").map_err(resp_err_kverr).map_err( - |err| { + stream_id: require_str(resp, "stream_id") + .map_err(resp_err_kverr) + .map_err(|err| { invalid_transfer_rpc_response(format!( "transfer read stream open response missing stream_id: err={}", transfer_rpc_response_err_text(&err) )) - }, - )?, - size: require_i64(resp, "size").map_err(resp_err_kverr).map_err(|err| { - invalid_transfer_rpc_response(format!( - "transfer read stream open response missing size: err={}", - transfer_rpc_response_err_text(&err) - )) - })?, + })?, + size: require_i64(resp, "size") + .map_err(resp_err_kverr) + .map_err(|err| { + invalid_transfer_rpc_response(format!( + "transfer read stream open response missing size: err={}", + transfer_rpc_response_err_text(&err) + )) + })?, }) } @@ -1980,11 +2005,15 @@ fn is_relpath_skipped(skip_entries: &[FluxonFsTransferSkipEntryWire], relpath: & } fn file_name_from_relpath(relpath: &str) -> Result<&str, FlatDict> { - relpath.rsplit('/').next().filter(|v| !v.is_empty()).ok_or_else(|| { - resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("relpath must contain file name: {}", relpath), - })) - }) + relpath + .rsplit('/') + .next() + .filter(|v| !v.is_empty()) + .ok_or_else(|| { + resp_err_kverr(KvError::Api(ApiError::InvalidArgument { + detail: format!("relpath must contain file name: {}", relpath), + })) + }) } fn transfer_staging_dir_for_file(staging_prefix: &str, relpath: &str) -> String { @@ -2106,8 +2135,12 @@ where match attempt() { Ok(value) => Ok(value), Err(initial_err) if initial_err.kind() == ErrorKind::PermissionDenied => { - let repair_dir = - repair_permission_denied_dir_for_retry(repair_anchor, op, target_path, &initial_err)?; + let repair_dir = repair_permission_denied_dir_for_retry( + repair_anchor, + op, + target_path, + &initial_err, + )?; attempt().map_err(|retry_err| { resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( @@ -2334,10 +2367,11 @@ fn collect_transfer_tree_with_deadline( continue; } }; - out.symlink_notices.push(FluxonFsTransferSymlinkNoticeEntryWire { - relpath: child_rel, - link_target: link_target.to_string_lossy().to_string(), - }); + out.symlink_notices + .push(FluxonFsTransferSymlinkNoticeEntryWire { + relpath: child_rel, + link_target: link_target.to_string_lossy().to_string(), + }); continue; } if md.is_dir() { @@ -2361,7 +2395,8 @@ fn collect_transfer_tree_with_deadline( } out.files.sort_by(|a, b| a.relpath.cmp(&b.relpath)); out.empty_dirs.sort(); - out.symlink_notices.sort_by(|a, b| a.relpath.cmp(&b.relpath)); + out.symlink_notices + .sort_by(|a, b| a.relpath.cmp(&b.relpath)); Ok(out) } @@ -2574,10 +2609,8 @@ fn build_transfer_scan_events_for_result( event_seq_no_start: i64, result: FluxonFsTransferScanResultWire, ) -> (Vec, bool, i64) { - let (child_scan_units, continue_locally) = split_same_root_continuation_from_child_scan_units( - assignment, - result.child_scan_units, - ); + let (child_scan_units, continue_locally) = + split_same_root_continuation_from_child_scan_units(assignment, result.child_scan_units); if continue_locally { let event = build_transfer_scan_event( assignment, @@ -2588,11 +2621,7 @@ fn build_transfer_scan_events_for_result( result.full_dir_batches, String::new(), ); - return ( - vec![event], - true, - event_seq_no_start.saturating_add(1), - ); + return (vec![event], true, event_seq_no_start.saturating_add(1)); } let mut next_event_seq_no = event_seq_no_start; let mut events = Vec::new(); @@ -2620,11 +2649,7 @@ fn build_transfer_scan_events_for_result( Vec::new(), String::new(), )); - ( - events, - false, - next_event_seq_no.saturating_add(1), - ) + (events, false, next_event_seq_no.saturating_add(1)) } fn send_transfer_scan_event_once( @@ -2637,10 +2662,7 @@ fn send_transfer_scan_event_once( } let event_json = serde_json::to_string(event) .map_err(|e| format!("serialize transfer scan event failed: {}", e))?; - let payload = FlatDict::from([( - "scan_event_json".to_string(), - FlatValue::String(event_json), - )]); + let payload = FlatDict::from([("scan_event_json".to_string(), FlatValue::String(event_json))]); let resp = api .rpc_client() .call( @@ -2793,30 +2815,28 @@ fn run_transfer_scan_background_task( } let mut next_event_seq_no = 1_i64; loop { - let result = match build_transfer_scan_result_for_root_dir_abs( - root_dir_abs.as_str(), - &assignment, - ) { - Ok(v) => v, - Err(resp) => { - let failed = build_transfer_scan_event( - &assignment, - next_event_seq_no, - FluxonFsTransferScanEventKindWire::Failed, - Vec::new(), - Vec::new(), - Vec::new(), - transfer_rpc_response_err_text(&resp), - ); - let _ = send_transfer_scan_event_with_retry( - api.as_ref(), - master_id.as_str(), - &mut assignment, - &failed, - ); - break; - } - }; + let result = + match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { + Ok(v) => v, + Err(resp) => { + let failed = build_transfer_scan_event( + &assignment, + next_event_seq_no, + FluxonFsTransferScanEventKindWire::Failed, + Vec::new(), + Vec::new(), + Vec::new(), + transfer_rpc_response_err_text(&resp), + ); + let _ = send_transfer_scan_event_with_retry( + api.as_ref(), + master_id.as_str(), + &mut assignment, + &failed, + ); + break; + } + }; let (events, continue_locally, next_seq_no_after_events) = build_transfer_scan_events_for_result(&assignment, next_event_seq_no, result); next_event_seq_no = next_seq_no_after_events; @@ -2922,17 +2942,14 @@ impl TransferScanRegistryHandle { let assignment2 = assignment.clone(); let thread_name = format!("fluxon_fs_transfer_scan_{}", assignment.scan_task_id); match thread::Builder::new().name(thread_name).spawn(move || { - run_transfer_scan_background_task( - registry, - api2, - master_id2, - exports2, - assignment2, - ); + run_transfer_scan_background_task(registry, api2, master_id2, exports2, assignment2); }) { Ok(_) => Ok(FluxonFsTransferScanLaunchResultWire::started()), Err(err) => { - self.state.lock().tasks.remove(assignment.scan_task_id.as_str()); + self.state + .lock() + .tasks + .remove(assignment.scan_task_id.as_str()); Err(resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( "spawn transfer scan thread failed: scan_task_id={} err={}", @@ -3006,22 +3023,16 @@ impl TransferWorkerRegistryHandle { let master_id2 = master_id.to_string(); let exports2 = exports.clone(); let assignment2 = assignment.clone(); - let thread_name = format!( - "fluxon_fs_transfer_worker_{}", - assignment.worker_task_id - ); + let thread_name = format!("fluxon_fs_transfer_worker_{}", assignment.worker_task_id); match thread::Builder::new().name(thread_name).spawn(move || { - run_transfer_worker_background_task( - registry, - api2, - master_id2, - exports2, - assignment2, - ); + run_transfer_worker_background_task(registry, api2, master_id2, exports2, assignment2); }) { Ok(_) => Ok(FluxonFsTransferWorkerLaunchResultWire::started()), Err(err) => { - self.state.lock().tasks.remove(assignment.worker_task_id.as_str()); + self.state + .lock() + .tasks + .remove(assignment.worker_task_id.as_str()); Err(resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( "spawn transfer worker thread failed: worker_task_id={} err={}", @@ -3298,10 +3309,7 @@ fn send_transfer_worker_result_once( detail: format!("serialize transfer worker result failed: {}", e), }))) })?; - let payload = FlatDict::from([( - "result_json".to_string(), - FlatValue::String(result_json), - )]); + let payload = FlatDict::from([("result_json".to_string(), FlatValue::String(result_json))]); let resp = api .rpc_client() .call( @@ -3396,7 +3404,10 @@ fn open_transfer_read_stream_via_rpc_once( "relpath".to_string(), FlatValue::String(file.relpath.clone()), ), - ("initial_offset".to_string(), FlatValue::Int64(initial_offset)), + ( + "initial_offset".to_string(), + FlatValue::Int64(initial_offset), + ), ]); let resp = api .rpc_client() @@ -3518,25 +3529,25 @@ impl TransferWorkerRemoteControl { loop { self.before_heartbeat_retry_attempt()?; let current_materialized_empty_dirs = self.progress.total_materialized_empty_dirs(); - match self - .heartbeat - .ensure_continue( - force, - current_materialized_empty_dirs, - |heartbeat_unix_ms, _heartbeat_detail| { - let progress_snapshot = - self.progress.snapshot(chrono::Utc::now().timestamp_millis()); - let telemetry = - Some(transfer_worker_telemetry_from_progress_snapshot(&progress_snapshot)); - send_transfer_worker_heartbeat_once( - self.api.as_ref(), - self.master_id.as_str(), - &self.assignment, - heartbeat_unix_ms, - telemetry, - ) - }, - ) { + match self.heartbeat.ensure_continue( + force, + current_materialized_empty_dirs, + |heartbeat_unix_ms, _heartbeat_detail| { + let progress_snapshot = self + .progress + .snapshot(chrono::Utc::now().timestamp_millis()); + let telemetry = Some(transfer_worker_telemetry_from_progress_snapshot( + &progress_snapshot, + )); + send_transfer_worker_heartbeat_once( + self.api.as_ref(), + self.master_id.as_str(), + &self.assignment, + heartbeat_unix_ms, + telemetry, + ) + }, + ) { Ok(()) => return Ok(()), Err(TransferWorkerHeartbeatGateError::Terminal(err)) => return Err(err), Err(TransferWorkerHeartbeatGateError::Retryable { @@ -3631,10 +3642,7 @@ impl TransferWorkerRemoteControl { ) } - fn close_stream_with_retry( - &self, - stream_id: &str, - ) -> Result<(), TransferWorkerExecutionError> { + fn close_stream_with_retry(&self, stream_id: &str) -> Result<(), TransferWorkerExecutionError> { let api = self.api.clone(); let assignment = self.assignment.clone(); let op_detail = format!( @@ -3744,9 +3752,9 @@ impl TransferWorkerRemoteControl { if ack.accepted { return Ok(()); } - Err(TransferWorkerExecutionError::Stop(stop_reason_or_superseded( - ack.stop_reason, - ))) + Err(TransferWorkerExecutionError::Stop( + stop_reason_or_superseded(ack.stop_reason), + )) } } @@ -3825,10 +3833,12 @@ impl TransferWorkerHeartbeatGate { mut heartbeat_op: HeartbeatOp, ) -> Result<(), TransferWorkerHeartbeatGateError> where - HeartbeatOp: FnMut( - i64, - &'static str, - ) -> Result, + HeartbeatOp: + FnMut( + i64, + &'static str, + ) + -> Result, { loop { let (heartbeat_unix_ms, heartbeat_detail) = { @@ -3875,15 +3885,13 @@ impl TransferWorkerHeartbeatGate { state.heartbeat_inflight = false; let result = match heartbeat_result { Ok(heartbeat_result) if heartbeat_result.continue_running => { - state.last_heartbeat_completed_unix_ms = - chrono::Utc::now().timestamp_millis(); + state.last_heartbeat_completed_unix_ms = chrono::Utc::now().timestamp_millis(); state.last_heartbeat_materialized_empty_dirs = current_materialized_empty_dirs; state.granted_lease_expire_unix_ms = heartbeat_result.lease_expire_unix_ms; Ok(()) } Ok(heartbeat_result) => { - state.last_heartbeat_completed_unix_ms = - chrono::Utc::now().timestamp_millis(); + state.last_heartbeat_completed_unix_ms = chrono::Utc::now().timestamp_millis(); state.last_heartbeat_materialized_empty_dirs = current_materialized_empty_dirs; let reason = stop_reason_or_superseded(heartbeat_result.stop_reason); state.terminal_state = @@ -3933,7 +3941,8 @@ fn run_transfer_worker_background_task( )); let dedup_expire_unix_ms = match control.ensure_continue(true) { Ok(()) => { - let dst_export_root = match exports.export_root_dir_abs(assignment.dst_export.as_str()) { + let dst_export_root = match exports.export_root_dir_abs(assignment.dst_export.as_str()) + { Ok(v) => v, Err(err) => { tracing::warn!( @@ -3996,7 +4005,9 @@ fn run_transfer_worker_background_task( }, { let control = control.clone(); - move |file, read_offset, length| control.read_chunk_with_retry(file, read_offset, length) + move |file, read_offset, length| { + control.read_chunk_with_retry(file, read_offset, length) + } }, ) { Ok(result) => { @@ -4004,34 +4015,38 @@ fn run_transfer_worker_background_task( if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("before_result_submit", &assignment, &resp); - } - match control.submit_result_with_retry(&result) { - Ok(()) => control.dedup_expire_unix_ms(), - Err(TransferWorkerExecutionError::Stop(reason)) => { - tracing::info!( - "transfer worker result submission stopped: job_id={} batch_id={} worker_id={} worker_task_id={} reason={:?}", - assignment.job_id, - assignment.batch_id, - assignment.worker_id, - assignment.worker_task_id, - reason + log_transfer_worker_cleanup_failure( + "before_result_submit", + &assignment, + &resp, ); - control.dedup_expire_unix_ms() } - Err(TransferWorkerExecutionError::Fatal(resp)) => { - tracing::warn!( - "transfer worker result submission failed: job_id={} batch_id={} worker_id={} worker_task_id={} resp={:?}", - assignment.job_id, - assignment.batch_id, - assignment.worker_id, - assignment.worker_task_id, - resp - ); - control.dedup_expire_unix_ms() + match control.submit_result_with_retry(&result) { + Ok(()) => control.dedup_expire_unix_ms(), + Err(TransferWorkerExecutionError::Stop(reason)) => { + tracing::info!( + "transfer worker result submission stopped: job_id={} batch_id={} worker_id={} worker_task_id={} reason={:?}", + assignment.job_id, + assignment.batch_id, + assignment.worker_id, + assignment.worker_task_id, + reason + ); + control.dedup_expire_unix_ms() + } + Err(TransferWorkerExecutionError::Fatal(resp)) => { + tracing::warn!( + "transfer worker result submission failed: job_id={} batch_id={} worker_id={} worker_task_id={} resp={:?}", + assignment.job_id, + assignment.batch_id, + assignment.worker_id, + assignment.worker_task_id, + resp + ); + control.dedup_expire_unix_ms() + } } } - } Err(TransferWorkerExecutionError::Stop(reason)) => { control.close_all_streams(); if let Err(resp) = @@ -4054,10 +4069,13 @@ fn run_transfer_worker_background_task( if let Err(cleanup_resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("after_fatal", &assignment, &cleanup_resp); + log_transfer_worker_cleanup_failure( + "after_fatal", + &assignment, + &cleanup_resp, + ); } - if let Some((fatal_kind, fatal_message)) = - classify_transfer_worker_fatal(&resp) + if let Some((fatal_kind, fatal_message)) = classify_transfer_worker_fatal(&resp) { match report_transfer_worker_fatal_once( control.api.as_ref(), @@ -4107,15 +4125,25 @@ fn run_transfer_worker_background_task( } } Err(TransferWorkerExecutionError::Stop(reason)) => { - let dst_export_root = exports.export_root_dir_abs(assignment.dst_export.as_str()).ok(); + let dst_export_root = exports + .export_root_dir_abs(assignment.dst_export.as_str()) + .ok(); let dst_root = dst_export_root.and_then(|dst_export_root| { - safe_join_root(dst_export_root.as_str(), assignment.dst_root_relpath.as_str()) - .ok() - .map(PathBuf::from) + safe_join_root( + dst_export_root.as_str(), + assignment.dst_root_relpath.as_str(), + ) + .ok() + .map(PathBuf::from) }); if let Some(dst_root) = dst_root { - if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("before_execution_stop", &assignment, &resp); + if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) + { + log_transfer_worker_cleanup_failure( + "before_execution_stop", + &assignment, + &resp, + ); } } tracing::info!( @@ -4129,11 +4157,16 @@ fn run_transfer_worker_background_task( control.dedup_expire_unix_ms() } Err(TransferWorkerExecutionError::Fatal(resp)) => { - let dst_export_root = exports.export_root_dir_abs(assignment.dst_export.as_str()).ok(); + let dst_export_root = exports + .export_root_dir_abs(assignment.dst_export.as_str()) + .ok(); let dst_root = dst_export_root.and_then(|dst_export_root| { - safe_join_root(dst_export_root.as_str(), assignment.dst_root_relpath.as_str()) - .ok() - .map(PathBuf::from) + safe_join_root( + dst_export_root.as_str(), + assignment.dst_root_relpath.as_str(), + ) + .ok() + .map(PathBuf::from) }); if let Some(dst_root) = dst_root { if let Err(cleanup_resp) = @@ -4509,7 +4542,9 @@ fn plan_transfer_subtree_batches( total_bytes: 0, root_is_empty: true, mergeable_empty_dir_count: 1, - mergeable_empty_dir_estimated_bytes: estimate_empty_dir_manifest_entry_bytes(root_relpath), + mergeable_empty_dir_estimated_bytes: estimate_empty_dir_manifest_entry_bytes( + root_relpath, + ), direct_files_only_batches: Vec::new(), full_dir_batches: Vec::new(), child_scan_units: Vec::new(), @@ -4540,8 +4575,7 @@ fn plan_transfer_subtree_batches( let child_empty_dir_count = child_plan.mergeable_empty_dir_count; let child_empty_dir_estimated_bytes = child_plan.mergeable_empty_dir_estimated_bytes; - if mergeable_empty_dir_count - .saturating_add(child_empty_dir_count) + if mergeable_empty_dir_count.saturating_add(child_empty_dir_count) > TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET || mergeable_empty_dir_estimated_bytes .saturating_add(child_empty_dir_estimated_bytes) @@ -4716,7 +4750,11 @@ fn build_root_direct_files_only_batch_from_entries( } fn sort_transfer_scan_batches(batches: &mut [FluxonFsTransferScanBatchWire]) { - batches.sort_by(|a, b| a.root_relpath.cmp(&b.root_relpath).then(a.batch_id.cmp(&b.batch_id))); + batches.sort_by(|a, b| { + a.root_relpath + .cmp(&b.root_relpath) + .then(a.batch_id.cmp(&b.batch_id)) + }); } fn build_full_dir_batch_for_mergeable_subtree( @@ -4750,14 +4788,12 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( root_dir_abs: &str, assignment: &FluxonFsTransferScanAssignmentWire, ) -> Result { - let Some(mut session) = take_transfer_subtree_streaming_session(root_dir_abs, assignment)? else { + let Some(mut session) = take_transfer_subtree_streaming_session(root_dir_abs, assignment)? + else { return Ok(build_finished_empty_subtree_stream_result(assignment)); }; loop { - if session - .dir_stack - .is_empty() - { + if session.dir_stack.is_empty() { let mut full_dir_batches = Vec::new(); if let Some(batch) = flush_pending_subtree_stream_batch(assignment, &mut session)? { full_dir_batches.push(batch); @@ -4776,7 +4812,9 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( finished: true, }); } - if TransferScanDeadline::from_assignment(assignment).is_some_and(|deadline| deadline.reached()) { + if TransferScanDeadline::from_assignment(assignment) + .is_some_and(|deadline| deadline.reached()) + { let mut full_dir_batches = Vec::new(); if let Some(batch) = flush_pending_subtree_stream_batch(assignment, &mut session)? { full_dir_batches.push(batch); @@ -4808,7 +4846,10 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( if should_flush_subtree_stream_batch( assignment.batch_ready_bytes, session.pending_bytes, - session.pending_files.len().saturating_add(session.pending_symlink_notices.len()), + session + .pending_files + .len() + .saturating_add(session.pending_symlink_notices.len()), session.pending_empty_dirs.len(), ) { let batch = flush_pending_subtree_stream_batch(assignment, &mut session)?.unwrap(); @@ -4865,22 +4906,30 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( }); } else if md.is_dir() { frame.saw_visible_child = true; - session.dir_stack.push(open_transfer_subtree_streaming_dir_frame( - child_path, - child_relpath, - )?); + session + .dir_stack + .push(open_transfer_subtree_streaming_dir_frame( + child_path, + child_relpath, + )?); } else if md.is_file() { frame.saw_visible_child = true; let size = md.len().min(i64::MAX as u64) as i64; session.pending_bytes = session.pending_bytes.saturating_add(size); session .pending_files - .push(FluxonFsTransferScanFrontierEntry { relpath: child_relpath, size }); + .push(FluxonFsTransferScanFrontierEntry { + relpath: child_relpath, + size, + }); } if should_flush_subtree_stream_batch( assignment.batch_ready_bytes, session.pending_bytes, - session.pending_files.len().saturating_add(session.pending_symlink_notices.len()), + session + .pending_files + .len() + .saturating_add(session.pending_symlink_notices.len()), session.pending_empty_dirs.len(), ) { let batch = flush_pending_subtree_stream_batch(assignment, &mut session)?.unwrap(); @@ -4928,11 +4977,12 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( ); } let deadline = TransferScanDeadline::from_assignment(assignment); - let root_listing = match collect_transfer_root_dir_listing_slice(root_dir_abs, assignment, deadline)? { - TransferRootDirListingOutcome::Complete(v) => v, - TransferRootDirListingOutcome::Finished(result) => return Ok(result), - TransferRootDirListingOutcome::Partial(result) => return Ok(result), - }; + let root_listing = + match collect_transfer_root_dir_listing_slice(root_dir_abs, assignment, deadline)? { + TransferRootDirListingOutcome::Complete(v) => v, + TransferRootDirListingOutcome::Finished(result) => return Ok(result), + TransferRootDirListingOutcome::Partial(result) => return Ok(result), + }; let mut direct_files = root_listing.direct_files; let mut direct_symlink_notices = root_listing.direct_symlink_notices; let mut direct_empty_dirs = root_listing.direct_empty_dirs; @@ -4976,7 +5026,10 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( if (!direct_files.is_empty() || !direct_symlink_notices.is_empty() || !direct_empty_dirs.is_empty()) - && !direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()) + && !direct_files_only_disposition_covers_root( + assignment, + assignment.root_relpath.as_str(), + ) { let mut next_partition_index = root_listing.emitted_direct_files_batch_count; direct_files_only_batches.extend(build_partitioned_root_direct_files_only_batches( @@ -4987,17 +5040,15 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( direct_empty_dirs.clone(), )?); } - child_scan_units.extend( - direct_dirs[delegated_child_scan_unit_count..] - .iter() - .map(|entry| { - new_child_scan_unit( - entry.relpath.clone(), - assignment.generation + 1, - delegated_child_scan_mode(), - ) - }), - ); + child_scan_units.extend(direct_dirs[delegated_child_scan_unit_count..].iter().map( + |entry| { + new_child_scan_unit( + entry.relpath.clone(), + assignment.generation + 1, + delegated_child_scan_mode(), + ) + }, + )); child_scan_units.sort_by(|a, b| a.root_relpath.cmp(&b.root_relpath)); sort_transfer_scan_batches(&mut direct_files_only_batches); return Ok(FluxonFsTransferScanResultWire { @@ -5027,7 +5078,8 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( let mut root_partitioned = root_listing.emitted_direct_files_batch_count > 0 || direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()); let mut mergeable_empty_dir_count = direct_empty_dirs.len(); - let mut mergeable_empty_dir_estimated_bytes = estimate_empty_dir_manifest_bytes(&direct_empty_dirs); + let mut mergeable_empty_dir_estimated_bytes = + estimate_empty_dir_manifest_bytes(&direct_empty_dirs); for child_relpath in direct_dirs.iter().map(|entry| entry.relpath.clone()) { let child_plan = plan_transfer_subtree_batches( root_dir_abs, @@ -5043,8 +5095,7 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( let child_empty_dir_count = child_plan.mergeable_empty_dir_count; let child_empty_dir_estimated_bytes = child_plan.mergeable_empty_dir_estimated_bytes; - if mergeable_empty_dir_count - .saturating_add(child_empty_dir_count) + if mergeable_empty_dir_count.saturating_add(child_empty_dir_count) > TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET || mergeable_empty_dir_estimated_bytes .saturating_add(child_empty_dir_estimated_bytes) @@ -5083,7 +5134,10 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( if (!direct_files.is_empty() || !direct_symlink_notices.is_empty() || !mergeable_empty_child_relpaths.is_empty()) - && !direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()) + && !direct_files_only_disposition_covers_root( + assignment, + assignment.root_relpath.as_str(), + ) { let mut next_partition_index = root_listing.emitted_direct_files_batch_count; direct_empty_dirs.extend(mergeable_empty_child_relpaths); @@ -5212,10 +5266,11 @@ fn handle_transfer_scan_assignment( assignment.generation, assignment.known_dispositions.len(), ); - let result = match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { - Ok(v) => v, - Err(resp) => return resp, - }; + let result = + match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { + Ok(v) => v, + Err(resp) => return resp, + }; encode_transfer_scan_result(&result, "transfer scan result") } @@ -5228,8 +5283,11 @@ fn prepare_transfer_file_streaming( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { let staging_relpath = transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) @@ -5239,9 +5297,12 @@ where .map_err(TransferWorkerExecutionError::fatal)?; ensure_transfer_parent_dirs(dst_root, final_relpath.as_str()) .map_err(TransferWorkerExecutionError::fatal)?; - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_relpath.as_str()) - .map_err(resp_err_kverr) - .map_err(TransferWorkerExecutionError::fatal)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr) + .map_err(TransferWorkerExecutionError::fatal)?; let mut dst_file = open_create_file_with_parent_dir_chmod_retry(&staging_abs) .map_err(TransferWorkerExecutionError::fatal)?; dst_file @@ -5254,14 +5315,14 @@ where let remaining = file.size.saturating_sub(copied); let chunk = coordinator.read_chunk(file, copied, remaining.min(CHUNK_BYTES as i64))?; if chunk.is_empty() { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( - ApiError::InvalidArgument { + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { detail: format!( "transfer worker source ended before expected size: relpath={} expected={} copied={}", file.relpath, file.size, copied ), - }, - )))); + }), + ))); } dst_file .write_all(&chunk) @@ -5271,13 +5332,14 @@ where coordinator.record_written_bytes(chunk.len() as i64); } if copied != file.size { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( - ApiError::InvalidArgument { - detail: format!( - "transfer worker size mismatch before staging completion: relpath={} expected={} actual={}", - file.relpath, file.size, copied - ), - })))); + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "transfer worker size mismatch before staging completion: relpath={} expected={} actual={}", + file.relpath, file.size, copied + ), + }), + ))); } // The staged file is still invisible at this point, so one more checkpoint // keeps supersession able to stop the worker before any later visible @@ -5302,48 +5364,57 @@ fn execute_transfer_single_file( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { coordinator.checkpoint_continue()?; - let prepared = match prepare_transfer_file_streaming(dst_root, staging_prefix, file, coordinator) { - Ok(v) => v, - Err(TransferWorkerExecutionError::Fatal(resp)) => { - if let Some(failed) = classify_transfer_failed_file(file, &resp) { - let staging_relpath = - transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) - .map_err(TransferWorkerExecutionError::fatal)?; - let staging_abs = safe_join_root( - dst_root.to_string_lossy().as_ref(), - staging_relpath.as_str(), - ) - .map_err(resp_err_kverr) - .map_err(TransferWorkerExecutionError::fatal)?; - match fs::remove_file(&staging_abs) { - Ok(()) => {} - Err(err) if err.kind() == ErrorKind::NotFound => {} - Err(err) => return Err(TransferWorkerExecutionError::fatal(resp_err_io(err))), + let prepared = + match prepare_transfer_file_streaming(dst_root, staging_prefix, file, coordinator) { + Ok(v) => v, + Err(TransferWorkerExecutionError::Fatal(resp)) => { + if let Some(failed) = classify_transfer_failed_file(file, &resp) { + let staging_relpath = + transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) + .map_err(TransferWorkerExecutionError::fatal)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr) + .map_err(TransferWorkerExecutionError::fatal)?; + match fs::remove_file(&staging_abs) { + Ok(()) => {} + Err(err) if err.kind() == ErrorKind::NotFound => {} + Err(err) => { + return Err(TransferWorkerExecutionError::fatal(resp_err_io(err))); + } + } + return Ok(TransferWorkerLaneOutcome::Failed( + TransferWorkerLaneFailedFileResult { result: failed }, + )); } - return Ok(TransferWorkerLaneOutcome::Failed( - TransferWorkerLaneFailedFileResult { result: failed }, - )); + return Err(TransferWorkerExecutionError::Fatal(resp)); } - return Err(TransferWorkerExecutionError::Fatal(resp)); - } - Err(err) => return Err(err), - }; + Err(err) => return Err(err), + }; coordinator.checkpoint_continue()?; - let result = promote_prepared_transfer_file(dst_root, PreparedTransferFile { - staging_relpath: prepared.staging_relpath.clone(), - final_relpath: prepared.final_relpath.clone(), - visible_size: prepared.visible_size, - }) + let result = promote_prepared_transfer_file( + dst_root, + PreparedTransferFile { + staging_relpath: prepared.staging_relpath.clone(), + final_relpath: prepared.final_relpath.clone(), + visible_size: prepared.visible_size, + }, + ) .map_err(TransferWorkerExecutionError::fatal); match result { - Ok(result) => Ok(TransferWorkerLaneOutcome::Visible(TransferWorkerLaneFileResult { - result, - })), + Ok(result) => Ok(TransferWorkerLaneOutcome::Visible( + TransferWorkerLaneFileResult { result }, + )), Err(TransferWorkerExecutionError::Fatal(resp)) => { if let Some(failed) = classify_transfer_failed_file(file, &resp) { let staging_abs = safe_join_root( @@ -5374,8 +5445,11 @@ fn execute_transfer_empty_dir( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { coordinator.checkpoint_continue()?; @@ -5393,11 +5467,14 @@ fn execute_transfer_worker_assignment_with_policy( read_chunk: ReadChunkFn, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { let policy = policy.normalized(); @@ -5424,11 +5501,14 @@ fn execute_transfer_worker_assignment_with_policy_and_progress Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { create_dir_all_with_parent_dir_chmod_retry(dst_root) @@ -5436,9 +5516,11 @@ where let manifest = FluxonFsTransferManifestWire::decode_from_blob(assignment.manifest_blob.as_slice()) .map_err(|e| { - TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("decode transfer worker manifest failed: {}", e), - }))) + TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( + ApiError::InvalidArgument { + detail: format!("decode transfer worker manifest failed: {}", e), + }, + ))) })?; if transfer_manifest_is_empty_dirs_only_batch(&manifest, assignment.collect_infos.as_slice()) { // Empty-dir-only batches never generate byte-based ramp-up signals, so @@ -5664,10 +5746,16 @@ fn promote_prepared_transfer_file( dst_root: &PathBuf, file: PreparedTransferFile, ) -> Result { - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), file.staging_relpath.as_str()) - .map_err(resp_err_kverr)?; - let final_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), file.final_relpath.as_str()) - .map_err(resp_err_kverr)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + file.staging_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; + let final_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + file.final_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; rename_with_dst_parent_dir_chmod_retry(&staging_abs, &final_abs)?; Ok(FluxonFsTransferWorkerFileResultWire { relpath: file.final_relpath.clone(), @@ -5694,15 +5782,15 @@ fn prepare_transfer_collect_info_materialization( ), })) })?; - let staging_relpath = transfer_collect_info_staging_relpath( - batch_id, - worker_task_id, - collect_info.collect_kind, - )?; + let staging_relpath = + transfer_collect_info_staging_relpath(batch_id, worker_task_id, collect_info.collect_kind)?; ensure_transfer_parent_dirs(dst_root, staging_relpath.as_str())?; ensure_transfer_parent_dirs(dst_root, output_relpath.as_str())?; - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_relpath.as_str()) - .map_err(resp_err_kverr)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; let mut dst_file = open_create_file_with_parent_dir_chmod_retry(&staging_abs)?; dst_file.set_len(0).map_err(resp_err_io)?; dst_file @@ -5723,14 +5811,15 @@ fn transfer_collect_info_staging_relpath( worker_task_id: &str, collect_kind: FluxonFsTransferCollectInfoKind, ) -> Result { - let output_relpath = transfer_collect_info_output_relpath(batch_id, collect_kind).map_err(|detail| { - resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!( - "build transfer collect info output relpath failed: batch_id={} err={}", - batch_id, detail - ), - })) - })?; + let output_relpath = + transfer_collect_info_output_relpath(batch_id, collect_kind).map_err(|detail| { + resp_err_kverr(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "build transfer collect info output relpath failed: batch_id={} err={}", + batch_id, detail + ), + })) + })?; Ok(format!("{}.{}.fluxon.part", output_relpath, worker_task_id)) } @@ -5750,9 +5839,12 @@ fn prune_empty_parent_dirs(mut current: PathBuf, root: &PathBuf) -> Result<(), F Ok(()) } -fn cleanup_attempt_staging_prefix(dst_root: &PathBuf, staging_prefix: &str) -> Result<(), FlatDict> { - let staging_abs = - safe_join_root(dst_root.to_string_lossy().as_ref(), staging_prefix).map_err(resp_err_kverr)?; +fn cleanup_attempt_staging_prefix( + dst_root: &PathBuf, + staging_prefix: &str, +) -> Result<(), FlatDict> { + let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_prefix) + .map_err(resp_err_kverr)?; match fs::remove_dir_all(&staging_abs) { Ok(()) => {} Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), @@ -5865,7 +5957,8 @@ pub(crate) fn read_transfer_chunk_from_root_dir_abs( return Ok(Vec::new()); } let to_read = std::cmp::min(length, size - offset) as usize; - f.seek(SeekFrom::Start(offset as u64)).map_err(resp_err_io)?; + f.seek(SeekFrom::Start(offset as u64)) + .map_err(resp_err_io)?; let mut buf = vec![0u8; to_read]; f.read_exact(&mut buf).map_err(resp_err_io)?; Ok(buf) @@ -5881,11 +5974,14 @@ pub(crate) fn execute_transfer_worker_assignment( read_chunk: ReadChunkFn, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { execute_transfer_worker_assignment_with_policy( @@ -5943,12 +6039,19 @@ pub(super) fn handle_transfer_read(exports: &AgentExportsHandle, payload: FlatDi Ok(v) => v, Err(e) => return resp_err_kverr(e), }; - let buf = - match read_transfer_chunk_from_root_dir_abs(root_dir_abs.as_str(), relpath.as_str(), offset, length) { - Ok(v) => v, - Err(resp) => return resp, - }; - resp_ok(BTreeMap::from([("data".to_string(), FlatValue::Bytes(buf))])) + let buf = match read_transfer_chunk_from_root_dir_abs( + root_dir_abs.as_str(), + relpath.as_str(), + offset, + length, + ) { + Ok(v) => v, + Err(resp) => return resp, + }; + resp_ok(BTreeMap::from([( + "data".to_string(), + FlatValue::Bytes(buf), + )])) } pub(super) fn handle_transfer_stream_open( @@ -6010,7 +6113,9 @@ pub(super) fn handle_transfer_worker( Err(resp) => return resp, }; match registry.launch_task(api, master_id, exports, assignment) { - Ok(result) => encode_transfer_worker_launch_result(&result, "transfer worker launch result"), + Ok(result) => { + encode_transfer_worker_launch_result(&result, "transfer worker launch result") + } Err(resp) => resp, } } @@ -6086,15 +6191,15 @@ mod tests { .collect() } - fn assert_all_child_scan_units_are_subtree_streaming( - result: &FluxonFsTransferScanResultWire, - ) { - assert!(result - .child_scan_units - .iter() - .all(|child| child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming)); - } - + fn assert_all_child_scan_units_are_subtree_streaming(result: &FluxonFsTransferScanResultWire) { + assert!( + result + .child_scan_units + .iter() + .all(|child| child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming) + ); + } + fn ok_bool(resp: &FlatDict) -> bool { matches!(resp.get("ok"), Some(FlatValue::Bool(true))) } @@ -6106,10 +6211,7 @@ mod tests { panic!("unexpected open result fatal decode error: {:?}", other) } Err(TransferWorkerRpcFailure::Retryable { detail }) => { - panic!( - "unexpected open result retryable decode error: {}", - detail - ) + panic!("unexpected open result retryable decode error: {}", detail) } } } @@ -6121,10 +6223,7 @@ mod tests { panic!("unexpected next result fatal decode error: {:?}", other) } Err(TransferWorkerRpcFailure::Retryable { detail }) => { - panic!( - "unexpected next result retryable decode error: {}", - detail - ) + panic!("unexpected next result retryable decode error: {}", detail) } } } @@ -6140,10 +6239,7 @@ mod tests { .collect() } - fn test_worker_assignment( - relpath: &str, - size: i64, - ) -> FluxonFsTransferWorkerAssignmentWire { + fn test_worker_assignment(relpath: &str, size: i64) -> FluxonFsTransferWorkerAssignmentWire { FluxonFsTransferWorkerAssignmentWire { job_id: "job".to_string(), batch_id: "batch".to_string(), @@ -6158,12 +6254,13 @@ mod tests { root_relpath: ".".to_string(), staging_prefix: ".fluxon.stage/job/batch".to_string(), lease_expire_unix_ms: 0, - manifest_blob: build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { + manifest_blob: build_transfer_manifest_blob( + vec![FluxonFsTransferScanFrontierEntry { relpath: relpath.to_string(), size, - }, - ], Vec::new()) + }], + Vec::new(), + ) .unwrap(), collect_infos: Vec::new(), } @@ -6183,7 +6280,11 @@ mod tests { read_chunk: ReadChunkFn, ) -> TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { let policy = Arc::new(TransferWorkerLanePolicy::production_default()); @@ -6217,16 +6318,19 @@ mod tests { #[test] fn build_transfer_manifest_blob_round_trips_entries() { - let blob = build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { - relpath: "a".to_string(), - size: 1, - }, - FluxonFsTransferScanFrontierEntry { - relpath: "b/c".to_string(), - size: 2, - }, - ], vec!["empty".to_string()]) + let blob = build_transfer_manifest_blob( + vec![ + FluxonFsTransferScanFrontierEntry { + relpath: "a".to_string(), + size: 1, + }, + FluxonFsTransferScanFrontierEntry { + relpath: "b/c".to_string(), + size: 2, + }, + ], + vec!["empty".to_string()], + ) .unwrap(); let manifest = FluxonFsTransferManifestWire::decode_from_blob(&blob).unwrap(); assert_eq!(manifest.entry_count, 2); @@ -6250,11 +6354,12 @@ mod tests { #[test] fn materialize_transfer_collect_info_writes_task_scoped_staging_then_output_file() { let root = TempDir::new().unwrap(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "root/link-file.bin".to_string(), - link_target: "target/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "root/link-file.bin".to_string(), + link_target: "target/file.bin".to_string(), + }]) + .unwrap(); let prepared = prepare_transfer_collect_info_materialization( &root.path().to_path_buf(), "batch-1", @@ -6383,7 +6488,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(2)), ("length".to_string(), FlatValue::Int64(3)), ]), @@ -6398,7 +6506,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(6)), ("length".to_string(), FlatValue::Int64(1)), ]), @@ -6423,7 +6534,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(1)), ("length".to_string(), FlatValue::Int64(3)), ]), @@ -6446,9 +6560,15 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(0)), - ("length".to_string(), FlatValue::Int64(CHUNK_BYTES as i64 + 1)), + ( + "length".to_string(), + FlatValue::Int64(CHUNK_BYTES as i64 + 1), + ), ]), ); assert!(matches!(resp.get("ok"), Some(FlatValue::Bool(false)))); @@ -6470,7 +6590,10 @@ mod tests { FlatValue::String("task-0".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(0)), ]), ); @@ -6551,7 +6674,10 @@ mod tests { FlatValue::String("task-1".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(0)), ]), ); @@ -6584,7 +6710,10 @@ mod tests { ("length".to_string(), FlatValue::Int64(2)), ]), ); - assert!(matches!(invalid_resp.get("ok"), Some(FlatValue::Bool(false)))); + assert!(matches!( + invalid_resp.get("ok"), + Some(FlatValue::Bool(false)) + )); } #[test] @@ -6603,7 +6732,10 @@ mod tests { FlatValue::String("task-2".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(3)), ]), ); @@ -6647,7 +6779,10 @@ mod tests { FlatValue::String("task-3".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(3)), ]), ); @@ -6718,7 +6853,10 @@ mod tests { let result = decode_result_json(&resp); assert!(result.finished); assert!(result.direct_files_only_batches.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -6781,7 +6919,8 @@ mod tests { } #[test] - fn handle_transfer_scan_assignment_groups_empty_children_into_direct_batch_without_direct_files() { + fn handle_transfer_scan_assignment_groups_empty_children_into_direct_batch_without_direct_files() + { let root = TempDir::new().unwrap(); write_file(&root, "root/big/data.bin", b"12345"); fs::create_dir_all(root.path().join("root/empty-a")).unwrap(); @@ -6814,9 +6953,10 @@ mod tests { assert_eq!(child_scan_unit_roots(&result), vec!["root/big".to_string()]); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); - let direct_manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let direct_manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert!(direct_manifest.entries.is_empty()); assert_eq!( direct_manifest.empty_dir_relpaths, @@ -6854,16 +6994,17 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.full_dir_batches.len(), 0); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/huge".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/huge".to_string() + ); let manifest = FluxonFsTransferManifestWire::decode_from_blob( &result.direct_files_only_batches[0].manifest_blob, ) .unwrap(); assert!(manifest.entries.is_empty()); assert!(!manifest.empty_dir_relpaths.is_empty()); - assert!( - manifest.empty_dir_relpaths.len() <= TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET - ); + assert!(manifest.empty_dir_relpaths.len() <= TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET); assert!( estimate_empty_dir_manifest_bytes(&manifest.empty_dir_relpaths) <= TRANSFER_MERGEABLE_EMPTY_DIR_ESTIMATED_BYTES_BUDGET @@ -6879,10 +7020,10 @@ mod tests { )) + 1; for idx in 0..child_count { - fs::create_dir_all(root.path().join(format!( - "root/branch-{idx:05}/{}", - "x".repeat(200) - ))) + fs::create_dir_all( + root.path() + .join(format!("root/branch-{idx:05}/{}", "x".repeat(200))), + ) .unwrap(); } let result = build_transfer_scan_result_for_root_dir_abs( @@ -6908,10 +7049,12 @@ mod tests { assert!(!result.finished); assert!(result.direct_files_only_batches.is_empty()); assert!(!result.child_scan_units.is_empty()); - assert!(result - .child_scan_units - .iter() - .any(|child| child.scan_mode == FluxonFsTransferScanMode::FullTree)); + assert!( + result + .child_scan_units + .iter() + .any(|child| child.scan_mode == FluxonFsTransferScanMode::FullTree) + ); assert!(result.child_scan_units.iter().all(|child| { child.scan_mode == FluxonFsTransferScanMode::FullTree || child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming @@ -6963,10 +7106,16 @@ mod tests { assert!(!continue_locally); assert_eq!(next_event_seq_no, 9); assert_eq!(events.len(), 2); - assert_eq!(events[0].event_kind, FluxonFsTransferScanEventKindWire::Append); + assert_eq!( + events[0].event_kind, + FluxonFsTransferScanEventKindWire::Append + ); assert_eq!(events[0].event_seq_no, 7); assert_eq!(events[0].full_dir_batches.len(), 1); - assert_eq!(events[1].event_kind, FluxonFsTransferScanEventKindWire::Finished); + assert_eq!( + events[1].event_kind, + FluxonFsTransferScanEventKindWire::Finished + ); assert_eq!(events[1].event_seq_no, 8); assert!(events[1].direct_files_only_batches.is_empty()); assert!(events[1].child_scan_units.is_empty()); @@ -6997,17 +7146,21 @@ mod tests { skip_entries: Vec::new(), }; - let first = build_transfer_scan_result_for_root_dir_abs( - root.path().to_str().unwrap(), - &assignment, - ) - .unwrap(); + let first = + build_transfer_scan_result_for_root_dir_abs(root.path().to_str().unwrap(), &assignment) + .unwrap(); assert!(!first.finished); assert!(!first.direct_files_only_batches.is_empty()); assert!(first.full_dir_batches.is_empty()); assert_eq!(first.child_scan_units.len(), 1); - assert_eq!(first.child_scan_units[0].scan_unit_id, assignment.scan_unit_id); - assert_eq!(first.child_scan_units[0].root_relpath, assignment.root_relpath); + assert_eq!( + first.child_scan_units[0].scan_unit_id, + assignment.scan_unit_id + ); + assert_eq!( + first.child_scan_units[0].root_relpath, + assignment.root_relpath + ); assert_eq!(first.child_scan_units[0].generation, assignment.generation); let first_entry_count = first .direct_files_only_batches @@ -7019,7 +7172,10 @@ mod tests { .len() }) .sum::(); - assert_eq!(first_entry_count, TRANSFER_SCAN_ROOT_LISTING_SLICE_ENTRY_LIMIT); + assert_eq!( + first_entry_count, + TRANSFER_SCAN_ROOT_LISTING_SLICE_ENTRY_LIMIT + ); let second_assignment = FluxonFsTransferScanAssignmentWire { scan_task_id: "task-2".to_string(), @@ -7086,7 +7242,8 @@ mod tests { } #[test] - fn build_transfer_scan_result_root_direct_fanout_only_emits_child_scan_units_without_recursing() { + fn build_transfer_scan_result_root_direct_fanout_only_emits_child_scan_units_without_recursing() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/child/payload.bin", b"xyz"); @@ -7114,14 +7271,18 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); assert!(result.full_dir_batches.is_empty()); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree ); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7133,7 +7294,8 @@ mod tests { } #[test] - fn build_transfer_scan_result_directory_direct_fanout_only_emits_child_scan_units_without_recursing() { + fn build_transfer_scan_result_directory_direct_fanout_only_emits_child_scan_units_without_recursing() + { let root = TempDir::new().unwrap(); write_file(&root, "root/child/direct.bin", b"abc"); write_file(&root, "root/child/grand/payload.bin", b"xyz"); @@ -7161,14 +7323,18 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); assert!(result.full_dir_batches.is_empty()); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child/grand".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child/grand".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree ); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7206,10 +7372,14 @@ mod tests { .unwrap(); assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); - assert_eq!(result.direct_files_only_batches[0].root_relpath, "root/child"); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.direct_files_only_batches[0].root_relpath, + "root/child" + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7218,7 +7388,10 @@ mod tests { }] ); assert!(manifest.empty_dir_relpaths.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child/grand".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child/grand".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7253,10 +7426,14 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child".to_string()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child".to_string() + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7330,9 +7507,10 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7407,7 +7585,10 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child-b".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child-b".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree @@ -7447,9 +7628,10 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7490,9 +7672,10 @@ mod tests { assert!(result.direct_files_only_batches.is_empty()); assert!(result.child_scan_units.is_empty()); assert_eq!(result.full_dir_batches.len(), 1); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.full_dir_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.full_dir_batches[0].manifest_blob, + ) + .unwrap(); assert!(manifest.entries.is_empty()); assert_eq!(manifest.empty_dir_relpaths, vec!["root".to_string()]); } @@ -7533,7 +7716,8 @@ mod tests { } #[test] - fn handle_transfer_scan_assignment_does_not_reaggregate_root_when_descendant_batch_is_durable() { + fn handle_transfer_scan_assignment_does_not_reaggregate_root_when_descendant_batch_is_durable() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/big/data.bin", b"12345"); @@ -7579,21 +7763,29 @@ mod tests { size: 3, }] ); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root")); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root/big")); - assert_eq!(child_scan_unit_roots(&result), vec!["root/small".to_string()]); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root") + ); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root/big") + ); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/small".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } #[test] - fn handle_transfer_scan_assignment_honors_cross_generation_descendant_full_dir_during_restart() { + fn handle_transfer_scan_assignment_honors_cross_generation_descendant_full_dir_during_restart() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/big/data.bin", b"12345"); @@ -7639,21 +7831,29 @@ mod tests { size: 3, }] ); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root")); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root/big")); - assert_eq!(child_scan_unit_roots(&result), vec!["root/small".to_string()]); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root") + ); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root/big") + ); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/small".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } #[test] - fn handle_transfer_scan_assignment_replays_descendant_current_layer_when_only_partial_descendant_direct_files_batch_is_durable() { + fn handle_transfer_scan_assignment_replays_descendant_current_layer_when_only_partial_descendant_direct_files_batch_is_durable() + { let root = TempDir::new().unwrap(); write_file(&root, "root/child/a.bin", b"ab"); write_file(&root, "root/child/b.bin", b"cd"); @@ -7688,10 +7888,14 @@ mod tests { assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); assert_eq!(result.direct_files_only_batches.len(), 1); - assert_eq!(result.direct_files_only_batches[0].root_relpath, "root/child"); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.direct_files_only_batches[0].root_relpath, + "root/child" + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![ @@ -7737,7 +7941,10 @@ mod tests { assert!(result.finished); assert!(result.direct_files_only_batches.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/parent".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/parent".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7772,9 +7979,10 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7782,7 +7990,10 @@ mod tests { size: 10, }] ); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7820,7 +8031,10 @@ mod tests { assert!(ok_bool(&resp)); assert!(result.finished); - assert_eq!(child_scan_unit_roots(&result), vec!["root/blocked".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/blocked".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7870,7 +8084,11 @@ mod tests { ); assert_eq!( child_scan_unit_roots(&result), - vec!["root/a".to_string(), "root/b".to_string(), "root/c".to_string()] + vec![ + "root/a".to_string(), + "root/b".to_string(), + "root/c".to_string() + ] ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); @@ -7905,7 +8123,10 @@ mod tests { ); let result = decode_result_json(&resp); assert_eq!(result.full_dir_batches.len(), 1); - assert_eq!(result.full_dir_batches[0].batch_kind, FluxonFsTransferBatchKind::SubtreeSlice); + assert_eq!( + result.full_dir_batches[0].batch_kind, + FluxonFsTransferBatchKind::SubtreeSlice + ); assert_eq!(result.full_dir_batches[0].collect_infos.len(), 1); assert_eq!( decode_symlink_notice_collect_blob( @@ -7957,11 +8178,15 @@ mod tests { assert!(result.direct_files_only_batches.is_empty()); assert!(result.child_scan_units.is_empty()); assert_eq!(result.full_dir_batches.len(), 1); - assert_eq!(result.full_dir_batches[0].batch_kind, FluxonFsTransferBatchKind::SubtreeSlice); + assert_eq!( + result.full_dir_batches[0].batch_kind, + FluxonFsTransferBatchKind::SubtreeSlice + ); assert_eq!(result.full_dir_batches[0].root_relpath, "root".to_string()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.full_dir_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.full_dir_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![ @@ -7982,7 +8207,9 @@ mod tests { FluxonFsTransferCollectInfoKind::SymlinkNotice ); let mut notices = decode_symlink_notice_collect_blob( - direct_files_only_batch.collect_infos[0].collect_blob.as_slice() + direct_files_only_batch.collect_infos[0] + .collect_blob + .as_slice(), ); notices.sort_by(|a, b| a.relpath.cmp(&b.relpath)); assert_eq!( @@ -8004,16 +8231,17 @@ mod tests { fn prepare_transfer_file_from_chunks_promotes_staged_file_to_final_path() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"ab".to_vec(), b"cde".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![ + b"ab".to_vec(), + b"cde".to_vec(), + Vec::new(), + ])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let prepared = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8038,32 +8266,31 @@ mod tests { fs::read(root.path().join("dir/file.bin")).unwrap(), b"abcde".to_vec() ); - assert!(!root - .path() - .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part") - .exists()); + assert!( + !root + .path() + .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part") + .exists() + ); } #[test] fn prepare_transfer_file_from_chunks_truncates_existing_staging_file() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let stale_staging = - root.path() - .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part"); + let stale_staging = root + .path() + .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part"); fs::create_dir_all(stale_staging.parent().unwrap()).unwrap(); fs::write(&stale_staging, b"stale-data").unwrap(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let prepared = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8076,23 +8303,23 @@ mod tests { .unwrap(); promote_prepared_transfer_file(&dst_root, prepared).unwrap(); - assert_eq!(fs::read(root.path().join("dir/file.bin")).unwrap(), b"xy".to_vec()); + assert_eq!( + fs::read(root.path().join("dir/file.bin")).unwrap(), + b"xy".to_vec() + ); } #[test] fn prepare_transfer_file_from_chunks_rejects_size_mismatch_and_keeps_staging_file() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let err = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8275,15 +8502,10 @@ mod tests { let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, _read_offset, _length| Ok(file_bytes.clone()) - }, - ) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, _read_offset, _length| Ok(file_bytes.clone()) + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8303,7 +8525,10 @@ mod tests { create_dir_all_with_parent_dir_chmod_retry(&target).unwrap(); assert!(target.is_dir()); - assert_eq!(fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, 0o777); + assert_eq!( + fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, + 0o777 + ); } #[cfg(unix)] @@ -8318,21 +8543,19 @@ mod tests { let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, _read_offset, _length| Ok(file_bytes.clone()) - }, - ) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, _read_offset, _length| Ok(file_bytes.clone()) + }) .unwrap(); assert_eq!(result.file_results.len(), 1); assert!(dst_root.is_dir()); assert_eq!(fs::read(dst_root.join("dir/file.bin")).unwrap(), file_bytes); - assert_eq!(fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, 0o777); + assert_eq!( + fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, + 0o777 + ); } #[test] @@ -8349,47 +8572,48 @@ mod tests { let assignment = assignment.clone(); let heartbeat_attempts = heartbeat_attempts.clone(); move || { - retry_transfer_worker_rpc_with_backoff( - &assignment, - "checkpoint", - "test-checkpoint", - BackoffConfig { - initial_secs: 0, - max_secs: 0, - }, - WarnConfig { - warn_interval_secs: 0, - }, - || { - let attempt = - heartbeat_attempts.fetch_add(1, Ordering::SeqCst) + 1; - if attempt < 3 { - return Err(TransferWorkerRpcFailure::Retryable { - detail: format!( - "transient heartbeat failure attempt={}", - attempt - ), - }); - } - Ok(()) - }, - ) - .map_err(TransferWorkerExecutionError::fatal) - } + retry_transfer_worker_rpc_with_backoff( + &assignment, + "checkpoint", + "test-checkpoint", + BackoffConfig { + initial_secs: 0, + max_secs: 0, + }, + WarnConfig { + warn_interval_secs: 0, + }, + || { + let attempt = heartbeat_attempts.fetch_add(1, Ordering::SeqCst) + 1; + if attempt < 3 { + return Err(TransferWorkerRpcFailure::Retryable { + detail: format!( + "transient heartbeat failure attempt={}", + attempt + ), + }); + } + Ok(()) + }, + ) + .map_err(TransferWorkerExecutionError::fatal) + } }, { let file_bytes = file_bytes.clone(); move |file, read_offset, _length| { - if file.relpath != "dir/file.bin" { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("unexpected file relpath: {}", file.relpath), - })))); - } - if read_offset == 0 { - return Ok(file_bytes.clone()); + if file.relpath != "dir/file.bin" { + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!("unexpected file relpath: {}", file.relpath), + }), + ))); + } + if read_offset == 0 { + return Ok(file_bytes.clone()); + } + Ok(Vec::new()) } - Ok(Vec::new()) - } }, ) .unwrap(); @@ -8409,24 +8633,20 @@ mod tests { let file_bytes = b"payload".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); let read_attempts = Arc::new(AtomicUsize::new(0)); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let assignment = assignment.clone(); - let file_bytes = file_bytes.clone(); - let read_attempts = read_attempts.clone(); - move |file, read_offset, _length| { + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let assignment = assignment.clone(); + let file_bytes = file_bytes.clone(); + let read_attempts = read_attempts.clone(); + move |file, read_offset, _length| { if file.relpath != "dir/file.bin" { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("unexpected file relpath: {}", file.relpath), - })))); + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!("unexpected file relpath: {}", file.relpath), + }), + ))); } - let op_detail = format!( - "test-read relpath={} offset={}", - file.relpath, read_offset - ); + let op_detail = + format!("test-read relpath={} offset={}", file.relpath, read_offset); retry_transfer_worker_rpc_with_backoff( &assignment, "read_chunk", @@ -8443,10 +8663,7 @@ mod tests { let attempt = read_attempts.fetch_add(1, Ordering::SeqCst) + 1; if attempt < 3 { return Err(TransferWorkerRpcFailure::Retryable { - detail: format!( - "transient read failure attempt={}", - attempt - ), + detail: format!("transient read failure attempt={}", attempt), }); } return Ok(file_bytes.clone()); @@ -8456,8 +8673,7 @@ mod tests { ) .map_err(TransferWorkerExecutionError::fatal) } - }, - ) + }) .unwrap(); assert_eq!(read_attempts.load(Ordering::SeqCst), 3); @@ -8481,23 +8697,23 @@ mod tests { { let checkpoint_calls = checkpoint_calls.clone(); move || { - let calls = checkpoint_calls.fetch_add(1, Ordering::SeqCst) + 1; - if calls >= 4 { - return Err(TransferWorkerExecutionError::Stop( - FluxonFsTransferWorkerStopReasonWire::Superseded, - )); + let calls = checkpoint_calls.fetch_add(1, Ordering::SeqCst) + 1; + if calls >= 4 { + return Err(TransferWorkerExecutionError::Stop( + FluxonFsTransferWorkerStopReasonWire::Superseded, + )); + } + Ok(()) } - Ok(()) - } }, { let file_bytes = file_bytes.clone(); move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); + if read_offset == 0 { + return Ok(file_bytes.clone()); + } + Ok(Vec::new()) } - Ok(Vec::new()) - } }, ); assert!(matches!( @@ -8567,12 +8783,7 @@ mod tests { break; } if max_in_flight - .compare_exchange( - observed, - current, - Ordering::SeqCst, - Ordering::SeqCst, - ) + .compare_exchange(observed, current, Ordering::SeqCst, Ordering::SeqCst) .is_ok() { break; @@ -8588,8 +8799,14 @@ mod tests { assert_eq!(result.file_results.len(), 2); assert!(max_in_flight.load(Ordering::SeqCst) >= 2); - assert_eq!(fs::read(root.path().join("dir/a.bin")).unwrap(), b"xxx".to_vec()); - assert_eq!(fs::read(root.path().join("dir/b.bin")).unwrap(), b"xxx".to_vec()); + assert_eq!( + fs::read(root.path().join("dir/a.bin")).unwrap(), + b"xxx".to_vec() + ); + assert_eq!( + fs::read(root.path().join("dir/b.bin")).unwrap(), + b"xxx".to_vec() + ); } #[test] @@ -8597,29 +8814,25 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "dir/link.bin".to_string(), - link_target: "dir/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "dir/link.bin".to_string(), + link_target: "dir/file.bin".to_string(), + }]) + .unwrap(); let assignment = FluxonFsTransferWorkerAssignmentWire { collect_infos: collect_infos.clone(), ..test_worker_assignment("dir/file.bin", file_bytes.len() as i64) }; - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); - } - Ok(Vec::new()) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, read_offset, _length| { + if read_offset == 0 { + return Ok(file_bytes.clone()); } - }, - ) + Ok(Vec::new()) + } + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8633,7 +8846,11 @@ mod tests { "fluxon_collect_info/batches/batch/symlinks.jsonl" ); assert_eq!( - fs::read(root.path().join("fluxon_collect_info/batches/batch/symlinks.jsonl")).unwrap(), + fs::read( + root.path() + .join("fluxon_collect_info/batches/batch/symlinks.jsonl") + ) + .unwrap(), collect_infos[0].collect_blob ); } @@ -8643,16 +8860,19 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let assignment = FluxonFsTransferWorkerAssignmentWire { - manifest_blob: build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { - relpath: "dir/good.bin".to_string(), - size: 5, - }, - FluxonFsTransferScanFrontierEntry { - relpath: "dir/bad.bin".to_string(), - size: 5, - }, - ], Vec::new()) + manifest_blob: build_transfer_manifest_blob( + vec![ + FluxonFsTransferScanFrontierEntry { + relpath: "dir/good.bin".to_string(), + size: 5, + }, + FluxonFsTransferScanFrontierEntry { + relpath: "dir/bad.bin".to_string(), + size: 5, + }, + ], + Vec::new(), + ) .unwrap(), ..test_worker_assignment("dir/good.bin", 5) }; @@ -8816,20 +9036,16 @@ mod tests { .unwrap(); let progress_heartbeat_count = Arc::new(AtomicUsize::new(0)); - gate.ensure_continue( - false, - TRANSFER_WORKER_HEARTBEAT_EMPTY_DIR_PROGRESS_COUNT, - { - let progress_heartbeat_count = progress_heartbeat_count.clone(); - move |_heartbeat_unix_ms, heartbeat_detail| { - assert_eq!(heartbeat_detail, "empty_dir_progress"); - progress_heartbeat_count.fetch_add(1, Ordering::SeqCst); - Ok(FluxonFsTransferWorkerHeartbeatResultWire::continue_running( - chrono::Utc::now().timestamp_millis() + 60_000, - )) - } - }, - ) + gate.ensure_continue(false, TRANSFER_WORKER_HEARTBEAT_EMPTY_DIR_PROGRESS_COUNT, { + let progress_heartbeat_count = progress_heartbeat_count.clone(); + move |_heartbeat_unix_ms, heartbeat_detail| { + assert_eq!(heartbeat_detail, "empty_dir_progress"); + progress_heartbeat_count.fetch_add(1, Ordering::SeqCst); + Ok(FluxonFsTransferWorkerHeartbeatResultWire::continue_running( + chrono::Utc::now().timestamp_millis() + 60_000, + )) + } + }) .unwrap(); gate.ensure_continue( @@ -8927,20 +9143,15 @@ mod tests { let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); - } - Ok(Vec::new()) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, read_offset, _length| { + if read_offset == 0 { + return Ok(file_bytes.clone()); } - }, - ) + Ok(Vec::new()) + } + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8948,7 +9159,10 @@ mod tests { cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment).unwrap(); - assert_eq!(fs::read(root.path().join("dir/file.bin")).unwrap(), file_bytes); + assert_eq!( + fs::read(root.path().join("dir/file.bin")).unwrap(), + file_bytes + ); assert!(!root.path().join(".fluxon.stage").exists()); } @@ -8957,11 +9171,12 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "root/link-file.bin".to_string(), - link_target: "target/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "root/link-file.bin".to_string(), + link_target: "target/file.bin".to_string(), + }]) + .unwrap(); let assignment = FluxonFsTransferWorkerAssignmentWire { collect_infos: collect_infos.clone(), ..test_worker_assignment("dir/file.bin", file_bytes.len() as i64) @@ -8976,9 +9191,11 @@ mod tests { let result = execute_transfer_worker_assignment( &assignment, &dst_root, - || Err(TransferWorkerExecutionError::Stop( - FluxonFsTransferWorkerStopReasonWire::Superseded, - )), + || { + Err(TransferWorkerExecutionError::Stop( + FluxonFsTransferWorkerStopReasonWire::Superseded, + )) + }, { let file_bytes = file_bytes.clone(); move |_file, read_offset, _length| { @@ -8996,11 +9213,20 @@ mod tests { FluxonFsTransferWorkerStopReasonWire::Superseded )) )); - assert!(root.path().join(prepared_collect.staging_relpath.as_str()).exists()); + assert!( + root.path() + .join(prepared_collect.staging_relpath.as_str()) + .exists() + ); cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment).unwrap(); assert!(!root.path().join(".fluxon.stage").exists()); - assert!(!root.path().join(prepared_collect.staging_relpath.as_str()).exists()); + assert!( + !root + .path() + .join(prepared_collect.staging_relpath.as_str()) + .exists() + ); } } diff --git a/fluxon_rs/fluxon_fs/src/cache_controller.rs b/fluxon_rs/fluxon_fs/src/cache_controller.rs index 8a0845c..77bba18 100644 --- a/fluxon_rs/fluxon_fs/src/cache_controller.rs +++ b/fluxon_rs/fluxon_fs/src/cache_controller.rs @@ -6,6 +6,7 @@ use dashmap::DashMap; use tokio::sync::mpsc; use fluxon_fs_core::config::{FluxonFsCacheControllerConfig, FluxonFsRequestIdentity}; +use fluxon_util::run_async_from_sync::spawn_blocking_allow_sync_async_bridge; /// Key of a single cacheable piece of an S3-gateway object. See design §4.1. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -128,7 +129,6 @@ impl CacheController { let rx = shared_rx.clone(); let inflight_clone = inflight.clone(); let fn_clone = stage_piece_fn.clone(); - let rt = rt_handle.clone(); let queue_depth_clone = queue_depth.clone(); let stage_success_count_clone = stage_success_count.clone(); let stage_fail_count_clone = stage_fail_count.clone(); @@ -142,7 +142,6 @@ impl CacheController { inflight_clone, fn_clone, stage_piece_range_fn_clone, - rt, queue_depth_clone, stage_success_count_clone, stage_fail_count_clone, @@ -280,7 +279,6 @@ async fn stage_worker_loop( inflight: Arc>, stage_piece_fn: StagePieceFn, stage_piece_range_fn: StagePieceRangeFn, - rt_handle: tokio::runtime::Handle, queue_depth: Arc, stage_success_count: Arc, stage_fail_count: Arc, @@ -355,20 +353,19 @@ async fn stage_worker_loop( let identity_clone = task.identity.clone(); let fn_clone = stage_piece_fn.clone(); let fn_range_clone = stage_piece_range_fn.clone(); - let result = rt_handle - .spawn_blocking(move || { - if staged_piece_count_clone > 1 { - stage_piece_range_fn_invoke( - &fn_range_clone, - &key_clone, - staged_piece_count_clone, - identity_clone.as_ref(), - ) - } else { - stage_piece_fn_invoke(&fn_clone, &key_clone, identity_clone.as_ref()) - } - }) - .await; + let result = spawn_blocking_allow_sync_async_bridge(move || { + if staged_piece_count_clone > 1 { + stage_piece_range_fn_invoke( + &fn_range_clone, + &key_clone, + staged_piece_count_clone, + identity_clone.as_ref(), + ) + } else { + stage_piece_fn_invoke(&fn_clone, &key_clone, identity_clone.as_ref()) + } + }) + .await; // Always release inflight. for piece_key in &staged_piece_keys { @@ -429,8 +426,9 @@ fn now_ms() -> i64 { #[cfg(test)] mod tests { use super::*; - use std::sync::mpsc; + use fluxon_util::run_async_from_sync::SyncAsyncBridge; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::mpsc; use std::sync::{Condvar, Mutex}; use tokio::time::{Duration, sleep}; @@ -483,6 +481,54 @@ mod tests { assert_eq!(snapshot.queue_depth, 0); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn stage_worker_allows_sync_async_bridge_in_blocking_callback() { + let handle = tokio::runtime::Handle::current(); + let (stage_result_tx, stage_result_rx) = mpsc::sync_channel(1); + let stage_piece_fn: StagePieceFn = Arc::new(move |_key, _identity| { + let value = handle + .run_async_from_sync(async { 7_usize }) + .map_err(|err| err.to_string())?; + let _ = stage_result_tx.send(value); + Ok(()) + }); + let stage_piece_range_fn: StagePieceRangeFn = + Arc::new(move |_key, _count, _identity| Ok(())); + let ctrl = CacheController::start( + FluxonFsCacheControllerConfig { + stage_worker_count: 1, + ..FluxonFsCacheControllerConfig::default() + }, + stage_piece_fn, + stage_piece_range_fn, + tokio::runtime::Handle::current(), + ); + + assert_eq!( + ctrl.handle_suggest(sample_key(), None), + SuggestOutcome::Enqueued + ); + assert_eq!( + stage_result_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("stage worker did not run within 5s"), + 7 + ); + + for _ in 0..50 { + let snapshot = ctrl.stats_snapshot(); + if snapshot.stage_success_count == 1 && snapshot.stage_fail_count == 0 { + break; + } + sleep(Duration::from_millis(20)).await; + } + + let snapshot = ctrl.stats_snapshot(); + assert_eq!(snapshot.stage_success_count, 1); + assert_eq!(snapshot.stage_fail_count, 0); + assert_eq!(snapshot.stage_panic_count, 0); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn suggest_dedupes_while_inflight() { let stage_calls = Arc::new(AtomicUsize::new(0)); diff --git a/fluxon_rs/fluxon_fs/src/local_transfer_checker.rs b/fluxon_rs/fluxon_fs/src/local_transfer_checker.rs index cd1c1e6..1c343c5 100644 --- a/fluxon_rs/fluxon_fs/src/local_transfer_checker.rs +++ b/fluxon_rs/fluxon_fs/src/local_transfer_checker.rs @@ -1725,9 +1725,9 @@ mod tests { const LOCAL_TRANSFER_TEST_CHILD_TRANSFER_STATE_STORE_JSON_ENV: &str = "FLUXON_FS_TEST_CHILD_TRANSFER_STATE_STORE_JSON"; const LOCAL_TRANSFER_TEST_IO_ROOT: &str = - "/mnt/nvme0/fluxon_fs_transfer_tikv/rust_local_checker_io"; + "/tmp/fluxon-example/fluxon_fs_transfer_tikv/rust_local_checker_io"; const LOCAL_TRANSFER_TEST_TIKV_WORK_ROOT: &str = - "/mnt/nvme0/fluxon_fs_transfer_tikv/rust_local_checker"; + "/tmp/fluxon-example/fluxon_fs_transfer_tikv/rust_local_checker"; const LOCAL_TRANSFER_TEST_TIKV_READY_TIMEOUT_SECS: u64 = 180; const LOCAL_TRANSFER_TEST_TIKV_PD_LEASE_SECS: u64 = 60; const LOCAL_TRANSFER_TEST_TIKV_UNIFIED_READPOOL_MAX_THREADS: u64 = 4; diff --git a/fluxon_rs/fluxon_fs/src/master_http.rs b/fluxon_rs/fluxon_fs/src/master_http.rs index 387f7ac..64b6cdf 100644 --- a/fluxon_rs/fluxon_fs/src/master_http.rs +++ b/fluxon_rs/fluxon_fs/src/master_http.rs @@ -36,6 +36,7 @@ use fluxon_kv::{ ClusterEvent, ClusterMember, ConfigArg, MembershipEventReceiver, run_client_with_startup_member_metadata, }; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use fluxon_util::run_async_from_sync::spawn_blocking_allow_sync_async_bridge; use fluxon_util::{FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2}; use futures::future::BoxFuture; @@ -972,6 +973,7 @@ async fn async_main( .cluster_manager_view() .cluster_manager() .etcd_endpoints(); + let etcd_backend = ManagedEtcdClient::acquire(EtcdEndpointSet::new(endpoints.clone())?); let rt_handle = rt.handle().clone(); let api = Arc::new( @@ -1001,7 +1003,8 @@ async fn async_main( let s3_backend = Arc::new(FsS3BackendAgent::new(s3_agent, s3_kv_miss_policy)); let etcd2 = Arc::new(tokio::sync::Mutex::new( - EtcdClient::connect(endpoints.clone(), None) + etcd_backend + .client() .await .with_context(|| "connect etcd (panel runtime)")?, )); @@ -1075,7 +1078,7 @@ async fn async_main( register_mount_registry_rpc( rt_handle.clone(), &cluster_name, - endpoints.clone(), + etcd_backend.clone(), api.clone(), s3_state.clone(), )?; @@ -1085,11 +1088,12 @@ async fn async_main( api.clone(), export_registry.clone(), cluster_name.clone(), - endpoints.clone(), + etcd_backend.clone(), s3_state.clone(), )?; - let mut etcd = EtcdClient::connect(endpoints.clone(), None) + let mut etcd = etcd_backend + .client() .await .with_context(|| "connect etcd")?; @@ -1100,7 +1104,7 @@ async fn async_main( kv_framework.clone(), export_registry.clone(), cluster_name.clone(), - endpoints.clone(), + etcd_backend, s3_state.clone(), FLUXON_FS_CONTROL_SCHEMA_VERSION, pull_interval_ms, @@ -1468,7 +1472,7 @@ fn build_mount_exports_from_registry_locked( fn register_mount_registry_rpc( rt: Handle, cluster_name: &str, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, api: Arc, s3_state: Arc, ) -> KvResult<()> { @@ -1559,7 +1563,7 @@ fn register_mount_registry_rpc( // the filesystem data-plane mount itself. // // Therefore we spawn the etcd write asynchronously and return ok immediately. - let endpoints2 = etcd_endpoints.clone(); + let etcd_backend = etcd_backend.clone(); let key2 = key.clone(); let value2 = value.clone(); let s3_state2 = s3_state.clone(); @@ -1575,7 +1579,7 @@ fn register_mount_registry_rpc( } }); rt.spawn(async move { - let mut etcd = match EtcdClient::connect(endpoints2, None).await { + let mut etcd = match etcd_backend.client().await { Ok(v) => v, Err(e) => { tracing::warn!( @@ -1692,7 +1696,7 @@ fn register_agent_exports_push_rpc( api: Arc, export_registry: Arc>, cluster_name: String, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, s3_state: Arc, ) -> KvResult<()> { let path = FS_MASTER_AGENT_EXPORTS_PUSH_RPC_PATH.to_string(); @@ -1791,7 +1795,7 @@ fn register_agent_exports_push_rpc( export_registry.clone(), FLUXON_FS_CONTROL_SCHEMA_VERSION, cluster_name.clone(), - etcd_endpoints.clone(), + etcd_backend.clone(), ); let mut out: FlatDict = FlatDict::new(); @@ -1882,7 +1886,7 @@ fn start_export_registry_manager( framework: Arc, export_registry: Arc>, cluster_name: String, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, s3_state: Arc, schema_version: i64, pull_interval_ms: u64, @@ -1909,7 +1913,7 @@ fn start_export_registry_manager( framework.clone(), export_registry.clone(), cluster_name.clone(), - etcd_endpoints.clone(), + etcd_backend.clone(), s3_state.clone(), schema_version, pull_interval_ms, @@ -1941,7 +1945,7 @@ fn start_export_registry_manager( framework.clone(), export_registry.clone(), cluster_name.clone(), - etcd_endpoints.clone(), + etcd_backend.clone(), s3_state.clone(), schema_version, pull_interval_ms, @@ -1963,7 +1967,7 @@ fn start_export_registry_manager( &member_id, schema_version, &cluster_name, - &etcd_endpoints, + &etcd_backend, ); } } @@ -1986,7 +1990,7 @@ fn remove_agent_exports_and_persist_snapshot( agent_instance_key: &str, schema_version: i64, cluster_name: &str, - etcd_endpoints: &[String], + etcd_backend: &ManagedEtcdClient, ) { { let mut reg = export_registry.lock(); @@ -1997,7 +2001,7 @@ fn remove_agent_exports_and_persist_snapshot( export_registry, schema_version, cluster_name.to_string(), - etcd_endpoints.to_vec(), + etcd_backend.clone(), ); if let Err(e) = s3_state.delete_fs_export_registry_for_agent(agent_instance_key) { tracing::warn!( @@ -2014,7 +2018,7 @@ fn spawn_sync_agent_exports_until_ready( framework: Arc, export_registry: Arc>, cluster_name: String, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, s3_state: Arc, schema_version: i64, pull_interval_ms: u64, @@ -2136,7 +2140,7 @@ fn spawn_sync_agent_exports_until_ready( export_registry.clone(), schema_version, cluster_name.clone(), - etcd_endpoints.clone(), + etcd_backend.clone(), ); break; } @@ -2264,7 +2268,7 @@ fn spawn_persist_export_registry_snapshot_to_etcd( export_registry: Arc>, schema_version: i64, cluster_name: String, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, ) { let snapshot = { let reg = export_registry.lock(); @@ -2292,7 +2296,7 @@ fn spawn_persist_export_registry_snapshot_to_etcd( let key = export_registry_snapshot_etcd_key(&cluster_name); rt.spawn(async move { - let mut etcd = match EtcdClient::connect(etcd_endpoints, None).await { + let mut etcd = match etcd_backend.client().await { Ok(v) => v, Err(e) => { tracing::warn!( diff --git a/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs b/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs index 827bb23..c199366 100644 --- a/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs +++ b/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs @@ -5344,10 +5344,9 @@ mod tests { }; use crate::transfer::encode_transfer_manifest_blob_with_empty_dirs; use fluxon_fs_core::config::{ - FS_CACHE_DEFAULT_WRITE_SESSION_TARGET_INFLIGHT_BYTES_V1, - FS_EXPORT_DEFAULT_INLINE_BYTES_MAX_BYTES_V1, - FS_EXPORT_DEFAULT_METADATA_CACHE_TTL_MS_V1, FLUXON_FS_LOCAL_TRANSFER_CHECK_DST_EXPORT, FLUXON_FS_LOCAL_TRANSFER_CHECK_SRC_EXPORT, + FS_CACHE_DEFAULT_WRITE_SESSION_TARGET_INFLIGHT_BYTES_V1, + FS_EXPORT_DEFAULT_INLINE_BYTES_MAX_BYTES_V1, FS_EXPORT_DEFAULT_METADATA_CACHE_TTL_MS_V1, FluxonFsAccessModel, FluxonFsAccessUser, FluxonFsExport, FluxonFsExportRoutingMode, FluxonFsGlobalConfig, FluxonFsLocalTransferCheckJobSpecWire, FluxonFsRequestIdentity, FluxonFsS3GatewayConfig, FluxonFsS3KvMissPolicy, FluxonFsS3PermissionAccount, @@ -5364,7 +5363,7 @@ mod tests { use std::sync::OnceLock; use uuid::Uuid; - const TEST_TIKV_WORK_ROOT: &str = "/mnt/nvme0/fluxon_fs_transfer_tikv/rust_gateway"; + const TEST_TIKV_WORK_ROOT: &str = "/tmp/fluxon-example/fluxon_fs_transfer_tikv/rust_gateway"; const TEST_TIKV_READY_TIMEOUT_SECS: u64 = 180; const TEST_TIKV_PD_LEASE_SECS: u64 = 60; const TEST_TIKV_UNIFIED_READPOOL_MAX_THREADS: u64 = 4; @@ -6753,29 +6752,29 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" #[test] fn test_build_fs_master_admin_managed_agent_exports_uses_runtime_and_overlay_per_agent() { let online_agent_ids = BTreeSet::from([ - "fluxon_fs_agent_infra44-ThinkStation-PX".to_string(), - "fluxon_fs_agent_infra46-ThinkStation-PX".to_string(), + "fluxon_fs_agent_example-node-a".to_string(), + "fluxon_fs_agent_example-node-b".to_string(), ]); let runtime_agent_exports = vec![ super::FsMasterAdminRuntimeAgentExports { - agent_instance_key: "fluxon_fs_agent_infra44-ThinkStation-PX".to_string(), + agent_instance_key: "fluxon_fs_agent_example-node-a".to_string(), runtime_exports: vec![ super::FsMasterAdminRuntimeExportRecord { - export_name: "deployer-runtime-infra44".to_string(), - remote_root_dir_abs: "/mnt/nvme0/store_team_dev/fluxon_deploy".to_string(), + export_name: "deployer-runtime-example-node-a".to_string(), + remote_root_dir_abs: "/tmp/fluxon-example/deploy".to_string(), }, super::FsMasterAdminRuntimeExportRecord { export_name: "fluxon-release".to_string(), - remote_root_dir_abs: - "/mnt/nvme0/store_team_dev/fluxon_deploy/fluxon_release".to_string(), + remote_root_dir_abs: "/tmp/fluxon-example/deploy/fluxon_release" + .to_string(), }, ], }, super::FsMasterAdminRuntimeAgentExports { - agent_instance_key: "fluxon_fs_agent_infra46-ThinkStation-PX".to_string(), + agent_instance_key: "fluxon_fs_agent_example-node-b".to_string(), runtime_exports: vec![super::FsMasterAdminRuntimeExportRecord { - export_name: "deployer-runtime-infra46".to_string(), - remote_root_dir_abs: "/mnt/nvme0/store_team_dev/fluxon_deploy".to_string(), + export_name: "deployer-runtime-example-node-b".to_string(), + remote_root_dir_abs: "/tmp/fluxon-example/deploy".to_string(), }], }, ]; @@ -6793,7 +6792,7 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" .iter() .map(|record| record.export_name.as_str()) .collect::>(), - vec!["deployer-runtime-infra44", "fluxon-release"] + vec!["deployer-runtime-example-node-a", "fluxon-release"] ); assert_eq!( managed_agent_exports[1] @@ -6801,7 +6800,7 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" .iter() .map(|record| record.export_name.as_str()) .collect::>(), - vec!["deployer-runtime-infra46"] + vec!["deployer-runtime-example-node-b"] ); } @@ -6809,12 +6808,10 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" async fn test_snapshot_fs_master_admin_filters_offline_agents() { let fs_master_admin_backend = StaticFsMasterAdminBackend { members: vec![ - test_agent_member("fluxon_fs_agent_infra44-ThinkStation-PX"), - test_agent_member("fluxon_fs_agent_infra46-ThinkStation-PX"), + test_agent_member("fluxon_fs_agent_example-node-a"), + test_agent_member("fluxon_fs_agent_example-node-b"), ], - online_member_ids: BTreeSet::from([ - "fluxon_fs_agent_infra44-ThinkStation-PX".to_string() - ]), + online_member_ids: BTreeSet::from(["fluxon_fs_agent_example-node-a".to_string()]), }; let st = GatewayState::new( "test_cluster".to_string(), @@ -6838,28 +6835,28 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" ) .unwrap(); st.replace_fs_export_registry_for_agent( - "fluxon_fs_agent_infra44-ThinkStation-PX", + "fluxon_fs_agent_example-node-a", &[super::FsExportRegistryRecord { - export_name: "deployer-runtime-infra44".to_string(), - agent_instance_key: "fluxon_fs_agent_infra44-ThinkStation-PX".to_string(), - remote_root_dir_abs: "/mnt/nvme0/store_team_dev/fluxon_deploy".to_string(), + export_name: "deployer-runtime-example-node-a".to_string(), + agent_instance_key: "fluxon_fs_agent_example-node-a".to_string(), + remote_root_dir_abs: "/tmp/fluxon-example/deploy".to_string(), export: super::agent_registry_export_for_name_and_root_v1( - "deployer-runtime-infra44", - "/mnt/nvme0/store_team_dev/fluxon_deploy", + "deployer-runtime-example-node-a", + "/tmp/fluxon-example/deploy", ), updated_unix_ms: 1, }], ) .unwrap(); st.replace_fs_export_registry_for_agent( - "fluxon_fs_agent_infra46-ThinkStation-PX", + "fluxon_fs_agent_example-node-b", &[super::FsExportRegistryRecord { - export_name: "deployer-runtime-infra46".to_string(), - agent_instance_key: "fluxon_fs_agent_infra46-ThinkStation-PX".to_string(), - remote_root_dir_abs: "/mnt/nvme0/store_team_dev/fluxon_deploy".to_string(), + export_name: "deployer-runtime-example-node-b".to_string(), + agent_instance_key: "fluxon_fs_agent_example-node-b".to_string(), + remote_root_dir_abs: "/tmp/fluxon-example/deploy".to_string(), export: super::agent_registry_export_for_name_and_root_v1( - "deployer-runtime-infra46", - "/mnt/nvme0/store_team_dev/fluxon_deploy", + "deployer-runtime-example-node-b", + "/tmp/fluxon-example/deploy", ), updated_unix_ms: 1, }], @@ -6871,12 +6868,12 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" assert_eq!(snapshot.runtime_agent_exports.len(), 1); assert_eq!( snapshot.runtime_agent_exports[0].agent_instance_key, - "fluxon_fs_agent_infra44-ThinkStation-PX" + "fluxon_fs_agent_example-node-a" ); assert_eq!(snapshot.managed_agent_exports.len(), 1); assert_eq!( snapshot.managed_agent_exports[0].agent_instance_key, - "fluxon_fs_agent_infra44-ThinkStation-PX" + "fluxon_fs_agent_example-node-a" ); } @@ -10679,13 +10676,13 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" #[tokio::test] async fn test_bucket_root_put_is_noop_success_for_existing_export() { let backend = Arc::new(ObjectBackend::default()); - let st = test_state_with_buckets(backend, &["deployer-runtime-infra44"]); + let st = test_state_with_buckets(backend, &["deployer-runtime-example-node-a"]); let resp = super::handle_any_authed_impl( st.clone(), test_auth_ctx(&st), Method::PUT, - "deployer-runtime-infra44".to_string(), - "/deployer-runtime-infra44".parse().unwrap(), + "deployer-runtime-example-node-a".to_string(), + "/deployer-runtime-example-node-a".parse().unwrap(), HeaderMap::new(), Body::empty(), ) @@ -10741,14 +10738,14 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" #[tokio::test] async fn test_put_object_after_bucket_root_put_writes_object() { let backend = Arc::new(ObjectBackend::default()); - let st = test_state_with_buckets(backend.clone(), &["deployer-runtime-infra44"]); + let st = test_state_with_buckets(backend.clone(), &["deployer-runtime-example-node-a"]); let bucket_resp = super::handle_any_authed_impl( st.clone(), test_auth_ctx(&st), Method::PUT, - "deployer-runtime-infra44".to_string(), - "/deployer-runtime-infra44".parse().unwrap(), + "deployer-runtime-example-node-a".to_string(), + "/deployer-runtime-example-node-a".parse().unwrap(), HeaderMap::new(), Body::empty(), ) @@ -10760,8 +10757,8 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" st.clone(), test_auth_ctx(&st), Method::PUT, - "deployer-runtime-infra44/download/rclone-e2e.txt".to_string(), - "/deployer-runtime-infra44/download/rclone-e2e.txt" + "deployer-runtime-example-node-a/download/rclone-e2e.txt".to_string(), + "/deployer-runtime-example-node-a/download/rclone-e2e.txt" .parse() .unwrap(), HeaderMap::new(), @@ -10773,7 +10770,7 @@ max-background-jobs = {TEST_TIKV_RAFTDB_MAX_BACKGROUND_JOBS}\n" assert_eq!(object_resp.status(), StatusCode::OK); assert_eq!( backend - .get("deployer-runtime-infra44", "download/rclone-e2e.txt") + .get("deployer-runtime-example-node-a", "download/rclone-e2e.txt") .unwrap(), b"fluxon-rclone-e2e\n".to_vec() ); diff --git a/fluxon_rs/fluxon_kv/Cargo.toml b/fluxon_rs/fluxon_kv/Cargo.toml index 22ff136..83b4993 100644 --- a/fluxon_rs/fluxon_kv/Cargo.toml +++ b/fluxon_rs/fluxon_kv/Cargo.toml @@ -75,6 +75,8 @@ bytes = "1" pprof = { version = "0.15", features = ["flamegraph"] } hex = "0.4" sha2 = "0.10" +io-uring = "0.7" +foyer = { version = "0.22.3", features = ["serde"] } tokio-tungstenite = { version = "0.21", default-features = false, features = ["connect", "handshake"], optional = true } sockudo-ws = { version = "^1.7.4", default-features = false, features = ["tokio-runtime", "fastrand"], optional = true } diff --git a/fluxon_rs/fluxon_kv/examples/cluster_example.rs b/fluxon_rs/fluxon_kv/examples/cluster_example.rs index 7fd158d..b4c9d66 100644 --- a/fluxon_rs/fluxon_kv/examples/cluster_example.rs +++ b/fluxon_rs/fluxon_kv/examples/cluster_example.rs @@ -46,7 +46,7 @@ struct Args { instance_name: String, /// etcd endpoint - #[arg(long, default_value = "127.0.0.1:2379")] + #[arg(long, default_value = "http://127.0.0.1:2379")] etcd_endpoint: String, /// Cluster name diff --git a/fluxon_rs/fluxon_kv/examples/p2p_rpc_example.rs b/fluxon_rs/fluxon_kv/examples/p2p_rpc_example.rs index 465eb04..567599f 100644 --- a/fluxon_rs/fluxon_kv/examples/p2p_rpc_example.rs +++ b/fluxon_rs/fluxon_kv/examples/p2p_rpc_example.rs @@ -332,7 +332,7 @@ struct Args { instance_name: String, #[arg(short, long, default_value_t = String::from("test-cluster"))] cluster_name: String, - #[arg(long, default_value_t = String::from("http://10.126.126.235:25579"))] + #[arg(long, default_value_t = String::from("http://192.0.2.10:2379"))] etcd_endpoint: String, } diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs index 8e37608..b0d4d1e 100755 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs @@ -124,6 +124,13 @@ pub async fn handle_batch_delete_client_kv_meta_cache( delete_item.put_version ); + // A remote GET holds this per-key lock until it has installed the + // MemoryInfo returned by GetDone. Waiting here prevents an eviction + // from arriving before that install and leaving a stale holder cached. + let get_lock = client_inner + .get_remote_kv_lock + .get_lock(delete_item.key.clone()); + let _get_guard = get_lock.lock().await; remove_local_cached_info_for_delete(client_inner, delete_item); deleted_count += 1; diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs index 60986be..997686a 100755 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs @@ -10,6 +10,7 @@ use crate::cluster_manager::NodeIDString; use crate::cluster_manager::{ META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME, }; +use crate::master_kv_router::msg_pack::GetSourceKind; use crate::memholder::ExternalMemHolderInfo; use crate::memholder::MemholderManagerTrait; use crate::memholder::NodeHolderKey; @@ -340,7 +341,7 @@ impl HandlerForExternalClient for ClientKvApi { self.validate_requester_owner_status_updated(req.started_time)?; // dummy implementation, tmp owner user memholder for temporary holding to make self memholder - let (memholder, _) = match inner.get(&req.key).await? { + let (memholder, get_info) = match inner.get(&req.key).await? { Some(holder) => holder, None => { return Ok(ExternalGetResp { @@ -368,6 +369,9 @@ impl HandlerForExternalClient for ClientKvApi { offset: memholder.get_offset(), len: memholder.get_length() as u32, holder_id: memholder.holder_id(), + source_kind: get_info + .map(|info| info.source_kind()) + .unwrap_or(GetSourceKind::Memory), }; Ok(ExternalGetResp { external_memholder_info: Some(external_memholder_info), @@ -396,6 +400,7 @@ impl HandlerForExternalClient for ClientKvApi { &req.key, req.len as u32, req.reject_if_inflight_same_key, + req.reject_if_exists, req.preferred_sub_cluster.as_deref(), source_node_id, ) diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs index f309dd0..8cf0b28 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs @@ -13,7 +13,7 @@ use crate::{ cluster_manager::NodeID, master_kv_router::msg_pack::{ GetAllocationMode, GetDoneReq, GetDoneResp, GetMetaReq, GetMetaResp, GetRevokeReq, - GetStartReq, GetStartResp, + GetSourceKind, GetStartReq, GetStartResp, }, p2p::msg_pack::MsgPack, rpcresp_kvresult_convert::msg_and_error::codes_api, @@ -22,23 +22,33 @@ use crate::{ use chrono::Utc; use std::sync::Arc; +const MAX_STALE_SSD_ROUTE_RETRIES: usize = 3; + #[derive(Debug, Clone)] pub struct RemoteGetInfo { get_id: u64, data_len: usize, + source_kind: GetSourceKind, src_addr: u64, target_addr: u64, node_id: NodeID, peer_is_src_or_target: bool, } +impl RemoteGetInfo { + pub fn source_kind(&self) -> GetSourceKind { + self.source_kind + } +} + impl std::fmt::Display for RemoteGetInfo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "GetInfo{{ get_id: {}, data_len: {} bytes, src_addr: {:#x}, target_addr: {:#x}, node_id: {:?}, remote_transfer: {} }}", + "GetInfo{{ get_id: {}, data_len: {} bytes, source_kind: {:?}, src_addr: {:#x}, target_addr: {:#x}, node_id: {:?}, remote_transfer: {} }}", self.get_id, self.data_len, + self.source_kind, self.src_addr, self.target_addr, self.node_id, @@ -116,234 +126,348 @@ impl ClientKvApiInner { } obe_get_cache_miss(&metrics, &client_id, &node_role, key); - let t1 = Utc::now().timestamp_micros(); - let resp = { - match self.get_start(key).await { - Ok(resp) => resp, - Err(err) => { - obe_get_start_error_rpc(&metrics, &client_id, &node_role, key); - return Err(err); + let mut stale_ssd_route_retries = 0usize; + 'remote_get: loop { + let t1 = Utc::now().timestamp_micros(); + let resp = { + match self.get_start(key).await { + Ok(resp) => resp, + Err(err) => { + obe_get_start_error_rpc(&metrics, &client_id, &node_role, key); + return Err(err); + } } + }; + let start_handle_us = resp.server_process_us; + let t2 = Utc::now().timestamp_micros(); + // start stage success + // Note: only record timestamps; no scope begin/end + // errors handled above and below + if resp.error_code != OK { + if resp.error_code == codes_api::API_KEY_NOT_FOUND { + obe_get_start_not_found(&metrics, &client_id, &node_role, key); + return Ok(None); + } + obe_get_start_error_status(&metrics, &client_id, &node_role, key); + crate::rpcresp_kvresult_convert::try_from_code( + resp.error_code, + resp.error_json.clone(), + )?; + unreachable!("try_from_code should have returned Err for non-OK, unreachable"); } - }; - let start_handle_us = resp.server_process_us; - let t2 = Utc::now().timestamp_micros(); - // start stage success - // Note: only record timestamps; no scope begin/end - // errors handled above and below - if resp.error_code != OK { - if resp.error_code == codes_api::API_KEY_NOT_FOUND { - obe_get_start_not_found(&metrics, &client_id, &node_role, key); - return Ok(None); - } - obe_get_start_error_status(&metrics, &client_id, &node_role, key); - crate::rpcresp_kvresult_convert::try_from_code( - resp.error_code, - resp.error_json.clone(), - )?; - unreachable!("try_from_code should have returned Err for non-OK, unreachable"); - } - obe_get_start_success(&metrics, &client_id, &node_role, key, t1, t2); - - let put_id = resp.put_id; - let get_id = resp.get_id; - let data_len = resp.len as usize; - - let abs_src = resp.src_addr; - let abs_target = resp.target_addr; - - // debug get slice from src_addr and len - tracing::debug!( - "kv get src addr {:#x} to target addr {:#x}", - abs_src, - abs_target - ); - - let peer_id = if &*resp.node_id == &*self.view.cluster_manager().get_self_info().id { - None - } else { - Some(resp.node_id.clone()) - }; + obe_get_start_success(&metrics, &client_id, &node_role, key, t1, t2); - #[cfg(test)] - { - self.test_record.add_transfering_get( - get_id, - key.to_string(), - data_len as u32, - abs_target, - resp.node_id.to_string(), - peer_id.is_some(), - ); - } + let put_id = resp.put_id; + let get_id = resp.get_id; + let data_len = resp.len as usize; - // transfer data (skip if local and src==target to avoid redundant copy) - if peer_id.is_none() && abs_src == abs_target { - tracing::debug!( - "kv get local no-op: src==target {:#x}, len={} (skip transfer)", - abs_target, - data_len - ); - } else { - // tracing::debug!( - // "kv get transfer in transfer engine path from {}", - // peer_id.as_ref().map(|v| &**v).unwrap_or("self") - // ); + let abs_src = resp.src_addr; + let abs_target = resp.target_addr; + + // debug get slice from src_addr and len tracing::debug!( - "p2p get transfer: key={}, remote_src={:#x} -> local_target={:#x}, len={}, peer={:?}", - key, + "kv get src addr {:#x} to target addr {:#x}", abs_src, - abs_target, - data_len, - peer_id + abs_target ); - if let Err(e) = self - .view - .client_transfer_engine() - .transfer_data_no_copy( - peer_id.clone(), - true, - abs_src, - abs_target, - data_len as u64, - None, - ) - .await + + let peer_id = if &*resp.node_id == &*self.view.cluster_manager().get_self_info().id { + None + } else { + Some(resp.node_id.clone()) + }; + + #[cfg(test)] { - tracing::warn!("transfer data failed: {:?}", e); + self.test_record.add_transfering_get( + get_id, + key.to_string(), + data_len as u32, + abs_target, + resp.node_id.to_string(), + peer_id.is_some(), + ); + } - #[cfg(test)] - { - self.test_record.remove_transfering_get(get_id); + let mut ssd_done_resp = None; + if resp.source_kind == GetSourceKind::Ssd { + let ssd_stage_len = resp.ssd_stage_len; + if ssd_stage_len < data_len as u64 { + #[cfg(test)] + { + self.test_record.remove_transfering_get(get_id); + } + + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "invalid ssd stage len for key={} get_id={} data_len={} ssd_stage_len={}", + key, get_id, data_len, ssd_stage_len + ), + }); + self.get_revoke_ssd_source(get_id).await?; + if stale_ssd_route_retries < MAX_STALE_SSD_ROUTE_RETRIES { + stale_ssd_route_retries += 1; + tracing::warn!( + "Retrying get after invalid SSD route metadata: key={} retry={}/{}", + key, + stale_ssd_route_retries, + MAX_STALE_SSD_ROUTE_RETRIES + ); + continue 'remote_get; + } + return Err(err); } + let done_resp = match self + .stage_kv_from_ssd_source( + &resp.node_id, + key, + put_id, + get_id, + abs_src, + abs_target, + data_len as u64, + ssd_stage_len, + ) + .await + { + Ok(done_resp) => done_resp, + Err(err) => { + tracing::warn!( + "kv get ssd stage failed: key={}, source_node={}, stage={:#x}, target={:#x}, len={}, ssd_stage_len={}, err={}", + key, + resp.node_id, + abs_src, + abs_target, + data_len, + ssd_stage_len, + err + ); + + #[cfg(test)] + { + self.test_record.remove_transfering_get(get_id); + } + + obe_get_transfer_error( + &metrics, + &client_id, + &node_role, + key, + data_len as u64, + ); + self.get_revoke_ssd_source(get_id).await?; + if stale_ssd_route_retries < MAX_STALE_SSD_ROUTE_RETRIES { + stale_ssd_route_retries += 1; + tracing::warn!( + "Retrying get after failed SSD source cleanup: key={} retry={}/{}", + key, + stale_ssd_route_retries, + MAX_STALE_SSD_ROUTE_RETRIES + ); + continue 'remote_get; + } + return Err(err); + } + }; + ssd_done_resp = Some(done_resp); + tracing::debug!( + "kv get ssd staged and pushed: key={}, source_node={}, stage={:#x}, target={:#x}, len={}, ssd_stage_len={}", + key, + resp.node_id, + abs_src, + abs_target, + data_len, + ssd_stage_len + ); + } - obe_get_transfer_error(&metrics, &client_id, &node_role, key, data_len as u64); - self.get_revoke(get_id).await?; - return Err(KvError::Api(ApiError::Transfer { - from_addr: abs_src, - to_addr: abs_target, - len: data_len as u64, - error: e.to_string(), - })); + // transfer data (skip if local and src==target to avoid redundant copy) + if resp.source_kind == GetSourceKind::Ssd { + tracing::debug!( + "kv get ssd owner push complete: key={}, target={:#x}, len={} (skip requester transfer)", + key, + abs_target, + data_len + ); + } else if peer_id.is_none() && abs_src == abs_target { + tracing::debug!( + "kv get local no-op: src==target {:#x}, len={} (skip transfer)", + abs_target, + data_len + ); } else { + // tracing::debug!( + // "kv get transfer in transfer engine path from {}", + // peer_id.as_ref().map(|v| &**v).unwrap_or("self") + // ); tracing::debug!( - "get_transfer success key={}, src_addr={:#x}, target_addr={:#x}, len={}, peer_id={:?}", + "p2p get transfer: key={}, remote_src={:#x} -> local_target={:#x}, len={}, peer={:?}", key, abs_src, abs_target, data_len, peer_id ); + if let Err(e) = self + .view + .client_transfer_engine() + .transfer_data_no_copy( + peer_id.clone(), + true, + abs_src, + abs_target, + data_len as u64, + None, + ) + .await + { + tracing::warn!("transfer data failed: {:?}", e); + + #[cfg(test)] + { + self.test_record.remove_transfering_get(get_id); + } + + obe_get_transfer_error(&metrics, &client_id, &node_role, key, data_len as u64); + self.get_revoke(get_id).await?; + return Err(KvError::Api(ApiError::Transfer { + from_addr: abs_src, + to_addr: abs_target, + len: data_len as u64, + error: e.to_string(), + })); + } else { + tracing::debug!( + "get_transfer success key={}, src_addr={:#x}, target_addr={:#x}, len={}, peer_id={:?}", + key, + abs_src, + abs_target, + data_len, + peer_id + ); + } } - } - let t3 = Utc::now().timestamp_micros(); - obe_get_transfer_success( - &metrics, - &client_id, - &node_role, - key, - data_len as u64, - t2, - t3, - ); - - // Removed post-transfer zero-header verification per request. - - // Complete the get operation and get holder_id - let done_resp = match self.get_done(get_id).await { - Ok(resp) => resp, - Err(err) => { - obe_get_end_error_rpc(&metrics, &client_id, &node_role, key, data_len as u64); - return Err(err); + let t3 = Utc::now().timestamp_micros(); + obe_get_transfer_success( + &metrics, + &client_id, + &node_role, + key, + data_len as u64, + t2, + t3, + ); + + // Removed post-transfer zero-header verification per request. + + // Complete the get operation and get holder_id. SSD source already called + // get_done after pushing into the requester target. + let done_resp = if let Some(done_resp) = ssd_done_resp { + done_resp + } else { + match self.get_done(get_id).await { + Ok(resp) => resp, + Err(err) => { + obe_get_end_error_rpc( + &metrics, + &client_id, + &node_role, + key, + data_len as u64, + ); + return Err(err); + } + } + }; + let end_handle_us = done_resp.server_process_us; + let t4 = Utc::now().timestamp_micros(); + if done_resp.error_code != OK { + obe_get_done_error_status(&metrics, &client_id, &node_role, key, data_len as u64); + #[cfg(test)] + { + self.test_record.remove_transfering_get(get_id); + } + + crate::rpcresp_kvresult_convert::try_from_code( + done_resp.error_code, + done_resp.error_json.clone(), + )?; + unreachable!("error path should have returned above"); } - }; - let end_handle_us = done_resp.server_process_us; - let t4 = Utc::now().timestamp_micros(); - if done_resp.error_code != OK { - obe_get_done_error_status(&metrics, &client_id, &node_role, key, data_len as u64); + // end/done stage success and push detailed metrics + obe_get_done_success( + &metrics, + &client_id, + &node_role, + key, + data_len as u64, + get_id, + t1, + t2, + t3, + t4, + start_handle_us, + end_handle_us, + ); + #[cfg(test)] { self.test_record.remove_transfering_get(get_id); } - crate::rpcresp_kvresult_convert::try_from_code( - done_resp.error_code, - done_resp.error_json.clone(), - )?; - unreachable!("error path should have returned above"); - } - // end/done stage success and push detailed metrics - obe_get_done_success( - &metrics, - &client_id, - &node_role, - key, - data_len as u64, - get_id, - t1, - t2, - t3, - t4, - start_handle_us, - end_handle_us, - ); - - #[cfg(test)] - { - self.test_record.remove_transfering_get(get_id); - } - - // pulses and network bytes emitted inside obe_get_done_success - - let holder_id = done_resp.holder_id; - let expose_kind = if done_resp.allocation_mode == GetAllocationMode::Temporary { - UserMemHolderExposeKind::OwnedCopy - } else { - UserMemHolderExposeKind::SegPtr - }; - let master_node_id: NodeID = self - .view - .cluster_manager() - .find_or_wait_master_node() - .await? - .into(); - - // Create MemHolder with keep alive functionality - // Convert target_addr to offset using base address from master response - let offset = resp.target_addr - resp.target_base_addr; - let memory_info = Arc::new( - MemoryInfo::new( - offset, - data_len as u32, - holder_id, - key.to_string(), - master_node_id, - self.view.clone(), - ) - .await, - ); - // Create GetInfo with information from the response - let get_info = RemoteGetInfo { - get_id, - data_len, - src_addr: abs_src, - target_addr: abs_target, - node_id: resp.node_id.into(), - peer_is_src_or_target: true, - }; + // pulses and network bytes emitted inside obe_get_done_success - if done_resp.allocation_mode != GetAllocationMode::Temporary { - self.cache_local_replica_after_get(key, put_id, memory_info.clone()); - metrics.observe_cache_value_size(&client_id, node_role.as_str(), data_len as u64); + let holder_id = done_resp.holder_id; + let expose_kind = if done_resp.allocation_mode == GetAllocationMode::Temporary { + UserMemHolderExposeKind::OwnedCopy + } else { + UserMemHolderExposeKind::SegPtr + }; + let master_node_id: NodeID = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await? + .into(); + + // Create MemHolder with keep alive functionality + // Convert target_addr to offset using base address from master response + let offset = resp.target_addr - resp.target_base_addr; + let memory_info = Arc::new( + MemoryInfo::new( + offset, + data_len as u32, + holder_id, + key.to_string(), + master_node_id, + self.view.clone(), + ) + .await, + ); + // Create GetInfo with information from the response + let get_info = RemoteGetInfo { + get_id, + data_len, + source_kind: resp.source_kind, + src_addr: abs_src, + target_addr: abs_target, + node_id: resp.node_id.into(), + peer_is_src_or_target: true, + }; + + if done_resp.allocation_mode != GetAllocationMode::Temporary { + self.cache_local_replica_after_get(key, put_id, memory_info.clone()); + metrics.observe_cache_value_size(&client_id, node_role.as_str(), data_len as u64); + } + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info, + self.get_or_init_all_memholder_refcount(), + expose_kind, + )); + // let partial_hex=&user_mem_holder.bytes()[..std::cmp::min(16, user_mem_holder.bytes().len())]; + // tracing::debug!("external get done, key={}, partial_hex={:?}", key, partial_hex); + return Ok(Some((user_mem_holder, Some(get_info)))); } - let user_mem_holder = Arc::new(UserMemHolder::new( - memory_info, - self.get_or_init_all_memholder_refcount(), - expose_kind, - )); - // let partial_hex=&user_mem_holder.bytes()[..std::cmp::min(16, user_mem_holder.bytes().len())]; - // tracing::debug!("external get done, key={}, partial_hex={:?}", key, partial_hex); - Ok(Some((user_mem_holder, Some(get_info)))) } pub async fn is_exist(&self, key: &str) -> KvResult { @@ -435,8 +559,19 @@ impl ClientKvApiInner { /// 撤销 Get 操作，释放已分配的资源 pub async fn get_revoke(&self, get_id: u64) -> KvResult<()> { + self.get_revoke_inner(get_id, false).await + } + + async fn get_revoke_ssd_source(&self, get_id: u64) -> KvResult<()> { + self.get_revoke_inner(get_id, true).await + } + + async fn get_revoke_inner(&self, get_id: u64, drop_ssd_source: bool) -> KvResult<()> { let req = MsgPack { - serialize_part: GetRevokeReq { get_id }, + serialize_part: GetRevokeReq { + get_id, + drop_ssd_source, + }, raw_bytes: Vec::new(), }; diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs index dec19f5..5efd48f 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs @@ -3,13 +3,20 @@ use crate::client_kv_api::msg_pack::{ ExternalDeleteAckReq, ExternalDeleteAckResp, ExternalDeleteReq, ExternalDeleteResp, ExternalGetReq, ExternalGetResp, ExternalIsExistReq, ExternalIsExistResp, ExternalPutCommitReq, ExternalPutCommitResp, ExternalPutRevokeReq, ExternalPutRevokeResp, ExternalPutStartReq, - ExternalPutStartResp, ExternalPutTransferEndReq, ExternalPutTransferEndResp, SyncKvToFileReq, - SyncKvToFileResp, TestPutPhaseTrace, + ExternalPutStartResp, ExternalPutTransferEndReq, ExternalPutTransferEndResp, + SsdReplicaPersistReq, SsdReplicaPersistResp, SsdStageReadReq, SsdStageReadResp, + SyncKvToFileReq, SyncKvToFileResp, TestPutPhaseTrace, }; use crate::cluster_manager::NodeIDString; +use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::config::TestSpecConfig; +use crate::kv_ssd_storage::{ + DEFAULT_READ_TRANSFER_PIPELINE_CHUNK_BYTES, DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT, + KvSsdPersistGuard, KvSsdStorage, KvSsdStorageInit, SsdLoadedChunk, +}; use crate::master_kv_router::msg_pack::{ - BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, DeleteClientKvMetaCacheItem, + BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, BatchSsdReplicaEvictReq, + DeleteClientKvMetaCacheItem, SsdReplicaEviction, }; use crate::master_lease_manager::msg_pack::{AllocateClientLeaseReq, ClientLeaseKeepaliveReq}; use crate::memholder::{AllMemholderRefCount, MemoryInfo, UserMemHolder}; @@ -22,8 +29,8 @@ use crate::{ client_transfer_engine::{ClientTransferEngine, ClientTransferEngineAccessTrait}, cluster_manager::{ClusterEvent, ClusterManager, ClusterManagerAccessTrait}, master_kv_router::msg_pack::{ - DeleteReq, GetDoneReq, GetMetaReq, GetRevokeReq, GetStartReq, PutDoneReq, PutRevokeReq, - PutStartReq, + DeleteReq, GetDoneReq, GetDoneResp, GetMetaReq, GetRevokeReq, GetStartReq, PutDoneReq, + PutRevokeReq, PutStartReq, SsdReplicaCommitReq, }, metric_reporter::{MetricReporter, MetricReporterAccessTrait}, metrics::{MetricsHandle, OperationKind, RequestStage}, @@ -37,13 +44,19 @@ use async_trait::async_trait; use dashmap::DashMap; use fluxon_framework::{LogicalModule, define_module}; use fluxon_util::map_lock::AMapLock; +use futures::stream::{FuturesUnordered, StreamExt}; use limit_thirdparty::tokio; use parking_lot::Mutex; +use std::collections::VecDeque; use std::sync::Weak; use std::sync::{Arc, OnceLock}; use std::time::Duration; use tracing::warn; +const SSD_EVICTION_BATCH_MAX_ITEMS: usize = 256; +const SSD_EVICTION_BATCH_MAX_DELAY: Duration = Duration::from_millis(10); +const SSD_EVICTION_BATCH_RETRY_DELAY: Duration = Duration::from_millis(100); + /// Information about a memholder held by external client #[derive(Clone)] pub struct ExternalHoldingGetInfo { @@ -79,6 +92,8 @@ pub enum PutOptionalArg { LeaseId(u64), /// Ask the master to fail-fast when the same key already has an inflight put. RejectIfInflightSameKey, + /// Ask the master to reject the put when the key exists or is being created. + RejectIfExists, /// Prefer placing the target allocation on a kvclient within this sub_cluster. PreferredSubCluster(String), /// Hidden test-only side-channel for collecting per-put phase timings. @@ -98,6 +113,7 @@ impl PutOptionalArgs { self.0.iter().rev().find_map(|a| match a { PutOptionalArg::LeaseId(id) => Some(*id), PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExists | PutOptionalArg::PreferredSubCluster(_) | PutOptionalArg::TestObservePutPhases(_) => None, }) @@ -109,12 +125,19 @@ impl PutOptionalArgs { .any(|arg| matches!(arg, PutOptionalArg::RejectIfInflightSameKey)) } + pub fn reject_if_exists(&self) -> bool { + self.0 + .iter() + .any(|arg| matches!(arg, PutOptionalArg::RejectIfExists)) + } + /// Get the last provided preferred_sub_cluster if any. pub fn preferred_sub_cluster(&self) -> Option<&str> { self.0.iter().rev().find_map(|a| match a { PutOptionalArg::PreferredSubCluster(sc) => Some(sc.as_str()), PutOptionalArg::LeaseId(_) | PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExists | PutOptionalArg::TestObservePutPhases(_) => None, }) } @@ -124,6 +147,7 @@ impl PutOptionalArgs { PutOptionalArg::TestObservePutPhases(sink) => Some(sink.clone()), PutOptionalArg::LeaseId(_) | PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExists | PutOptionalArg::PreferredSubCluster(_) => None, }) } @@ -451,6 +475,91 @@ async fn handle_external_put_revoke( } } +async fn handle_ssd_stage_read( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let inner = view.client_kv_api().inner(); + let done_resp = match inner + .load_and_push_kv_from_ssd( + &req.key, + req.put_id, + req.stage_addr, + req.stage_len, + &req.target_node_id, + req.target_addr, + req.len, + ) + .await + { + Ok(()) => inner.get_done(req.get_id).await, + Err(err) => Err(err), + }; + + match done_resp { + Ok(done_resp) => MsgPack { + serialize_part: SsdStageReadResp { + done_holder_id: done_resp.holder_id, + done_allocation_mode: done_resp.allocation_mode, + done_error_code: done_resp.error_code, + done_error_json: done_resp.error_json, + done_server_process_us: done_resp.server_process_us, + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + }, + Err(err) => MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }, + } +} + +async fn handle_ssd_replica_persist( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let inner = view.client_kv_api().inner(); + let persist_guard = match inner + .persist_local_kv_to_ssd(&req.key, req.put_id, req.target_addr, req.len) + .await + { + Ok(persist_guard) => persist_guard, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + }; + + let persisted = persist_guard.is_some(); + if persisted { + if let Err(err) = inner + .commit_ssd_replica_to_master(&req.key, req.put_id, req.len) + .await + { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + } + drop(persist_guard); + + MsgPack { + serialize_part: SsdReplicaPersistResp { + persisted, + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + async fn handle_external_delete_ack( view: &ClientKvApiView, msg: &MsgPack, @@ -729,6 +838,7 @@ define_module!( #[derive(Clone, Debug)] pub struct ClientKvApiNewArg { pub test_spec_config: TestSpecConfig, + pub ssd_storage: Option, } pub struct ClientKvApi(ClientKvApiInner); @@ -772,9 +882,84 @@ impl std::ops::Deref for ClientKvApiViewHolder { } } +fn spawn_ssd_replica_eviction_batch( + view: ClientKvApiView, + mut rx: ::tokio::sync::mpsc::Receiver>, +) { + let view_for_task = view.clone(); + let _ = view.spawn("ssd_replica_eviction_batch", async move { + let mut pending = VecDeque::new(); + let mut rx_closed = false; + let mut shutdown_waiter = view_for_task.register_shutdown_waiter(); + + 'actor: loop { + if pending.is_empty() { + let received = tokio::select! { + received = rx.recv() => received, + _ = shutdown_waiter.wait() => break 'actor, + }; + match received { + Some(replicas) => pending.extend(replicas), + None => break, + } + } + + let deadline = ::tokio::time::Instant::now() + SSD_EVICTION_BATCH_MAX_DELAY; + while pending.len() < SSD_EVICTION_BATCH_MAX_ITEMS && !rx_closed { + let received = tokio::select! { + received = ::tokio::time::timeout_at(deadline, rx.recv()) => received, + _ = shutdown_waiter.wait() => break 'actor, + }; + match received { + Ok(Some(replicas)) => pending.extend(replicas), + Ok(None) => rx_closed = true, + Err(_) => break, + } + } + + let batch_len = pending.len().min(SSD_EVICTION_BATCH_MAX_ITEMS); + let batch = pending.drain(..batch_len).collect::>(); + let retry_batch = batch.clone(); + let send_result = tokio::select! { + result = view_for_task + .client_kv_api() + .inner() + .evict_ssd_replicas_from_master(batch) => result, + _ = shutdown_waiter.wait() => break 'actor, + }; + match send_result { + Ok(removed_count) => tracing::debug!( + "Sent SSD replica eviction batch: requested={} removed={}", + batch_len, + removed_count + ), + Err(err) => { + tracing::warn!( + "Failed to send SSD replica eviction batch; retaining it for retry: requested={} err={}", + batch_len, + err + ); + for eviction in retry_batch.into_iter().rev() { + pending.push_front(eviction); + } + tokio::select! { + _ = ::tokio::time::sleep(SSD_EVICTION_BATCH_RETRY_DELAY) => {} + _ = shutdown_waiter.wait() => break 'actor, + } + } + } + + if rx_closed && pending.is_empty() { + break; + } + } + }); +} + pub struct ClientKvApiInner { view: ClientKvApiViewHolder, test_spec_config: TestSpecConfig, + ssd_storage: Option>, metrics: OnceLock>, /// make sure each remote kv get run in order @@ -818,6 +1003,9 @@ pub struct ClientKvApiInner { rpc_caller_external_put_commit: RPCCaller, rpc_caller_external_put_revoke: RPCCaller, rpc_caller_resolve_side_transfer_lane: RPCCaller, + rpc_caller_ssd_stage_read: RPCCaller, + rpc_caller_ssd_replica_commit: RPCCaller, + rpc_caller_batch_ssd_replica_evict: RPCCaller, /// Default lease id recorded for inspection/convenience, but NOT auto-applied. /// Callers must explicitly pass `Some(lease_id)` to attach a put to a lease. @@ -900,6 +1088,275 @@ impl ClientKvApiInner { pub(crate) fn skip_put_end_commit_enabled(&self) -> bool { self.test_spec_config.skip_put_end_commit } + + pub(crate) async fn persist_local_kv_to_ssd( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + abs_addr: u64, + len: u64, + ) -> KvResult> { + let Some(store) = self.ssd_storage.as_ref() else { + return Ok(None); + }; + let guard = store.persist_from_addr(key, put_id, abs_addr, len).await?; + Ok(Some(guard)) + } + + pub(crate) fn kv_ssd_storage_usage_snapshot( + &self, + ) -> Vec { + self.ssd_storage + .as_ref() + .map(|store| store.device_usage_snapshot()) + .unwrap_or_default() + } + + pub(crate) async fn commit_ssd_replica_to_master( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + len: u64, + ) -> KvResult<()> { + let node_id = self.view.cluster_manager().get_self_info().id.clone(); + let req = MsgPack { + serialize_part: SsdReplicaCommitReq { + key: key.to_string(), + put_id, + node_id, + len, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_ssd_replica_commit + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json, + ) + } + + async fn evict_ssd_replicas_from_master( + &self, + evictions: Vec, + ) -> KvResult { + let req = MsgPack { + serialize_part: BatchSsdReplicaEvictReq { evictions }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_ssd_replica_evict + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(Duration::from_secs(10)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json, + )?; + Ok(resp.serialize_part.removed_count) + } + + pub(crate) async fn load_and_push_kv_from_ssd( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + stage_addr: u64, + stage_len: u64, + target_node_id: &NodeIDString, + target_addr: u64, + len: u64, + ) -> KvResult<()> { + let Some(store) = self.ssd_storage.as_ref() else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage is not enabled on this owner".to_string(), + })); + }; + + let self_node_id = &self.view.cluster_manager().get_self_info().id; + let peer_id = if target_node_id == self_node_id { + None + } else { + Some(target_node_id.clone()) + }; + if peer_id.is_none() && stage_addr == target_addr { + tracing::debug!( + "Loading local SSD replica directly into target: key={} target={:#x} len={} capacity={}", + key, + target_addr, + len, + stage_len + ); + return store + .load_into_addr(key, put_id, target_addr, len, stage_len) + .await; + } + + let (chunk_tx, chunk_rx) = ::tokio::sync::mpsc::channel( + DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT + .saturating_mul(2) + .max(1), + ); + let producer = store.load_into_addr_chunks( + key, + put_id, + stage_addr, + len, + stage_len, + DEFAULT_READ_TRANSFER_PIPELINE_CHUNK_BYTES, + DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT, + chunk_tx, + ); + let consumer = self.transfer_loaded_ssd_chunks(peer_id, target_addr, chunk_rx); + let (producer_res, consumer_res) = ::tokio::join!(producer, consumer); + match (producer_res, consumer_res) { + (Ok(()), Ok(())) => Ok(()), + (_, Err(err)) => Err(err), + (Err(err), _) => Err(err), + } + } + + async fn transfer_loaded_ssd_chunks( + &self, + peer_id: Option, + target_addr: u64, + mut chunk_rx: ::tokio::sync::mpsc::Receiver, + ) -> KvResult<()> { + let mut inflight = FuturesUnordered::new(); + let mut rx_open = true; + + loop { + tokio::select! { + maybe_chunk = chunk_rx.recv(), if rx_open && inflight.len() < DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT => { + match maybe_chunk { + Some(chunk) => { + let chunk_target_addr = target_addr.checked_add(chunk.offset).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd transfer target addr overflow: target_addr={:#x} offset={}", + target_addr, + chunk.offset + ), + }) + })?; + let transfer_engine = self.view.client_transfer_engine(); + let peer_id = peer_id.clone(); + inflight.push(async move { + transfer_engine + .transfer_data_no_copy( + peer_id, + false, + chunk.stage_addr, + chunk_target_addr, + chunk.len, + None, + ) + .await?; + Ok::<(), KvError>(()) + }); + } + None => { + rx_open = false; + } + } + } + Some(result) = inflight.next(), if !inflight.is_empty() => { + result?; + } + else => { + if !rx_open && inflight.is_empty() { + break; + } + } + } + } + Ok(()) + } + + pub(crate) async fn stage_kv_from_ssd_source( + &self, + source_node_id: &NodeIDString, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + get_id: u64, + stage_addr: u64, + target_addr: u64, + len: u64, + stage_len: u64, + ) -> KvResult { + let self_node_id = self.view.cluster_manager().get_self_info().id.clone(); + if source_node_id == &self_node_id { + self.load_and_push_kv_from_ssd( + key, + put_id, + stage_addr, + stage_len, + &self_node_id, + target_addr, + len, + ) + .await?; + return self.get_done(get_id).await; + } + + let req = MsgPack { + serialize_part: SsdStageReadReq { + key: key.to_string(), + put_id, + get_id, + stage_addr, + stage_len, + target_node_id: self_node_id, + target_addr, + len, + }, + raw_bytes: Vec::new(), + }; + let resp = self + .rpc_caller_ssd_stage_read + .call( + self.view.p2p_module(), + source_node_id.clone().into(), + req, + Some(Duration::from_secs(60)), + 0, + ) + .await + .map_err(KvError::from)?; + let resp = resp.serialize_part; + crate::rpcresp_kvresult_convert::try_from_code(resp.error_code, resp.error_json)?; + Ok(GetDoneResp { + holder_id: resp.done_holder_id, + allocation_mode: resp.done_allocation_mode, + error_code: resp.done_error_code, + error_json: resp.done_error_json, + server_process_us: resp.done_server_process_us, + }) + } } #[derive(Debug, Clone)] @@ -1512,16 +1969,27 @@ impl ClientKvApi { &self.0 } + pub(crate) fn kv_ssd_storage_usage_snapshot( + &self, + ) -> Vec { + self.inner().kv_ssd_storage_usage_snapshot() + } + pub fn attach_view(&self, view: ClientKvApiView) { self.0.view.attach(view); } pub async fn construct(arg: ClientKvApiNewArg) -> Result { tracing::info!("Constructing ClientKvApi in Client mode (PreView)"); + let ssd_storage = match arg.ssd_storage { + Some(init) => Some(Arc::new(KvSsdStorage::new(init).await?)), + None => None, + }; let inner = ClientKvApiInner { view: ClientKvApiViewHolder::new(), test_spec_config: arg.test_spec_config, + ssd_storage, metrics: OnceLock::new(), all_memholder_refcount: OnceLock::new(), get_remote_kv_lock: AMapLock::new(Duration::from_secs(60)), @@ -1554,6 +2022,9 @@ impl ClientKvApi { rpc_caller_external_put_commit: RPCCaller::new(), rpc_caller_external_put_revoke: RPCCaller::new(), rpc_caller_resolve_side_transfer_lane: RPCCaller::new(), + rpc_caller_ssd_stage_read: RPCCaller::new(), + rpc_caller_ssd_replica_commit: RPCCaller::new(), + rpc_caller_batch_ssd_replica_evict: RPCCaller::new(), default_lease_id: parking_lot::RwLock::new(None), }; Ok(Self(inner)) @@ -1587,6 +2058,15 @@ impl ClientKvApi { inner .rpc_caller_resolve_side_transfer_lane .regist(inner.view.p2p_module()); + inner + .rpc_caller_ssd_stage_read + .regist(inner.view.p2p_module()); + inner + .rpc_caller_ssd_replica_commit + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_ssd_replica_evict + .regist(inner.view.p2p_module()); crate::key_prefix::init_for_p2p_owner(inner.view.p2p_module()); crate::kvlease::init_for_p2p_owner(inner.view.p2p_module()); // Register master-only metric RPC callers @@ -1686,6 +2166,31 @@ impl ClientKvApi { }, ); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_ssd_stage_read", async move { + let result = handle_ssd_stage_read(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_ssd_replica_persist", async move { + let result = handle_ssd_replica_persist(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); @@ -1773,6 +2278,13 @@ impl ClientKvApi { .expect("delete_ack_batch rx already taken, that's impossible"); delete::spawn_owner_delete_ack_batch(inner.view.clone_view(), delete_ack_batch_rx); + if let Some(store) = inner.ssd_storage.as_ref() { + let eviction_rx = store + .take_eviction_rx() + .expect("SSD eviction receiver must only be taken during client initialization"); + spawn_ssd_replica_eviction_batch(inner.view.clone_view(), eviction_rx); + } + // Spawn cluster listener to clean up get_holding when external_client leaves let view = inner.view.clone_view(); let view2 = view.clone(); diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs index 55f0970..565a916 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs @@ -1,8 +1,10 @@ +use crate::master_kv_router::msg_pack::GetAllocationMode; use crate::master_kv_router::put::PutIDForAKey; use crate::p2p::msg_pack::{MsgPackSerializePart, RPCReq}; use crate::rpcresp_kvresult_convert::msg_and_error::ErrorCode; use bitcode::{Decode, Encode}; +use crate::cluster_manager::NodeIDString; use crate::memholder::ExternalMemHolderInfo; #[derive(Default, Debug, Clone, Encode, Decode)] @@ -89,6 +91,76 @@ impl MsgPackSerializePart for ExternalGetResp { } } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageReadReq { + pub key: String, + pub put_id: PutIDForAKey, + pub get_id: u64, + pub stage_addr: u64, + pub stage_len: u64, + pub target_node_id: NodeIDString, + pub target_addr: u64, + pub len: u64, +} + +impl MsgPackSerializePart for SsdStageReadReq { + fn msg_id(&self) -> u32 { + 4020 + } +} + +impl RPCReq for SsdStageReadReq { + type Resp = SsdStageReadResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageReadResp { + pub done_holder_id: u64, + pub done_allocation_mode: GetAllocationMode, + pub done_error_code: ErrorCode, + pub done_error_json: String, + pub done_server_process_us: i64, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for SsdStageReadResp { + fn msg_id(&self) -> u32 { + 4021 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdReplicaPersistReq { + pub key: String, + pub put_id: PutIDForAKey, + pub target_addr: u64, + pub len: u64, +} + +impl MsgPackSerializePart for SsdReplicaPersistReq { + fn msg_id(&self) -> u32 { + 4022 + } +} + +impl RPCReq for SsdReplicaPersistReq { + type Resp = SsdReplicaPersistResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdReplicaPersistResp { + pub persisted: bool, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for SsdReplicaPersistResp { + fn msg_id(&self) -> u32 { + 4023 + } +} + // #[derive(Default, Debug, Clone, Encode, Decode)] // pub struct ExternalPutReq { // pub key: String, @@ -114,6 +186,7 @@ pub struct ExternalPutStartReq { pub key: String, pub len: u64, pub reject_if_inflight_same_key: bool, + pub reject_if_exists: bool, /// Prefer placing the target allocation on any kvclient within this sub_cluster. pub preferred_sub_cluster: Option, /// Owner node_start_time observed by external when request starts diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs index 374ed71..584df80 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs @@ -35,6 +35,7 @@ impl ClientKvApiInner { payload_len: u64, len_for_start: u32, reject_if_inflight_same_key: bool, + reject_if_exists: bool, preferred_sub_cluster: Option<&str>, lease_id: Option, _test_payload_len_u32: u32, @@ -57,6 +58,7 @@ impl ClientKvApiInner { key, len_for_start, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster, ) .await @@ -243,6 +245,7 @@ impl ClientKvApiInner { ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exists = opts.reject_if_exists(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let payload_len = calc_flat_dict_encoded_len(&ptrs)?; @@ -251,6 +254,7 @@ impl ClientKvApiInner { payload_len, payload_len as u32, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster.as_deref(), lease_id, payload_len as u32, @@ -276,6 +280,7 @@ impl ClientKvApiInner { ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exists = opts.reject_if_exists(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let payload_len = value.len() as u64; self.put_common( @@ -283,6 +288,7 @@ impl ClientKvApiInner { payload_len, value.len() as u32, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster.as_deref(), lease_id, value.len() as u32, @@ -401,6 +407,7 @@ impl ClientKvApiInner { key: &str, len: u32, reject_if_inflight_same_key: bool, + reject_if_exists: bool, preferred_sub_cluster: Option<&str>, source_node_id: Option, ) -> KvResult<(PutStartResp, i64)> { @@ -409,6 +416,7 @@ impl ClientKvApiInner { key: key.to_string(), len: len as u64, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), source_node_id, }, @@ -458,12 +466,14 @@ impl ClientKvApiInner { key: &str, len: u32, reject_if_inflight_same_key: bool, + reject_if_exists: bool, preferred_sub_cluster: Option<&str>, ) -> KvResult<(PutStartResp, i64)> { self.put_start_with_source_node( key, len, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster, None, ) diff --git a/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs b/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs index 1aa6954..8c7cc78 100644 --- a/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs +++ b/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs @@ -237,10 +237,7 @@ impl ClientSegPool { std::path::Path::new(share_mem_path).join(SIDE_TRANSFER_PEERS_DIRNAME) } - pub fn side_transfer_peer_file_path( - share_mem_path: &str, - side_id: &str, - ) -> std::path::PathBuf { + pub fn side_transfer_peer_file_path(share_mem_path: &str, side_id: &str) -> std::path::PathBuf { Self::side_transfer_peers_dir(share_mem_path).join(format!("{side_id}.json")) } @@ -399,17 +396,13 @@ impl ClientSegPool { crate::rpcresp_kvresult_convert::msg_and_error::SharedMemError::MappingFailed { path: String::new(), len: map_len as u64, - detail: "share_mem_path is empty; explicit configuration required" - .to_string(), + detail: "share_mem_path is empty; explicit configuration required".to_string(), }, )); } let base_path = &share_mem_path; - tracing::info!( - "Using share_mem_path: {} for memory-mapped file", - base_path - ); + tracing::info!("Using share_mem_path: {} for memory-mapped file", base_path); std::fs::create_dir_all(base_path).map_err(|e| { KvError::SharedMem( crate::rpcresp_kvresult_convert::msg_and_error::SharedMemError::MappingFailed { diff --git a/fluxon_rs/fluxon_kv/src/cluster_manager/cluster_manager_test.rs b/fluxon_rs/fluxon_kv/src/cluster_manager/cluster_manager_test.rs index 2a668c5..1b937a3 100644 --- a/fluxon_rs/fluxon_kv/src/cluster_manager/cluster_manager_test.rs +++ b/fluxon_rs/fluxon_kv/src/cluster_manager/cluster_manager_test.rs @@ -4,7 +4,7 @@ use fluxon_framework::{AnyResult, LogicalModule, define_framework, define_module use limit_thirdparty::tokio; use limit_thirdparty::tokio::time::sleep; use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use std::{collections::HashMap, sync::Arc}; use super::{ @@ -37,6 +37,29 @@ async fn clean_etcd_members(cluster_name: &str) { }; } +async fn wait_for_member_ids( + cluster_manager: &ClusterManager, + expected_ids: &[&str], +) -> Vec { + let deadline = Instant::now() + Duration::from_secs(20); + let mut expected_ids: Vec = expected_ids.iter().map(|id| (*id).to_string()).collect(); + expected_ids.sort(); + + loop { + let members = cluster_manager.get_members(); + let mut actual_ids: Vec = members.iter().map(|member| member.id.clone()).collect(); + actual_ids.sort(); + if actual_ids == expected_ids { + return members; + } + assert!( + Instant::now() < deadline, + "timed out waiting for cluster members: expected={expected_ids:?}, actual={actual_ids:?}" + ); + sleep(Duration::from_millis(50)).await; + } +} + /// 创建一个测试模块来包装 ClusterManager pub struct TestClusterModule { view: std::sync::OnceLock, @@ -493,12 +516,11 @@ async fn test_cluster_manager_multiple_members() { let test_module3 = framework3.0.test_cluster_module.get().unwrap(); let cluster_manager3 = test_module3.cluster_manager(); - sleep(Duration::from_millis(10000)).await; - - // 验证所有成员都能看到彼此 - let members1 = cluster_manager1.get_members(); - let members2 = cluster_manager2.get_members(); - let members3 = cluster_manager3.get_members(); + // Wait for the observable contract instead of sampling after a fixed delay. + let expected_ids = ["test-member-1", "test-member-2", "test-member-3"]; + let members1 = wait_for_member_ids(cluster_manager1, &expected_ids).await; + let members2 = wait_for_member_ids(cluster_manager2, &expected_ids).await; + let members3 = wait_for_member_ids(cluster_manager3, &expected_ids).await; assert_eq!(members1.len(), 3); assert_eq!(members2.len(), 3); diff --git a/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs b/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs index 97a5758..6d2789c 100644 --- a/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs +++ b/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs @@ -5,6 +5,15 @@ pub use fluxon_commu::{ META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME, }; +pub(crate) const META_KEY_KV_SSD_STORAGE: &str = "kv_ssd_storage"; + +pub(crate) fn member_has_kv_ssd_storage(member: &ClusterMember) -> bool { + member + .metadata + .get(META_KEY_KV_SSD_STORAGE) + .is_some_and(|value| value == "true") +} + pub mod app_logic_ext; #[cfg(test)] diff --git a/fluxon_rs/fluxon_kv/src/config.rs b/fluxon_rs/fluxon_kv/src/config.rs index f9c7691..5c80d46 100644 --- a/fluxon_rs/fluxon_kv/src/config.rs +++ b/fluxon_rs/fluxon_kv/src/config.rs @@ -25,6 +25,13 @@ pub enum YamlNullable { Value(T), } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(untagged)] +pub enum SizeBytesYaml { + Bytes(u64), + Text(String), +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct GreptimeOtlpLogConfigYaml { @@ -97,6 +104,32 @@ pub enum SideTransferRole { Worker, } +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum KvSsdUringMode { + SingleBuffer, + Iovec, +} + +impl Default for KvSsdUringMode { + fn default() -> Self { + Self::SingleBuffer + } +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum KvSsdStorageBackend { + Native, + Foyer, +} + +impl Default for KvSsdStorageBackend { + fn default() -> Self { + Self::Native + } +} + const TEST_SPEC_TCP_THREAD_REACTOR_SHARD_COUNT_MIN: u8 = 1; const TEST_SPEC_TCP_THREAD_REACTOR_SHARD_COUNT_MAX: u8 = 16; const TEST_SPEC_TCP_THREAD_BULK_LANE_COUNT_MIN: u8 = 1; @@ -155,6 +188,10 @@ pub struct TestSpecConfig { pub side_transfer_worker_p2p_port_base: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub side_transfer_role: Option, + #[serde(default)] + pub kv_ssd_uring_mode: KvSsdUringMode, + #[serde(default)] + pub kv_ssd_storage_backend: KvSsdStorageBackend, } impl Default for TestSpecConfig { @@ -182,10 +219,24 @@ impl Default for TestSpecConfig { side_transfer_worker_count: 0, side_transfer_worker_p2p_port_base: None, side_transfer_role: None, + kv_ssd_uring_mode: KvSsdUringMode::default(), + kv_ssd_storage_backend: KvSsdStorageBackend::default(), } } } +fn validate_test_spec_kv_ssd_backend(test_spec_config: &TestSpecConfig) -> KvResult<()> { + if test_spec_config.kv_ssd_storage_backend == KvSsdStorageBackend::Foyer + && test_spec_config.kv_ssd_uring_mode != KvSsdUringMode::SingleBuffer + { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.kv_ssd_uring_mode is only valid with test_spec_config.kv_ssd_storage_backend=native".to_string(), + } + .into_kverror()); + } + Ok(()) +} + fn resolve_enable_transfer_rpc_fast_path( default_enabled: bool, test_spec_config: Option<&TestSpecConfig>, @@ -566,8 +617,8 @@ pub struct MasterConfig { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct ContributeToClusterPoolSizeYaml { - pub dram: u64, // bytes - pub vram: HashMap, // gpu_id -> bytes + pub dram: SizeBytesYaml, // bytes + pub vram: HashMap, // gpu_id -> bytes } /// Configuration for Fluxon KV backend specifications @@ -581,6 +632,8 @@ pub struct FluxonKvSpecYaml { #[serde(skip_serializing_if = "Option::is_none")] pub large_file_paths: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub large_limit_size: Option>>, + #[serde(skip_serializing_if = "Option::is_none")] pub p2p_listen_port: Option, #[serde(skip_serializing_if = "Option::is_none")] pub redis_compat: Option>, @@ -592,6 +645,11 @@ pub struct FluxonKvSpecYaml { #[serde(transparent)] pub struct LargeFilePathsYaml(pub Vec); +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct KvSsdStorageConfig { + pub large_limit_size: Vec, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct RedisCompatConfigYaml { @@ -682,6 +740,34 @@ impl LargeFilePaths { .into_kverror()) } + fn resolve_all_usable_root_subdirs( + &self, + relative_dir: &Path, + target_name: &str, + ) -> KvResult> { + self.require_configured_paths()?; + let mut out = Vec::new(); + let mut errors = Vec::new(); + for root in &self.paths { + let candidate = Path::new(root).join(relative_dir); + match fs::create_dir_all(&candidate) { + Ok(()) => out.push(candidate), + Err(err) => errors.push(format!("{} ({})", candidate.display(), err)), + } + } + if out.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "large_file_paths contains no usable root for {}; tried: {}", + target_name, + errors.join(", ") + ), + } + .into_kverror()); + } + Ok(out) + } + pub fn kv_logs_dir(&self, cluster_name: &str) -> KvResult { let relative_dir = PathBuf::from(format!("{cluster_name}_cluster_kv_logs")); self.resolve_preferred_root_subdir(&relative_dir, "kv logs") @@ -714,6 +800,60 @@ impl LargeFilePaths { "fluxon fs disk cache", ) } + + pub fn kv_ssd_storage_dirs( + &self, + cluster_name: &str, + instance_key: &str, + ) -> KvResult> { + let relative_dir = PathBuf::from(format!( + "{cluster_name}_cluster_kv_ssd_storage/{}", + crate::kv_ssd_storage::safe_path_component(instance_key) + )); + self.resolve_all_usable_root_subdirs(&relative_dir, "kv ssd storage") + } + + pub fn kv_ssd_storage_root_limits( + &self, + cluster_name: &str, + instance_key: &str, + large_limit_size: &[u64], + ) -> KvResult> { + self.require_configured_paths()?; + if large_limit_size.len() != self.paths.len() { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "large_limit_size length must match large_file_paths length: large_limit_size={} large_file_paths={}", + large_limit_size.len(), + self.paths.len() + ), + } + .into_kverror()); + } + let relative_dir = PathBuf::from(format!( + "{cluster_name}_cluster_kv_ssd_storage/{}", + crate::kv_ssd_storage::safe_path_component(instance_key) + )); + let mut out = Vec::new(); + let mut errors = Vec::new(); + for (root, limit_bytes) in self.paths.iter().zip(large_limit_size.iter().copied()) { + let candidate = Path::new(root).join(&relative_dir); + match fs::create_dir_all(&candidate) { + Ok(()) => out.push((candidate, limit_bytes)), + Err(err) => errors.push(format!("{} ({})", candidate.display(), err)), + } + } + if out.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "large_file_paths contains no usable root for kv ssd storage; tried: {}", + errors.join(", ") + ), + } + .into_kverror()); + } + Ok(out) + } } /// KV client backend types supported by the system @@ -733,13 +873,131 @@ pub struct ClientConfig { pub pprof_duration_seconds: Option, pub redis_compat_listen_addr: Option, pub fluxonkv_spec: FluxonKvSpec, - pub share_mem_path: String, // Mandatory shared bundle path + pub share_mem_path: String, // Mandatory shared bundle path pub large_file_paths: LargeFilePaths, // Mandatory large-file roots for logs and caches + pub ssd_storage: Option, pub test_spec_config: TestSpecConfig, } const CAPACITY_ALIGNMENT_BYTES: u64 = 16 * 1024 * 1024; +fn parse_size_bytes(value: &SizeBytesYaml, field: &str) -> KvResult { + match value { + SizeBytesYaml::Bytes(bytes) => Ok(*bytes), + SizeBytesYaml::Text(text) => parse_size_bytes_text(text, field), + } +} + +fn parse_size_bytes_text(raw: &str, field: &str) -> KvResult { + let input = raw.trim(); + if input.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field} size must not be empty"), + } + .into_kverror()); + } + + let mut number_end = 0usize; + let mut saw_digit = false; + let mut saw_dot = false; + for (idx, ch) in input.char_indices() { + if ch.is_ascii_digit() { + saw_digit = true; + number_end = idx + ch.len_utf8(); + continue; + } + if ch == '.' && !saw_dot { + saw_dot = true; + number_end = idx + ch.len_utf8(); + continue; + } + break; + } + let number = input[..number_end].trim(); + let unit = input[number_end..].trim(); + if !saw_digit || number == "." { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field} size must start with a number, got: {raw}"), + } + .into_kverror()); + } + if unit.chars().any(|ch| !ch.is_ascii_alphabetic()) { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field} size unit must be alphabetic, got: {raw}"), + } + .into_kverror()); + } + let multiplier = match unit.to_ascii_lowercase().as_str() { + "" | "b" | "byte" | "bytes" => 1u128, + "kb" | "kib" => 1024u128, + "mb" | "mib" => 1024u128.pow(2), + "gb" | "gib" => 1024u128.pow(3), + "tb" | "tib" => 1024u128.pow(4), + _ => { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "{field} size unit must be one of B, KB, MB, GB, TB, KiB, MiB, GiB, TiB; got: {raw}" + ), + } + .into_kverror()); + } + }; + + let (whole, frac) = match number.split_once('.') { + Some((whole, frac)) => { + if frac.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field} size has empty fractional part: {raw}"), + } + .into_kverror()); + } + (whole, frac) + } + None => (number, ""), + }; + let digits = format!("{whole}{frac}"); + let mut numerator = 0u128; + for ch in digits.chars() { + let digit = ch.to_digit(10).ok_or_else(|| { + ConfigError::InvalidClientConfig { + detail: format!("{field} size has invalid number: {raw}"), + } + .into_kverror() + })? as u128; + numerator = numerator + .checked_mul(10) + .and_then(|value| value.checked_add(digit)) + .ok_or_else(|| { + ConfigError::InvalidClientConfig { + detail: format!("{field} size is too large: {raw}"), + } + .into_kverror() + })?; + } + let scale = (0..frac.len()).try_fold(1u128, |acc, _| { + acc.checked_mul(10).ok_or_else(|| { + ConfigError::InvalidClientConfig { + detail: format!("{field} size has too many decimal places: {raw}"), + } + .into_kverror() + }) + })?; + let bytes_num = numerator.checked_mul(multiplier).ok_or_else(|| { + ConfigError::InvalidClientConfig { + detail: format!("{field} size is too large: {raw}"), + } + .into_kverror() + })?; + let bytes = bytes_num / scale; + if bytes > u64::MAX as u128 { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field} size exceeds u64 bytes: {raw}"), + } + .into_kverror()); + } + Ok(bytes as u64) +} + fn _validate_host_port_no_scheme(value: &str, field: &str) -> KvResult<()> { let s = value.trim(); if s.is_empty() { @@ -787,17 +1045,8 @@ fn _validate_host_port_no_scheme(value: &str, field: &str) -> KvResult<()> { Ok(()) } -fn _validate_capacity_multiple_of_alignment(value: u64, field: &str) -> KvResult<()> { - if value % CAPACITY_ALIGNMENT_BYTES != 0 { - return Err(ConfigError::InvalidClientConfig { - detail: format!( - "{} must be multiple of {} bytes, got: {}", - field, CAPACITY_ALIGNMENT_BYTES, value - ), - } - .into_kverror()); - } - Ok(()) +fn align_capacity_down(value: u64) -> u64 { + value / CAPACITY_ALIGNMENT_BYTES * CAPACITY_ALIGNMENT_BYTES } pub fn normalize_etcd_addresses(addresses: &[String]) -> KvResult> { @@ -914,6 +1163,7 @@ impl ClientConfigYaml { materialize_default_test_spec_transport_mode(&mut test_spec_config); validate_required_transfer_rpc_fast_path_ready_timeout(&test_spec_config)?; validate_test_spec_tcp_thread_tuning(&test_spec_config)?; + validate_test_spec_kv_ssd_backend(&test_spec_config)?; // Role selection contract: // - Missing contribute_to_cluster_pool_size means "zero-contribution" mode. @@ -930,49 +1180,39 @@ impl ClientConfigYaml { }, ), Some(c) => { - _validate_capacity_multiple_of_alignment( - c.dram, + let dram = align_capacity_down(parse_size_bytes( + &c.dram, "contribute_to_cluster_pool_size.dram", - )?; + )?); + let mut vram = HashMap::new(); for (gpu_id, size) in c.vram.iter() { - _validate_capacity_multiple_of_alignment( - *size, - &format!("contribute_to_cluster_pool_size.vram.{gpu_id}"), - )?; + vram.insert( + gpu_id.clone(), + align_capacity_down(parse_size_bytes( + size, + &format!("contribute_to_cluster_pool_size.vram.{gpu_id}"), + )?), + ); } - let vram_is_zero = c.vram.values().all(|&v| v == 0); - if c.dram == 0 && !vram_is_zero { + let vram_is_zero = vram.values().all(|&v| v == 0); + if dram == 0 && !vram_is_zero { return Err(ConfigError::InvalidClientConfig { detail: "contribute_to_cluster_pool_size is partially zero: dram=0 but vram has non-zero values".to_string(), } .into_kverror()); } - let is_zero = c.dram == 0 && vram_is_zero; - - ( - is_zero, - ContributeToClusterPoolSize { - dram: c.dram, - vram: c.vram.clone(), - }, - ) + let is_zero = dram == 0 && vram_is_zero; + + (is_zero, ContributeToClusterPoolSize { dram, vram }) } }; - if !is_external { - let Some(contrib_yaml) = &self.contribute_to_cluster_pool_size else { - return Err(ConfigError::InvalidClientConfig { - detail: "contribute_to_cluster_pool_size is required for owner mode (non-zero contribution)".to_string(), - } - .into_kverror()); - }; - if contrib_yaml.dram == 0 { - return Err(ConfigError::InvalidClientConfig { - detail: "owner mode requires non-zero contribute_to_cluster_pool_size.dram" - .to_string(), - } - .into_kverror()); + if !is_external && contribute_to_cluster_pool_size.dram == 0 { + return Err(ConfigError::InvalidClientConfig { + detail: "owner mode requires non-zero contribute_to_cluster_pool_size.dram" + .to_string(), } + .into_kverror()); } let is_side_transfer_worker = matches!( @@ -1028,6 +1268,13 @@ impl ClientConfigYaml { } .into_kverror()); } + if self.fluxonkv_spec.large_limit_size.is_some() { + return Err(ConfigError::InvalidClientConfig { + detail: "fluxonkv_spec.large_limit_size is forbidden in zero-contribution mode" + .to_string(), + } + .into_kverror()); + } } // Preserve historical behavior for configs that omit `protocol`, but allow @@ -1170,13 +1417,15 @@ impl ClientConfigYaml { } else { let Some(large_file_paths_yaml) = self.fluxonkv_spec.large_file_paths.as_ref() else { return Err(ConfigError::InvalidClientConfig { - detail: "fluxonkv_spec.large_file_paths is required for owner mode" - .to_string(), + detail: "fluxonkv_spec.large_file_paths is required for owner mode".to_string(), } .into_kverror()); }; LargeFilePaths { - paths: verify_non_empty_root_path_list(&large_file_paths_yaml.0, "large_file_paths")?, + paths: verify_non_empty_root_path_list( + &large_file_paths_yaml.0, + "large_file_paths", + )?, } }; @@ -1204,6 +1453,45 @@ impl ClientConfigYaml { } }; + let ssd_storage = if is_external { + None + } else { + match std::mem::take(&mut self.fluxonkv_spec.large_limit_size) { + None | Some(YamlNullable::Null) => None, + Some(YamlNullable::Value(raw_large_limit_size)) => { + let mut large_limit_size = Vec::with_capacity(raw_large_limit_size.len()); + for (idx, limit) in raw_large_limit_size.iter().enumerate() { + large_limit_size.push(parse_size_bytes( + limit, + &format!("fluxonkv_spec.large_limit_size[{idx}]"), + )?); + } + if large_limit_size.len() != large_file_paths.paths.len() { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "fluxonkv_spec.large_limit_size length must match fluxonkv_spec.large_file_paths length: large_limit_size={} large_file_paths={}", + large_limit_size.len(), + large_file_paths.paths.len() + ), + } + .into_kverror()); + } + for (idx, limit) in large_limit_size.iter().enumerate() { + if *limit < crate::kv_ssd_storage::SSD_ALIGNMENT as u64 { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "fluxonkv_spec.large_limit_size[{idx}] must be >= {}", + crate::kv_ssd_storage::SSD_ALIGNMENT + ), + } + .into_kverror()); + } + } + Some(KvSsdStorageConfig { large_limit_size }) + } + } + }; + Ok(ClientConfig { cluster_name: fluxonkv_spec.cluster_name.clone(), etcd_addresses_raw, @@ -1215,6 +1503,7 @@ impl ClientConfigYaml { fluxonkv_spec, share_mem_path, large_file_paths, + ssd_storage, test_spec_config, }) } @@ -1499,6 +1788,7 @@ impl MasterConfigYaml { materialize_default_test_spec_transport_mode(&mut test_spec_config); validate_required_transfer_rpc_fast_path_ready_timeout(&test_spec_config)?; validate_test_spec_tcp_thread_tuning(&test_spec_config)?; + validate_test_spec_kv_ssd_backend(&test_spec_config)?; let protocol = apply_test_spec_rdma_device_names_to_protocol( self.protocol.unwrap_or(ProtocolConfig { protocol_type: ProtocolType::Rdma, @@ -1647,7 +1937,241 @@ fluxonkv_spec: .unwrap(); let err = cfg.verify().unwrap_err(); let text = format!("{err}"); - assert!(text.contains("fluxonkv_spec.large_file_paths is forbidden in zero-contribution mode")); + assert!( + text.contains("fluxonkv_spec.large_file_paths is forbidden in zero-contribution mode") + ); + } + + #[test] + fn client_config_owner_accepts_large_limit_size() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [1048576] + sub_cluster: rack-a +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + assert_eq!( + verified + .ssd_storage + .as_ref() + .map(|cfg| cfg.large_limit_size.as_slice()), + Some([1048576].as_slice()) + ); + assert_eq!( + verified.test_spec_config.kv_ssd_uring_mode, + KvSsdUringMode::SingleBuffer + ); + assert_eq!( + verified.test_spec_config.kv_ssd_storage_backend, + KvSsdStorageBackend::Native + ); + } + + #[test] + fn client_config_accepts_size_strings() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 1.1GB + vram: + "0": 0.5GB +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [512.9B] + sub_cluster: rack-a +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + assert_eq!( + verified.contribute_to_cluster_pool_size.dram, + 70 * CAPACITY_ALIGNMENT_BYTES + ); + assert_eq!( + verified + .contribute_to_cluster_pool_size + .vram + .get("0") + .copied(), + Some(512 * 1024 * 1024) + ); + assert_eq!( + verified + .ssd_storage + .as_ref() + .map(|cfg| cfg.large_limit_size.as_slice()), + Some([512].as_slice()) + ); + } + + #[test] + fn size_string_truncates_fractional_bytes() { + assert_eq!(parse_size_bytes_text("0.9B", "test_size").unwrap(), 0); + } + + #[test] + fn client_test_spec_config_accepts_kv_ssd_iovec_ablation_mode() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [1048576] + sub_cluster: rack-a +test_spec_config: + kv_ssd_uring_mode: iovec +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + assert_eq!( + verified.test_spec_config.kv_ssd_uring_mode, + KvSsdUringMode::Iovec + ); + } + + #[test] + fn client_test_spec_config_accepts_foyer_kv_ssd_backend() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [134217728] + sub_cluster: rack-a +test_spec_config: + kv_ssd_storage_backend: foyer +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + assert_eq!( + verified.test_spec_config.kv_ssd_storage_backend, + KvSsdStorageBackend::Foyer + ); + } + + #[test] + fn client_test_spec_config_rejects_foyer_with_iovec_mode() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [134217728] + sub_cluster: rack-a +test_spec_config: + kv_ssd_storage_backend: foyer + kv_ssd_uring_mode: iovec +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + assert!(format!("{err}").contains("kv_ssd_uring_mode is only valid")); + } + + #[test] + fn client_config_owner_rejects_large_limit_size_length_mismatch() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large_a, /tmp/test_owner_large_b] + large_limit_size: [1048576] + sub_cluster: rack-a +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + let text = format!("{err}"); + assert!( + text.contains("fluxonkv_spec.large_limit_size length must match fluxonkv_spec.large_file_paths length"), + "{text}" + ); + } + + #[test] + fn client_config_owner_rejects_too_small_large_limit_size() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner + large_file_paths: [/tmp/test_owner_large] + large_limit_size: [1] + sub_cluster: rack-a +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + let text = format!("{err}"); + assert!( + text.contains("fluxonkv_spec.large_limit_size[0] must be >= 512"), + "{text}" + ); + } + + #[test] + fn client_config_zero_contribution_rejects_large_limit_size() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_external +fluxonkv_spec: + cluster_name: test_cluster + share_mem_path: /tmp/test_external + large_limit_size: [1048576] +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + let text = format!("{err}"); + assert!( + text.contains("fluxonkv_spec.large_limit_size is forbidden in zero-contribution mode") + ); } #[test] @@ -1667,7 +2191,9 @@ fluxonkv_spec: let logs_dir = large_file_paths.kv_logs_dir("test_cluster").unwrap(); assert_eq!( logs_dir, - first_root.join("child").join("test_cluster_cluster_kv_logs") + first_root + .join("child") + .join("test_cluster_cluster_kv_logs") ); assert!(logs_dir.exists()); @@ -1683,6 +2209,32 @@ fluxonkv_spec: assert!(third_party_logs_dir.exists()); } + #[test] + fn large_file_paths_uses_all_usable_roots_for_kv_ssd_storage() { + let tempdir = new_test_dir("fluxon_large_paths_uses_all_usable_roots_for_kv_ssd_storage"); + let first_root = tempdir.join("first_root"); + let second_root = tempdir.join("second_root"); + + let large_file_paths = LargeFilePaths { + paths: vec![ + first_root.to_string_lossy().into_owned(), + second_root.to_string_lossy().into_owned(), + ], + }; + + let dirs = large_file_paths + .kv_ssd_storage_dirs("test_cluster", "owner/a:b") + .unwrap(); + assert_eq!( + dirs, + vec![ + first_root.join("test_cluster_cluster_kv_ssd_storage/owner_a_b"), + second_root.join("test_cluster_cluster_kv_ssd_storage/owner_a_b"), + ] + ); + assert!(dirs.iter().all(|dir| dir.exists())); + } + #[test] fn client_test_spec_config_accepts_explicit_rdma_device_names() { let cfg = ClientConfigYaml::from_str( diff --git a/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs b/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs index da701cd..630a8ea 100644 --- a/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs +++ b/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs @@ -89,6 +89,7 @@ fn new_client_config( large_file_paths: LargeFilePaths { paths: vec![format!("{}_large", shm_path)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -130,6 +131,7 @@ fn new_zero_contribution_client_config( }, share_mem_path: shm_path.to_string(), large_file_paths: LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } diff --git a/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs b/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs index 9cb291f..a2e844c 100644 --- a/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs +++ b/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs @@ -18,6 +18,7 @@ use crate::{ cluster_manager::{ ClusterManager, ClusterManagerAccessTrait, IpcBandwidthAttributorHandle, NodeRole, }, + master_kv_router::msg_pack::GetSourceKind, master_lease_manager::msg_pack::{AllocateClientLeaseReq, ClientLeaseKeepaliveReq}, memholder::ExternalMemHolder, p2p::{ @@ -865,8 +866,7 @@ impl ExternalInner { return Ok(false); } - self.finish_owner_recover(&share_mem_path, payload) - .await?; + self.finish_owner_recover(&share_mem_path, payload).await?; Ok(true) } @@ -1610,11 +1610,29 @@ impl ExternalInner { } } - /// External Get operation (outer): retry + wait wrapper around get_inner + /// External Get operation with the selected storage source for this call. + pub(crate) async fn get_with_source( + &self, + key: &str, + ) -> KvResult, GetSourceKind)>> { + self.get_with_source_inner(key).await + } + + /// External Get operation (outer): retry + wait wrapper around get_inner. pub async fn get( &self, key: &str, ) -> KvResult>> { + Ok(self + .get_with_source_inner(key) + .await? + .map(|(holder, _source_kind)| holder)) + } + + async fn get_with_source_inner( + &self, + key: &str, + ) -> KvResult, GetSourceKind)>> { tracing::debug!("External get request for key: {}", key); // Ensure external mode configured; if not, block until owner is ready once @@ -1630,7 +1648,7 @@ impl ExternalInner { // 1) Fast path: try weak-index lookup first if let Some(h) = self.try_get_from_weak_cache(key).await { - return Ok(Some(h)); + return Ok(Some((h, GetSourceKind::Memory))); } // 2) Ensure only one inflight get() per key using a keyed semaphore (permits=1) @@ -1647,7 +1665,7 @@ impl ExternalInner { key ); drop(permit); - return Ok(Some(h)); + return Ok(Some((h, GetSourceKind::Memory))); } let mut recover_attempts: usize = 0; @@ -1660,7 +1678,7 @@ impl ExternalInner { match self.get_inner(key, prev_owner_start_time).await { Ok(v) => { // Update weak index on success if Some - if let Some(ref h) = v { + if let Some((ref h, _source_kind)) = v { // let hex= &h.bytes()[..std::cmp::min(16, h.len as usize)]; // tracing::info!("external get done, key={}, partial_hex={:?}", key, hex); self.key_weak_memholder_index @@ -1708,7 +1726,7 @@ key={}, attempt={}/{}, err={}", &self, key: &str, started_time: i64, - ) -> KvResult>> { + ) -> KvResult, GetSourceKind)>> { // Ensure external mode configured and compute base address let base_ptr = self.base_ptr_ro().await.expect( "ExternalClientApi.get_inner called in non-external mode (no shared memory configured)", @@ -1751,6 +1769,7 @@ key={}, attempt={}/{}, err={}", ); match result { Some(info) => { + let source_kind = info.source_kind; // Attribute external<->owner shared-memory payload bytes to the owner topology edge. // // Causal chain: @@ -1788,7 +1807,7 @@ key={}, attempt={}/{}, err={}", info.len, info.holder_id ); - Ok(Some(external_memholder)) + Ok(Some((external_memholder, source_kind))) } None => Ok(None), } @@ -1803,6 +1822,7 @@ key={}, attempt={}/{}, err={}", ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exists = opts.reject_if_exists(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let observe_sink = opts.test_observe_put_phases(); let observe_enabled = true; @@ -1838,6 +1858,7 @@ key={}, attempt={}/{}, err={}", base_addr, lease_id, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster.as_deref(), observe_enabled, ) @@ -1890,6 +1911,7 @@ key={}, attempt={}/{}, err={}", ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exists = opts.reject_if_exists(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let observe_sink = opts.test_observe_put_phases(); let observe_enabled = true; @@ -1924,6 +1946,7 @@ key={}, attempt={}/{}, err={}", base_addr, lease_id, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster.as_deref(), observe_enabled, ) @@ -1972,6 +1995,7 @@ key={}, attempt={}/{}, err={}", base_addr: usize, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, preferred_sub_cluster: Option<&str>, observe_enabled: bool, ) -> KvResult { @@ -1981,6 +2005,7 @@ key={}, attempt={}/{}, err={}", key: key.to_string(), len: payload_len, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), started_time, test_observe_put_phases: true, @@ -2129,6 +2154,7 @@ key={}, attempt={}/{}, err={}", base_addr: usize, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, preferred_sub_cluster: Option<&str>, observe_enabled: bool, ) -> KvResult { @@ -2139,6 +2165,7 @@ key={}, attempt={}/{}, err={}", key: key.to_string(), len: value.len() as u64, reject_if_inflight_same_key, + reject_if_exists, preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), started_time, test_observe_put_phases: true, diff --git a/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs b/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs new file mode 100644 index 0000000..588f503 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs @@ -0,0 +1,3696 @@ +use crate::config::{KvSsdStorageBackend, KvSsdUringMode}; +use crate::kv_ssd_storage_foyer::{FoyerKvSsdPersistGuard, FoyerKvSsdStorage}; +use crate::master_kv_router::msg_pack::SsdReplicaEviction; +use crate::master_kv_router::put::PutIDForAKey; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}; +use ::tokio::{ + sync::{Notify, mpsc as tokio_mpsc, oneshot}, + task, +}; +use futures::stream::{FuturesUnordered, StreamExt}; +use io_uring::{IoUring, opcode, types::Fd}; +use parking_lot::Mutex; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::fs::{self, OpenOptions}; +use std::io; +use std::os::fd::{AsRawFd, RawFd}; +use std::os::unix::fs::MetadataExt; +use std::os::unix::fs::OpenOptionsExt; +use std::path::{Path, PathBuf}; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::thread::JoinHandle; + +pub(crate) const SSD_ALIGNMENT: usize = 512; +const DEFAULT_SHARDS_PER_OWNER: usize = 4; +const DEFAULT_URING_THREADS: usize = 16; +const DEFAULT_URING_IO_DEPTH: usize = 128; +const DEFAULT_URING_READ_WEIGHT: usize = 2; +const DEFAULT_WRITE_QUEUE_DEPTH: usize = 8; +const DEFAULT_READ_QUEUE_DEPTH: usize = 16; +const DEFAULT_WRITE_INFLIGHT: usize = 2; +const DEFAULT_READ_INFLIGHT: usize = 16; +const DEFAULT_EVICTION_QUEUE_DEPTH: usize = 1024; +pub(crate) const DEFAULT_READ_TRANSFER_PIPELINE_CHUNK_BYTES: u64 = 4 * 1024 * 1024; +pub(crate) const DEFAULT_READ_TRANSFER_PIPELINE_INFLIGHT: usize = 4; + +#[derive(Clone, Debug)] +pub struct KvSsdStorageInit { + pub root_limits: Vec, + pub uring_mode: KvSsdUringMode, + pub backend: KvSsdStorageBackend, +} + +#[derive(Clone, Debug)] +pub struct KvSsdStorageRootLimit { + pub root_dir: PathBuf, + pub limit_bytes: u64, +} + +#[derive(Debug)] +pub struct KvSsdStorage { + root_dirs: Vec, + devices: Vec, + shard_to_device: Vec, + next_write_device: AtomicUsize, + inner: Arc>, + space_notify: Arc, + eviction_rx: Mutex>>>, + _eviction_tx_guard: Option>>, + foyer: Option, +} + +#[derive(Debug)] +struct SsdDeviceWorker { + device_id: u64, + root_dir: PathBuf, + shard_ids: Vec, + max_shard_capacity: u64, + _files: Vec, + _io: Arc, + write_tx: tokio_mpsc::Sender, + read_tx: tokio_mpsc::Sender, +} + +#[derive(Clone, Debug)] +struct SsdDeviceRoot { + device_id: u64, + root_dir: PathBuf, + limit_bytes: u64, +} + +struct OpenedSsdShard { + shard_id: usize, + device_idx: usize, + capacity: u64, + file: std::fs::File, +} + +#[derive(Clone, Debug)] +struct SsdShardSpec { + shard_id: usize, + device_idx: usize, + capacity: u64, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct SsdLoadedChunk { + pub offset: u64, + pub stage_addr: u64, + pub len: u64, +} + +#[derive(Clone, Debug)] +pub(crate) struct KvSsdStorageDeviceUsage { + pub device: String, + pub capacity_bytes: u64, + pub used_bytes: u64, +} + +#[derive(Debug)] +struct KvSsdStorageInner { + ring: SsdRingBuffer, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct KvSsdKey { + key: String, + put_id: PutIDForAKey, +} + +#[derive(Clone, Debug)] +struct SsdIndexEntry { + shard_id: usize, + begin: u64, + len: u64, + aligned_len: u64, + file_offset: u64, +} + +#[derive(Clone, Debug)] +struct SsdReadPinInfo { + entry: SsdIndexEntry, + count: usize, +} + +#[derive(Clone, Debug)] +enum SsdEntryState { + Writing(SsdIndexEntry), + Committed(SsdIndexEntry), +} + +impl SsdEntryState { + fn entry(&self) -> &SsdIndexEntry { + match self { + Self::Writing(entry) | Self::Committed(entry) => entry, + } + } +} + +#[derive(Debug)] +struct SsdShardRing { + capacity: u64, + head: u64, + tail: u64, + order: VecDeque, +} + +#[derive(Debug)] +struct SsdRingBuffer { + shards: Vec, + next_shard: usize, + entries: HashMap, + read_pins: HashMap, +} + +#[derive(Debug)] +enum SsdPreparedWrite { + Ready { + entry: SsdIndexEntry, + evicted: Vec, + }, + Existing, + BlockedByBusyIo, +} + +#[derive(Debug)] +enum SsdAllocation { + Ready { + begin: u64, + file_offset: u64, + evicted: Vec, + }, + BlockedByBusyIo, + TooLarge, +} + +impl SsdRingBuffer { + fn new(shard_capacities: Vec) -> Self { + assert!(!shard_capacities.is_empty()); + Self { + shards: shard_capacities + .into_iter() + .map(|capacity| SsdShardRing { + capacity, + head: 0, + tail: 0, + order: VecDeque::new(), + }) + .collect(), + next_shard: 0, + entries: HashMap::new(), + read_pins: HashMap::new(), + } + } + + #[cfg(test)] + fn get(&self, key: &KvSsdKey) -> Option { + match self.entries.get(key) { + Some(SsdEntryState::Committed(entry)) if self.is_offset_valid(entry) => { + Some(entry.clone()) + } + _ => None, + } + } + + fn pin_read(&mut self, key: &KvSsdKey) -> Option { + let entry = match self.entries.get(key) { + Some(SsdEntryState::Committed(entry)) if self.is_offset_valid(entry) => entry.clone(), + _ => return None, + }; + let pin = self + .read_pins + .entry(key.clone()) + .or_insert_with(|| SsdReadPinInfo { + entry: entry.clone(), + count: 0, + }); + pin.count += 1; + Some(entry) + } + + fn unpin_read(&mut self, key: &KvSsdKey) { + match self.read_pins.get_mut(key) { + Some(pin) if pin.count > 1 => pin.count -= 1, + Some(_) => { + self.read_pins.remove(key); + } + None => debug_assert!(false, "missing kv ssd read pin for key={key:?}"), + } + } + + #[cfg(test)] + fn prepare_write(&mut self, key: KvSsdKey, len: u64) -> KvResult { + let allowed_shards = (0..self.shards.len()).collect::>(); + self.prepare_write_on_shards(key, len, &allowed_shards) + } + + fn prepare_write_on_shards( + &mut self, + key: KvSsdKey, + len: u64, + allowed_shards: &[usize], + ) -> KvResult { + if allowed_shards.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd device has no shards".to_string(), + })); + } + if let Some(state) = self.entries.get(&key) { + return Ok(match state { + SsdEntryState::Writing(_) => SsdPreparedWrite::BlockedByBusyIo, + SsdEntryState::Committed(entry) if entry.len == len => SsdPreparedWrite::Existing, + SsdEntryState::Committed(entry) => { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd duplicate persist length mismatch: key={} put_id=({},{}) existing_len={} requested_len={}", + key.key, key.put_id.0, key.put_id.1, entry.len, len + ), + })); + } + }); + } + let aligned_len = align_up_u64(len, SSD_ALIGNMENT as u64)?; + let max_capacity = self + .shards + .iter() + .enumerate() + .filter(|(idx, _)| allowed_shards.contains(idx)) + .map(|(_, shard)| shard.capacity) + .max() + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd device has invalid shard set: {allowed_shards:?}"), + }) + })?; + if aligned_len > max_capacity { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd value len={} aligned_len={} exceeds shard capacity={}", + len, aligned_len, max_capacity + ), + })); + } + + let shard_count = self.shards.len(); + for offset in 0..shard_count { + let shard_id = (self.next_shard + offset) % shard_count; + if !allowed_shards.contains(&shard_id) { + continue; + } + let (begin, file_offset, evicted) = match self + .allocate_contiguous(shard_id, aligned_len) + { + SsdAllocation::Ready { + begin, + file_offset, + evicted, + } => (begin, file_offset, evicted), + SsdAllocation::BlockedByBusyIo => continue, + SsdAllocation::TooLarge => unreachable!("aligned_len was checked against capacity"), + }; + self.next_shard = (shard_id + 1) % shard_count; + + let entry = SsdIndexEntry { + shard_id, + begin, + len, + aligned_len, + file_offset, + }; + self.entries + .insert(key.clone(), SsdEntryState::Writing(entry.clone())); + self.shards[shard_id].order.push_back(key); + return Ok(SsdPreparedWrite::Ready { entry, evicted }); + } + + Ok(SsdPreparedWrite::BlockedByBusyIo) + } + + fn allocate_contiguous(&mut self, shard_id: usize, size: u64) -> SsdAllocation { + let shard = &self.shards[shard_id]; + if size > shard.capacity { + return SsdAllocation::TooLarge; + } + let capacity = shard.capacity; + let mut head = shard.head; + let phys = head % capacity; + let space_until_end = capacity - phys; + if size > space_until_end { + head += space_until_end; + } + let begin = head; + let new_head = head + size; + let new_tail = new_head.saturating_sub(capacity); + if self.has_busy_entries_before(shard_id, new_tail) { + return SsdAllocation::BlockedByBusyIo; + } + + self.shards[shard_id].head = new_head; + let evicted = self.advance_tail(shard_id, new_tail); + SsdAllocation::Ready { + begin, + file_offset: begin % capacity, + evicted, + } + } + + fn advance_tail(&mut self, shard_id: usize, new_tail: u64) -> Vec { + if new_tail <= self.shards[shard_id].tail { + return Vec::new(); + } + debug_assert!(!self.has_busy_entries_before(shard_id, new_tail)); + self.shards[shard_id].tail = new_tail; + let mut evicted = Vec::new(); + + while let Some(key) = self.shards[shard_id].order.front() { + match self.entries.get(key) { + Some(state) if state.entry().begin >= new_tail => break, + _ => { + let key = self.shards[shard_id] + .order + .pop_front() + .expect("front key exists"); + if matches!(self.entries.remove(&key), Some(SsdEntryState::Committed(_))) { + evicted.push(key); + } + } + } + } + evicted + } + + fn commit(&mut self, key: &KvSsdKey, success: bool) -> bool { + let Some(state) = self.entries.get(key) else { + return false; + }; + let entry = match state { + SsdEntryState::Writing(entry) => entry.clone(), + SsdEntryState::Committed(_) => return true, + }; + if !self.is_offset_valid(&entry) || !success { + self.entries.remove(key); + return false; + } + self.entries + .insert(key.clone(), SsdEntryState::Committed(entry)); + true + } + + fn remove(&mut self, key: &KvSsdKey) { + self.entries.remove(key); + } + + fn is_offset_valid(&self, entry: &SsdIndexEntry) -> bool { + self.shards + .get(entry.shard_id) + .is_some_and(|shard| entry.begin >= shard.tail) + } + + fn used_bytes_by_shard(&self) -> Vec { + let mut out = vec![0u64; self.shards.len()]; + for state in self.entries.values() { + let entry = state.entry(); + if self.is_offset_valid(entry) { + out[entry.shard_id] = out[entry.shard_id].saturating_add(entry.aligned_len); + } + } + out + } + + fn has_busy_entries_before(&self, shard_id: usize, new_tail: u64) -> bool { + if new_tail <= self.shards[shard_id].tail { + return false; + } + let writing_busy = self.entries.values().any(|state| match state { + SsdEntryState::Writing(entry) => entry.shard_id == shard_id && entry.begin < new_tail, + SsdEntryState::Committed(_) => false, + }); + if writing_busy { + return true; + } + self.read_pins + .values() + .any(|pin| pin.entry.shard_id == shard_id && pin.entry.begin < new_tail) + } +} + +struct SsdReadPin { + inner: Arc>, + space_notify: Arc, + key: KvSsdKey, +} + +pub(crate) struct KvSsdPersistGuard { + _native_pin: Option, + _foyer_guard: Option, +} + +impl Drop for SsdReadPin { + fn drop(&mut self) { + self.inner.lock().ring.unpin_read(&self.key); + self.space_notify.notify_one(); + } +} + +struct WriteCommand { + key: KvSsdKey, + entry_len: u64, + data: AlignedBuffer, + done_tx: oneshot::Sender>, +} + +struct ReadCommand { + key: KvSsdKey, + entry: SsdIndexEntry, + file_offset: u64, + target: ReadTarget, + _read_pin: Option, + done_tx: oneshot::Sender>, +} + +struct WriteTask { + key: KvSsdKey, + entry: SsdIndexEntry, + data: AlignedBuffer, + done_tx: oneshot::Sender>, +} + +struct ReadTask { + key: KvSsdKey, + entry: SsdIndexEntry, + file_offset: u64, + target: ReadTarget, + _read_pin: Option, + done_tx: oneshot::Sender>, +} + +struct WriteCompletion { + key: KvSsdKey, + success: bool, + result: KvResult<()>, + done_tx: oneshot::Sender>, +} + +struct ReadCompletion { + key: KvSsdKey, + entry: SsdIndexEntry, + result: KvResult, + _read_pin: Option, + done_tx: oneshot::Sender>, +} + +enum ReadTarget { + Scratch(AlignedBuffer), + Direct { target_addr: u64, len: usize }, +} + +enum ReadOutput { + Scratch(AlignedBuffer), + Direct, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SsdReadPath { + Scratch, + Direct, +} + +pub fn safe_path_component(raw: &str) -> String { + let mut out = String::with_capacity(raw.len().max(1)); + for ch in raw.chars() { + if ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.') { + out.push(ch); + } else { + out.push('_'); + } + } + if out.is_empty() { + "unnamed".to_string() + } else { + out + } +} + +impl KvSsdStorage { + pub async fn new(init: KvSsdStorageInit) -> KvResult { + if init.root_limits.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage root_limits must contain at least one path".to_string(), + })); + } + for (idx, root_limit) in init.root_limits.iter().enumerate() { + if root_limit.limit_bytes < SSD_ALIGNMENT as u64 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd storage root_limits[{idx}].limit_bytes must be >= {}", + SSD_ALIGNMENT + ), + })); + } + } + + if init.backend == KvSsdStorageBackend::Foyer { + if init.root_limits.len() != 1 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd backend requires exactly one root, got {}", + init.root_limits.len() + ), + })); + } + tracing::warn!( + "Using test-only Foyer KV SSD backend; proactive SSD eviction notifications are unavailable" + ); + let root_limit = init + .root_limits + .into_iter() + .next() + .expect("foyer root count checked above"); + let dummy_capacity = root_limit.limit_bytes; + let foyer = FoyerKvSsdStorage::new(root_limit).await?; + let root_dirs = vec![foyer.root_dir().to_path_buf()]; + let (eviction_tx, eviction_rx) = tokio_mpsc::channel(DEFAULT_EVICTION_QUEUE_DEPTH); + return Ok(Self { + root_dirs, + devices: Vec::new(), + shard_to_device: Vec::new(), + next_write_device: AtomicUsize::new(0), + inner: Arc::new(Mutex::new(KvSsdStorageInner { + ring: SsdRingBuffer::new(vec![dummy_capacity]), + })), + space_notify: Arc::new(Notify::new()), + eviction_rx: Mutex::new(Some(eviction_rx)), + _eviction_tx_guard: Some(eviction_tx), + foyer: Some(foyer), + }); + } + + let device_roots = deduplicate_device_roots(&init.root_limits)?; + let effective_root_dirs = device_roots + .iter() + .map(|root| root.root_dir.clone()) + .collect::>(); + let shard_specs = build_shard_specs(&device_roots)?; + let opened_shards = open_cache_files(&device_roots, &shard_specs)?; + let inner = Arc::new(Mutex::new(KvSsdStorageInner { + ring: SsdRingBuffer::new( + shard_specs + .iter() + .map(|spec| spec.capacity) + .collect::>(), + ), + })); + let space_notify = Arc::new(Notify::new()); + let (eviction_tx, eviction_rx) = tokio_mpsc::channel(DEFAULT_EVICTION_QUEUE_DEPTH); + let shard_count = shard_specs.len(); + let mut shard_to_device = vec![0usize; shard_count]; + let mut device_shards = device_roots + .iter() + .map(|root| (root.clone(), Vec::<(usize, std::fs::File)>::new())) + .collect::>(); + for opened in opened_shards { + debug_assert_eq!( + inner.lock().ring.shards[opened.shard_id].capacity, + opened.capacity + ); + shard_to_device[opened.shard_id] = opened.device_idx; + device_shards[opened.device_idx] + .1 + .push((opened.shard_id, opened.file)); + } + + let mut devices = Vec::with_capacity(device_shards.len()); + for (device_root, shard_files) in device_shards { + let shard_ids = shard_files + .iter() + .map(|(shard_id, _)| *shard_id) + .collect::>(); + let max_shard_capacity = shard_ids + .iter() + .filter_map(|shard_id| shard_specs.get(*shard_id)) + .map(|spec| spec.capacity) + .max() + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd device has no shards: device_id={} root_dir={}", + device_root.device_id, + device_root.root_dir.display() + ), + }) + })?; + let fds = shard_files + .iter() + .map(|(shard_id, file)| (*shard_id, file.as_raw_fd())) + .collect::>(); + let io = Arc::new(UringIoEngine::new_multi( + fds, + UringConfig { + threads: DEFAULT_URING_THREADS, + io_depth: DEFAULT_URING_IO_DEPTH, + mode: init.uring_mode, + }, + )?); + let (write_tx, write_rx) = tokio_mpsc::channel(DEFAULT_WRITE_QUEUE_DEPTH); + let (read_tx, read_rx) = tokio_mpsc::channel(DEFAULT_READ_QUEUE_DEPTH); + + task::spawn(ssd_writer_loop( + Arc::clone(&inner), + write_rx, + Arc::clone(&io), + Arc::clone(&space_notify), + DEFAULT_WRITE_INFLIGHT, + shard_ids.clone(), + eviction_tx.clone(), + )); + task::spawn(ssd_reader_loop( + Arc::clone(&inner), + read_rx, + Arc::clone(&io), + DEFAULT_READ_INFLIGHT, + )); + + devices.push(SsdDeviceWorker { + device_id: device_root.device_id, + root_dir: device_root.root_dir, + shard_ids, + max_shard_capacity, + _files: shard_files + .into_iter() + .map(|(_, file)| file) + .collect::>(), + _io: io, + write_tx, + read_tx, + }); + } + + Ok(Self { + root_dirs: effective_root_dirs, + devices, + shard_to_device, + next_write_device: AtomicUsize::new(0), + inner, + space_notify, + eviction_rx: Mutex::new(Some(eviction_rx)), + _eviction_tx_guard: None, + foyer: None, + }) + } + + pub fn root_dirs(&self) -> &[PathBuf] { + &self.root_dirs + } + + pub(crate) fn device_usage_snapshot(&self) -> Vec { + if let Some(foyer) = self.foyer.as_ref() { + return vec![KvSsdStorageDeviceUsage { + device: format!("foyer:{}", foyer.root_dir().display()), + capacity_bytes: foyer.capacity_bytes(), + used_bytes: foyer.logical_used_bytes(), + }]; + } + let inner = self.inner.lock(); + let used_by_shard = inner.ring.used_bytes_by_shard(); + self.devices + .iter() + .enumerate() + .map(|(device_idx, device)| { + let mut capacity_bytes = 0u64; + let mut used_bytes = 0u64; + for shard_id in &device.shard_ids { + if self + .shard_to_device + .get(*shard_id) + .is_some_and(|mapped| *mapped == device_idx) + { + let Some(shard) = inner.ring.shards.get(*shard_id) else { + continue; + }; + capacity_bytes = capacity_bytes.saturating_add(shard.capacity); + used_bytes = used_bytes + .saturating_add(used_by_shard.get(*shard_id).copied().unwrap_or(0)); + } + } + KvSsdStorageDeviceUsage { + device: format!("dev:{}:{}", device.device_id, device.root_dir.display()), + capacity_bytes, + used_bytes, + } + }) + .collect() + } + + pub(crate) fn take_eviction_rx(&self) -> Option>> { + self.eviction_rx.lock().take() + } + + fn next_write_tx(&self, entry_len: u64) -> KvResult> { + if self.devices.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage has no active device".to_string(), + })); + } + let aligned_len = align_ssd_io_len(entry_len)?; + let capacities = self + .devices + .iter() + .map(|device| device.max_shard_capacity) + .collect::>(); + loop { + let start_idx = self.next_write_device.load(Ordering::Relaxed); + let Some(device_idx) = select_write_device(&capacities, aligned_len, start_idx) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd value len={} aligned_len={} exceeds every device shard capacity: {:?}", + entry_len, aligned_len, capacities + ), + })); + }; + let next_idx = (device_idx + 1) % self.devices.len(); + if self + .next_write_device + .compare_exchange_weak(start_idx, next_idx, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + return Ok(self.devices[device_idx].write_tx.clone()); + } + } + } + + fn read_tx_for_shard(&self, shard_id: usize) -> KvResult> { + let Some(device_idx) = self.shard_to_device.get(shard_id).copied() else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd invalid shard id for read: {}", shard_id), + })); + }; + let Some(device) = self.devices.get(device_idx) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd invalid device index for read: shard_id={} device_idx={}", + shard_id, device_idx + ), + })); + }; + if !device.shard_ids.contains(&shard_id) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd shard/device route mismatch: shard_id={} device_idx={} device_id={} root_dir={}", + shard_id, + device_idx, + device.device_id, + device.root_dir.display() + ), + })); + } + Ok(device.read_tx.clone()) + } + + pub(crate) async fn persist_from_addr( + &self, + key: &str, + put_id: PutIDForAKey, + addr: u64, + len: u64, + ) -> KvResult { + validate_key(key)?; + if let Some(foyer) = self.foyer.as_ref() { + return foyer + .persist_from_addr(key, put_id, addr, len) + .await + .map(|guard| KvSsdPersistGuard { + _native_pin: None, + _foyer_guard: Some(guard), + }); + } + let len_usize = usize::try_from(len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd persist len does not fit usize: {}", len), + }) + })?; + let aligned_len = align_up_usize(len_usize, SSD_ALIGNMENT)?; + let data = unsafe { AlignedBuffer::copy_from_addr(addr, len_usize, aligned_len)? }; + self.persist_buffer(key, put_id, len, data).await + } + + pub async fn persist(&self, key: &str, put_id: PutIDForAKey, data: &[u8]) -> KvResult<()> { + validate_key(key)?; + if let Some(foyer) = self.foyer.as_ref() { + let guard = foyer.persist(key, put_id, data).await?; + drop(guard); + return Ok(()); + } + let aligned_len = align_up_usize(data.len(), SSD_ALIGNMENT)?; + let mut buffer = AlignedBuffer::zeroed(aligned_len)?; + unsafe { + std::ptr::copy_nonoverlapping(data.as_ptr(), buffer.as_mut_ptr(), data.len()); + } + let guard = self + .persist_buffer(key, put_id, data.len() as u64, buffer) + .await?; + drop(guard); + Ok(()) + } + + async fn persist_buffer( + &self, + key: &str, + put_id: PutIDForAKey, + entry_len: u64, + data: AlignedBuffer, + ) -> KvResult { + let (done_tx, done_rx) = oneshot::channel(); + let write_tx = self.next_write_tx(entry_len)?; + write_tx + .send(WriteCommand { + key: KvSsdKey { + key: key.to_string(), + put_id, + }, + entry_len, + data, + done_tx, + }) + .await + .map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd write queue closed: {}", err), + }) + })?; + done_rx.await.map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd write completion closed: {}", err), + }) + })? + } + + pub async fn load_into_addr( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_len: u64, + ) -> KvResult<()> { + validate_key(key)?; + if let Some(foyer) = self.foyer.as_ref() { + return foyer + .load_into_addr(key, put_id, target_addr, len, target_len) + .await; + } + if target_len < len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd target capacity too small for key={} put_id=({},{}) len={} target_len={}", + key, put_id.0, put_id.1, len, target_len + ), + })); + } + let key = KvSsdKey { + key: key.to_string(), + put_id, + }; + let (entry, read_pin) = { + let mut inner = self.inner.lock(); + let Some(entry) = inner.ring.pin_read(&key) else { + return Err(KvError::Api(ApiError::KeyNotFound { + key: key.key.clone(), + })); + }; + ( + entry, + SsdReadPin { + inner: Arc::clone(&self.inner), + space_notify: Arc::clone(&self.space_notify), + key: key.clone(), + }, + ) + }; + if entry.len != len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd length mismatch for key={} put_id=({},{}) expected={} actual={}", + key.key, put_id.0, put_id.1, len, entry.len + ), + })); + } + + let len_usize = usize::try_from(len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd load len does not fit usize: {}", len), + }) + })?; + let aligned_len_usize = usize::try_from(entry.aligned_len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd aligned load len does not fit usize: {}", + entry.aligned_len + ), + }) + })?; + let target = match choose_read_path(&entry, target_addr, len, target_len) { + SsdReadPath::Direct => ReadTarget::Direct { + target_addr, + len: aligned_len_usize, + }, + SsdReadPath::Scratch => ReadTarget::Scratch(AlignedBuffer::zeroed(aligned_len_usize)?), + }; + let output = self + .submit_read_command( + key, + entry.clone(), + entry.file_offset, + target, + Some(read_pin), + ) + .await?; + if let ReadOutput::Scratch(buffer) = output { + unsafe { + std::ptr::copy_nonoverlapping(buffer.as_ptr(), target_addr as *mut u8, len_usize); + } + } + Ok(()) + } + + pub(crate) async fn load_into_addr_chunks( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_len: u64, + chunk_bytes: u64, + max_read_inflight: usize, + ready_tx: tokio_mpsc::Sender, + ) -> KvResult<()> { + validate_key(key)?; + if let Some(foyer) = self.foyer.as_ref() { + return foyer + .load_into_addr_chunks( + key, + put_id, + target_addr, + len, + target_len, + chunk_bytes, + max_read_inflight, + ready_tx, + ) + .await; + } + if target_len < len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd target capacity too small for chunked load: key={} put_id=({},{}) len={} target_len={}", + key, put_id.0, put_id.1, len, target_len + ), + })); + } + let chunk_bytes = align_up_u64(chunk_bytes.max(1), SSD_ALIGNMENT as u64)?; + let key = KvSsdKey { + key: key.to_string(), + put_id, + }; + let (entry, _read_pin) = { + let mut inner = self.inner.lock(); + let Some(entry) = inner.ring.pin_read(&key) else { + return Err(KvError::Api(ApiError::KeyNotFound { + key: key.key.clone(), + })); + }; + ( + entry, + SsdReadPin { + inner: Arc::clone(&self.inner), + space_notify: Arc::clone(&self.space_notify), + key: key.clone(), + }, + ) + }; + if entry.len != len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd length mismatch for chunked load: key={} put_id=({},{}) expected={} actual={}", + key.key, put_id.0, put_id.1, len, entry.len + ), + })); + } + + let mut next_offset = 0u64; + let mut inflight = FuturesUnordered::new(); + let max_read_inflight = max_read_inflight.max(1); + + loop { + while next_offset < len && inflight.len() < max_read_inflight { + let payload_len = chunk_bytes.min(len - next_offset); + let stage_addr = checked_add_u64(target_addr, next_offset, "chunk stage addr")?; + let remaining_target_len = target_len - next_offset; + inflight.push(self.load_entry_range_into_addr( + key.clone(), + entry.clone(), + next_offset, + payload_len, + stage_addr, + remaining_target_len, + )); + next_offset += payload_len; + } + + let Some(chunk) = inflight.next().await else { + break; + }; + let chunk = chunk?; + ready_tx.send(chunk).await.map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd chunk ready queue closed: {}", err), + }) + })?; + } + Ok(()) + } + + async fn load_entry_range_into_addr( + &self, + key: KvSsdKey, + entry: SsdIndexEntry, + offset: u64, + payload_len: u64, + target_addr: u64, + target_len: u64, + ) -> KvResult { + if payload_len == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd chunk payload len must be positive".to_string(), + })); + } + let payload_end = checked_add_u64(offset, payload_len, "chunk payload end")?; + if payload_end > entry.len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd chunk exceeds entry len: offset={} len={} entry_len={}", + offset, payload_len, entry.len + ), + })); + } + let read_len = align_up_u64(payload_len, SSD_ALIGNMENT as u64)?; + let read_end = checked_add_u64(offset, read_len, "chunk read end")?; + if read_end > entry.aligned_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd aligned chunk exceeds entry aligned len: offset={} read_len={} aligned_len={}", + offset, read_len, entry.aligned_len + ), + })); + } + if target_len < read_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd chunk target capacity too small: offset={} read_len={} target_len={}", + offset, read_len, target_len + ), + })); + } + let file_offset = checked_add_u64(entry.file_offset, offset, "chunk file offset")?; + let read_len_usize = usize::try_from(read_len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd chunk read len does not fit usize: {}", read_len), + }) + })?; + let payload_len_usize = usize::try_from(payload_len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd chunk payload len does not fit usize: {}", + payload_len + ), + }) + })?; + let target = match choose_chunk_read_path(target_addr, read_len, target_len, file_offset) { + SsdReadPath::Direct => ReadTarget::Direct { + target_addr, + len: read_len_usize, + }, + SsdReadPath::Scratch => ReadTarget::Scratch(AlignedBuffer::zeroed(read_len_usize)?), + }; + let output = self + .submit_read_command(key, entry, file_offset, target, None) + .await?; + if let ReadOutput::Scratch(buffer) = output { + unsafe { + std::ptr::copy_nonoverlapping( + buffer.as_ptr(), + target_addr as *mut u8, + payload_len_usize, + ); + } + } + Ok(SsdLoadedChunk { + offset, + stage_addr: target_addr, + len: payload_len, + }) + } + + async fn submit_read_command( + &self, + key: KvSsdKey, + entry: SsdIndexEntry, + file_offset: u64, + target: ReadTarget, + read_pin: Option, + ) -> KvResult { + let (done_tx, done_rx) = oneshot::channel(); + let read_tx = self.read_tx_for_shard(entry.shard_id)?; + read_tx + .send(ReadCommand { + key, + entry, + file_offset, + target, + _read_pin: read_pin, + done_tx, + }) + .await + .map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd read queue closed: {}", err), + }) + })?; + done_rx.await.map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd read completion closed: {}", err), + }) + })? + } + + #[cfg(test)] + async fn has_entry(&self, key: &str, put_id: PutIDForAKey) -> bool { + if let Some(foyer) = self.foyer.as_ref() { + return foyer.contains(key, put_id); + } + let key = KvSsdKey { + key: key.to_string(), + put_id, + }; + self.inner.lock().ring.get(&key).is_some() + } +} + +fn prepare_ssd_write( + inner: &Arc>, + space_notify: &Arc, + key: &KvSsdKey, + entry_len: u64, + shard_ids: &[usize], +) -> (KvResult, Option) { + let (prepared, existing_pinned) = { + let mut inner = inner.lock(); + let prepared = inner + .ring + .prepare_write_on_shards(key.clone(), entry_len, shard_ids); + let existing_pinned = matches!(&prepared, Ok(SsdPreparedWrite::Existing)) + && inner.ring.pin_read(key).is_some(); + (prepared, existing_pinned) + }; + let guard = existing_pinned.then(|| KvSsdPersistGuard { + _native_pin: Some(SsdReadPin { + inner: Arc::clone(inner), + space_notify: Arc::clone(space_notify), + key: key.clone(), + }), + _foyer_guard: None, + }); + (prepared, guard) +} + +async fn ssd_writer_loop( + inner: Arc>, + mut rx: tokio_mpsc::Receiver, + io: Arc, + space_notify: Arc, + write_inflight: usize, + shard_ids: Vec, + eviction_tx: tokio_mpsc::Sender>, +) { + let mut pending: VecDeque = VecDeque::new(); + let mut inflight = FuturesUnordered::new(); + let max_inflight = write_inflight.max(1); + + loop { + while inflight.len() < max_inflight { + let Some(cmd) = pending.pop_front() else { + break; + }; + let (prepared, existing_guard) = + prepare_ssd_write(&inner, &space_notify, &cmd.key, cmd.entry_len, &shard_ids); + match prepared { + Ok(SsdPreparedWrite::Ready { entry, evicted }) => { + publish_ssd_evictions(&eviction_tx, evicted); + inflight.push(execute_write( + WriteTask { + key: cmd.key, + entry, + data: cmd.data, + done_tx: cmd.done_tx, + }, + Arc::clone(&io), + )); + } + Ok(SsdPreparedWrite::Existing) => { + let result = existing_guard.ok_or_else(|| { + KvError::Api(ApiError::KeyNotFound { + key: cmd.key.key.clone(), + }) + }); + let _ = cmd.done_tx.send(result); + } + Ok(SsdPreparedWrite::BlockedByBusyIo) => { + pending.push_front(cmd); + break; + } + Err(err) => { + let _ = cmd.done_tx.send(Err(err)); + } + } + } + + tokio::select! { + Some(completion) = inflight.next(), if !inflight.is_empty() => { + finish_write_completion(&inner, &space_notify, completion); + } + Some(cmd) = rx.recv() => { + pending.push_back(cmd); + } + _ = space_notify.notified(), if !pending.is_empty() => { + // Retry pending commands after an active read/write releases a ring position. + } + else => { + if pending.is_empty() && inflight.is_empty() { + break; + } + }, + } + } + + while !pending.is_empty() || !inflight.is_empty() { + while inflight.len() < max_inflight { + let Some(cmd) = pending.pop_front() else { + break; + }; + let (prepared, existing_guard) = + prepare_ssd_write(&inner, &space_notify, &cmd.key, cmd.entry_len, &shard_ids); + match prepared { + Ok(SsdPreparedWrite::Ready { entry, evicted }) => { + publish_ssd_evictions(&eviction_tx, evicted); + inflight.push(execute_write( + WriteTask { + key: cmd.key, + entry, + data: cmd.data, + done_tx: cmd.done_tx, + }, + Arc::clone(&io), + )); + } + Ok(SsdPreparedWrite::Existing) => { + let result = existing_guard.ok_or_else(|| { + KvError::Api(ApiError::KeyNotFound { + key: cmd.key.key.clone(), + }) + }); + let _ = cmd.done_tx.send(result); + } + Ok(SsdPreparedWrite::BlockedByBusyIo) => { + pending.push_front(cmd); + break; + } + Err(err) => { + let _ = cmd.done_tx.send(Err(err)); + } + } + } + + if let Some(completion) = inflight.next().await { + finish_write_completion(&inner, &space_notify, completion); + } else if !pending.is_empty() { + space_notify.notified().await; + } + } +} + +fn publish_ssd_evictions( + eviction_tx: &tokio_mpsc::Sender>, + evicted: Vec, +) { + if evicted.is_empty() { + return; + } + let replicas = evicted + .into_iter() + .map(|key| SsdReplicaEviction { + key: key.key, + put_id: key.put_id, + }) + .collect(); + if let Err(err) = eviction_tx.try_send(replicas) { + tracing::warn!( + "Dropping SSD eviction notification because its queue is unavailable: {err}" + ); + } +} + +fn finish_write_completion( + inner: &Arc>, + space_notify: &Arc, + completion: WriteCompletion, +) { + let (committed, commit_pinned) = { + let mut inner = inner.lock(); + let committed = inner.ring.commit(&completion.key, completion.success); + let commit_pinned = committed && inner.ring.pin_read(&completion.key).is_some(); + (committed, commit_pinned) + }; + space_notify.notify_one(); + let result = match completion.result { + Err(err) => Err(err), + Ok(()) if committed && commit_pinned => Ok(KvSsdPersistGuard { + _native_pin: Some(SsdReadPin { + inner: Arc::clone(inner), + space_notify: Arc::clone(space_notify), + key: completion.key.clone(), + }), + _foyer_guard: None, + }), + Ok(()) => Err(KvError::Api(ApiError::KeyNotFound { + key: completion.key.key.clone(), + })), + }; + let _ = completion.done_tx.send(result); +} + +async fn execute_write(task: WriteTask, io: Arc) -> WriteCompletion { + let WriteTask { + key, + entry, + data, + done_tx, + } = task; + let data_len = data.len(); + let shard_id = entry.shard_id; + let file_offset = entry.file_offset; + let result = async move { + let rx = { + let data_ptr = data.as_ptr(); + io.write_at_async(shard_id, data_ptr, data_len, file_offset)? + }; + let written = rx + .await + .map_err(|_| io::Error::other("kv ssd write completion dropped"))??; + if written != data_len { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + format!("short kv ssd write: {} != {}", written, data_len), + ) + .into()); + } + Ok(()) + } + .await; + let result = result.map_err(|err| file_error_for_entry(&key, file_offset, err)); + WriteCompletion { + key, + success: result.is_ok(), + result, + done_tx, + } +} + +async fn ssd_reader_loop( + inner: Arc>, + mut rx: tokio_mpsc::Receiver, + io: Arc, + read_inflight: usize, +) { + let mut pending = VecDeque::new(); + let mut inflight = FuturesUnordered::new(); + let max_inflight = read_inflight.max(1); + + loop { + while inflight.len() < max_inflight { + let Some(task) = pending.pop_front() else { + break; + }; + inflight.push(execute_read(task, Arc::clone(&io))); + } + + tokio::select! { + Some(completion) = inflight.next(), if !inflight.is_empty() => { + let valid = inner.lock().ring.is_offset_valid(&completion.entry); + let result = if valid { + completion.result + } else { + inner.lock().ring.remove(&completion.key); + Err(KvError::Api(ApiError::KeyNotFound { + key: completion.key.key.clone(), + })) + }; + let _ = completion.done_tx.send(result); + } + Some(cmd) = rx.recv() => { + pending.push_back(ReadTask { + key: cmd.key, + entry: cmd.entry, + file_offset: cmd.file_offset, + target: cmd.target, + _read_pin: cmd._read_pin, + done_tx: cmd.done_tx, + }); + } + else => break, + } + } + + while let Some(completion) = inflight.next().await { + let valid = inner.lock().ring.is_offset_valid(&completion.entry); + let result = if valid { + completion.result + } else { + inner.lock().ring.remove(&completion.key); + Err(KvError::Api(ApiError::KeyNotFound { + key: completion.key.key.clone(), + })) + }; + let _ = completion.done_tx.send(result); + } +} + +async fn execute_read(task: ReadTask, io: Arc) -> ReadCompletion { + let ReadTask { + key, + entry, + file_offset, + target, + _read_pin, + done_tx, + } = task; + let shard_id = entry.shard_id; + let result = async move { + match target { + ReadTarget::Scratch(mut buffer) => { + let buffer_len = buffer.len(); + let rx = { + let buffer_ptr = buffer.as_mut_ptr(); + io.read_at_async(shard_id, buffer_ptr, buffer_len, file_offset)? + }; + let read = rx + .await + .map_err(|_| io::Error::other("kv ssd read completion dropped"))??; + if read != buffer_len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("short kv ssd read: {} != {}", read, buffer_len), + )); + } + Ok(ReadOutput::Scratch(buffer)) + } + ReadTarget::Direct { target_addr, len } => { + let rx = io.read_at_async(shard_id, target_addr as *mut u8, len, file_offset)?; + let read = rx + .await + .map_err(|_| io::Error::other("kv ssd read completion dropped"))??; + if read != len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("short kv ssd direct read: {} != {}", read, len), + )); + } + Ok(ReadOutput::Direct) + } + } + } + .await + .map_err(|err| file_error_for_entry(&key, file_offset, err)); + ReadCompletion { + key, + entry, + result, + _read_pin, + done_tx, + } +} + +#[derive(Clone, Copy)] +struct UringConfig { + threads: usize, + io_depth: usize, + mode: KvSsdUringMode, +} + +#[derive(Clone, Copy)] +enum IoType { + Read, + Write, + Readv, + Writev, +} + +struct IoCtx { + io_type: IoType, + fd: RawFd, + len: usize, + offset: u64, + complete: oneshot::Sender>, + buffer: Option<(*mut u8, usize)>, + iovecs: Option>, +} + +unsafe impl Send for IoCtx {} + +struct UringShard { + read_rx: crossbeam::channel::Receiver, + write_rx: crossbeam::channel::Receiver, + uring: Option, + io_depth: usize, + read_weight: usize, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SubmitWaitAction { + Retry, + DrainCompletions, + FailWorker, +} + +fn submit_wait_error_action(err: &io::Error) -> SubmitWaitAction { + if err.kind() == io::ErrorKind::Interrupted { + SubmitWaitAction::Retry + } else if err.raw_os_error() == Some(libc::EBUSY) { + SubmitWaitAction::DrainCompletions + } else { + SubmitWaitAction::FailWorker + } +} + +impl UringShard { + fn run(mut self) { + let mut read_inflight = 0usize; + let mut write_inflight = 0usize; + let mut read_closed = false; + let mut write_closed = false; + let mut submitted = HashMap::new(); + let mut next_token = 1u64; + + loop { + let mut inflight = read_inflight + write_inflight; + while inflight < self.io_depth && !(read_closed && write_closed) { + let next = self.try_recv_weighted( + &mut read_closed, + &mut write_closed, + read_inflight, + write_inflight, + ); + let Some(ctx) = next else { + break; + }; + self.submit_ctx( + ctx, + &mut submitted, + &mut next_token, + &mut read_inflight, + &mut write_inflight, + ); + inflight = read_inflight + write_inflight; + } + + if read_closed && write_closed && inflight == 0 { + debug_assert!(submitted.is_empty()); + return; + } + if inflight == 0 { + let Some(ctx) = self.recv_blocking(&mut read_closed, &mut write_closed) else { + continue; + }; + self.submit_ctx( + ctx, + &mut submitted, + &mut next_token, + &mut read_inflight, + &mut write_inflight, + ); + continue; + } + let submit_result = self + .uring + .as_mut() + .expect("io_uring must exist while its worker is running") + .submit_and_wait(1); + if let Err(err) = submit_result { + match submit_wait_error_action(&err) { + SubmitWaitAction::Retry | SubmitWaitAction::DrainCompletions => { + self.drain_completions( + &mut submitted, + &mut read_inflight, + &mut write_inflight, + ); + continue; + } + SubmitWaitAction::FailWorker => { + let error_kind = err.kind(); + let detail = format!("io_uring submit failed: {err}"); + drop(self.uring.take()); + for (_, ctx) in submitted.drain() { + let _ = ctx + .complete + .send(Err(io::Error::new(error_kind, detail.clone()))); + } + return; + } + } + } + + self.drain_completions(&mut submitted, &mut read_inflight, &mut write_inflight); + } + } + + fn drain_completions( + &mut self, + submitted: &mut HashMap, + read_inflight: &mut usize, + write_inflight: &mut usize, + ) { + let uring = self + .uring + .as_mut() + .expect("io_uring must exist while draining completions"); + for cqe in uring.completion() { + let token = cqe.user_data(); + if token == 0 { + continue; + } + let Some(ctx) = submitted.remove(&token) else { + tracing::warn!("Ignoring io_uring completion with unknown token {token}"); + continue; + }; + match ctx.io_type { + IoType::Read | IoType::Readv => *read_inflight = read_inflight.saturating_sub(1), + IoType::Write | IoType::Writev => { + *write_inflight = write_inflight.saturating_sub(1) + } + } + let result = cqe.result(); + let send_result = if result < 0 { + Err(io::Error::from_raw_os_error(-result)) + } else { + Ok(result as usize) + }; + let _ = ctx.complete.send(send_result); + } + } + + fn try_recv_weighted( + &self, + read_closed: &mut bool, + write_closed: &mut bool, + read_inflight: usize, + write_inflight: usize, + ) -> Option { + let prefer_read = read_inflight <= write_inflight.saturating_mul(self.read_weight); + if prefer_read { + self.try_recv_read(read_closed) + .or_else(|| self.try_recv_write(write_closed)) + } else { + self.try_recv_write(write_closed) + .or_else(|| self.try_recv_read(read_closed)) + } + } + + fn try_recv_read(&self, read_closed: &mut bool) -> Option { + if *read_closed { + return None; + } + match self.read_rx.try_recv() { + Ok(ctx) => Some(ctx), + Err(crossbeam::channel::TryRecvError::Empty) => None, + Err(crossbeam::channel::TryRecvError::Disconnected) => { + *read_closed = true; + None + } + } + } + + fn try_recv_write(&self, write_closed: &mut bool) -> Option { + if *write_closed { + return None; + } + match self.write_rx.try_recv() { + Ok(ctx) => Some(ctx), + Err(crossbeam::channel::TryRecvError::Empty) => None, + Err(crossbeam::channel::TryRecvError::Disconnected) => { + *write_closed = true; + None + } + } + } + + fn recv_blocking(&self, read_closed: &mut bool, write_closed: &mut bool) -> Option { + loop { + match (!*read_closed, !*write_closed) { + (true, true) => { + crossbeam::channel::select! { + recv(self.read_rx) -> msg => match msg { + Ok(ctx) => return Some(ctx), + Err(_) => *read_closed = true, + }, + recv(self.write_rx) -> msg => match msg { + Ok(ctx) => return Some(ctx), + Err(_) => *write_closed = true, + }, + } + } + (true, false) => match self.read_rx.recv() { + Ok(ctx) => return Some(ctx), + Err(_) => *read_closed = true, + }, + (false, true) => match self.write_rx.recv() { + Ok(ctx) => return Some(ctx), + Err(_) => *write_closed = true, + }, + (false, false) => return None, + } + } + } + + fn submit_ctx( + &mut self, + ctx: IoCtx, + submitted: &mut HashMap, + next_token: &mut u64, + read_inflight: &mut usize, + write_inflight: &mut usize, + ) { + let fd = Fd(ctx.fd); + let buffer = ctx.buffer; + let iovecs_ptr = ctx + .iovecs + .as_ref() + .map(|iovecs| iovecs.as_ptr()) + .unwrap_or(std::ptr::null()); + let sqe = match ctx.io_type { + IoType::Read => { + let Some((ptr, len)) = buffer else { + let _ = ctx.complete.send(Err(io::Error::other( + "single-buffer read missing buffer pointer", + ))); + return; + }; + opcode::Read::new(fd, ptr, len as _) + .offset(ctx.offset) + .build() + } + IoType::Write => { + let Some((ptr, len)) = buffer else { + let _ = ctx.complete.send(Err(io::Error::other( + "single-buffer write missing buffer pointer", + ))); + return; + }; + opcode::Write::new(fd, ptr as *const u8, len as _) + .offset(ctx.offset) + .build() + } + IoType::Readv => opcode::Readv::new(fd, iovecs_ptr, ctx.len as _) + .offset(ctx.offset) + .build(), + IoType::Writev => opcode::Writev::new(fd, iovecs_ptr, ctx.len as _) + .offset(ctx.offset) + .build(), + }; + let io_type = ctx.io_type; + let token = next_submission_token(next_token, submitted); + let sqe = sqe.user_data(token); + let Some(uring) = self.uring.as_mut() else { + let _ = ctx + .complete + .send(Err(io::Error::other("io_uring worker is closed"))); + return; + }; + let push_result = unsafe { uring.submission().push(&sqe) }; + if push_result.is_err() { + let _ = ctx + .complete + .send(Err(io::Error::other("submission queue full"))); + return; + } + let replaced = submitted.insert(token, ctx); + debug_assert!( + replaced.is_none(), + "io_uring submission token must be unique" + ); + match io_type { + IoType::Read | IoType::Readv => *read_inflight += 1, + IoType::Write | IoType::Writev => *write_inflight += 1, + } + } +} + +fn next_submission_token(next_token: &mut u64, submitted: &HashMap) -> u64 { + loop { + if *next_token == 0 { + *next_token = 1; + } + let token = *next_token; + *next_token = token.wrapping_add(1); + if !submitted.contains_key(&token) { + return token; + } + } +} + +#[derive(Debug)] +struct UringIoEngine { + fds: HashMap, + read_txs: Vec>, + write_txs: Vec>, + handles: Vec>, + mode: KvSsdUringMode, +} + +impl UringIoEngine { + fn new_multi(shard_fds: Vec<(usize, RawFd)>, cfg: UringConfig) -> io::Result { + if cfg.threads == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "threads must be > 0", + )); + } + if shard_fds.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "at least one fd is required", + )); + } + let fds = shard_fds.into_iter().collect::>(); + let mut read_txs = Vec::with_capacity(cfg.threads); + let mut write_txs = Vec::with_capacity(cfg.threads); + let mut handles = Vec::with_capacity(cfg.threads); + for idx in 0..cfg.threads { + let (read_tx, read_rx) = crossbeam::channel::bounded(cfg.io_depth * 2); + let (write_tx, write_rx) = crossbeam::channel::bounded(cfg.io_depth * 2); + let uring = IoUring::builder().build(cfg.io_depth as u32)?; + let handle = std::thread::Builder::new() + .name(format!("fluxon-kv-ssd-uring-{idx}")) + .spawn(move || { + UringShard { + read_rx, + write_rx, + uring: Some(uring), + io_depth: cfg.io_depth, + read_weight: DEFAULT_URING_READ_WEIGHT, + } + .run() + })?; + read_txs.push(read_tx); + write_txs.push(write_tx); + handles.push(handle); + } + Ok(Self { + fds, + read_txs, + write_txs, + handles, + mode: cfg.mode, + }) + } + + fn read_at_async( + &self, + shard_id: usize, + ptr: *mut u8, + len: usize, + offset: u64, + ) -> io::Result>> { + match self.mode { + KvSsdUringMode::SingleBuffer => { + self.submit_buffer(IoType::Read, shard_id, ptr, len, offset) + } + KvSsdUringMode::Iovec => self.readv_at_async(shard_id, vec![(ptr, len)], offset), + } + } + + fn write_at_async( + &self, + shard_id: usize, + ptr: *const u8, + len: usize, + offset: u64, + ) -> io::Result>> { + match self.mode { + KvSsdUringMode::SingleBuffer => { + self.submit_buffer(IoType::Write, shard_id, ptr as *mut u8, len, offset) + } + KvSsdUringMode::Iovec => self.writev_at_async(shard_id, vec![(ptr, len)], offset), + } + } + + fn readv_at_async( + &self, + shard_id: usize, + iovecs: Vec<(*mut u8, usize)>, + offset: u64, + ) -> io::Result>> { + self.submit_iovecs(IoType::Readv, shard_id, iovecs, offset) + } + + fn writev_at_async( + &self, + shard_id: usize, + iovecs: Vec<(*const u8, usize)>, + offset: u64, + ) -> io::Result>> { + let iovecs = iovecs + .into_iter() + .map(|(ptr, len)| (ptr as *mut u8, len)) + .collect(); + self.submit_iovecs(IoType::Writev, shard_id, iovecs, offset) + } + + fn submit_buffer( + &self, + io_type: IoType, + shard_id: usize, + ptr: *mut u8, + len: usize, + offset: u64, + ) -> io::Result>> { + if !matches!(io_type, IoType::Read | IoType::Write) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "single-buffer submit requires read/write io type", + )); + } + validate_direct_io([(ptr as usize, len)], offset)?; + let (tx, rx) = oneshot::channel(); + let ctx = IoCtx { + io_type, + fd: self.fd(shard_id)?, + len, + offset, + complete: tx, + buffer: Some((ptr, len)), + iovecs: None, + }; + self.pick_tx(io_type, shard_id).send(ctx).map_err(|err| { + io::Error::new( + io::ErrorKind::BrokenPipe, + format!("io_uring send failed: {}", err), + ) + })?; + Ok(rx) + } + + fn submit_iovecs( + &self, + io_type: IoType, + shard_id: usize, + iovecs: Vec<(*mut u8, usize)>, + offset: u64, + ) -> io::Result>> { + if iovecs.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "readv/writev requires at least one iovec", + )); + } + validate_direct_io( + iovecs.iter().map(|(ptr, len)| (*ptr as usize, *len)), + offset, + )?; + let iovecs_libc = iovecs + .iter() + .map(|(ptr, len)| libc::iovec { + iov_base: *ptr as *mut libc::c_void, + iov_len: *len, + }) + .collect::>() + .into_boxed_slice(); + let (tx, rx) = oneshot::channel(); + let ctx = IoCtx { + io_type, + fd: self.fd(shard_id)?, + len: iovecs_libc.len(), + offset, + complete: tx, + buffer: None, + iovecs: Some(iovecs_libc), + }; + self.pick_tx(io_type, shard_id).send(ctx).map_err(|err| { + io::Error::new( + io::ErrorKind::BrokenPipe, + format!("io_uring send failed: {}", err), + ) + })?; + Ok(rx) + } + + fn fd(&self, shard_id: usize) -> io::Result { + self.fds.get(&shard_id).copied().ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("invalid SSD shard id {shard_id}"), + ) + }) + } + + fn pick_tx(&self, io_type: IoType, shard_id: usize) -> &crossbeam::channel::Sender { + match io_type { + IoType::Read => &self.read_txs[shard_id % self.read_txs.len()], + IoType::Write => &self.write_txs[shard_id % self.write_txs.len()], + IoType::Readv => &self.read_txs[shard_id % self.read_txs.len()], + IoType::Writev => &self.write_txs[shard_id % self.write_txs.len()], + } + } +} + +impl Drop for UringIoEngine { + fn drop(&mut self) { + self.read_txs.clear(); + self.write_txs.clear(); + for handle in self.handles.drain(..) { + let _ = handle.join(); + } + } +} + +struct AlignedBuffer { + ptr: NonNull, + len: usize, +} + +unsafe impl Send for AlignedBuffer {} + +impl AlignedBuffer { + fn zeroed(len: usize) -> KvResult { + if len == 0 || !len.is_multiple_of(SSD_ALIGNMENT) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "aligned buffer len must be positive and {}-byte aligned: {}", + SSD_ALIGNMENT, len + ), + })); + } + let mut raw = std::ptr::null_mut(); + let rc = unsafe { libc::posix_memalign(&mut raw, SSD_ALIGNMENT, len) }; + if rc != 0 || raw.is_null() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("posix_memalign failed with rc={}", rc), + })); + } + unsafe { + std::ptr::write_bytes(raw as *mut u8, 0, len); + } + Ok(Self { + ptr: NonNull::new(raw as *mut u8).expect("posix_memalign returned non-null"), + len, + }) + } + + unsafe fn copy_from_addr(addr: u64, actual_len: usize, aligned_len: usize) -> KvResult { + let mut buffer = Self::zeroed(aligned_len)?; + unsafe { + std::ptr::copy_nonoverlapping(addr as *const u8, buffer.as_mut_ptr(), actual_len); + } + Ok(buffer) + } + + fn as_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.ptr.as_ptr() + } + + fn len(&self) -> usize { + self.len + } +} + +impl Drop for AlignedBuffer { + fn drop(&mut self) { + unsafe { + libc::free(self.ptr.as_ptr() as *mut libc::c_void); + } + } +} + +fn validate_key(key: &str) -> KvResult<()> { + if key.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage key must be non-empty".to_string(), + })); + } + Ok(()) +} + +fn choose_read_path( + entry: &SsdIndexEntry, + target_addr: u64, + len: u64, + target_len: u64, +) -> SsdReadPath { + if len == 0 || entry.len != len { + return SsdReadPath::Scratch; + } + if target_addr.is_multiple_of(SSD_ALIGNMENT as u64) + && target_len >= entry.aligned_len + && entry.file_offset.is_multiple_of(SSD_ALIGNMENT as u64) + { + SsdReadPath::Direct + } else { + SsdReadPath::Scratch + } +} + +fn choose_chunk_read_path( + target_addr: u64, + read_len: u64, + target_len: u64, + file_offset: u64, +) -> SsdReadPath { + if read_len != 0 + && target_addr.is_multiple_of(SSD_ALIGNMENT as u64) + && read_len.is_multiple_of(SSD_ALIGNMENT as u64) + && target_len >= read_len + && file_offset.is_multiple_of(SSD_ALIGNMENT as u64) + { + SsdReadPath::Direct + } else { + SsdReadPath::Scratch + } +} + +fn choose_device_shard_count(limit_bytes: u64) -> usize { + let max_aligned_shards = (limit_bytes / SSD_ALIGNMENT as u64).max(1) as usize; + DEFAULT_SHARDS_PER_OWNER.min(max_aligned_shards).max(1) +} + +fn select_write_device( + max_shard_capacities: &[u64], + aligned_len: u64, + start_idx: usize, +) -> Option { + (0..max_shard_capacities.len()) + .map(|offset| start_idx.wrapping_add(offset) % max_shard_capacities.len()) + .find(|device_idx| max_shard_capacities[*device_idx] >= aligned_len) +} + +fn aligned_shard_capacity(capacity_bytes: u64, shard_count: usize) -> KvResult { + let raw = capacity_bytes / shard_count as u64; + let capacity = raw / SSD_ALIGNMENT as u64 * SSD_ALIGNMENT as u64; + if capacity == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage capacity is too small for shard count".to_string(), + })); + } + Ok(capacity) +} + +fn build_shard_specs(device_roots: &[SsdDeviceRoot]) -> KvResult> { + if device_roots.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage root_limits must contain at least one path".to_string(), + })); + } + let mut shard_specs = Vec::new(); + for (device_idx, root) in device_roots.iter().enumerate() { + let shard_count = choose_device_shard_count(root.limit_bytes); + let shard_capacity = aligned_shard_capacity(root.limit_bytes, shard_count)?; + for _ in 0..shard_count { + let shard_id = shard_specs.len(); + shard_specs.push(SsdShardSpec { + shard_id, + device_idx, + capacity: shard_capacity, + }); + } + } + Ok(shard_specs) +} + +fn deduplicate_device_roots(root_limits: &[KvSsdStorageRootLimit]) -> KvResult> { + if root_limits.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage root_limits must contain at least one path".to_string(), + })); + } + let mut seen_devices = HashSet::new(); + let mut device_roots = Vec::new(); + for root_limit in root_limits { + fs::create_dir_all(&root_limit.root_dir) + .map_err(|err| file_error(&root_limit.root_dir, 0, err))?; + let metadata = fs::metadata(&root_limit.root_dir) + .map_err(|err| file_error(&root_limit.root_dir, 0, err))?; + let device_id = metadata.dev(); + if seen_devices.insert(device_id) { + device_roots.push(SsdDeviceRoot { + device_id, + root_dir: root_limit.root_dir.clone(), + limit_bytes: root_limit.limit_bytes, + }); + } + } + if device_roots.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage root_limits contains no usable device".to_string(), + })); + } + Ok(device_roots) +} + +fn open_cache_files( + device_roots: &[SsdDeviceRoot], + shard_specs: &[SsdShardSpec], +) -> KvResult> { + if device_roots.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "kv ssd storage root_limits must contain at least one path".to_string(), + })); + } + let mut files = Vec::with_capacity(shard_specs.len()); + for spec in shard_specs { + let Some(device_root) = device_roots.get(spec.device_idx) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd shard spec references invalid device index: shard_id={} device_idx={}", + spec.shard_id, spec.device_idx + ), + })); + }; + let root_dir = &device_root.root_dir; + let shards_dir = root_dir.join("shards"); + fs::create_dir_all(&shards_dir).map_err(|err| file_error(&shards_dir, 0, err))?; + let path = shards_dir.join(format!("shard-{:06}.dat", spec.shard_id)); + let file = OpenOptions::new() + .create(true) + .truncate(true) + .read(true) + .write(true) + .custom_flags(libc::O_DIRECT) + .open(&path) + .map_err(|err| file_error(&path, 0, err))?; + file.set_len(spec.capacity) + .map_err(|err| file_error(&path, 0, err))?; + files.push(OpenedSsdShard { + shard_id: spec.shard_id, + device_idx: spec.device_idx, + capacity: spec.capacity, + file, + }); + } + Ok(files) +} + +fn align_up_usize(value: usize, alignment: usize) -> KvResult { + value + .checked_add(alignment - 1) + .map(|v| v / alignment * alignment) + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("alignment overflow for value={}", value), + }) + }) +} + +fn align_up_u64(value: u64, alignment: u64) -> KvResult { + value + .checked_add(alignment - 1) + .map(|v| v / alignment * alignment) + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("alignment overflow for value={}", value), + }) + }) +} + +pub(crate) fn align_ssd_io_len(len: u64) -> KvResult { + align_up_u64(len, SSD_ALIGNMENT as u64) +} + +fn checked_add_u64(lhs: u64, rhs: u64, label: &str) -> KvResult { + lhs.checked_add(rhs).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd {label} overflow: {lhs} + {rhs}"), + }) + }) +} + +fn validate_direct_io( + iovecs: impl IntoIterator, + offset: u64, +) -> io::Result<()> { + ensure_aligned("offset", offset as usize)?; + for (addr, len) in iovecs { + ensure_aligned("buffer address", addr)?; + ensure_aligned("iovec length", len)?; + } + Ok(()) +} + +fn ensure_aligned(name: &str, value: usize) -> io::Result<()> { + if value.is_multiple_of(SSD_ALIGNMENT) { + Ok(()) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("O_DIRECT {name} {value:#x} is not {SSD_ALIGNMENT}-byte aligned"), + )) + } +} + +fn file_error_for_entry(key: &KvSsdKey, offset: u64, err: io::Error) -> KvError { + KvError::Api(ApiError::FileWriteError { + path: format!("kv-ssd://{}@({},{})", key.key, key.put_id.0, key.put_id.1), + offset, + detail: err.to_string(), + }) +} + +fn file_error(path: &Path, offset: u64, err: io::Error) -> KvError { + KvError::Api(ApiError::FileWriteError { + path: path.to_string_lossy().to_string(), + offset, + detail: err.to_string(), + }) +} + +impl From for KvError { + fn from(err: io::Error) -> Self { + KvError::Api(ApiError::FileWriteError { + path: "kv-ssd://io".to_string(), + offset: 0, + detail: err.to_string(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use uuid::Uuid; + + fn new_root() -> PathBuf { + std::env::current_dir() + .unwrap() + .join("target") + .join("fluxon_kv_ssd_tests") + .join(Uuid::new_v4().to_string()) + } + + async fn new_store_with_mode(max_bytes: u64, uring_mode: KvSsdUringMode) -> KvSsdStorage { + KvSsdStorage::new(KvSsdStorageInit { + root_limits: vec![KvSsdStorageRootLimit { + root_dir: new_root(), + limit_bytes: max_bytes, + }], + uring_mode, + backend: KvSsdStorageBackend::Native, + }) + .await + .unwrap() + } + + async fn new_store(max_bytes: u64) -> KvSsdStorage { + new_store_with_mode(max_bytes, KvSsdUringMode::SingleBuffer).await + } + + async fn new_foyer_store(max_bytes: u64) -> KvSsdStorage { + KvSsdStorage::new(KvSsdStorageInit { + root_limits: vec![KvSsdStorageRootLimit { + root_dir: new_root(), + limit_bytes: max_bytes, + }], + uring_mode: KvSsdUringMode::SingleBuffer, + backend: KvSsdStorageBackend::Foyer, + }) + .await + .unwrap() + } + + #[::tokio::test] + async fn foyer_backend_persists_and_loads_from_disk_without_memory_admission() { + let store = new_foyer_store(128 * 1024 * 1024).await; + let data = (0..8193).map(|idx| (idx % 251) as u8).collect::>(); + let put_id = (91, 2); + + store.persist("foyer", put_id, &data).await.unwrap(); + let foyer = store.foyer.as_ref().unwrap(); + assert_eq!(foyer.memory_usage(), 0); + + let mut loaded = vec![0u8; data.len()]; + store + .load_into_addr( + "foyer", + put_id, + loaded.as_mut_ptr() as u64, + data.len() as u64, + loaded.len() as u64, + ) + .await + .unwrap(); + assert_eq!(loaded, data); + assert_eq!(foyer.source_counts(), (0, 1, 0)); + assert_eq!(foyer.memory_usage(), 0); + + let mut chunked = vec![0u8; data.len()]; + let (ready_tx, mut ready_rx) = tokio_mpsc::channel(8); + store + .load_into_addr_chunks( + "foyer", + put_id, + chunked.as_mut_ptr() as u64, + data.len() as u64, + chunked.len() as u64, + 4096, + 4, + ready_tx, + ) + .await + .unwrap(); + let mut chunks = Vec::new(); + while let Some(chunk) = ready_rx.recv().await { + chunks.push((chunk.offset, chunk.len)); + } + assert_eq!(chunked, data); + assert_eq!(chunks, vec![(0, 4096), (4096, 4096), (8192, 1)]); + assert_eq!(foyer.source_counts(), (0, 2, 0)); + assert_eq!(foyer.memory_usage(), 0); + } + + #[::tokio::test] + async fn foyer_backend_backpressures_concurrent_persists_until_durable() { + const ENTRY_COUNT: usize = 24; + const ENTRY_BYTES: usize = 1024 * 1024; + + let store = new_foyer_store(128 * 1024 * 1024).await; + let keys = (0..ENTRY_COUNT) + .map(|idx| format!("foyer-concurrent-{idx}")) + .collect::>(); + let values = (0..ENTRY_COUNT) + .map(|idx| vec![(idx % 251) as u8; ENTRY_BYTES]) + .collect::>(); + let persists = keys + .iter() + .zip(values.iter()) + .enumerate() + .map(|(idx, (key, value))| store.persist(key, (100 + idx as u64, 0), value)); + for result in futures::future::join_all(persists).await { + result.unwrap(); + } + + let foyer = store.foyer.as_ref().unwrap(); + assert_eq!(foyer.memory_usage(), 0); + for (idx, (key, expected)) in keys.iter().zip(values.iter()).enumerate() { + let mut loaded = vec![0u8; expected.len()]; + store + .load_into_addr( + key, + (100 + idx as u64, 0), + loaded.as_mut_ptr() as u64, + loaded.len() as u64, + loaded.len() as u64, + ) + .await + .unwrap(); + assert_eq!(&loaded, expected); + } + assert_eq!(foyer.memory_usage(), 0); + } + + #[::tokio::test] + async fn foyer_backend_rejects_multiple_roots() { + let err = KvSsdStorage::new(KvSsdStorageInit { + root_limits: vec![ + KvSsdStorageRootLimit { + root_dir: new_root(), + limit_bytes: 128 * 1024 * 1024, + }, + KvSsdStorageRootLimit { + root_dir: new_root(), + limit_bytes: 128 * 1024 * 1024, + }, + ], + uring_mode: KvSsdUringMode::SingleBuffer, + backend: KvSsdStorageBackend::Foyer, + }) + .await + .unwrap_err(); + assert!(format!("{err}").contains("requires exactly one root")); + } + + fn test_key(key: &str, version: u64) -> KvSsdKey { + KvSsdKey { + key: key.to_string(), + put_id: (version, 0), + } + } + + fn prepare_ready(ring: &mut SsdRingBuffer, key: &KvSsdKey) -> SsdIndexEntry { + match ring.prepare_write(key.clone(), 500).unwrap() { + SsdPreparedWrite::Ready { entry, .. } => entry, + other => panic!("expected ready SSD write, got {other:?}"), + } + } + + #[test] + fn shard_specs_preserve_per_device_limits() { + let device_roots = vec![ + SsdDeviceRoot { + device_id: 1, + root_dir: PathBuf::from("/tmp/fluxon-test-ssd-a"), + limit_bytes: 4 * SSD_ALIGNMENT as u64, + }, + SsdDeviceRoot { + device_id: 2, + root_dir: PathBuf::from("/tmp/fluxon-test-ssd-b"), + limit_bytes: 8 * SSD_ALIGNMENT as u64, + }, + ]; + + let shard_specs = build_shard_specs(&device_roots).unwrap(); + let mut capacity_by_device = [0u64; 2]; + for spec in shard_specs { + capacity_by_device[spec.device_idx] = + capacity_by_device[spec.device_idx].saturating_add(spec.capacity); + } + + assert_eq!(capacity_by_device[0], 4 * SSD_ALIGNMENT as u64); + assert_eq!(capacity_by_device[1], 8 * SSD_ALIGNMENT as u64); + } + + #[test] + fn write_device_selection_skips_devices_with_undersized_shards() { + let capacities = [1024, 4096]; + + assert_eq!(select_write_device(&capacities, 2048, 0), Some(1)); + assert_eq!(select_write_device(&capacities, 2048, 1), Some(1)); + assert_eq!(select_write_device(&capacities, 8192, 0), None); + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum ReferenceSqeShape { + Read, + Write, + Readv, + Writev, + } + + fn fluxon_current_sqe_shape(io_type: IoType, iovec_count: usize) -> ReferenceSqeShape { + assert!(iovec_count > 0); + match io_type { + IoType::Read => ReferenceSqeShape::Read, + IoType::Write => ReferenceSqeShape::Write, + IoType::Readv => ReferenceSqeShape::Readv, + IoType::Writev => ReferenceSqeShape::Writev, + } + } + + fn ringline_reference_direct_io_sqe_shape( + io_type: IoType, + iovec_count: usize, + ) -> Option { + assert!(iovec_count > 0); + match (io_type, iovec_count) { + (IoType::Readv, 1) => Some(ReferenceSqeShape::Read), + (IoType::Writev, 1) => Some(ReferenceSqeShape::Write), + (_, _) => None, + } + } + + fn ringline_reference_direct_io_allowed(buf_addr: u64, len: u64, offset: u64) -> bool { + len != 0 + && buf_addr.is_multiple_of(SSD_ALIGNMENT as u64) + && len.is_multiple_of(SSD_ALIGNMENT as u64) + && offset.is_multiple_of(SSD_ALIGNMENT as u64) + } + + #[::tokio::test] + async fn persist_and_load_roundtrip() { + let store = new_store(1024 * 1024).await; + let data = b"hello from ssd"; + let put_id = (10, 1); + store.persist("k", put_id, data).await.unwrap(); + + let mut out = vec![0u8; data.len()]; + store + .load_into_addr( + "k", + put_id, + out.as_mut_ptr() as u64, + out.len() as u64, + out.len() as u64, + ) + .await + .unwrap(); + assert_eq!(out, data); + } + + #[::tokio::test] + async fn persist_and_load_roundtrip_supports_iovec_ablation_mode() { + let store = new_store_with_mode(1024 * 1024, KvSsdUringMode::Iovec).await; + let data = b"hello from ssd through iovec"; + let put_id = (10, 2); + store.persist("k-iovec", put_id, data).await.unwrap(); + + let mut out = vec![0u8; data.len()]; + store + .load_into_addr( + "k-iovec", + put_id, + out.as_mut_ptr() as u64, + out.len() as u64, + out.len() as u64, + ) + .await + .unwrap(); + assert_eq!(out, data); + } + + #[::tokio::test] + async fn aligned_load_roundtrip_uses_direct_target() { + let store = new_store(1024 * 1024).await; + let data = (0..4096).map(|idx| (idx % 251) as u8).collect::>(); + let put_id = (11, 1); + store.persist("aligned", put_id, &data).await.unwrap(); + + let mut out = AlignedBuffer::zeroed(data.len()).unwrap(); + let target_addr = out.as_mut_ptr() as u64; + let entry = { + let key = KvSsdKey { + key: "aligned".to_string(), + put_id, + }; + store.inner.lock().ring.get(&key).unwrap() + }; + assert_eq!( + choose_read_path(&entry, target_addr, data.len() as u64, data.len() as u64), + SsdReadPath::Direct + ); + + store + .load_into_addr( + "aligned", + put_id, + target_addr, + data.len() as u64, + data.len() as u64, + ) + .await + .unwrap(); + + let out_slice = unsafe { std::slice::from_raw_parts(out.as_ptr(), data.len()) }; + assert_eq!(out_slice, data.as_slice()); + } + + #[::tokio::test] + async fn chunked_load_roundtrip_streams_ready_chunks() { + let store = new_store(1024 * 1024).await; + let data = (0..2500).map(|idx| (idx % 251) as u8).collect::>(); + let put_id = (13, 1); + store.persist("chunked", put_id, &data).await.unwrap(); + + let mut out = + AlignedBuffer::zeroed(align_ssd_io_len(data.len() as u64).unwrap() as usize).unwrap(); + let target_addr = out.as_mut_ptr() as u64; + let (tx, mut rx) = ::tokio::sync::mpsc::channel(2); + let producer = store.load_into_addr_chunks( + "chunked", + put_id, + target_addr, + data.len() as u64, + out.len() as u64, + 1024, + 2, + tx, + ); + let consumer = async { + let mut chunks = Vec::new(); + while let Some(chunk) = rx.recv().await { + chunks.push((chunk.offset, chunk.len)); + } + chunks + }; + let (producer_res, mut chunks) = ::tokio::join!(producer, consumer); + producer_res.unwrap(); + chunks.sort_unstable(); + assert_eq!(chunks, vec![(0, 1024), (1024, 1024), (2048, 452)]); + + let out_slice = unsafe { std::slice::from_raw_parts(out.as_ptr(), data.len()) }; + assert_eq!(out_slice, data.as_slice()); + } + + #[test] + fn read_path_uses_direct_for_aligned_target_with_enough_capacity() { + let aligned = SsdIndexEntry { + shard_id: 0, + begin: 0, + len: 4096, + aligned_len: 4096, + file_offset: 0, + }; + assert_eq!( + choose_read_path(&aligned, 4096, 4096, 4096), + SsdReadPath::Direct + ); + assert_eq!( + choose_read_path(&aligned, 4097, 4096, 4096), + SsdReadPath::Scratch + ); + + let unaligned_len = SsdIndexEntry { + len: 500, + aligned_len: 512, + ..aligned + }; + assert_eq!( + choose_read_path(&unaligned_len, 4096, 500, 512), + SsdReadPath::Direct + ); + assert_eq!( + choose_read_path(&unaligned_len, 4096, 500, 500), + SsdReadPath::Scratch + ); + } + + #[test] + fn chunk_read_path_matches_ringline_alignment_plus_stage_capacity() { + let cases = [ + (4096, 512, 512, 0, SsdReadPath::Direct), + (4097, 512, 512, 0, SsdReadPath::Scratch), + (4096, 500, 512, 0, SsdReadPath::Scratch), + (4096, 512, 511, 0, SsdReadPath::Scratch), + (4096, 512, 512, 1, SsdReadPath::Scratch), + (4096, 0, 512, 0, SsdReadPath::Scratch), + ]; + + for (target_addr, read_len, target_len, file_offset, expected) in cases { + let ringline_direct = + ringline_reference_direct_io_allowed(target_addr, read_len, file_offset); + assert_eq!( + choose_chunk_read_path(target_addr, read_len, target_len, file_offset), + expected + ); + assert_eq!( + expected == SsdReadPath::Direct, + ringline_direct && target_len >= read_len + ); + } + } + + #[test] + fn full_read_path_adds_payload_contract_to_ringline_alignment() { + let entry = SsdIndexEntry { + shard_id: 0, + begin: 0, + len: 500, + aligned_len: 512, + file_offset: 0, + }; + let target_addr = 4096; + + assert!(ringline_reference_direct_io_allowed( + target_addr, + entry.aligned_len, + entry.file_offset + )); + assert_eq!( + choose_read_path(&entry, target_addr, entry.len, entry.aligned_len), + SsdReadPath::Direct + ); + assert_eq!( + choose_read_path(&entry, target_addr, entry.len - 1, entry.aligned_len), + SsdReadPath::Scratch + ); + assert_eq!( + choose_read_path(&entry, target_addr, entry.len, entry.len), + SsdReadPath::Scratch + ); + } + + #[test] + fn ringline_reference_uses_single_buffer_direct_io_opcodes() { + assert_eq!( + fluxon_current_sqe_shape(IoType::Readv, 1), + ReferenceSqeShape::Readv + ); + assert_eq!( + ringline_reference_direct_io_sqe_shape(IoType::Readv, 1), + Some(ReferenceSqeShape::Read) + ); + assert_eq!( + fluxon_current_sqe_shape(IoType::Writev, 1), + ReferenceSqeShape::Writev + ); + assert_eq!( + ringline_reference_direct_io_sqe_shape(IoType::Writev, 1), + Some(ReferenceSqeShape::Write) + ); + assert_eq!( + ringline_reference_direct_io_sqe_shape(IoType::Readv, 2), + None + ); + } + + #[test] + fn submit_wait_error_action_classifies_transient_errors() { + let interrupted = io::Error::from(io::ErrorKind::Interrupted); + let busy = io::Error::from_raw_os_error(libc::EBUSY); + let invalid = io::Error::from_raw_os_error(libc::EINVAL); + + assert_eq!( + submit_wait_error_action(&interrupted), + SubmitWaitAction::Retry + ); + assert_eq!( + submit_wait_error_action(&busy), + SubmitWaitAction::DrainCompletions + ); + assert_eq!( + submit_wait_error_action(&invalid), + SubmitWaitAction::FailWorker + ); + } + + #[derive(Clone, Copy, Debug)] + enum PerfOpcode { + FluxonReadv, + FluxonWritev, + RinglineRead, + RinglineWrite, + } + + #[derive(Clone, Copy, Debug)] + struct PerfResult { + opcode: PerfOpcode, + ops: usize, + bytes_per_op: usize, + elapsed: std::time::Duration, + } + + impl PerfResult { + fn ns_per_op(&self) -> f64 { + self.elapsed.as_nanos() as f64 / self.ops as f64 + } + + fn mib_per_sec(&self) -> f64 { + let bytes = self.ops as f64 * self.bytes_per_op as f64; + bytes / self.elapsed.as_secs_f64() / 1024.0 / 1024.0 + } + } + + fn best_perf_result(results: &[PerfResult]) -> PerfResult { + results + .iter() + .min_by_key(|result| result.elapsed) + .copied() + .unwrap() + } + + fn median_perf_result(results: &[PerfResult]) -> PerfResult { + let mut sorted = results.to_vec(); + sorted.sort_by_key(|result| result.elapsed); + sorted[sorted.len() / 2] + } + + fn print_perf_pair_result( + stat: &str, + label: &str, + fluxon: PerfResult, + ringline: PerfResult, + rounds: usize, + ) { + println!( + "uring opcode perf {stat}-of-{rounds} {label}: {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + fluxon.opcode, + fluxon.ops, + fluxon.bytes_per_op, + fluxon.elapsed, + fluxon.ns_per_op(), + fluxon.mib_per_sec() + ); + println!( + "uring opcode perf {stat}-of-{rounds} {label}: {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + ringline.opcode, + ringline.ops, + ringline.bytes_per_op, + ringline.elapsed, + ringline.ns_per_op(), + ringline.mib_per_sec() + ); + println!( + "uring opcode perf delta {stat}-of-{rounds} {label}: ringline-style relative to Fluxon = {:.2}%", + (fluxon.ns_per_op() - ringline.ns_per_op()) / fluxon.ns_per_op() * 100.0 + ); + } + + fn print_perf_pair( + label: &str, + fluxon_results: &[PerfResult], + ringline_results: &[PerfResult], + ) { + assert_eq!(fluxon_results.len(), ringline_results.len()); + let rounds = fluxon_results.len(); + print_perf_pair_result( + "best", + label, + best_perf_result(fluxon_results), + best_perf_result(ringline_results), + rounds, + ); + print_perf_pair_result( + "median", + label, + median_perf_result(fluxon_results), + median_perf_result(ringline_results), + rounds, + ); + } + + fn run_uring_io_perf( + opcode: PerfOpcode, + fd: RawFd, + buffer: &mut AlignedBuffer, + ops: usize, + bytes_per_op: usize, + offset_slots: usize, + ) -> io::Result { + assert!(offset_slots > 0); + let mut ring = IoUring::builder().build(64)?; + let start = std::time::Instant::now(); + for idx in 0..ops { + let offset = u64::try_from((idx % offset_slots) * bytes_per_op).unwrap(); + let iovec = libc::iovec { + iov_base: buffer.as_mut_ptr() as *mut libc::c_void, + iov_len: bytes_per_op, + }; + let sqe = match opcode { + PerfOpcode::FluxonReadv => { + opcode::Readv::new(Fd(fd), &iovec, 1).offset(offset).build() + } + PerfOpcode::FluxonWritev => opcode::Writev::new(Fd(fd), &iovec, 1) + .offset(offset) + .build(), + PerfOpcode::RinglineRead => { + opcode::Read::new(Fd(fd), buffer.as_mut_ptr(), bytes_per_op as _) + .offset(offset) + .build() + } + PerfOpcode::RinglineWrite => { + opcode::Write::new(Fd(fd), buffer.as_ptr(), bytes_per_op as _) + .offset(offset) + .build() + } + } + .user_data((idx + 1) as u64); + unsafe { + ring.submission() + .push(&sqe) + .map_err(|_| io::Error::other("submission queue full"))?; + } + ring.submit_and_wait(1)?; + let mut cq = ring.completion(); + let cqe: io_uring::cqueue::Entry = cq + .next() + .ok_or_else(|| io::Error::other("missing completion"))?; + if cqe.result() != bytes_per_op as i32 { + return Err(io::Error::other(format!( + "short uring perf completion: {} != {}", + cqe.result(), + bytes_per_op + ))); + } + } + Ok(PerfResult { + opcode, + ops, + bytes_per_op, + elapsed: start.elapsed(), + }) + } + + #[test] + #[ignore = "manual perf comparison for Fluxon Readv/Writev vs ringline-style Read/Write"] + fn perf_compare_fluxon_iovecs_with_ringline_single_buffer_ops() { + let dir = std::env::temp_dir().join(format!("fluxon-kv-ssd-perf-{}", std::process::id())); + fs::create_dir_all(&dir).unwrap(); + for (bytes_per_op, ops, offset_slots, rounds) in [ + (4096usize, 256usize, 128usize, 3usize), + (1024 * 1024usize, 12usize, 12usize, 3usize), + (10 * 1024 * 1024usize, 3usize, 3usize, 3usize), + ] { + assert!(bytes_per_op.is_multiple_of(SSD_ALIGNMENT)); + let file_len = bytes_per_op.checked_mul(offset_slots).unwrap(); + let path = dir.join(format!("direct-{bytes_per_op}.dat")); + let file = OpenOptions::new() + .create(true) + .truncate(true) + .read(true) + .write(true) + .custom_flags(libc::O_DIRECT) + .open(&path) + .unwrap(); + file.set_len(file_len as u64).unwrap(); + + let mut seed = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + unsafe { + std::ptr::write_bytes(seed.as_mut_ptr(), 0x5a, bytes_per_op); + } + for slot in 0..offset_slots { + let offset = u64::try_from(slot * bytes_per_op).unwrap(); + let written = unsafe { + libc::pwrite( + file.as_raw_fd(), + seed.as_ptr() as *const _, + bytes_per_op, + offset as libc::off_t, + ) + }; + assert_eq!(written, bytes_per_op as isize); + } + + let mut readv_buffer = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut read_buffer = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut writev_buffer = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut write_buffer = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + unsafe { + std::ptr::write_bytes(writev_buffer.as_mut_ptr(), 0xa5, bytes_per_op); + std::ptr::write_bytes(write_buffer.as_mut_ptr(), 0x3c, bytes_per_op); + } + + let warmup_ops = ops.min(16); + let _ = run_uring_io_perf( + PerfOpcode::FluxonReadv, + file.as_raw_fd(), + &mut readv_buffer, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_uring_io_perf( + PerfOpcode::RinglineRead, + file.as_raw_fd(), + &mut read_buffer, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_uring_io_perf( + PerfOpcode::FluxonWritev, + file.as_raw_fd(), + &mut writev_buffer, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_uring_io_perf( + PerfOpcode::RinglineWrite, + file.as_raw_fd(), + &mut write_buffer, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + + let mut readv_results = Vec::with_capacity(rounds); + let mut read_results = Vec::with_capacity(rounds); + let mut writev_results = Vec::with_capacity(rounds); + let mut write_results = Vec::with_capacity(rounds); + for round_idx in 0..rounds { + if round_idx % 2 == 0 { + readv_results.push( + run_uring_io_perf( + PerfOpcode::FluxonReadv, + file.as_raw_fd(), + &mut readv_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + read_results.push( + run_uring_io_perf( + PerfOpcode::RinglineRead, + file.as_raw_fd(), + &mut read_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } else { + read_results.push( + run_uring_io_perf( + PerfOpcode::RinglineRead, + file.as_raw_fd(), + &mut read_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + readv_results.push( + run_uring_io_perf( + PerfOpcode::FluxonReadv, + file.as_raw_fd(), + &mut readv_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } + + if round_idx % 2 == 0 { + writev_results.push( + run_uring_io_perf( + PerfOpcode::FluxonWritev, + file.as_raw_fd(), + &mut writev_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + write_results.push( + run_uring_io_perf( + PerfOpcode::RinglineWrite, + file.as_raw_fd(), + &mut write_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } else { + write_results.push( + run_uring_io_perf( + PerfOpcode::RinglineWrite, + file.as_raw_fd(), + &mut write_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + writev_results.push( + run_uring_io_perf( + PerfOpcode::FluxonWritev, + file.as_raw_fd(), + &mut writev_buffer, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } + } + + print_perf_pair("read", &readv_results, &read_results); + print_perf_pair("write", &writev_results, &write_results); + + fs::remove_file(&path).ok(); + } + fs::remove_dir(&dir).ok(); + } + + #[derive(Clone, Copy, Debug)] + enum StoragePerfOp { + Persist, + Load, + } + + #[derive(Clone, Copy, Debug)] + struct StoragePerfResult { + mode: KvSsdUringMode, + op: StoragePerfOp, + ops: usize, + bytes_per_op: usize, + elapsed: std::time::Duration, + } + + impl StoragePerfResult { + fn ns_per_op(&self) -> f64 { + self.elapsed.as_nanos() as f64 / self.ops as f64 + } + + fn mib_per_sec(&self) -> f64 { + let bytes = self.ops as f64 * self.bytes_per_op as f64; + bytes / self.elapsed.as_secs_f64() / 1024.0 / 1024.0 + } + } + + fn storage_perf_capacity(bytes_per_op: usize, ops: usize) -> u64 { + let bytes = bytes_per_op + .checked_mul(ops) + .and_then(|value| value.checked_mul(2)) + .unwrap(); + align_up_u64(bytes as u64, SSD_ALIGNMENT as u64).unwrap() + } + + fn storage_perf_data(bytes_per_op: usize) -> Vec { + (0..bytes_per_op) + .map(|idx| ((idx * 31 + idx / 251) % 251) as u8) + .collect() + } + + async fn run_storage_persist_perf( + mode: KvSsdUringMode, + bytes_per_op: usize, + ops: usize, + ) -> StoragePerfResult { + let store = new_store_with_mode(storage_perf_capacity(bytes_per_op, ops), mode).await; + let data = storage_perf_data(bytes_per_op); + let start = std::time::Instant::now(); + for idx in 0..ops { + store + .persist( + &format!("persist-{mode:?}-{idx}"), + (idx as u64, 0), + data.as_slice(), + ) + .await + .unwrap(); + } + StoragePerfResult { + mode, + op: StoragePerfOp::Persist, + ops, + bytes_per_op, + elapsed: start.elapsed(), + } + } + + async fn storage_with_seeded_values( + mode: KvSsdUringMode, + bytes_per_op: usize, + ops: usize, + data: &[u8], + ) -> KvSsdStorage { + let store = new_store_with_mode(storage_perf_capacity(bytes_per_op, ops), mode).await; + for idx in 0..ops { + store + .persist(&format!("load-{mode:?}-{idx}"), (idx as u64, 0), data) + .await + .unwrap(); + } + store + } + + async fn run_storage_load_perf( + mode: KvSsdUringMode, + bytes_per_op: usize, + ops: usize, + ) -> StoragePerfResult { + let data = storage_perf_data(bytes_per_op); + let store = storage_with_seeded_values(mode, bytes_per_op, ops, data.as_slice()).await; + let mut out = AlignedBuffer::zeroed(bytes_per_op).unwrap(); + let start = std::time::Instant::now(); + for idx in 0..ops { + store + .load_into_addr( + &format!("load-{mode:?}-{idx}"), + (idx as u64, 0), + out.as_mut_ptr() as u64, + bytes_per_op as u64, + out.len() as u64, + ) + .await + .unwrap(); + } + let elapsed = start.elapsed(); + let out_slice = unsafe { std::slice::from_raw_parts(out.as_ptr(), data.len()) }; + assert_eq!(out_slice, data.as_slice()); + StoragePerfResult { + mode, + op: StoragePerfOp::Load, + ops, + bytes_per_op, + elapsed, + } + } + + fn best_storage_perf_result(results: &[StoragePerfResult]) -> StoragePerfResult { + results + .iter() + .min_by_key(|result| result.elapsed) + .copied() + .unwrap() + } + + fn median_storage_perf_result(results: &[StoragePerfResult]) -> StoragePerfResult { + let mut sorted = results.to_vec(); + sorted.sort_by_key(|result| result.elapsed); + sorted[sorted.len() / 2] + } + + fn print_storage_perf_pair_result( + stat: &str, + label: &str, + iovec: StoragePerfResult, + single_buffer: StoragePerfResult, + rounds: usize, + ) { + println!( + "kv ssd storage perf {stat}-of-{rounds} {label}: {:?} {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + iovec.mode, + iovec.op, + iovec.ops, + iovec.bytes_per_op, + iovec.elapsed, + iovec.ns_per_op(), + iovec.mib_per_sec() + ); + println!( + "kv ssd storage perf {stat}-of-{rounds} {label}: {:?} {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + single_buffer.mode, + single_buffer.op, + single_buffer.ops, + single_buffer.bytes_per_op, + single_buffer.elapsed, + single_buffer.ns_per_op(), + single_buffer.mib_per_sec() + ); + println!( + "kv ssd storage perf delta {stat}-of-{rounds} {label}: single_buffer relative to iovec = {:.2}%", + (iovec.ns_per_op() - single_buffer.ns_per_op()) / iovec.ns_per_op() * 100.0 + ); + } + + fn print_storage_perf_pair( + label: &str, + iovec_results: &[StoragePerfResult], + single_buffer_results: &[StoragePerfResult], + ) { + assert_eq!(iovec_results.len(), single_buffer_results.len()); + let rounds = iovec_results.len(); + print_storage_perf_pair_result( + "best", + label, + best_storage_perf_result(iovec_results), + best_storage_perf_result(single_buffer_results), + rounds, + ); + print_storage_perf_pair_result( + "median", + label, + median_storage_perf_result(iovec_results), + median_storage_perf_result(single_buffer_results), + rounds, + ); + } + + #[::tokio::test] + #[ignore = "manual KvSsdStorage-level perf comparison for iovec vs single-buffer uring mode"] + async fn perf_compare_kv_ssd_storage_iovec_with_single_buffer_mode() { + for (bytes_per_op, ops, rounds) in [ + (1024 * 1024usize, 4usize, 3usize), + (10 * 1024 * 1024usize, 2usize, 3usize), + ] { + assert!(bytes_per_op.is_multiple_of(SSD_ALIGNMENT)); + + let mut iovec_persist_results = Vec::with_capacity(rounds); + let mut single_persist_results = Vec::with_capacity(rounds); + let mut iovec_load_results = Vec::with_capacity(rounds); + let mut single_load_results = Vec::with_capacity(rounds); + + for round_idx in 0..rounds { + if round_idx % 2 == 0 { + iovec_persist_results.push( + run_storage_persist_perf(KvSsdUringMode::Iovec, bytes_per_op, ops).await, + ); + single_persist_results.push( + run_storage_persist_perf(KvSsdUringMode::SingleBuffer, bytes_per_op, ops) + .await, + ); + iovec_load_results.push( + run_storage_load_perf(KvSsdUringMode::Iovec, bytes_per_op, ops).await, + ); + single_load_results.push( + run_storage_load_perf(KvSsdUringMode::SingleBuffer, bytes_per_op, ops) + .await, + ); + } else { + single_persist_results.push( + run_storage_persist_perf(KvSsdUringMode::SingleBuffer, bytes_per_op, ops) + .await, + ); + iovec_persist_results.push( + run_storage_persist_perf(KvSsdUringMode::Iovec, bytes_per_op, ops).await, + ); + single_load_results.push( + run_storage_load_perf(KvSsdUringMode::SingleBuffer, bytes_per_op, ops) + .await, + ); + iovec_load_results.push( + run_storage_load_perf(KvSsdUringMode::Iovec, bytes_per_op, ops).await, + ); + } + } + + let label = format!("{} bytes/op persist", bytes_per_op); + print_storage_perf_pair(&label, &iovec_persist_results, &single_persist_results); + let label = format!("{} bytes/op load", bytes_per_op); + print_storage_perf_pair(&label, &iovec_load_results, &single_load_results); + } + } + + #[::tokio::test] + async fn unaligned_payload_loads_direct_when_stage_capacity_is_aligned() { + let store = new_store(1024 * 1024).await; + let data = (0..500).map(|idx| (idx % 251) as u8).collect::>(); + let put_id = (12, 1); + store.persist("unaligned", put_id, &data).await.unwrap(); + + let mut out = AlignedBuffer::zeroed(SSD_ALIGNMENT).unwrap(); + let target_addr = out.as_mut_ptr() as u64; + let entry = { + let key = KvSsdKey { + key: "unaligned".to_string(), + put_id, + }; + store.inner.lock().ring.get(&key).unwrap() + }; + assert_eq!(entry.len, data.len() as u64); + assert_eq!(entry.aligned_len, SSD_ALIGNMENT as u64); + assert_eq!( + choose_read_path(&entry, target_addr, data.len() as u64, SSD_ALIGNMENT as u64), + SsdReadPath::Direct + ); + + store + .load_into_addr( + "unaligned", + put_id, + target_addr, + data.len() as u64, + SSD_ALIGNMENT as u64, + ) + .await + .unwrap(); + + let out_slice = unsafe { std::slice::from_raw_parts(out.as_ptr(), data.len()) }; + assert_eq!(out_slice, data.as_slice()); + } + + #[::tokio::test] + async fn storage_deduplicates_root_dirs_on_same_device() { + let root_a = new_root(); + let root_b = new_root(); + let store = KvSsdStorage::new(KvSsdStorageInit { + root_limits: vec![ + KvSsdStorageRootLimit { + root_dir: root_a.clone(), + limit_bytes: 4 * SSD_ALIGNMENT as u64, + }, + KvSsdStorageRootLimit { + root_dir: root_b.clone(), + limit_bytes: 8 * SSD_ALIGNMENT as u64, + }, + ], + uring_mode: KvSsdUringMode::SingleBuffer, + backend: KvSsdStorageBackend::Native, + }) + .await + .unwrap(); + + assert_eq!( + fs::metadata(&root_a).unwrap().dev(), + fs::metadata(&root_b).unwrap().dev() + ); + assert_eq!(store.root_dirs(), &[root_a.clone()]); + assert_eq!(store.devices.len(), 1); + assert_eq!(store.shard_to_device, vec![0, 0, 0, 0]); + assert!(root_a.join("shards/shard-000000.dat").exists()); + assert!(root_a.join("shards/shard-000001.dat").exists()); + assert!(root_a.join("shards/shard-000002.dat").exists()); + assert!(root_a.join("shards/shard-000003.dat").exists()); + assert!(!root_b.join("shards").exists()); + } + + #[test] + fn ring_prepare_write_on_shards_uses_only_allowed_shards() { + let mut ring = SsdRingBuffer::new(vec![1024, 1024, 1024, 1024]); + let mut allocated_shards = Vec::new(); + + for version in 0..4 { + let key = test_key("per-device", version); + let entry = match ring + .prepare_write_on_shards(key.clone(), 500, &[1, 3]) + .unwrap() + { + SsdPreparedWrite::Ready { entry, .. } => entry, + other => panic!("expected ready SSD write, got {other:?}"), + }; + allocated_shards.push(entry.shard_id); + assert!(ring.commit(&key, true)); + } + + assert_eq!(allocated_shards, vec![1, 3, 1, 3]); + } + + #[::tokio::test] + async fn ring_keeps_new_entry_and_expires_old() { + let store = new_store(1024).await; + let mut eviction_rx = store + .take_eviction_rx() + .expect("SSD eviction receiver must be available once"); + store.persist("old", (1, 0), &[1u8; 500]).await.unwrap(); + store.persist("filler", (2, 0), &[2u8; 500]).await.unwrap(); + store.persist("new", (3, 0), &[3u8; 500]).await.unwrap(); + + assert!(!store.has_entry("old", (1, 0)).await); + assert!(store.has_entry("filler", (2, 0)).await); + assert!(store.has_entry("new", (3, 0)).await); + let evicted = ::tokio::time::timeout(Duration::from_secs(1), eviction_rx.recv()) + .await + .expect("SSD eviction notification must arrive") + .expect("SSD eviction channel must stay open"); + assert_eq!( + evicted, + vec![SsdReplicaEviction { + key: "old".to_string(), + put_id: (1, 0), + }] + ); + } + + #[::tokio::test] + async fn persist_guard_blocks_eviction_until_route_commit_finishes() { + let store = Arc::new(new_store(512).await); + let mut eviction_rx = store + .take_eviction_rx() + .expect("SSD eviction receiver must be available once"); + let old_data = vec![1u8; 500]; + let guard = store + .persist_from_addr( + "old", + (1, 0), + old_data.as_ptr() as u64, + old_data.len() as u64, + ) + .await + .unwrap(); + + let store_for_write = Arc::clone(&store); + let mut new_write = + ::tokio::spawn( + async move { store_for_write.persist("new", (2, 0), &[2u8; 500]).await }, + ); + assert!( + ::tokio::time::timeout(Duration::from_millis(50), &mut new_write) + .await + .is_err(), + "new write must wait while the old route commit holds an entry pin" + ); + + drop(guard); + ::tokio::time::timeout(Duration::from_secs(1), &mut new_write) + .await + .expect("new write must resume after the route commit pin is dropped") + .expect("new write task must complete") + .expect("new write must succeed"); + let evicted = ::tokio::time::timeout(Duration::from_secs(1), eviction_rx.recv()) + .await + .expect("SSD eviction notification must arrive") + .expect("SSD eviction channel must stay open"); + assert_eq!(evicted[0].key, "old"); + assert_eq!(evicted[0].put_id, (1, 0)); + } + + #[test] + fn ring_read_pin_blocks_overwrite_until_unpinned() { + let mut ring = SsdRingBuffer::new(vec![1024]); + let old = test_key("old", 1); + let filler = test_key("filler", 2); + let new = test_key("new", 3); + + let old_entry = prepare_ready(&mut ring, &old); + assert_eq!(old_entry.begin, 0); + assert!(ring.commit(&old, true)); + prepare_ready(&mut ring, &filler); + assert!(ring.commit(&filler, true)); + + let pinned = ring.pin_read(&old).unwrap(); + assert_eq!(pinned.begin, old_entry.begin); + assert!(matches!( + ring.prepare_write(new.clone(), 500).unwrap(), + SsdPreparedWrite::BlockedByBusyIo + )); + assert!(ring.get(&old).is_some()); + + ring.unpin_read(&old); + let new_entry = prepare_ready(&mut ring, &new); + assert_eq!(new_entry.file_offset, 0); + assert!(ring.commit(&new, true)); + assert!(ring.get(&old).is_none()); + } + + #[test] + fn ring_writing_entry_blocks_overwrite_until_write_finishes() { + let mut ring = SsdRingBuffer::new(vec![1024]); + let old = test_key("old", 1); + let filler = test_key("filler", 2); + let new = test_key("new", 3); + + let old_entry = prepare_ready(&mut ring, &old); + assert_eq!(old_entry.begin, 0); + prepare_ready(&mut ring, &filler); + + assert!(matches!( + ring.prepare_write(new.clone(), 500).unwrap(), + SsdPreparedWrite::BlockedByBusyIo + )); + + assert!(ring.commit(&old, true)); + let new_entry = prepare_ready(&mut ring, &new); + assert_eq!(new_entry.file_offset, 0); + } + + #[test] + fn safe_component_replaces_path_separators() { + assert_eq!(safe_path_component("owner/a:b"), "owner_a_b"); + } +} diff --git a/fluxon_rs/fluxon_kv/src/kv_ssd_storage2.rs b/fluxon_rs/fluxon_kv/src/kv_ssd_storage2.rs new file mode 100644 index 0000000..b0bf792 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/kv_ssd_storage2.rs @@ -0,0 +1,441 @@ +use crate::kv_ssd_storage::SSD_ALIGNMENT; +use io_uring::{IoUring, opcode, types::Fd}; +use std::fs::{self, OpenOptions}; +use std::io; +use std::os::fd::{AsRawFd, RawFd}; +use std::os::unix::fs::OpenOptionsExt; +use std::ptr::NonNull; + +#[derive(Clone, Copy, Debug)] +enum DirectOpcode { + BaselineReadv, + BaselineWritev, + OptimizedRead, + OptimizedWrite, +} + +#[derive(Clone, Copy, Debug)] +struct DirectIoPerfResult { + opcode: DirectOpcode, + ops: usize, + bytes_per_op: usize, + elapsed: std::time::Duration, +} + +impl DirectIoPerfResult { + fn ns_per_op(&self) -> f64 { + self.elapsed.as_nanos() as f64 / self.ops as f64 + } + + fn mib_per_sec(&self) -> f64 { + let bytes = self.ops as f64 * self.bytes_per_op as f64; + bytes / self.elapsed.as_secs_f64() / 1024.0 / 1024.0 + } +} + +struct DirectAlignedBuffer { + ptr: NonNull, + len: usize, +} + +impl DirectAlignedBuffer { + fn zeroed(len: usize) -> io::Result { + if len == 0 || !len.is_multiple_of(SSD_ALIGNMENT) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "direct buffer len must be positive and {}-byte aligned: {}", + SSD_ALIGNMENT, len + ), + )); + } + let mut raw = std::ptr::null_mut(); + let rc = unsafe { libc::posix_memalign(&mut raw, SSD_ALIGNMENT, len) }; + if rc != 0 || raw.is_null() { + return Err(io::Error::other(format!("posix_memalign failed: rc={rc}"))); + } + unsafe { + std::ptr::write_bytes(raw as *mut u8, 0, len); + } + Ok(Self { + ptr: NonNull::new(raw as *mut u8).expect("posix_memalign returned non-null"), + len, + }) + } + + fn as_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.ptr.as_ptr() + } + + fn len(&self) -> usize { + self.len + } +} + +impl Drop for DirectAlignedBuffer { + fn drop(&mut self) { + unsafe { + libc::free(self.ptr.as_ptr() as *mut libc::c_void); + } + } +} + +fn baseline_current_sqe( + opcode: DirectOpcode, + fd: RawFd, + iovec: &libc::iovec, + offset: u64, +) -> io_uring::squeue::Entry { + match opcode { + DirectOpcode::BaselineReadv => opcode::Readv::new(Fd(fd), iovec, 1).offset(offset).build(), + DirectOpcode::BaselineWritev => { + opcode::Writev::new(Fd(fd), iovec, 1).offset(offset).build() + } + DirectOpcode::OptimizedRead | DirectOpcode::OptimizedWrite => { + unreachable!("optimized opcodes do not use iovec SQEs") + } + } +} + +fn optimized_single_buffer_sqe( + opcode: DirectOpcode, + fd: RawFd, + buffer: &mut DirectAlignedBuffer, + offset: u64, +) -> io_uring::squeue::Entry { + match opcode { + DirectOpcode::OptimizedRead => { + opcode::Read::new(Fd(fd), buffer.as_mut_ptr(), buffer.len() as _) + .offset(offset) + .build() + } + DirectOpcode::OptimizedWrite => { + opcode::Write::new(Fd(fd), buffer.as_ptr(), buffer.len() as _) + .offset(offset) + .build() + } + DirectOpcode::BaselineReadv | DirectOpcode::BaselineWritev => { + unreachable!("baseline opcodes use iovec SQEs") + } + } +} + +fn run_direct_io_perf( + opcode: DirectOpcode, + fd: RawFd, + buffer: &mut DirectAlignedBuffer, + ops: usize, + bytes_per_op: usize, + offset_slots: usize, +) -> io::Result { + assert!(offset_slots > 0); + assert_eq!(buffer.len(), bytes_per_op); + let mut ring = IoUring::builder().build(64)?; + let start = std::time::Instant::now(); + for idx in 0..ops { + let offset = u64::try_from((idx % offset_slots) * bytes_per_op).unwrap(); + let iovec = libc::iovec { + iov_base: buffer.as_mut_ptr() as *mut libc::c_void, + iov_len: bytes_per_op, + }; + let sqe = match opcode { + DirectOpcode::BaselineReadv | DirectOpcode::BaselineWritev => { + baseline_current_sqe(opcode, fd, &iovec, offset) + } + DirectOpcode::OptimizedRead | DirectOpcode::OptimizedWrite => { + optimized_single_buffer_sqe(opcode, fd, buffer, offset) + } + } + .user_data((idx + 1) as u64); + unsafe { + ring.submission() + .push(&sqe) + .map_err(|_| io::Error::other("submission queue full"))?; + } + ring.submit_and_wait(1)?; + let mut cq = ring.completion(); + let cqe: io_uring::cqueue::Entry = cq + .next() + .ok_or_else(|| io::Error::other("missing completion"))?; + if cqe.result() != bytes_per_op as i32 { + return Err(io::Error::other(format!( + "short direct I/O completion: {} != {}", + cqe.result(), + bytes_per_op + ))); + } + } + Ok(DirectIoPerfResult { + opcode, + ops, + bytes_per_op, + elapsed: start.elapsed(), + }) +} + +fn best_result(results: &[DirectIoPerfResult]) -> DirectIoPerfResult { + results + .iter() + .min_by_key(|result| result.elapsed) + .copied() + .unwrap() +} + +fn median_result(results: &[DirectIoPerfResult]) -> DirectIoPerfResult { + let mut sorted = results.to_vec(); + sorted.sort_by_key(|result| result.elapsed); + sorted[sorted.len() / 2] +} + +fn print_pair( + label: &str, + baseline_results: &[DirectIoPerfResult], + optimized_results: &[DirectIoPerfResult], +) { + assert_eq!(baseline_results.len(), optimized_results.len()); + let rounds = baseline_results.len(); + for (stat, baseline, optimized) in [ + ( + "best", + best_result(baseline_results), + best_result(optimized_results), + ), + ( + "median", + median_result(baseline_results), + median_result(optimized_results), + ), + ] { + println!( + "kv_ssd_storage2 {stat}-of-{rounds} {label} baseline: {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + baseline.opcode, + baseline.ops, + baseline.bytes_per_op, + baseline.elapsed, + baseline.ns_per_op(), + baseline.mib_per_sec() + ); + println!( + "kv_ssd_storage2 {stat}-of-{rounds} {label} optimized: {:?}: ops={} bytes/op={} elapsed={:?} ns/op={:.1} MiB/s={:.1}", + optimized.opcode, + optimized.ops, + optimized.bytes_per_op, + optimized.elapsed, + optimized.ns_per_op(), + optimized.mib_per_sec() + ); + println!( + "kv_ssd_storage2 delta {stat}-of-{rounds} {label}: optimized relative to baseline = {:.2}%", + (baseline.ns_per_op() - optimized.ns_per_op()) / baseline.ns_per_op() * 100.0 + ); + } +} + +fn seed_direct_file(fd: RawFd, bytes_per_op: usize, offset_slots: usize) -> io::Result<()> { + let mut seed = DirectAlignedBuffer::zeroed(bytes_per_op)?; + unsafe { + std::ptr::write_bytes(seed.as_mut_ptr(), 0x5a, seed.len()); + } + for slot in 0..offset_slots { + let offset = u64::try_from(slot * bytes_per_op).unwrap(); + let written = unsafe { + libc::pwrite( + fd, + seed.as_ptr() as *const _, + bytes_per_op, + offset as libc::off_t, + ) + }; + if written != bytes_per_op as isize { + return Err(io::Error::other(format!( + "short direct seed write: {written} != {bytes_per_op}" + ))); + } + } + Ok(()) +} + +#[test] +#[ignore = "manual perf comparison for kv_ssd_storage2 single-buffer direct I/O fast path"] +fn perf_compare_current_iovec_with_storage2_single_buffer_fast_path() { + let dir = std::env::temp_dir().join(format!("fluxon-kv-ssd-storage2-{}", std::process::id())); + fs::create_dir_all(&dir).unwrap(); + for (bytes_per_op, ops, offset_slots, rounds) in [ + (4096usize, 256usize, 128usize, 3usize), + (1024 * 1024usize, 12usize, 12usize, 3usize), + (10 * 1024 * 1024usize, 3usize, 3usize, 3usize), + ] { + assert!(bytes_per_op.is_multiple_of(SSD_ALIGNMENT)); + let file_len = bytes_per_op.checked_mul(offset_slots).unwrap(); + let path = dir.join(format!("direct-{bytes_per_op}.dat")); + let file = OpenOptions::new() + .create(true) + .truncate(true) + .read(true) + .write(true) + .custom_flags(libc::O_DIRECT) + .open(&path) + .unwrap(); + file.set_len(file_len as u64).unwrap(); + seed_direct_file(file.as_raw_fd(), bytes_per_op, offset_slots).unwrap(); + + let mut baseline_read = DirectAlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut optimized_read = DirectAlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut baseline_write = DirectAlignedBuffer::zeroed(bytes_per_op).unwrap(); + let mut optimized_write = DirectAlignedBuffer::zeroed(bytes_per_op).unwrap(); + unsafe { + std::ptr::write_bytes(baseline_write.as_mut_ptr(), 0xa5, baseline_write.len()); + std::ptr::write_bytes(optimized_write.as_mut_ptr(), 0x3c, optimized_write.len()); + } + + let warmup_ops = ops.min(16); + let _ = run_direct_io_perf( + DirectOpcode::BaselineReadv, + file.as_raw_fd(), + &mut baseline_read, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_direct_io_perf( + DirectOpcode::OptimizedRead, + file.as_raw_fd(), + &mut optimized_read, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_direct_io_perf( + DirectOpcode::BaselineWritev, + file.as_raw_fd(), + &mut baseline_write, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + let _ = run_direct_io_perf( + DirectOpcode::OptimizedWrite, + file.as_raw_fd(), + &mut optimized_write, + warmup_ops, + bytes_per_op, + offset_slots, + ) + .unwrap(); + + let mut baseline_read_results = Vec::with_capacity(rounds); + let mut optimized_read_results = Vec::with_capacity(rounds); + let mut baseline_write_results = Vec::with_capacity(rounds); + let mut optimized_write_results = Vec::with_capacity(rounds); + for round_idx in 0..rounds { + if round_idx % 2 == 0 { + baseline_read_results.push( + run_direct_io_perf( + DirectOpcode::BaselineReadv, + file.as_raw_fd(), + &mut baseline_read, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + optimized_read_results.push( + run_direct_io_perf( + DirectOpcode::OptimizedRead, + file.as_raw_fd(), + &mut optimized_read, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } else { + optimized_read_results.push( + run_direct_io_perf( + DirectOpcode::OptimizedRead, + file.as_raw_fd(), + &mut optimized_read, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + baseline_read_results.push( + run_direct_io_perf( + DirectOpcode::BaselineReadv, + file.as_raw_fd(), + &mut baseline_read, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } + + if round_idx % 2 == 0 { + baseline_write_results.push( + run_direct_io_perf( + DirectOpcode::BaselineWritev, + file.as_raw_fd(), + &mut baseline_write, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + optimized_write_results.push( + run_direct_io_perf( + DirectOpcode::OptimizedWrite, + file.as_raw_fd(), + &mut optimized_write, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } else { + optimized_write_results.push( + run_direct_io_perf( + DirectOpcode::OptimizedWrite, + file.as_raw_fd(), + &mut optimized_write, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + baseline_write_results.push( + run_direct_io_perf( + DirectOpcode::BaselineWritev, + file.as_raw_fd(), + &mut baseline_write, + ops, + bytes_per_op, + offset_slots, + ) + .unwrap(), + ); + } + } + + print_pair("read", &baseline_read_results, &optimized_read_results); + print_pair("write", &baseline_write_results, &optimized_write_results); + + fs::remove_file(&path).ok(); + } + fs::remove_dir(&dir).ok(); +} diff --git a/fluxon_rs/fluxon_kv/src/kv_ssd_storage_foyer.rs b/fluxon_rs/fluxon_kv/src/kv_ssd_storage_foyer.rs new file mode 100644 index 0000000..6ffb155 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/kv_ssd_storage_foyer.rs @@ -0,0 +1,421 @@ +use crate::kv_ssd_storage::{KvSsdStorageRootLimit, SsdLoadedChunk, align_ssd_io_len}; +use crate::master_kv_router::put::PutIDForAKey; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}; +use ::tokio::sync::{Mutex as TokioMutex, mpsc as tokio_mpsc}; +use foyer::{ + BlockEngineConfig, DeviceBuilder, FsDeviceBuilder, HybridCache, HybridCacheBuilder, + HybridCacheEntry, HybridCachePolicy, PsyncIoEngineConfig, Source, +}; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; + +const FOYER_MEMORY_CAPACITY_BYTES: usize = 1; +const FOYER_MEMORY_SHARDS: usize = 1; +const FOYER_BLOCK_SIZE_BYTES: usize = 64 * 1024 * 1024; + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +struct FoyerKvSsdKey { + key: String, + put_id: PutIDForAKey, +} + +type FoyerEntry = HybridCacheEntry>; + +pub(crate) struct FoyerKvSsdPersistGuard { + _entry: FoyerEntry, +} + +#[derive(Debug)] +pub(crate) struct FoyerKvSsdStorage { + cache: HybridCache>, + root_dir: PathBuf, + capacity_bytes: u64, + // Serialize forced storage writes because Foyer drops submissions when its byte queue is full. + persist_gate: TokioMutex<()>, + entry_lengths: Mutex>, + logical_used_bytes: AtomicU64, + memory_hits: AtomicU64, + disk_hits: AtomicU64, + outer_hits: AtomicU64, +} + +impl FoyerKvSsdStorage { + pub(crate) async fn new(root_limit: KvSsdStorageRootLimit) -> KvResult { + if root_limit.limit_bytes < FOYER_BLOCK_SIZE_BYTES as u64 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd capacity must be at least {} bytes, got {}", + FOYER_BLOCK_SIZE_BYTES, root_limit.limit_bytes + ), + })); + } + let capacity_bytes = usize::try_from(root_limit.limit_bytes).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd capacity does not fit usize: {}", + root_limit.limit_bytes + ), + }) + })?; + + fs::create_dir_all(&root_limit.root_dir) + .map_err(|err| foyer_file_error(&root_limit.root_dir, err))?; + let foyer_root = root_limit.root_dir.join("foyer"); + match fs::remove_dir_all(&foyer_root) { + Ok(()) => {} + Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} + Err(err) => return Err(foyer_file_error(&foyer_root, err)), + } + fs::create_dir_all(&foyer_root).map_err(|err| foyer_file_error(&foyer_root, err))?; + + let device = FsDeviceBuilder::new(&foyer_root) + .with_capacity(capacity_bytes) + // Keep the ablation backend outside the kernel page cache, matching native SSD I/O. + .with_direct(true) + .build() + .map_err(|err| foyer_error("build filesystem device", err))?; + let engine = BlockEngineConfig::new(device).with_block_size(FOYER_BLOCK_SIZE_BYTES); + let cache = HybridCacheBuilder::new() + .with_name("fluxon_kv_ssd_foyer") + .with_policy(HybridCachePolicy::WriteOnInsertion) + .with_flush_on_close(false) + .memory(FOYER_MEMORY_CAPACITY_BYTES) + .with_shards(FOYER_MEMORY_SHARDS) + .with_weighter(|_key: &FoyerKvSsdKey, value: &Vec| value.len()) + .with_filter(|_key: &FoyerKvSsdKey, _value: &Vec| false) + .storage() + .with_io_engine_config(PsyncIoEngineConfig::new()) + .with_engine_config(engine) + .build() + .await + .map_err(|err| foyer_error("build hybrid cache", err))?; + + tracing::warn!( + root_dir = %root_limit.root_dir.display(), + capacity_bytes = root_limit.limit_bytes, + memory_capacity_bytes = FOYER_MEMORY_CAPACITY_BYTES, + memory_shards = FOYER_MEMORY_SHARDS, + block_size_bytes = FOYER_BLOCK_SIZE_BYTES, + storage_admission = "forced", + persist_concurrency = 1, + direct_io = true, + "Initialized test-only Foyer KV SSD backend with memory admission disabled" + ); + + Ok(Self { + cache, + root_dir: root_limit.root_dir, + capacity_bytes: root_limit.limit_bytes, + persist_gate: TokioMutex::new(()), + entry_lengths: Mutex::new(HashMap::new()), + logical_used_bytes: AtomicU64::new(0), + memory_hits: AtomicU64::new(0), + disk_hits: AtomicU64::new(0), + outer_hits: AtomicU64::new(0), + }) + } + + pub(crate) fn root_dir(&self) -> &Path { + &self.root_dir + } + + pub(crate) fn capacity_bytes(&self) -> u64 { + self.capacity_bytes + } + + pub(crate) fn logical_used_bytes(&self) -> u64 { + self.logical_used_bytes + .load(Ordering::Relaxed) + .min(self.capacity_bytes) + } + + #[cfg(test)] + pub(crate) fn contains(&self, key: &str, put_id: PutIDForAKey) -> bool { + self.cache.contains(&FoyerKvSsdKey { + key: key.to_string(), + put_id, + }) + } + + pub(crate) async fn persist_from_addr( + &self, + key: &str, + put_id: PutIDForAKey, + addr: u64, + len: u64, + ) -> KvResult { + let len_usize = usize::try_from(len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("foyer kv ssd persist len does not fit usize: {len}"), + }) + })?; + if len_usize == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "foyer kv ssd persist len must be positive".to_string(), + })); + } + let _persist_permit = self.persist_gate.lock().await; + let data = unsafe { std::slice::from_raw_parts(addr as *const u8, len_usize) }.to_vec(); + self.persist_vec_locked(key, put_id, data).await + } + + pub(crate) async fn persist( + &self, + key: &str, + put_id: PutIDForAKey, + data: &[u8], + ) -> KvResult { + if data.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "foyer kv ssd persist len must be positive".to_string(), + })); + } + let _persist_permit = self.persist_gate.lock().await; + self.persist_vec_locked(key, put_id, data.to_vec()).await + } + + async fn persist_vec_locked( + &self, + key: &str, + put_id: PutIDForAKey, + data: Vec, + ) -> KvResult { + let cache_key = FoyerKvSsdKey { + key: key.to_string(), + put_id, + }; + let len = data.len() as u64; + if let Some(existing_len) = self.entry_lengths.lock().get(&cache_key).copied() { + if existing_len != len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd duplicate persist length mismatch: key={} put_id=({},{}) existing_len={} requested_len={}", + key, put_id.0, put_id.1, existing_len, len + ), + })); + } + } + + let entry = self + .cache + .storage_writer(cache_key.clone()) + .force() + .insert(data) + .ok_or_else(|| { + KvError::Api(ApiError::FileWriteError { + path: self.root_dir.display().to_string(), + offset: 0, + detail: format!( + "foyer rejected forced storage admission for key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }) + })?; + self.cache.storage().wait().await; + if !self.cache.storage().may_contains(&cache_key) { + return Err(KvError::Api(ApiError::FileWriteError { + path: self.root_dir.display().to_string(), + offset: 0, + detail: format!( + "foyer did not commit key={} put_id=({},{}) to storage", + key, put_id.0, put_id.1 + ), + })); + } + + let mut entry_lengths = self.entry_lengths.lock(); + if entry_lengths.insert(cache_key, len).is_none() { + self.logical_used_bytes.fetch_add(len, Ordering::Relaxed); + } + drop(entry_lengths); + Ok(FoyerKvSsdPersistGuard { _entry: entry }) + } + + pub(crate) async fn load_into_addr( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_len: u64, + ) -> KvResult<()> { + if target_len < len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd target capacity too small for key={} put_id=({},{}) len={} target_len={}", + key, put_id.0, put_id.1, len, target_len + ), + })); + } + let entry = self.load_entry(key, put_id, len).await?; + let len_usize = usize::try_from(len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("foyer kv ssd load len does not fit usize: {len}"), + }) + })?; + unsafe { + std::ptr::copy_nonoverlapping( + entry.value().as_ptr(), + target_addr as *mut u8, + len_usize, + ); + } + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn load_into_addr_chunks( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_len: u64, + chunk_bytes: u64, + _max_read_inflight: usize, + ready_tx: tokio_mpsc::Sender, + ) -> KvResult<()> { + if target_len < len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd target capacity too small for chunked load: key={} put_id=({},{}) len={} target_len={}", + key, put_id.0, put_id.1, len, target_len + ), + })); + } + let chunk_bytes = align_ssd_io_len(chunk_bytes.max(1))?; + let entry = self.load_entry(key, put_id, len).await?; + let mut offset = 0u64; + while offset < len { + let payload_len = chunk_bytes.min(len - offset); + let end = offset.checked_add(payload_len).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd chunk range overflow: offset={offset} len={payload_len}" + ), + }) + })?; + let stage_addr = target_addr.checked_add(offset).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd chunk address overflow: target_addr={target_addr} offset={offset}" + ), + }) + })?; + let offset_usize = usize::try_from(offset).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("foyer kv ssd chunk offset does not fit usize: {offset}"), + }) + })?; + let end_usize = usize::try_from(end).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("foyer kv ssd chunk end does not fit usize: {end}"), + }) + })?; + unsafe { + std::ptr::copy_nonoverlapping( + entry.value()[offset_usize..end_usize].as_ptr(), + stage_addr as *mut u8, + end_usize - offset_usize, + ); + } + ready_tx + .send(SsdLoadedChunk { + offset, + stage_addr, + len: payload_len, + }) + .await + .map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("foyer kv ssd chunk ready queue closed: {err}"), + }) + })?; + offset = end; + } + Ok(()) + } + + async fn load_entry( + &self, + key: &str, + put_id: PutIDForAKey, + expected_len: u64, + ) -> KvResult { + let cache_key = FoyerKvSsdKey { + key: key.to_string(), + put_id, + }; + let entry = self + .cache + .get(&cache_key) + .await + .map_err(|err| foyer_error("load entry", err))? + .ok_or_else(|| { + self.forget_entry(&cache_key); + KvError::Api(ApiError::KeyNotFound { + key: key.to_string(), + }) + })?; + self.record_source(entry.source()); + if entry.value().len() as u64 != expected_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "foyer kv ssd length mismatch for key={} put_id=({},{}) expected={} actual={}", + key, + put_id.0, + put_id.1, + expected_len, + entry.value().len() + ), + })); + } + Ok(entry) + } + + fn forget_entry(&self, key: &FoyerKvSsdKey) { + if let Some(len) = self.entry_lengths.lock().remove(key) { + self.logical_used_bytes.fetch_sub(len, Ordering::Relaxed); + } + } + + fn record_source(&self, source: Source) { + match source { + Source::Memory => self.memory_hits.fetch_add(1, Ordering::Relaxed), + Source::Disk => self.disk_hits.fetch_add(1, Ordering::Relaxed), + Source::Outer => self.outer_hits.fetch_add(1, Ordering::Relaxed), + }; + } + + #[cfg(test)] + pub(crate) fn memory_usage(&self) -> usize { + self.cache.memory().usage() + } + + #[cfg(test)] + pub(crate) fn source_counts(&self) -> (u64, u64, u64) { + ( + self.memory_hits.load(Ordering::Relaxed), + self.disk_hits.load(Ordering::Relaxed), + self.outer_hits.load(Ordering::Relaxed), + ) + } +} + +fn foyer_error(operation: &str, err: impl std::fmt::Display) -> KvError { + KvError::Api(ApiError::Unknown { + detail: format!("foyer kv ssd {operation} failed: {err}"), + }) +} + +fn foyer_file_error(path: &Path, err: std::io::Error) -> KvError { + KvError::Api(ApiError::FileWriteError { + path: path.display().to_string(), + offset: 0, + detail: err.to_string(), + }) +} diff --git a/fluxon_rs/fluxon_kv/src/kv_test.rs b/fluxon_rs/fluxon_kv/src/kv_test.rs index 5f0a9e2..41f7b5b 100644 --- a/fluxon_rs/fluxon_kv/src/kv_test.rs +++ b/fluxon_rs/fluxon_kv/src/kv_test.rs @@ -11,9 +11,11 @@ use crate::cluster_manager::ClusterManagerRdmaControlInit; use crate::config::{ - ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, LargeFilePaths, MasterConfig, MonitoringConfig, - ProtocolConfig, ProtocolType, TestSpecConfig, TestSpecTransportMode, TransferEngineType, + ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, KvSsdStorageConfig, LargeFilePaths, + MasterConfig, MonitoringConfig, ProtocolConfig, ProtocolType, TestSpecConfig, + TestSpecTransportMode, TransferEngineType, }; +use crate::master_kv_router::msg_pack::GetSourceKind; use crate::run_master_with_test_overrides; use crate::{ClientRunTestOverrides, MasterRunTestOverrides, run_client_with_test_overrides}; // external client runs via run_client when contribution is zero @@ -38,6 +40,8 @@ const CLIENT_COMMUNICATION_VALUE: &[u8] = b"message_from_client1_to_client2"; const TRANSFER_DATA_PROBE_VALUE_LEN: usize = 256 * 1024; const KV_TEST_TRANSFER_PROBE_IO_TIMEOUT_SECS: u64 = 10; const KV_TEST_SHUTDOWN_TIMEOUT_SECS: u64 = 60; +const KV_TEST_SSD_STORAGE_BYTES: u64 = 64 * 1024 * 1024; +const KV_TEST_STORAGE_PROFILE_SSD_ROUTE_TIMEOUT_SECS: u64 = 30; fn kv_test_run_scope() -> &'static str { static RUN_SCOPE: OnceLock = OnceLock::new(); @@ -610,6 +614,7 @@ struct KvTestClientOptions { enable_transfer_rpc_fast_path: Option, contribute_to_cluster_pool_size: Option, share_mem_path: Option, + ssd_storage: Option, etcd_mode: Option, } @@ -642,6 +647,10 @@ impl KvTestClientOptions { .share_mem_path .clone() .or_else(|| self.share_mem_path.clone()), + ssd_storage: overrides + .ssd_storage + .clone() + .or_else(|| self.ssd_storage.clone()), etcd_mode: overrides .etcd_mode .clone() @@ -650,6 +659,40 @@ impl KvTestClientOptions { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum KvTestStorageProfile { + Memory, + Ssd, + MemorySsd, +} + +impl KvTestStorageProfile { + fn round_suffix(self) -> &'static str { + match self { + Self::Memory => "", + Self::Ssd => "_ssd", + Self::MemorySsd => "_memory_ssd", + } + } + + fn ssd_storage(self) -> Option { + match self { + Self::Memory => None, + Self::Ssd | Self::MemorySsd => Some(KvSsdStorageConfig { + large_limit_size: vec![KV_TEST_SSD_STORAGE_BYTES], + }), + } + } + + fn requires_memory_source(self) -> bool { + matches!(self, Self::Memory | Self::MemorySsd) + } + + fn requires_ssd_source(self) -> bool { + matches!(self, Self::Ssd | Self::MemorySsd) + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum KvTestRoundProfile { P2pOnly, @@ -760,6 +803,7 @@ fn kv_test_round_test_spec_config(round_profile: KvTestRoundProfile) -> TestSpec #[derive(Clone, Debug)] struct KvTestRoundOptions { round_profile: KvTestRoundProfile, + storage_profile: KvTestStorageProfile, round_name: String, cluster_name: String, master_port: Option, @@ -803,6 +847,9 @@ impl KvTestRoundOptions { ) } + fn owner_sub_cluster(&self) -> String { + format!("{}_owners", self.round_name) + } } #[derive(Clone, Debug)] @@ -842,8 +889,7 @@ fn default_client_large_file_paths( instance_key: &str, contribute_to_cluster_pool_size: &ContributeToClusterPoolSize, ) -> LargeFilePaths { - if contribute_to_cluster_pool_size.dram == 0 - && contribute_to_cluster_pool_size.vram.is_empty() + if contribute_to_cluster_pool_size.dram == 0 && contribute_to_cluster_pool_size.vram.is_empty() { return LargeFilePaths { paths: Vec::new() }; } @@ -852,7 +898,10 @@ fn default_client_large_file_paths( } } -fn default_owner_test_client_options(round_profile: KvTestRoundProfile) -> KvTestClientOptions { +fn default_owner_test_client_options( + round_profile: KvTestRoundProfile, + storage_profile: KvTestStorageProfile, +) -> KvTestClientOptions { KvTestClientOptions { protocol_config: Some(round_profile.protocol_config()), transfer_engine: Some(round_profile.owner_transfer_engine()), @@ -861,6 +910,7 @@ fn default_owner_test_client_options(round_profile: KvTestRoundProfile) -> KvTes enable_transfer_rpc_fast_path: Some(round_profile.enable_transfer_rpc_fast_path()), contribute_to_cluster_pool_size: Some(default_owner_contribute_to_cluster_pool_size()), share_mem_path: None, + ssd_storage: storage_profile.ssd_storage(), etcd_mode: Some(KvTestEtcdMode::Enabled), } } @@ -874,6 +924,7 @@ fn default_master_test_client_options(round_profile: KvTestRoundProfile) -> KvTe enable_transfer_rpc_fast_path: Some(round_profile.enable_transfer_rpc_fast_path()), contribute_to_cluster_pool_size: None, share_mem_path: None, + ssd_storage: None, etcd_mode: None, } } @@ -887,22 +938,31 @@ fn default_external_test_client_options() -> KvTestClientOptions { enable_transfer_rpc_fast_path: Some(false), contribute_to_cluster_pool_size: Some(default_external_contribute_to_cluster_pool_size()), share_mem_path: None, + ssd_storage: None, etcd_mode: Some(KvTestEtcdMode::Disabled), } } -fn new_kv_test_round(round_profile: KvTestRoundProfile) -> KvTestRoundOptions { - let round_name = round_profile.round_name(); +fn new_kv_test_round( + round_profile: KvTestRoundProfile, + storage_profile: KvTestStorageProfile, +) -> KvTestRoundOptions { + let round_name = format!( + "{}{}", + round_profile.round_name(), + storage_profile.round_suffix() + ); KvTestRoundOptions { round_profile, - round_name: round_name.to_string(), + storage_profile, + round_name: round_name.clone(), // Keep each process run on its own cluster namespace so a crashed/aborted previous run // cannot poison the next rerun with stale members. cluster_name: format!("test_cluster_{}_{}", round_name, kv_test_run_scope()), master_port: None, step8_master_port: None, master_options: default_master_test_client_options(round_profile), - owner_client_options: default_owner_test_client_options(round_profile), + owner_client_options: default_owner_test_client_options(round_profile, storage_profile), external_client_options: default_external_test_client_options(), } } @@ -919,15 +979,35 @@ fn default_kv_test_run_options() -> KvTestRunOptions { .filter(|item| !item.is_empty()) { let profile = match round_name { - "p2p_only" => KvTestRoundProfile::P2pOnly, + "p2p_only" => { + rounds.push(new_kv_test_round( + KvTestRoundProfile::P2pOnly, + KvTestStorageProfile::Memory, + )); + continue; + } + "p2p_only_ssd" => { + rounds.push(new_kv_test_round( + KvTestRoundProfile::P2pOnly, + KvTestStorageProfile::Ssd, + )); + continue; + } + "p2p_only_memory_ssd" => { + rounds.push(new_kv_test_round( + KvTestRoundProfile::P2pOnly, + KvTestStorageProfile::MemorySsd, + )); + continue; + } "rdma_transfer_only" => KvTestRoundProfile::RdmaTransferOnly, "rdma_transfer_with_rpc" => KvTestRoundProfile::RdmaTransferWithRpc, other => panic!( - "unsupported FLUXON_KV_TEST_ROUNDS entry '{}'; expected one of: p2p_only, rdma_transfer_only, rdma_transfer_with_rpc", + "unsupported FLUXON_KV_TEST_ROUNDS entry '{}'; expected one of: p2p_only, p2p_only_ssd, p2p_only_memory_ssd, rdma_transfer_only, rdma_transfer_with_rpc", other ), }; - rounds.push(new_kv_test_round(profile)); + rounds.push(new_kv_test_round(profile, KvTestStorageProfile::Memory)); } if rounds.is_empty() { panic!("FLUXON_KV_TEST_ROUNDS was set but produced no valid rounds"); @@ -937,9 +1017,17 @@ fn default_kv_test_run_options() -> KvTestRunOptions { KvTestRunOptions { rounds: vec![ - new_kv_test_round(KvTestRoundProfile::P2pOnly), - new_kv_test_round(KvTestRoundProfile::RdmaTransferOnly), - new_kv_test_round(KvTestRoundProfile::RdmaTransferWithRpc), + new_kv_test_round(KvTestRoundProfile::P2pOnly, KvTestStorageProfile::Memory), + new_kv_test_round(KvTestRoundProfile::P2pOnly, KvTestStorageProfile::Ssd), + new_kv_test_round(KvTestRoundProfile::P2pOnly, KvTestStorageProfile::MemorySsd), + new_kv_test_round( + KvTestRoundProfile::RdmaTransferOnly, + KvTestStorageProfile::Memory, + ), + new_kv_test_round( + KvTestRoundProfile::RdmaTransferWithRpc, + KvTestStorageProfile::Memory, + ), ], } } @@ -1022,6 +1110,8 @@ fn build_client_launch( let contribute_to_cluster_pool_size = options .contribute_to_cluster_pool_size .unwrap_or(default_owner_contribute_to_cluster_pool_size()); + let is_external = contribute_to_cluster_pool_size.dram == 0 + && contribute_to_cluster_pool_size.vram.is_empty(); let share_mem_path = options .share_mem_path .unwrap_or_else(|| format!("/tmp/kvcache_shared_memory/{}", instance_key)); @@ -1043,7 +1133,11 @@ fn build_client_launch( enable_transfer_rpc_fast_path: options .enable_transfer_rpc_fast_path .expect("kv_test requires enable_transfer_rpc_fast_path to be set explicitly"), - sub_cluster: None, + sub_cluster: if is_external { + None + } else { + Some(round.owner_sub_cluster()) + }, }, // English note: // kv_test uses a per-instance shared memory path by default so each owner/external share @@ -1054,6 +1148,7 @@ fn build_client_launch( &instance_key, &contribute_to_cluster_pool_size, ), + ssd_storage: options.ssd_storage, // Mirror round intent into the generated config so logs and runtime behavior // agree on whether this launch is transfer_only vs transfer_with_rpc. test_spec_config: kv_test_round_test_spec_config(round.round_profile), @@ -1381,7 +1476,10 @@ async fn key_meta_cache_check( } } - tracing::info!("🔍 Starting PUT and GET in parallel: {}", parallel_unique_key); + tracing::info!( + "🔍 Starting PUT and GET in parallel: {}", + parallel_unique_key + ); for i in 0..10 { let (put_client, other_client) = if i % 2 == 0 { (client, client2) @@ -1420,7 +1518,9 @@ async fn key_meta_cache_check( } assert!( - put_client.client_kv_api().has_cached_key(parallel_unique_key), + put_client + .client_kv_api() + .has_cached_key(parallel_unique_key), "put client should have immediate local cache metadata for key {} after put time {}", parallel_unique_key, i @@ -1577,6 +1677,321 @@ async fn shutdown_framework_with_timeout(label: &str, framework: &crate::Framewo } } +fn build_storage_profile_probe_value(tag: &str) -> Vec { + const STORAGE_PROFILE_PROBE_VALUE_LEN: usize = 64 * 1024; + build_storage_profile_probe_value_with_len(tag, STORAGE_PROFILE_PROBE_VALUE_LEN) +} + +fn build_storage_profile_probe_value_with_len(tag: &str, len: usize) -> Vec { + let pattern = format!("kv_test_storage_profile:{tag}:").into_bytes(); + let mut value = Vec::with_capacity(len); + while value.len() < len { + value.extend_from_slice(pattern.as_slice()); + } + value.truncate(len); + value +} + +async fn force_evict_memory_replicas_for_storage_probe( + master_framework: &crate::Framework, + key: &str, +) { + let master_view = master_framework.master_kv_router_view(); + let deadline = + Instant::now() + Duration::from_secs(KV_TEST_STORAGE_PROFILE_SSD_ROUTE_TIMEOUT_SECS); + let (put_id, memory_replica_nodes) = loop { + if let Some(route) = master_view.master_kv_router().inner().kv_routes.get(key) { + let put_id = route.put_id; + let memory_replica_nodes = route + .node_replicas + .read() + .iter() + .filter_map(|(node_id, replicas)| { + replicas.memory.is_some().then(|| node_id.clone()) + }) + .collect::>(); + let ssd_replica_count = route + .node_replicas + .read() + .values() + .filter(|replicas| replicas.ssd.is_some()) + .count(); + if ssd_replica_count > 0 { + break (put_id, memory_replica_nodes); + } + } + + if Instant::now() >= deadline { + panic!( + "storage profile probe expected at least one SSD replica before memory eviction: key={} timeout={}s", + key, KV_TEST_STORAGE_PROFILE_SSD_ROUTE_TIMEOUT_SECS + ); + } + sleep(Duration::from_millis(50)).await; + }; + + for node_id in memory_replica_nodes { + crate::master_kv_router::delete::evict_one_kv_replica_for_node( + &master_view, + key.to_string(), + node_id.clone(), + put_id, + ) + .unwrap_or_else(|code| { + panic!( + "storage profile probe failed to evict memory replica: key={} node={} put_id=({},{}) code={}", + key, node_id, put_id.0, put_id.1, code + ) + }); + } + + let Some(route) = master_view.master_kv_router().inner().kv_routes.get(key) else { + panic!("storage profile probe route disappeared after memory replicas eviction: key={key}"); + }; + assert!( + route + .node_replicas + .read() + .values() + .all(|replicas| replicas.memory.is_none()), + "storage profile probe memory replicas still exist after eviction: key={}", + key + ); + assert!( + route + .node_replicas + .read() + .values() + .any(|replicas| replicas.ssd.is_some()), + "storage profile probe SSD replica disappeared after memory replicas eviction: key={}", + key + ); +} + +async fn assert_owner_get_source_kind( + reader_framework: &crate::Framework, + key: &str, + expected_value: &[u8], + expected_source_kind: GetSourceKind, +) { + let reader_view = reader_framework.client_kv_api_view().clone(); + let reader_api = reader_view.client_kv_api(); + let (mem_holder, get_info) = reader_api + .inner() + .get(key) + .await + .unwrap_or_else(|err| { + panic!( + "storage profile probe get failed: key={} expected_source={:?} err={}", + key, expected_source_kind, err + ) + }) + .unwrap_or_else(|| { + panic!( + "storage profile probe get returned None: key={} expected_source={:?}", + key, expected_source_kind + ) + }); + assert_eq!( + mem_holder.bytes(), + expected_value, + "storage profile probe value mismatch for key={key}" + ); + let Some(get_info) = get_info else { + panic!( + "storage profile probe expected remote get info for key={} source={:?}", + key, expected_source_kind + ); + }; + assert_eq!( + get_info.source_kind(), + expected_source_kind, + "storage profile probe source kind mismatch for key={key}" + ); +} + +fn invalidate_ssd_route_len_for_storage_probe(master_framework: &crate::Framework, key: &str) { + let master_view = master_framework.master_kv_router_view(); + let route = master_view + .master_kv_router() + .inner() + .kv_routes + .get(key) + .unwrap_or_else(|| panic!("stale SSD route probe is missing key={key}")); + let mut node_replicas = route.node_replicas.write(); + for replicas in node_replicas.values_mut() { + if let Some(ssd) = replicas.ssd.as_mut() { + ssd.len = ssd + .len + .checked_add(1) + .expect("stale SSD route probe length must not overflow"); + return; + } + } + panic!("stale SSD route probe found no SSD replica for key={key}"); +} + +async fn assert_stale_ssd_route_retries_to_miss( + master_framework: &crate::Framework, + reader_framework: &crate::Framework, + key: &str, +) { + let reader_view = reader_framework.client_kv_api_view().clone(); + let result = reader_view + .client_kv_api() + .inner() + .get(key) + .await + .unwrap_or_else(|err| { + panic!("stale SSD route probe returned an error: key={key} err={err}") + }); + assert!( + result.is_none(), + "stale SSD route probe must retry and converge to a miss: key={key}" + ); + assert!( + !master_framework + .master_kv_router_view() + .master_kv_router() + .inner() + .kv_routes + .contains_key(key), + "stale SSD route probe must remove the empty route: key={key}" + ); +} + +async fn run_non_rdma_storage_profile_coverage( + round: &KvTestRoundOptions, + master_framework: &crate::Framework, + writer_framework: &crate::Framework, +) -> Option> { + if round.round_profile != KvTestRoundProfile::P2pOnly { + return None; + } + + info!( + "📋 Storage profile coverage: round={} storage={:?}", + round.round_name, round.storage_profile + ); + + let writer_view = writer_framework.client_kv_api_view().clone(); + let writer_api = writer_view.client_kv_api(); + let storage_probe_put_opts = || { + crate::client_kv_api::PutOptionalArgs(vec![ + crate::client_kv_api::PutOptionalArg::PreferredSubCluster(round.owner_sub_cluster()), + ]) + }; + + let memory_key = format!("storage_profile_memory_key_{}", round.round_name); + let memory_value = build_storage_profile_probe_value(&format!("{}:memory", round.round_name)); + if round.storage_profile.requires_memory_source() { + writer_api + .inner() + .put(&memory_key, &memory_value, storage_probe_put_opts()) + .await + .unwrap_or_else(|err| { + panic!( + "storage profile memory probe put failed: key={} err={}", + memory_key, err + ) + }); + } + + let ssd_key = format!("storage_profile_ssd_key_{}", round.round_name); + let ssd_value = build_storage_profile_probe_value_with_len( + &format!("{}:ssd", round.round_name), + 64 * 1024 + 123, + ); + let local_small_ssd_key = format!("storage_profile_local_small_ssd_key_{}", round.round_name); + let local_small_ssd_value = build_storage_profile_probe_value_with_len( + &format!("{}:local_small_ssd", round.round_name), + 123, + ); + let stale_ssd_key = format!("storage_profile_stale_ssd_key_{}", round.round_name); + let stale_ssd_value = build_storage_profile_probe_value_with_len( + &format!("{}:stale_ssd", round.round_name), + 64 * 1024 + 321, + ); + if round.storage_profile.requires_ssd_source() { + writer_api + .inner() + .put(&ssd_key, &ssd_value, storage_probe_put_opts()) + .await + .unwrap_or_else(|err| { + panic!( + "storage profile SSD probe put failed: key={} err={}", + ssd_key, err + ) + }); + force_evict_memory_replicas_for_storage_probe(master_framework, &ssd_key).await; + + writer_api + .inner() + .put( + &local_small_ssd_key, + &local_small_ssd_value, + storage_probe_put_opts(), + ) + .await + .unwrap_or_else(|err| { + panic!( + "storage profile local small SSD probe put failed: key={} err={}", + local_small_ssd_key, err + ) + }); + force_evict_memory_replicas_for_storage_probe(master_framework, &local_small_ssd_key).await; + + writer_api + .inner() + .put(&stale_ssd_key, &stale_ssd_value, storage_probe_put_opts()) + .await + .unwrap_or_else(|err| { + panic!( + "storage profile stale SSD probe put failed: key={} err={}", + stale_ssd_key, err + ) + }); + force_evict_memory_replicas_for_storage_probe(master_framework, &stale_ssd_key).await; + invalidate_ssd_route_len_for_storage_probe(master_framework, &stale_ssd_key); + } + + let reader_launch = new_client_launch(round, "test_storage_profile_reader", None); + let (reader_framework, _) = run_kv_test_client(reader_launch) + .await + .expect("Failed to start storage profile reader"); + + sleep(Duration::from_secs(10)).await; + + if round.storage_profile.requires_memory_source() { + assert_owner_get_source_kind( + &reader_framework, + &memory_key, + &memory_value, + GetSourceKind::Memory, + ) + .await; + } + if round.storage_profile.requires_ssd_source() { + assert_owner_get_source_kind( + writer_framework, + &local_small_ssd_key, + &local_small_ssd_value, + GetSourceKind::Ssd, + ) + .await; + assert_owner_get_source_kind(&reader_framework, &ssd_key, &ssd_value, GetSourceKind::Ssd) + .await; + assert_stale_ssd_route_retries_to_miss(master_framework, &reader_framework, &stale_ssd_key) + .await; + } + + info!( + "✅ Storage profile coverage passed: round={} storage={:?}", + round.round_name, round.storage_profile + ); + Some(reader_framework) +} + async fn run_kv_step8(round: &KvTestRoundOptions) { info!("📋 Step 8: Verifying external client blocking and recovery behavior"); @@ -2462,10 +2877,10 @@ async fn run_kv_round(round: &KvTestRoundOptions) { .kv_routes .get(CLIENT_COMMUNICATION_KEY) { - let replicas = one_kv_nodes_routes.nodes_replicas.read(); + let replicas = one_kv_nodes_routes.node_replicas.read(); let active_replica_count = replicas .iter() - .filter(|(_, kv_info)| !kv_info.tomb_tag.is_tomb()) + .filter(|(_, replicas)| !replicas.tomb_tag.is_tomb() && replicas.memory.is_some()) .count(); info!( @@ -2720,6 +3135,9 @@ async fn run_kv_round(round: &KvTestRoundOptions) { info!("✅ Key meta cache testing completed"); } + let storage_profile_reader_framework = + run_non_rdma_storage_profile_coverage(round, &master_framework, &client1_framework).await; + // 清理旧资源 { info!("🧹 Cleaning up resources"); @@ -2743,6 +3161,14 @@ async fn run_kv_round(round: &KvTestRoundOptions) { .unwrap_or_else(|e| panic!("Client 1 framework shutdown failed: {}", e)); info!("✅ Client 1 framework shutdown successfully"); + if let Some(storage_profile_reader_framework) = storage_profile_reader_framework { + shutdown_framework_with_timeout( + "storage profile reader", + &storage_profile_reader_framework, + ) + .await; + } + master_framework .shutdown() .await diff --git a/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs b/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs index c74b64a..43d3c09 100644 --- a/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs +++ b/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs @@ -148,6 +148,7 @@ fn new_client_config_with_cluster_and_dram( large_file_paths: crate::config::LargeFilePaths { paths: vec![format!("{}/large/{}", base, instance_key)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), }; println!("fluxonkv core created client config for test: {:?}", conf); diff --git a/fluxon_rs/fluxon_kv/src/lib.rs b/fluxon_rs/fluxon_kv/src/lib.rs index edaa386..117da53 100644 --- a/fluxon_rs/fluxon_kv/src/lib.rs +++ b/fluxon_rs/fluxon_kv/src/lib.rs @@ -7,6 +7,10 @@ pub mod external_client_api; pub mod panel_proxy; // #[cfg(test)] pub mod key_prefix; +pub mod kv_ssd_storage; +#[cfg(test)] +mod kv_ssd_storage2; +mod kv_ssd_storage_foyer; #[cfg(feature = "test_bins")] pub mod kv_test; pub mod kvlease; @@ -56,6 +60,7 @@ use crate::external_client_api::ExternalClientApiViewTrait; use crate::master_kv_router::MasterKvRouterAccessTrait; use crate::master_kv_router::MasterKvRouterView; use crate::master_kv_router::MasterKvRouterViewTrait; +use crate::master_kv_router::msg_pack::GetSourceKind; use crate::master_lease_manager::MasterLeaseManager; use crate::master_lease_manager::master_lease_manager::MasterLeaseManagerNewArg; use crate::master_lease_manager::{ @@ -76,7 +81,9 @@ use clap::{Parser, Subcommand}; use client_kv_api::ClientKvApiNewArg; use client_seg_pool::{ClientSegPool, ClientSegPoolNewArg}; use client_transfer_engine::ClientTransferEngine; -use cluster_manager::{ClusterManager, ClusterManagerNewArg, ClusterManagerRdmaControlInit}; +use cluster_manager::{ + ClusterManager, ClusterManagerNewArg, ClusterManagerRdmaControlInit, META_KEY_KV_SSD_STORAGE, +}; use config::{ ClientConfig, ClientConfigYaml, ContributeToClusterPoolSize, FluxonKvSpec, MasterConfig, MasterConfigYaml, ProtocolConfig, ProtocolType, SideTransferRole, TestSpecConfig, @@ -203,6 +210,24 @@ pub enum KvGetResult { External(Option>), } +/// Internal GET result that preserves the source selected for this operation. +#[derive(Clone)] +pub enum KvGetResultWithSource { + Owner(Option<(Arc, GetSourceKind)>), + External(Option<(Arc, GetSourceKind)>), +} + +impl KvGetResultWithSource { + fn without_source(self) -> KvGetResult { + match self { + Self::Owner(result) => KvGetResult::Owner(result.map(|(holder, _source_kind)| holder)), + Self::External(result) => { + KvGetResult::External(result.map(|(holder, _source_kind)| holder)) + } + } + } +} + /// A thin trait on `Framework` to route KV ops by role. /// /// Rules: @@ -245,8 +270,14 @@ pub trait KvClientTrait { opts: crate::client_kv_api::PutOptionalArgs, ) -> KvResult<()>; - /// Get by role; returns role-tagged holder - async fn kv_get(&self, key: &str) -> KvResult; + /// Get by role; returns role-tagged holder. + async fn kv_get(&self, key: &str) -> KvResult { + Ok(self.kv_get_with_source(key).await?.without_source()) + } + + /// Benchmark instrumentation path that also returns the selected storage source. + #[doc(hidden)] + async fn kv_get_with_source(&self, key: &str) -> KvResult; /// Delete by role async fn kv_delete(&self, key: &str) -> KvResult<()>; @@ -351,15 +382,15 @@ impl KvClientTrait for Framework { } } - async fn kv_get(&self, key: &str) -> KvResult { + async fn kv_get_with_source(&self, key: &str) -> KvResult { if self.is_external_mode() { let r = self .external_client_api_view() .external_client_api() .inner() - .get(key) + .get_with_source(key) .await?; - Ok(KvGetResult::External(r)) + Ok(KvGetResultWithSource::External(r)) } else { let r = self .client_kv_api_view() @@ -367,7 +398,12 @@ impl KvClientTrait for Framework { .inner() .get(key) .await?; - Ok(KvGetResult::Owner(r.map(|(h, _)| h))) + Ok(KvGetResultWithSource::Owner(r.map(|(holder, get_info)| { + let source_kind = get_info + .map(|info| info.source_kind()) + .unwrap_or(GetSourceKind::Memory); + (holder, source_kind) + }))) } } @@ -797,6 +833,7 @@ fn build_side_transfer_worker_config( }, share_mem_path: owner_config.share_mem_path.clone(), large_file_paths: owner_config.large_file_paths.clone(), + ssd_storage: None, test_spec_config, }) } @@ -841,6 +878,7 @@ fn build_side_transfer_worker_config_yaml( cluster_name: side_config.cluster_name, share_mem_path: side_config.share_mem_path, large_file_paths: None, + large_limit_size: None, p2p_listen_port: side_config.fluxonkv_spec.p2p_listen_port, redis_compat: None, sub_cluster: None, @@ -1915,6 +1953,9 @@ async fn run_client_impl( if is_side_transfer_worker { metadata.insert("side_transfer_worker".to_string(), "true".to_string()); } + if !is_external && !is_side_transfer_worker && config.ssd_storage.is_some() { + metadata.insert(META_KEY_KV_SSD_STORAGE.to_string(), "true".to_string()); + } // Local IPC routing requires both share-group owner id and the local IPC root. // The owner id is also published via a dedicated share-group key; we denormalize it into @@ -2004,6 +2045,33 @@ async fn run_client_impl( .await .map_err(|e| anyhow::anyhow!("Failed to initialize framework: {:#}", e))?; } else { + let ssd_storage = if is_side_transfer_worker { + None + } else if let Some(ssd_cfg) = config.ssd_storage.as_ref() { + let root_limits = config + .large_file_paths + .kv_ssd_storage_root_limits( + &config.cluster_name, + &config.instance_key, + &ssd_cfg.large_limit_size, + ) + .map_err(|err| anyhow::anyhow!("invalid kv ssd storage dirs: {}", err))?; + Some(crate::kv_ssd_storage::KvSsdStorageInit { + root_limits: root_limits + .into_iter() + .map( + |(root_dir, limit_bytes)| crate::kv_ssd_storage::KvSsdStorageRootLimit { + root_dir, + limit_bytes, + }, + ) + .collect(), + uring_mode: config.test_spec_config.kv_ssd_uring_mode, + backend: config.test_spec_config.kv_ssd_storage_backend, + }) + } else { + None + }; let init_args = InitArgsOwner { cluster_manager_arg: ClusterManagerNewArg { etcd_endpoints: config.fluxonkv_spec.etcd_addresses.clone(), @@ -2036,6 +2104,7 @@ async fn run_client_impl( }, client_kv_api_arg: ClientKvApiNewArg { test_spec_config: config.test_spec_config.clone(), + ssd_storage, }, client_seg_pool_arg: ClientSegPoolNewArg { contribute_size: config.contribute_to_cluster_pool_size.clone(), @@ -2468,6 +2537,7 @@ mod tests { large_file_paths: crate::config::LargeFilePaths { paths: vec!["/tmp/fluxon_side_transfer_test_large".to_string()], }, + ssd_storage: None, test_spec_config: TestSpecConfig { enable_side_transfer: true, side_transfer_worker_count: 4, @@ -2736,8 +2806,8 @@ mod tests { large_file_paths: crate::config::LargeFilePaths { paths: vec![owner_large_root.to_string_lossy().into_owned()], }, - protocol_version: - fluxon_util::git_version_build_record::get_current_git_commitid().unwrap(), + protocol_version: fluxon_util::git_version_build_record::get_current_git_commitid() + .unwrap(), write_ts: Some(chrono::Utc::now().timestamp_micros()), }; let shared_meta_json = serde_json::to_string(&shared_meta).unwrap(); @@ -2773,6 +2843,7 @@ mod tests { }, share_mem_path: share_mem_root.to_string_lossy().into_owned(), large_file_paths: crate::config::LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), }; diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs index 12a55ee..5d945ba 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs @@ -1,8 +1,9 @@ use super::{ MasterKvRouterView, msg_pack::{ - BatchDeleteAckReq, BatchDeleteAckResp, BatchDeleteClientKvMetaCacheReq, DeleteAckReq, - DeleteAckResp, DeleteClientKvMetaCacheItem, DeleteReq, DeleteResp, + BatchDeleteAckReq, BatchDeleteAckResp, BatchDeleteClientKvMetaCacheReq, + BatchSsdReplicaEvictReq, BatchSsdReplicaEvictResp, DeleteAckReq, DeleteAckResp, + DeleteClientKvMetaCacheItem, DeleteReq, DeleteResp, }, }; use crate::master_kv_router::OneKvNodesRoutes; @@ -63,8 +64,11 @@ pub fn do_delete_one_kv_all_replicas( } // Remove from all node caches that hold replicas of this key - let nodes_replicas = kv_route_info.nodes_replicas.read(); - for (node_id, _kv_info) in nodes_replicas.iter() { + let node_replicas = kv_route_info.node_replicas.read(); + for (node_id, replicas) in node_replicas.iter() { + if replicas.tomb_tag.is_tomb() || replicas.memory.is_none() { + continue; + } if let Some(cache) = view.master_kv_router().get_node_cache_controller(node_id) { let _ = cache.remove(&key_clone); @@ -115,10 +119,7 @@ pub fn evict_one_kv_replica_for_node( return Ok(()); } - let removed_replica = { - let mut nodes_replicas = route.nodes_replicas.write(); - nodes_replicas.remove(&node_id).is_some() - }; + let removed_replica = route.remove_memory_replica(&node_id); if !removed_replica { tracing::debug!( "Local replica eviction ignored because node replica is already absent: key={} node_id={} put_id=({},{})", @@ -130,13 +131,18 @@ pub fn evict_one_kv_replica_for_node( return Ok(()); } - let last_replica_gone = route.nodes_replicas.read().is_empty(); + let last_replica_gone = !route.has_live_replica(); if last_replica_gone { + let route_for_compare = route.clone(); let removed = view .master_kv_router() .inner() .kv_routes - .remove_if(&key, |_, current| current.put_id == put_id) + .remove_if(&key, |_, current| { + Arc::ptr_eq(current, &route_for_compare) + && current.put_id == put_id + && !current.has_live_replica() + }) .is_some(); if removed && view.master_kv_router().prefix_index_enabled() { let view_task = view.clone(); @@ -213,6 +219,138 @@ pub fn evict_one_kv_replica_for_node( Ok(()) } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct SsdReplicaRemoval { + pub replica_removed: bool, + pub route_removed: bool, +} + +fn remove_ssd_replica_if_version( + route: &OneKvNodesRoutes, + node_id: &NodeID, + put_id: PutIDForAKey, +) -> bool { + route.put_id == put_id && route.remove_ssd_replica(node_id) +} + +pub(crate) fn remove_one_ssd_replica_for_node( + view: &MasterKvRouterView, + key: &str, + node_id: &NodeID, + put_id: PutIDForAKey, +) -> SsdReplicaRemoval { + let route = { + let Some(route) = view.master_kv_router().inner().kv_routes.get(key) else { + return SsdReplicaRemoval::default(); + }; + route.clone() + }; + if !remove_ssd_replica_if_version(&route, node_id, put_id) { + return SsdReplicaRemoval::default(); + } + + let route_removed = if route.has_live_replica() { + false + } else { + let route_for_compare = route.clone(); + view.master_kv_router() + .inner() + .kv_routes + .remove_if(key, |_, current| { + Arc::ptr_eq(current, &route_for_compare) + && current.put_id == put_id + && !current.has_live_replica() + }) + .is_some() + }; + + SsdReplicaRemoval { + replica_removed: true, + route_removed, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::master_kv_router::{KvNodeReplicas, KvSsdReplicaInfo}; + use crate::master_seg_manager::NodeTombTag; + + #[test] + fn ssd_replica_cleanup_is_scoped_to_put_id() { + let route = OneKvNodesRoutes::new((10, 2), None); + let node_id: NodeID = "owner-a".to_string().into(); + route.node_replicas.write().insert( + node_id.clone(), + KvNodeReplicas { + tomb_tag: NodeTombTag::new(), + memory: None, + ssd: Some(KvSsdReplicaInfo { len: 4096 }), + }, + ); + + assert!(!remove_ssd_replica_if_version(&route, &node_id, (9, 7))); + assert!( + route + .node_replicas + .read() + .get(&node_id) + .is_some_and(|replicas| replicas.ssd.is_some()) + ); + + assert!(remove_ssd_replica_if_version(&route, &node_id, (10, 2))); + assert!(route.node_replicas.read().get(&node_id).is_none()); + } +} + +pub async fn handle_batch_ssd_replica_evict( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let requested_count = req.serialize_part.evictions.len(); + let mut removed_count = 0u32; + let mut removed_route_keys = Vec::new(); + + for replica in &req.serialize_part.evictions { + let removal = + remove_one_ssd_replica_for_node(&view, &replica.key, &req_node_id, replica.put_id); + if removal.replica_removed { + removed_count = removed_count.saturating_add(1); + } + if removal.route_removed { + removed_route_keys.push(replica.key.clone()); + } + } + + if !removed_route_keys.is_empty() && view.master_kv_router().prefix_index_enabled() { + let inner = view.master_kv_router().inner(); + let mut tree = inner.prefix_index.write().await; + for key in &removed_route_keys { + if !inner.kv_routes.contains_key(key) { + tree.remove(key); + } + } + } + + tracing::debug!( + "Handled SSD replica eviction batch: node={} requested={} removed={} routes_removed={}", + req_node_id, + requested_count, + removed_count, + removed_route_keys.len() + ); + + MsgPack { + serialize_part: BatchSsdReplicaEvictResp { + removed_count, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + pub async fn handle_delete( view: MasterKvRouterView, req: MsgPack, diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs index 8c17155..5179879 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs @@ -1,27 +1,192 @@ use super::{ - InflightGetInfo, KvRouteInfo, MasterKvRouterView, NodeValueReplicaDesc, OwnerHoldingGetInfo, + InflightGetInfo, MasterKvRouterView, NodeCacheCapacityReservation, NodeValueReplicaDesc, + OwnerHoldingGetInfo, msg_pack::{ GetAllocationMode, GetDoneReq, GetDoneResp, GetMetaReq, GetMetaResp, GetRevokeReq, - GetRevokeResp, GetStartReq, GetStartResp, + GetRevokeResp, GetSourceKind, GetStartReq, GetStartResp, }, }; +use crate::kv_ssd_storage::{SSD_ALIGNMENT, align_ssd_io_len}; use crate::master_kv_router::OneKvNodesRoutes; +use crate::master_kv_router::delete::remove_one_ssd_replica_for_node; use crate::master_kv_router::put::PutIDForAKey; use crate::memholder::MemholderManagerTrait; use crate::{ cluster_manager::NodeID, master_seg_manager::one_seg_allocator::Allocation, - p2p::msg_pack::MsgPack, rpcresp_kvresult_convert::msg_and_error, + master_seg_manager::one_seg_allocator::OneSegAllocator, p2p::msg_pack::MsgPack, + rpcresp_kvresult_convert::msg_and_error, }; use rand::Rng; use rand::seq::SliceRandom; use std::collections::HashSet; -use std::{ - collections::HashMap, - sync::{Arc, atomic::Ordering}, -}; +use std::future::Future; +use std::sync::{Arc, atomic::Ordering}; +use std::time::{Duration, Instant}; + +const GET_ALLOCATION_WAIT_TIMEOUT: Duration = Duration::from_secs(5); +const GET_ALLOCATION_RETRY_INTERVAL: Duration = Duration::from_millis(2); + +struct AllocationWaitResult { + allocation: Option, + attempts: usize, + elapsed: Duration, +} + +fn try_allocate_from_segments( + allocators: &[Arc], + len: u64, + attempts: &mut usize, +) -> Option { + for allocator in allocators { + *attempts += 1; + if let Ok(allocation) = allocator.allocate(len) { + return Some(allocation); + } + } + None +} + +async fn allocate_with_bounded_wait( + allocators: &[Arc], + len: u64, + timeout: Duration, + retry_interval: Duration, + mut advance_reclamation: F, +) -> AllocationWaitResult +where + F: FnMut() -> Fut, + Fut: Future, +{ + let started = Instant::now(); + let mut attempts = 0usize; + + loop { + if let Some(allocation) = try_allocate_from_segments(allocators, len, &mut attempts) { + return AllocationWaitResult { + allocation: Some(allocation), + attempts, + elapsed: started.elapsed(), + }; + } + + if started.elapsed() >= timeout { + return AllocationWaitResult { + allocation: None, + attempts, + elapsed: started.elapsed(), + }; + } + + // Moka eviction removes the route first and releases the allocation after + // the owner processes its invalidation RPC. Drive pending eviction work, + // then yield so that RPC and holder ACK tasks can make progress. + advance_reclamation().await; + + // In-process maintenance can release an allocation synchronously. Retry + // once before sleeping so this path does not add a fixed 2 ms penalty. + if let Some(allocation) = try_allocate_from_segments(allocators, len, &mut attempts) { + return AllocationWaitResult { + allocation: Some(allocation), + attempts, + elapsed: started.elapsed(), + }; + } + + tokio::time::sleep(retry_interval).await; + } +} + +async fn allocate_get_buffer_on_node( + view: &MasterKvRouterView, + node_id: &NodeID, + len: u64, + get_id: u64, + purpose: &str, +) -> Result, msg_and_error::KvError> { + let mut node_allocators = view.master_seg_manager().get_node_allocators(node_id); + if node_allocators.is_empty() { + tracing::info!( + "No allocators found for {} during get: {}, node is not ready", + purpose, + node_id + ); + return Err(msg_and_error::KvError::Unreachable( + msg_and_error::UnreachableError::OwnerNoSeg { detail: "config=0 initializes as external; non-zero initializes as owner; the owner must have memory space (segment)".to_string() } + )); + } + + node_allocators.shuffle(&mut rand::thread_rng()); + let cache = view + .master_kv_router() + .get_node_cache_controller(node_id.as_ref()); + let result = allocate_with_bounded_wait( + &node_allocators, + len, + GET_ALLOCATION_WAIT_TIMEOUT, + GET_ALLOCATION_RETRY_INTERVAL, + || { + if let Some(cache) = cache.as_ref() { + cache.run_pending_tasks(); + } + let inflight_gets = view.master_kv_router().inner().inflight_gets.clone(); + let inflight_puts = view.master_kv_router().inner().inflight_puts.clone(); + async move { + // Moka removes an in-flight entry from lookup immediately, but + // its maintenance queue can retain the RAII value briefly. Drain + // both queues under allocator pressure so completed PUT/GET + // staging and target allocations are actually released. + inflight_gets.run_pending_tasks().await; + inflight_puts.run_pending_tasks().await; + } + }, + ) + .await; + + if let Some(allocation) = result.allocation { + if result.attempts > node_allocators.len() { + tracing::debug!( + "{} allocation recovered after {} attempts and {:?} for get_id {} on node {}", + purpose, + result.attempts, + result.elapsed, + get_id, + node_id + ); + } + return Ok(Arc::new(allocation)); + } + + let allocator = node_allocators + .iter() + .max_by_key(|allocator| { + allocator + .total_size_bytes() + .saturating_sub(allocator.used_size_bytes()) + }) + .expect("non-empty allocator list must have a diagnostic allocator"); + let total = allocator.total_size_bytes(); + let used = allocator.used_size_bytes(); + let free = total.saturating_sub(used); + tracing::info!( + "{} allocation timed out after {} attempts and {:?} for get_id {} on node {}", + purpose, + result.attempts, + result.elapsed, + get_id, + node_id + ); + Err(msg_and_error::KvError::Api( + msg_and_error::ApiError::NoSpace { + node: node_id.as_ref().to_string(), + segment: allocator.seg_device_id.clone(), + total_capacity: total, + free_capacity: free, + }, + )) +} fn update_moka_for_node( - view: MasterKvRouterView, + view: &MasterKvRouterView, node_id: String, key: String, weight: u32, @@ -31,41 +196,75 @@ fn update_moka_for_node( if !view.master_kv_router().replica_cache_enabled() { return; } - let view_task = view.clone(); - let _ = view.spawn("update_moka_for_node", async move { - if let Some(cache) = view_task - .master_kv_router() - .get_node_cache_controller(&node_id) - { - if new_inserted { - cache.insert( - key.clone(), - NodeValueReplicaDesc { - weight_bytes: weight, - put_id, - }, - ); - tracing::debug!( - "Inserted key: {:?} into node cache: {}, weight={}", - key, - node_id, - weight - ); - } else { - let _ = cache.get(&key); - tracing::debug!( - "Touched key: {:?} on node cache: {} (TTL refresh)", - key, - node_id - ); - } + if let Some(cache) = view.master_kv_router().get_node_cache_controller(&node_id) { + if new_inserted { + cache.insert( + key.clone(), + NodeValueReplicaDesc { + weight_bytes: weight, + put_id, + }, + ); + cache.run_pending_tasks(); + tracing::debug!( + "Inserted key: {:?} into node cache: {}, weight={}", + key, + node_id, + weight + ); } else { - tracing::warn!( - "No cache controller found for node: {} when updating moka", + let _ = cache.get(&key); + tracing::debug!( + "Touched key: {:?} on node cache: {} (TTL refresh)", + key, node_id ); } - }); + } else { + tracing::warn!( + "No cache controller found for node: {} when updating moka", + node_id + ); + } +} + +fn reserve_durable_get_target( + view: &MasterKvRouterView, + route: &Arc, + node_id: &NodeID, + target_allocation: &mut Arc, +) -> (GetAllocationMode, Option>) { + if !route.try_reserve_get_durable_slot() { + return (GetAllocationMode::Temporary, None); + } + + let reservation = match view + .master_kv_router() + .reserve_node_cache_capacity(node_id.as_ref(), target_allocation.capcity()) + { + Ok(reservation) => reservation, + Err(err) => { + route.release_get_durable_slot(); + tracing::warn!( + "Falling back to a temporary GET target because cache capacity reservation failed: node={} bytes={} err={}", + node_id, + target_allocation.capcity(), + err + ); + return (GetAllocationMode::Temporary, None); + } + }; + + if route.lease_id.is_some() { + if let Some(reservation) = reservation { + Arc::get_mut(target_allocation) + .expect("new GET target allocation must be uniquely owned") + .set_on_drop(move || drop(reservation)); + } + (GetAllocationMode::DurableReplica, None) + } else { + (GetAllocationMode::DurableReplica, reservation.map(Arc::new)) + } } pub async fn handle_get_start( @@ -79,21 +278,25 @@ pub async fn handle_get_start( key: &str, ) { if let Some((tombs, put_id)) = tombs_and_put_id { - let mut remove_in_kv_routes = false; - if let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(key) { - one_kv_nodes_routes.clean_up_tomb_nodes_replicas(put_id, tombs, view); - if one_kv_nodes_routes.nodes_replicas.read().is_empty() { - remove_in_kv_routes = true; - } - } + let route_to_remove = view + .master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|route| route.value().clone()) + .filter(|route| { + route.remove_tombed_node_replicas(put_id, tombs) && !route.has_live_replica() + }); - if remove_in_kv_routes { - view.master_kv_router() - .inner() - .kv_routes - .remove_if(key, |_, one_kv_nodes_routes| { - one_kv_nodes_routes.put_id == put_id - }); + if let Some(route_to_remove) = route_to_remove { + view.master_kv_router().inner().kv_routes.remove_if( + key, + |_, one_kv_nodes_routes| { + Arc::ptr_eq(one_kv_nodes_routes, &route_to_remove) + && one_kv_nodes_routes.put_id == put_id + && !one_kv_nodes_routes.has_live_replica() + }, + ); } } } @@ -113,6 +316,16 @@ pub async fn handle_get_start( }, ) } + fn align_ssd_stage_addr(raw_addr: u64) -> Result { + raw_addr + .checked_add(SSD_ALIGNMENT as u64 - 1) + .map(|addr| addr / SSD_ALIGNMENT as u64 * SSD_ALIGNMENT as u64) + .ok_or_else(|| { + msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!("ssd source staging address alignment overflow: {raw_addr}"), + }) + }) + } tracing::debug!("Handling GetStartReq: {:?}", req.serialize_part); @@ -138,76 +351,53 @@ pub async fn handle_get_start( return failed_resp_err(err, None, &view, &req.serialize_part.key); }; - let replicas: HashMap = one_kv_nodes_routes.nodes_replicas.read().clone(); - // Currently we are holding the lock with `replicas` - // 选择一个replica (这里可以实现更复杂的选择逻辑) - let mut replica_keys = replicas.keys().collect::>(); + let replicas = one_kv_nodes_routes.node_replicas.read().clone(); + let mut replica_keys = replicas + .iter() + .filter_map(|(node_id, replicas)| replicas.memory.is_some().then_some(node_id)) + .collect::>(); let mut tombs = HashSet::new(); let mut target_allocations = None; let mut allocation_mode = GetAllocationMode::Temporary; - for _ in 0..replicas.len() { + let mut target_capacity_reservation = None; + while !replica_keys.is_empty() { let to_remove_idx = rand::thread_rng().gen_range(0..replica_keys.len()); let selected_replica_key = replica_keys.remove(to_remove_idx); - let selected_replica = replicas.get(&*selected_replica_key).unwrap(); + let selected_replica = replicas + .get(selected_replica_key) + .expect("selected memory replica node must exist"); if selected_replica.tomb_tag.is_tomb() { tombs.insert(selected_replica_key.to_owned()); continue; } - let src_allocation = selected_replica.allocation.clone(); - let src_node_id = selected_replica.node_id.clone(); + let src_allocation = selected_replica + .memory + .as_ref() + .expect("selected memory replica must exist") + .clone(); + let src_node_id = selected_replica_key.to_owned(); // 为get调用方分配接收内存作为传输target if target_allocations.is_none() { - target_allocations = if let Some(replica_on_recv_node) = replicas.get(&req_node_id) { + target_allocations = if let Some(replica_on_recv_node) = replicas + .get(&req_node_id) + .filter(|replicas| !replicas.tomb_tag.is_tomb()) + .and_then(|replicas| replicas.memory.as_ref()) + { allocation_mode = GetAllocationMode::ReuseReplica; - Some(replica_on_recv_node.allocation.clone()) + Some(replica_on_recv_node.clone()) } else { - let target_allocation = { - let req_node_allocators = - view.master_seg_manager().get_node_allocators(&req_node_id); - if req_node_allocators.is_empty() { - tracing::info!( - "No allocators found for requesting node: {}, node is not ready", - req_node_id - ); - let err = msg_and_error::KvError::Unreachable( - msg_and_error::UnreachableError::OwnerNoSeg { detail: "config=0 initializes as external; non-zero initializes as owner; the owner must have memory space (segment)".to_string() } - ); - return failed_resp_err( - err, - Some((tombs, one_kv_nodes_routes.put_id)), - &view, - &req.serialize_part.key, - ); - } - - let target_allocator = - req_node_allocators.choose(&mut rand::thread_rng()).unwrap(); - - let mut allocated_addr: Option = None; - for attempt in 1..=3 { - if let Ok(allocation) = target_allocator.allocate(src_allocation.size()) { - allocated_addr = Some(allocation); - break; - } else { - tracing::info!( - "Requesting node as target allocation attempt {}/3 failed for get_id {}", - attempt, - get_id - ); - } - } - if allocated_addr.is_none() { - tracing::info!("No space left for target(Requesting node) allocation"); - let total = target_allocator.total_size_bytes(); - let used = target_allocator.used_size_bytes(); - let free = total.saturating_sub(used); - let err = msg_and_error::KvError::Api(msg_and_error::ApiError::NoSpace { - node: req_node_id.as_ref().to_string(), - segment: target_allocator.seg_device_id.clone(), - total_capacity: total, - free_capacity: free, - }); + let mut target_allocation = match allocate_get_buffer_on_node( + &view, + &req_node_id, + src_allocation.size(), + get_id, + "requesting target", + ) + .await + { + Ok(allocation) => allocation, + Err(err) => { return failed_resp_err( err, Some((tombs, one_kv_nodes_routes.put_id)), @@ -215,14 +405,16 @@ pub async fn handle_get_start( &req.serialize_part.key, ); } - allocated_addr.unwrap() }; - if one_kv_nodes_routes.try_reserve_get_durable_slot() { - allocation_mode = GetAllocationMode::DurableReplica; - } else { - allocation_mode = GetAllocationMode::Temporary; - } - Some(Arc::new(target_allocation)) + let (mode, reservation) = reserve_durable_get_target( + &view, + &one_kv_nodes_routes, + &req_node_id, + &mut target_allocation, + ); + allocation_mode = mode; + target_capacity_reservation = reservation; + Some(target_allocation) }; } @@ -253,11 +445,13 @@ pub async fn handle_get_start( put_id: one_kv_nodes_routes.put_id, get_id, node_id: resp_node_id.clone().into(), + source_kind: GetSourceKind::Memory, src_addr: resp_src_addr, target_addr: resp_target_addr, src_base_addr: resp_src_base, target_base_addr: resp_target_base, len: src_allocation.size(), + ssd_stage_len: 0, error_code: msg_and_error::OK, error_json: String::new(), server_process_us: 0, @@ -270,8 +464,12 @@ pub async fn handle_get_start( req_node_id, len: src_allocation.size(), allocation: target_allocation, // 存储target allocation + // Keep the source alive if Moka evicts its route during the transfer. + source_allocation: Some(src_allocation), route: one_kv_nodes_routes.clone(), allocation_mode, + source_kind: GetSourceKind::Memory, + cache_capacity_reservation: target_capacity_reservation, }; view.master_kv_router() @@ -286,7 +484,7 @@ pub async fn handle_get_start( // unnecessary cache work. if one_kv_nodes_routes.lease_id.is_none() { update_moka_for_node( - view.clone(), + &view, src_node_id.to_string(), req.serialize_part.key.clone(), 0, @@ -308,6 +506,222 @@ pub async fn handle_get_start( }, ); } + + let mut ssd_replica_keys = replicas + .iter() + .filter_map(|(node_id, replicas)| replicas.ssd.is_some().then_some(node_id)) + .collect::>(); + while !ssd_replica_keys.is_empty() { + let to_remove_idx = rand::thread_rng().gen_range(0..ssd_replica_keys.len()); + let selected_ssd_key = ssd_replica_keys.remove(to_remove_idx); + let selected_node_replicas = replicas + .get(selected_ssd_key) + .expect("selected SSD replica node must exist"); + if selected_node_replicas.tomb_tag.is_tomb() { + tombs.insert(selected_ssd_key.to_owned()); + } else { + let ssd_replica = selected_node_replicas + .ssd + .as_ref() + .expect("selected SSD replica must exist"); + let ssd_stage_len = match align_ssd_io_len(ssd_replica.len) { + Ok(len) => len, + Err(err) => { + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + let mut target_allocation = match allocate_get_buffer_on_node( + &view, + &req_node_id, + ssd_replica.len, + get_id, + "requesting target", + ) + .await + { + Ok(allocation) => allocation, + Err(err) => { + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + let target_base = target_allocation.base_addr(); + let target_addr = match target_base.checked_add(target_allocation.addr()) { + Some(addr) => addr, + None => { + let err = + msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "requesting target address overflow: base={} offset={}", + target_base, + target_allocation.addr() + ), + }); + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + + let local_ssd_read = selected_ssd_key == &req_node_id; + let (source_allocation, source_base, source_addr, response_ssd_stage_len) = + if local_ssd_read { + let target_capacity = target_allocation.capcity(); + tracing::debug!( + "Using local SSD read for get_id {} on node {}: target={:#x} len={} capacity={}", + get_id, + selected_ssd_key, + target_addr, + ssd_replica.len, + target_capacity + ); + (None, target_base, target_addr, target_capacity) + } else { + let source_alloc_len = match ssd_stage_len.checked_add(SSD_ALIGNMENT as u64 - 1) + { + Some(len) => len, + None => { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidArgument { + detail: format!( + "ssd source staging allocation length overflow: {ssd_stage_len}" + ), + }, + ); + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + let source_allocation = match allocate_get_buffer_on_node( + &view, + selected_ssd_key, + source_alloc_len, + get_id, + "ssd source staging", + ) + .await + { + Ok(allocation) => allocation, + Err(err) => { + tracing::info!( + "Skipping SSD source for get_id {} on node {}: {}", + get_id, + selected_ssd_key, + err + ); + continue; + } + }; + let source_base = source_allocation.base_addr(); + let source_raw_addr = match source_base.checked_add(source_allocation.addr()) { + Some(addr) => addr, + None => { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidArgument { + detail: format!( + "ssd source staging raw address overflow: base={} offset={}", + source_base, + source_allocation.addr() + ), + }, + ); + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + let source_addr = match align_ssd_stage_addr(source_raw_addr) { + Ok(addr) => addr, + Err(err) => { + return failed_resp_err( + err, + Some((tombs, one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + ( + Some(source_allocation), + source_base, + source_addr, + ssd_stage_len, + ) + }; + let (allocation_mode, cache_capacity_reservation) = reserve_durable_get_target( + &view, + &one_kv_nodes_routes, + &req_node_id, + &mut target_allocation, + ); + let resp = GetStartResp { + put_id: one_kv_nodes_routes.put_id, + get_id, + node_id: selected_ssd_key.to_owned().into(), + source_kind: GetSourceKind::Ssd, + src_addr: source_addr, + target_addr, + src_base_addr: source_base, + target_base_addr: target_base, + len: ssd_replica.len, + ssd_stage_len: response_ssd_stage_len, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }; + let info = InflightGetInfo { + put_id: one_kv_nodes_routes.put_id, + src_node_id: selected_ssd_key.to_owned(), + key: req.serialize_part.key.clone(), + req_node_id, + len: ssd_replica.len, + allocation: target_allocation, + source_allocation, + route: one_kv_nodes_routes.clone(), + allocation_mode, + source_kind: GetSourceKind::Ssd, + cache_capacity_reservation, + }; + + view.master_kv_router() + .inner() + .inflight_gets + .insert(get_id, info) + .await; + + clean_up_tombs( + &view, + Some((tombs, one_kv_nodes_routes.put_id)), + &req.serialize_part.key, + ); + return ( + get_id, + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + } tracing::info!("Key not found: {}", req.serialize_part.key); { let err = msg_and_error::KvError::Api(msg_and_error::ApiError::KeyNotFound { @@ -322,6 +736,49 @@ pub async fn handle_get_start( } } +fn drop_failed_ssd_source(view: &MasterKvRouterView, inflight_info: &InflightGetInfo) { + if inflight_info.source_kind != GetSourceKind::Ssd { + tracing::warn!( + "Ignoring drop_ssd_source for non-SSD get: get_key={} put_id=({},{}) source_kind={:?}", + inflight_info.key, + inflight_info.put_id.0, + inflight_info.put_id.1, + inflight_info.source_kind + ); + return; + } + + let removal = remove_one_ssd_replica_for_node( + view, + &inflight_info.key, + &inflight_info.src_node_id, + inflight_info.put_id, + ); + if !removal.replica_removed { + return; + } + + tracing::warn!( + "Removed failed SSD replica: key={} node={} put_id=({},{})", + inflight_info.key, + inflight_info.src_node_id, + inflight_info.put_id.0, + inflight_info.put_id.1 + ); + + if removal.route_removed && view.master_kv_router().prefix_index_enabled() { + let view_task = view.clone(); + let key_for_prefix = inflight_info.key.clone(); + let _ = view.spawn("ssd_failure_remove_prefix_index", async move { + let inner = view_task.master_kv_router().inner(); + let mut tree = inner.prefix_index.write().await; + if !inner.kv_routes.contains_key(&key_for_prefix) { + tree.remove(&key_for_prefix); + } + }); + } +} + pub async fn handle_get_revoke( view: MasterKvRouterView, req: MsgPack, @@ -338,6 +795,9 @@ pub async fn handle_get_revoke( .remove(&get_id) .await { + if req.serialize_part.drop_ssd_source { + drop_failed_ssd_source(&view, &inflight_info); + } inflight_info.release_durable_slot_if_needed(); tracing::info!("Revoked get operation with get_id: {}", get_id); } else { @@ -361,13 +821,14 @@ pub async fn handle_get_done( let get_id = req.serialize_part.get_id; // Remove from inflight_gets and transfer to get_holding - if let Some(inflight_info) = view + if let Some(mut inflight_info) = view .master_kv_router() .inner() .inflight_gets .remove(&get_id) .await { + let mut cache_capacity_reservation = inflight_info.cache_capacity_reservation.take(); let mut allocation_mode = inflight_info.allocation_mode; let route = inflight_info.route.clone(); // clone req_node_id to avoid borrow/move conflict when inserting into kv_routes @@ -381,7 +842,6 @@ pub async fn handle_get_done( .next_holder_id .fetch_add(1, Ordering::Relaxed); - let src_node_id = inflight_info.src_node_id; let key = inflight_info.key; // Create holding info @@ -400,20 +860,22 @@ pub async fn handle_get_done( if allocation_mode == GetAllocationMode::DurableReplica { let mut promote_committed = false; - if let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(&key) { + let current_route = view + .master_kv_router() + .inner() + .kv_routes + .get(&key) + .map(|route| route.value().clone()); + if let Some(one_kv_nodes_routes) = current_route { if one_kv_nodes_routes.put_id == inflight_info.put_id { - let mut nodes_replicas = one_kv_nodes_routes.nodes_replicas.write(); if let Some(tomb_tag) = - view.master_seg_manager().get_node_tomb_tag(&src_node_id) + view.master_seg_manager().get_node_tomb_tag(&req_node_id) { if !tomb_tag.is_tomb() { - nodes_replicas.insert( + one_kv_nodes_routes.insert_memory_replica( inflight_info.req_node_id.clone(), - KvRouteInfo { - node_id: inflight_info.req_node_id, - allocation: inflight_info.allocation, - tomb_tag, - }, + inflight_info.allocation, + tomb_tag, ); promote_committed = true; // Read lease binding from route snapshot: for this put_id, @@ -433,8 +895,11 @@ pub async fn handle_get_done( } else { alloc_cap as u32 }; + // Move the target from the in-flight reservation into + // Moka in the same handler turn. + drop(cache_capacity_reservation.take()); update_moka_for_node( - view.clone(), + &view, req_node_id.to_string(), key.clone(), req_weight, @@ -613,19 +1078,36 @@ pub async fn handle_get_meta( .get(&req.serialize_part.key) { // lock and clone, release the lock quickly - let nodes_replicas: HashMap = - (*one_kv_nodes_routes.nodes_replicas.read()).clone(); + let node_replicas = one_kv_nodes_routes.node_replicas.read().clone(); - // Key exists, get metadata from the first replica - for (_, kv_info) in nodes_replicas.iter() { - if kv_info.tomb_tag.is_tomb() { + for replicas in node_replicas.values() { + if replicas.tomb_tag.is_tomb() { continue; } - let len = kv_info.allocation.size(); + let Some(allocation) = replicas.memory.as_ref() else { + continue; + }; return MsgPack { serialize_part: GetMetaResp { exists: true, - len, + len: allocation.size(), + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + }; + } + for replicas in node_replicas.values() { + if replicas.tomb_tag.is_tomb() { + continue; + } + let Some(ssd) = replicas.ssd.as_ref() else { + continue; + }; + return MsgPack { + serialize_part: GetMetaResp { + exists: true, + len: ssd.len, error_code: msg_and_error::OK, error_json: String::new(), }, @@ -684,3 +1166,102 @@ pub async fn handle_get_meta( } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::master_seg_manager::msg_pack::SegmentDeviceDescription; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[::tokio::test] + async fn allocation_waits_for_delayed_release() { + let allocator = Arc::new( + OneSegAllocator::new( + "get-wait-test".to_string(), + SegmentDeviceDescription::Cpu, + 0, + 4096, + ) + .expect("test allocator must be created"), + ); + let held = allocator + .allocate(4096) + .expect("test allocator must initially have capacity"); + let reclamation_attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_wait = Arc::clone(&reclamation_attempts); + let allocators = [Arc::clone(&allocator)]; + + let wait_for_allocation = allocate_with_bounded_wait( + &allocators, + 4096, + Duration::from_millis(500), + Duration::from_millis(1), + move || { + attempts_for_wait.fetch_add(1, Ordering::Relaxed); + async {} + }, + ); + let release_capacity = async move { + tokio::time::sleep(Duration::from_millis(20)).await; + drop(held); + }; + + let (result, ()) = tokio::join!(wait_for_allocation, release_capacity); + assert!(result.allocation.is_some()); + assert!(result.attempts > 1); + assert!(reclamation_attempts.load(Ordering::Relaxed) > 0); + } + + #[::tokio::test] + async fn allocation_wait_drains_removed_moka_value() { + let allocator = Arc::new( + OneSegAllocator::new( + "get-moka-drain-test".to_string(), + SegmentDeviceDescription::Cpu, + 0, + 4096, + ) + .expect("test allocator must be created"), + ); + let allocation = Arc::new( + allocator + .allocate(4096) + .expect("test allocator must initially have capacity"), + ); + let allocation_weak = Arc::downgrade(&allocation); + let cache = moka::future::Cache::builder() + .eviction_listener(|_key: Arc, _value: Arc, _cause| {}) + .build(); + cache.insert(1, allocation).await; + cache.run_pending_tasks().await; + + let removed = cache + .remove(&1) + .await + .expect("the test allocation must still be present"); + drop(removed); + assert!( + allocation_weak.upgrade().is_some(), + "Moka's pending remove should still retain the allocation before maintenance" + ); + + let allocators = [Arc::clone(&allocator)]; + let result = allocate_with_bounded_wait( + &allocators, + 4096, + Duration::from_millis(500), + Duration::from_millis(1), + || { + let cache = cache.clone(); + async move { + cache.run_pending_tasks().await; + } + }, + ) + .await; + + assert!(result.allocation.is_some()); + assert!(result.attempts > 1); + assert!(allocation_weak.upgrade().is_none()); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs index ee4ca2b..fe43d79 100644 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs @@ -8,20 +8,22 @@ mod count_prefix_index; use self::{ count_prefix_index::PrefixRadixTree, - delete::handle_batch_delete_ack, delete::handle_delete, delete::handle_delete_ack, + delete::{handle_batch_delete_ack, handle_batch_ssd_replica_evict}, get::{handle_get_done, handle_get_meta, handle_get_revoke, handle_get_start}, msg_pack::{ - BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, CountPrefixReq, CountPrefixResp, - DeleteAckReq, DeleteReq, GetAllocationMode, GetDoneReq, GetMetaReq, GetRevokeReq, - GetStartReq, PutDoneReq, PutRevokeReq, PutStartReq, + BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, BatchSsdReplicaEvictReq, + CountPrefixReq, CountPrefixResp, DeleteAckReq, DeleteReq, GetAllocationMode, GetDoneReq, + GetMetaReq, GetRevokeReq, GetSourceKind, GetStartReq, PutDoneReq, PutRevokeReq, + PutStartReq, SsdReplicaCommitReq, }, placement::{PlacementDefault, PlacementPolicy}, - put::{handle_put_done, handle_put_revoke, handle_put_start}, + put::{handle_put_done, handle_put_revoke, handle_put_start, handle_ssd_replica_commit}, }; use crate::ClientKvApiAccessTrait; use crate::client_kv_api::ClientKvApi; +use crate::client_kv_api::msg_pack::SsdReplicaPersistReq; use crate::cluster_manager::{ ClusterEvent, ClusterManager, ClusterManagerAccessTrait, NodeID, NodeIDString, }; @@ -42,7 +44,7 @@ use fluxon_framework::{LogicalModule, define_module}; use async_trait::async_trait; use chrono::Utc; -use dashmap::DashMap; +use dashmap::{DashMap, mapref::entry::Entry}; use limit_thirdparty::tokio::sync::ARwLock; use limit_thirdparty::tokio::{self, sync::ampsc}; use moka::notification::RemovalCause; @@ -104,6 +106,7 @@ pub struct InflightPutInfo { pub key: String, pub req_node_id: NodeID, pub len: u64, + pub persist_to_ssd: bool, pub src_target_allocation: Arc>>, } @@ -116,8 +119,11 @@ pub struct InflightGetInfo { pub req_node_id: NodeID, pub len: u64, pub allocation: Arc, + pub source_allocation: Option>, pub route: Arc, pub allocation_mode: GetAllocationMode, + pub source_kind: GetSourceKind, + pub(crate) cache_capacity_reservation: Option>, } impl InflightGetInfo { @@ -192,13 +198,90 @@ pub struct NodeValueReplicaDesc { pub put_id: PutIDForAKey, } -/// Information about a completed `put` operation that can be retrieved via `get`. -/// Now supports multiple replicas per key. +type NodeReplicaCache = moka::sync::SegmentedCache; + +pub(crate) struct NodeCacheCapacityReservation { + node_id: NodeIDString, + bytes: u64, + base_capacity: u64, + reserved_bytes: Arc>, + cache: Arc, +} + +impl NodeCacheCapacityReservation { + fn reserve( + node_id: NodeIDString, + bytes: u64, + base_capacity: u64, + reserved_bytes: Arc>, + cache: Arc, + ) -> Result { + let mut reserved = reserved_bytes.lock(); + let new_reserved = reserved.checked_add(bytes).ok_or_else(|| { + format!( + "cache capacity reservation overflow: node_id={} reserved={} bytes={}", + node_id, *reserved, bytes + ) + })?; + let new_capacity = base_capacity.saturating_sub(new_reserved); + cache.set_max_capacity(new_capacity).map_err(|err| { + format!( + "moka.set_max_capacity failed while reserving: node_id={} new_capacity={} err={}", + node_id, new_capacity, err + ) + })?; + *reserved = new_reserved; + drop(reserved); + + Ok(Self { + node_id, + bytes, + base_capacity, + reserved_bytes, + cache, + }) + } +} + +impl Drop for NodeCacheCapacityReservation { + fn drop(&mut self) { + let mut reserved = self.reserved_bytes.lock(); + let new_reserved = reserved.checked_sub(self.bytes).unwrap_or_else(|| { + panic!( + "cache capacity reservation underflow: node_id={} reserved={} bytes={}", + self.node_id, *reserved, self.bytes + ) + }); + let new_capacity = self.base_capacity.saturating_sub(new_reserved); + self.cache + .set_max_capacity(new_capacity) + .unwrap_or_else(|err| { + panic!( + "moka.set_max_capacity failed while releasing reservation: node_id={} new_capacity={} err={}", + self.node_id, new_capacity, err + ) + }); + *reserved = new_reserved; + } +} + #[derive(Clone, Debug)] -pub struct KvRouteInfo { - pub node_id: NodeID, - pub allocation: Arc, +pub struct KvSsdReplicaInfo { + pub len: u64, +} + +#[derive(Clone, Debug)] +pub struct KvNodeReplicas { pub tomb_tag: NodeTombTag, + pub memory: Option>, + pub ssd: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SsdReplicaCommitStatus { + Committed, + MissingMemory, + TombedNode, } #[derive(Debug)] @@ -228,26 +311,114 @@ pub struct OneKvNodesRoutes { /// nodes must not insert leased keys into moka caches. pub lease_id: Option, - /// node_id -> KvRouteInfo - pub nodes_replicas: RwLock>, + /// Per-node memory and SSD replicas for this key-version. + pub node_replicas: RwLock>, pub get_durable_slots_used: AtomicU32, } impl OneKvNodesRoutes { - fn clean_up_tomb_nodes_replicas( + fn new(put_id: PutIDForAKey, lease_id: Option) -> Self { + Self { + put_id, + lease_id, + node_replicas: RwLock::new(HashMap::new()), + get_durable_slots_used: AtomicU32::new(0), + } + } + + fn remove_tombed_node_replicas( &self, verify_put_id: PutIDForAKey, tombs: HashSet, - _view: &MasterKvRouterView, ) -> bool { if self.put_id != verify_put_id { return false; } - let mut nodes_replicas = self.nodes_replicas.write(); - nodes_replicas.retain(|_, kv_info| !tombs.contains(&kv_info.node_id)); + self.node_replicas + .write() + .retain(|node_id, replicas| !tombs.contains(node_id) || !replicas.tomb_tag.is_tomb()); + + true + } + + fn has_live_replica(&self) -> bool { + self.node_replicas.read().values().any(|replicas| { + !replicas.tomb_tag.is_tomb() && (replicas.memory.is_some() || replicas.ssd.is_some()) + }) + } + + fn has_memory_replica(&self, node_id: &NodeID) -> bool { + self.node_replicas + .read() + .get(node_id) + .is_some_and(|replicas| !replicas.tomb_tag.is_tomb() && replicas.memory.is_some()) + } - return true; + fn insert_memory_replica( + &self, + node_id: NodeID, + allocation: Arc, + tomb_tag: NodeTombTag, + ) { + let mut node_replicas = self.node_replicas.write(); + if let Some(replicas) = node_replicas.get_mut(&node_id) { + if !replicas.tomb_tag.is_tomb() { + replicas.memory = Some(allocation); + return; + } + } + + node_replicas.insert( + node_id, + KvNodeReplicas { + tomb_tag, + memory: Some(allocation), + ssd: None, + }, + ); + } + + fn commit_ssd_replica(&self, node_id: &NodeID, len: u64) -> SsdReplicaCommitStatus { + let mut node_replicas = self.node_replicas.write(); + let Some(replicas) = node_replicas.get_mut(node_id) else { + return SsdReplicaCommitStatus::MissingMemory; + }; + if replicas.memory.is_none() { + return SsdReplicaCommitStatus::MissingMemory; + } + if replicas.tomb_tag.is_tomb() { + return SsdReplicaCommitStatus::TombedNode; + } + + replicas.ssd = Some(KvSsdReplicaInfo { len }); + SsdReplicaCommitStatus::Committed + } + + fn remove_memory_replica(&self, node_id: &NodeID) -> bool { + let mut node_replicas = self.node_replicas.write(); + let Some(replicas) = node_replicas.get_mut(node_id) else { + return false; + }; + let removed = replicas.memory.take().is_some(); + let remove_node = replicas.ssd.is_none(); + if remove_node { + node_replicas.remove(node_id); + } + removed + } + + fn remove_ssd_replica(&self, node_id: &NodeID) -> bool { + let mut node_replicas = self.node_replicas.write(); + let Some(replicas) = node_replicas.get_mut(node_id) else { + return false; + }; + let removed = replicas.ssd.take().is_some(); + let remove_node = replicas.memory.is_none(); + if remove_node { + node_replicas.remove(node_id); + } + removed } fn try_reserve_get_durable_slot(&self) -> bool { @@ -275,16 +446,13 @@ impl OneKvNodesRoutes { mod tests { use super::*; use crate::cluster_manager::ClusterMember; + use crate::master_seg_manager::msg_pack::SegmentDeviceDescription; + use crate::master_seg_manager::one_seg_allocator::OneSegAllocator; use std::collections::HashMap; #[test] fn one_kv_nodes_routes_only_reserves_two_get_durable_slots() { - let routes = OneKvNodesRoutes { - put_id: (1, 0), - lease_id: None, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), - }; + let routes = OneKvNodesRoutes::new((1, 0), None); assert!(routes.try_reserve_get_durable_slot()); assert!(routes.try_reserve_get_durable_slot()); @@ -294,6 +462,82 @@ mod tests { assert!(routes.try_reserve_get_durable_slot()); } + #[test] + fn one_kv_nodes_routes_updates_memory_and_ssd_independently() { + let routes = OneKvNodesRoutes::new((1, 0), None); + let node_id: NodeID = "node-a".to_string().into(); + let allocator = Arc::new( + OneSegAllocator::new( + "route-test".to_string(), + SegmentDeviceDescription::Cpu, + 0, + 4096, + ) + .expect("test allocator must be created"), + ); + let allocation = Arc::new( + allocator + .allocate(512) + .expect("test allocation must be created"), + ); + + assert_eq!( + routes.commit_ssd_replica(&node_id, 512), + SsdReplicaCommitStatus::MissingMemory + ); + routes.insert_memory_replica(node_id.clone(), allocation, NodeTombTag::new()); + assert_eq!( + routes.commit_ssd_replica(&node_id, 512), + SsdReplicaCommitStatus::Committed + ); + + assert!(routes.remove_memory_replica(&node_id)); + { + let node_replicas = routes.node_replicas.read(); + let replicas = node_replicas + .get(&node_id) + .expect("SSD replica must keep the node entry alive"); + assert!(replicas.memory.is_none()); + assert_eq!(replicas.ssd.as_ref().map(|ssd| ssd.len), Some(512)); + } + + assert!(routes.remove_ssd_replica(&node_id)); + assert!(routes.node_replicas.read().is_empty()); + + let old_tomb_tag = NodeTombTag::new(); + let old_allocation = Arc::new( + allocator + .allocate(512) + .expect("old-incarnation test allocation must be created"), + ); + routes.insert_memory_replica(node_id.clone(), old_allocation, old_tomb_tag.clone()); + assert_eq!( + routes.commit_ssd_replica(&node_id, 512), + SsdReplicaCommitStatus::Committed + ); + old_tomb_tag.set_tomb(); + + let new_allocation = Arc::new( + allocator + .allocate(512) + .expect("new-incarnation test allocation must be created"), + ); + routes.insert_memory_replica(node_id.clone(), new_allocation, NodeTombTag::new()); + assert!( + routes.remove_tombed_node_replicas(routes.put_id, HashSet::from([node_id.clone()]),) + ); + assert!(routes.has_memory_replica(&node_id)); + assert!( + routes + .node_replicas + .read() + .get(&node_id) + .expect("new node incarnation must remain indexed") + .ssd + .is_none() + ); + } + fn new_test_member(metadata: HashMap) -> ClusterMember { ClusterMember { id: "node-a".to_string(), @@ -346,8 +590,8 @@ pub struct MasterKvRouterInner { /// (key, put_time_ms, put_version) -> inflight_put_info pub inflight_puts: moka::future::Cache<(String, u64, u32), InflightPutInfo>, - /// key -> inflight put count - pub inflight_put_key_counts: Arc>, + /// key -> inflight put admission state + pub(crate) inflight_put_key_counts: Arc>, pub inflight_gets: moka::future::Cache, /// Cache for holding get operations (owned, flattened by (node_id, holder_id)) @@ -371,12 +615,10 @@ pub struct MasterKvRouterInner { pub prefix_index: ARwLock, /// Support replicas: node_id -> key -> route_info - pub node_kv_cache_controller: - DashMap>>, + pub node_kv_cache_controller: DashMap>, - /// Per-node total bytes reserved for leased replicas. We subtract this from - /// the base max capacity of each node's moka cache. Acts like a fetch_sub/add counter. - pub lease_reserved_bytes: DashMap>, + /// Per-node bytes held outside Moka while still consuming allocator space. + pub cache_reserved_bytes: DashMap>>, /// Historical final put placement decisions by target node. pub put_target_decision_counts: DashMap>, @@ -393,6 +635,12 @@ pub struct MasterKvRouterInner { pub delete_broadcast: EnsureMemholderMgmtDeleteHandle, } +#[derive(Debug)] +pub(crate) struct InflightPutKeyAdmission { + inflight_count: u32, + create_only: bool, +} + impl MasterKvRouterInner { fn view(&self) -> &MasterKvRouterView { self.view.get().unwrap() @@ -476,7 +724,8 @@ impl MasterKvRouter { } else { INFLIGHT_PUT_TTL_SECONDS }; - let inflight_put_key_counts: Arc> = Arc::new(DashMap::new()); + let inflight_put_key_counts: Arc> = + Arc::new(DashMap::new()); let inflight_put_key_counts_for_listener = inflight_put_key_counts.clone(); let inflight_puts = moka::future::Cache::builder() .time_to_live(Duration::from_secs(inflight_put_ttl_seconds)) @@ -510,7 +759,7 @@ impl MasterKvRouter { kv_routes: DashMap::new(), prefix_index: ARwLock::new(PrefixRadixTree::new()), node_kv_cache_controller: DashMap::new(), - lease_reserved_bytes: DashMap::new(), + cache_reserved_bytes: DashMap::new(), put_target_decision_counts: DashMap::new(), put_requester_target_decision_counts: DashMap::new(), put_placement_mode_counts: DashMap::new(), @@ -572,13 +821,24 @@ impl MasterKvRouter { (put_time_ms, put_version) } - fn release_inflight_put_key_count_map(counts: &DashMap, key: &str) { - if let Some(mut entry) = counts.get_mut(key) { - if *entry <= 1 { - drop(entry); - counts.remove(key); + fn release_inflight_put_key_count_map( + counts: &DashMap, + key: &str, + ) { + if let Entry::Occupied(mut entry) = counts.entry(key.to_string()) { + let admission = entry.get_mut(); + assert!( + admission.inflight_count > 0, + "inflight put admission count must remain positive" + ); + if admission.inflight_count == 1 { + entry.remove(); } else { - *entry -= 1; + assert!( + !admission.create_only, + "create-only put admission must remain exclusive" + ); + admission.inflight_count -= 1; } } } @@ -587,17 +847,45 @@ impl MasterKvRouter { &self, key: &str, reject_if_inflight_same_key: bool, + reject_if_exists: bool, ) -> Result<(), KvError> { let counts = &self.inner().inflight_put_key_counts; - let mut entry = counts.entry(key.to_string()).or_insert(0); - if reject_if_inflight_same_key && *entry > 0 { - return Err(KvError::Api( - crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { - key: key.to_string(), - }, - )); + match counts.entry(key.to_string()) { + Entry::Occupied(mut entry) => { + if reject_if_exists { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyAlreadyExists { + key: key.to_string(), + }, + )); + } + if reject_if_inflight_same_key || entry.get().create_only { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { + key: key.to_string(), + }, + )); + } + let admission = entry.get_mut(); + admission.inflight_count = admission + .inflight_count + .checked_add(1) + .expect("inflight put admission count overflow"); + } + Entry::Vacant(entry) => { + if reject_if_exists && self.inner().kv_routes.contains_key(key) { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyAlreadyExists { + key: key.to_string(), + }, + )); + } + entry.insert(InflightPutKeyAdmission { + inflight_count: 1, + create_only: reject_if_exists, + }); + } } - *entry += 1; Ok(()) } @@ -607,6 +895,7 @@ impl MasterKvRouter { fn register_rpc_callers(&self) { RPCCaller::::new().regist(self.0.view().p2p_module()); + RPCCaller::::new().regist(self.0.view().p2p_module()); } fn register_rpc_handlers(&self) { @@ -766,6 +1055,36 @@ impl MasterKvRouter { Ok(()) }); + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_ssd_replica_commit", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_ssd_replica_commit(view_task, msg).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send SsdReplicaCommitResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_batch_ssd_replica_evict", async move { + let ack = handle_batch_ssd_replica_evict(view_task, msg, req_node_id).await; + if let Err(err) = resp.send_resp(ack).await { + error!("Failed to send BatchSsdReplicaEvictResp: {:?}", err); + } + }); + Ok(()) + }); + // --- MemHolder Handlers --- // let view = inner.view.clone(); // RPCHandler::::new().regist(p2p, move |resp, msg| { @@ -1376,10 +1695,7 @@ impl MasterKvRouter { }); } - pub fn get_node_cache_controller( - &self, - node_id: &str, - ) -> Option>> { + pub fn get_node_cache_controller(&self, node_id: &str) -> Option> { if !self.replica_cache_enabled() { return None; } @@ -1451,55 +1767,27 @@ impl MasterKvRouter { ) } - /// Atomically adjust a node's cache capacity reservation by `delta_bytes`. - /// Positive delta reserves capacity (fetch_sub from usable capacity), - /// negative delta releases reservation (fetch_add back to usable capacity). - pub fn adjust_node_cache_capacity_for_lease( + pub(crate) fn reserve_node_cache_capacity( &self, node_id: &str, - delta_bytes: i64, - ) -> crate::rpcresp_kvresult_convert::msg_and_error::KvResult<()> { + bytes: u64, + ) -> crate::rpcresp_kvresult_convert::msg_and_error::KvResult< + Option, + > { if !self.replica_cache_enabled() { - return Ok(()); - } - // Track per-node reserved bytes with an atomic counter - let reserved_counter = self - .inner() - .lease_reserved_bytes - .entry(node_id.to_string()) - .or_insert_with(|| Arc::new(AtomicU64::new(0))) - .value() - .clone(); - - // Apply delta to the counter with simple fetch ops per user's preference - if delta_bytes >= 0 { - reserved_counter.fetch_add(delta_bytes as u64, Ordering::Relaxed); - } else { - let sub = (-delta_bytes) as u64; - reserved_counter.fetch_sub(sub, Ordering::Relaxed); + return Ok(None); } - - // Recompute target capacity: base(=MOKA_CACHE_CAPACITY_RATIO * node_space) - reserved_total - let reserved_total = reserved_counter.load(Ordering::Relaxed); let node_space_size = self .inner() .view() .master_seg_manager() .get_node_space_size(node_id); if node_space_size == 0 { - // Node not ready: this should not happen in a successful put_done path. - // Revert the counter delta before returning error. - if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); - } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); - } return Err( crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { detail: format!( - "node_id={} has no segment (node_space_size=0) while adjusting cache capacity", + "node_id={} has no segment while reserving cache capacity", node_id ), }, @@ -1507,62 +1795,158 @@ impl MasterKvRouter { ); } let base_capacity = (node_space_size as f32 * MOKA_CACHE_CAPACITY_RATIO) as u64; - let new_capacity = base_capacity.saturating_sub(reserved_total); + let cache = self.get_node_cache_controller(node_id).ok_or_else(|| { + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { + detail: format!("node_id={} cache_controller not found", node_id), + }, + ) + })?; + let reserved_bytes = self + .inner() + .cache_reserved_bytes + .entry(node_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(0))) + .value() + .clone(); - if let Some(cache) = self.get_node_cache_controller(node_id) { - if let Err(e) = cache.set_max_capacity(new_capacity) { - // Revert counter and return error. - if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); - } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); - } - return Err(crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { - rpc_input_json: format!( - "moka.set_max_capacity failed: node_id={}, new_capacity={}, err={}", - node_id, new_capacity, e - ), - } - )); - } - Ok(()) - } else { - // Revert counter and return error. - if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); - } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); - } - Err( - crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { - detail: format!("node_id={} cache_controller not found", node_id), - }, - ), + NodeCacheCapacityReservation::reserve( + node_id.to_string(), + bytes, + base_capacity, + reserved_bytes, + cache, + ) + .map(Some) + .map_err(|detail| { + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { + detail, + }, ) - } + }) } - - // Note: no additional getters for reserved bytes; policy currently relies only on adjust calls. } +// moved to crate::metrics::client -impl MasterKvRouterView { - pub fn try_adjust_node_cache_capacity_for_lease( - &self, - node_id: &str, - delta_bytes: i64, - ) -> Option> { - let _view_guard = self.try_upgrade()?; - Some( - self.master_kv_router() - .adjust_node_cache_capacity_for_lease(node_id, delta_bytes), +#[cfg(test)] +mod cache_capacity_reservation_tests { + use super::*; + use crate::master_seg_manager::msg_pack::SegmentDeviceDescription; + use crate::master_seg_manager::one_seg_allocator::OneSegAllocator; + + fn new_cache(capacity: u64) -> Arc { + Arc::new( + moka::sync::SegmentedCache::builder(1) + .max_capacity(capacity) + .weigher(Box::new(|_key: &String, value: &NodeValueReplicaDesc| { + value.weight_bytes + })) + .build(), ) } + + #[test] + fn pending_persist_reservation_reduces_capacity_until_guard_drops() { + let cache = new_cache(800); + let reserved_bytes = Arc::new(Mutex::new(0)); + let first = NodeCacheCapacityReservation::reserve( + "node-a".to_string(), + 120, + 800, + reserved_bytes.clone(), + cache.clone(), + ) + .unwrap(); + let second = NodeCacheCapacityReservation::reserve( + "node-a".to_string(), + 80, + 800, + reserved_bytes.clone(), + cache.clone(), + ) + .unwrap(); + + assert_eq!(*reserved_bytes.lock(), 200); + assert_eq!(cache.policy().max_capacity(), Some(600)); + + drop(first); + assert_eq!(*reserved_bytes.lock(), 80); + assert_eq!(cache.policy().max_capacity(), Some(720)); + + drop(second); + assert_eq!(*reserved_bytes.lock(), 0); + assert_eq!(cache.policy().max_capacity(), Some(800)); + } + + #[test] + fn inflight_get_keeps_capacity_reserved_until_handoff_drops() { + let cache = new_cache(800); + let reserved_bytes = Arc::new(Mutex::new(0)); + let reservation = Arc::new( + NodeCacheCapacityReservation::reserve( + "node-a".to_string(), + 120, + 800, + reserved_bytes.clone(), + cache.clone(), + ) + .unwrap(), + ); + let allocator = Arc::new( + OneSegAllocator::new( + "inflight-get-reservation-test".to_string(), + SegmentDeviceDescription::Cpu, + 0, + 4096, + ) + .expect("test allocator must be created"), + ); + let allocation = Arc::new( + allocator + .allocate(512) + .expect("test allocation must be created"), + ); + let route = Arc::new(OneKvNodesRoutes::new((1, 0), None)); + let mut inflight = InflightGetInfo { + put_id: (1, 0), + src_node_id: "node-b".to_string().into(), + key: "key-a".to_string(), + req_node_id: "node-a".to_string().into(), + len: 512, + allocation, + source_allocation: None, + route, + allocation_mode: GetAllocationMode::DurableReplica, + source_kind: GetSourceKind::Memory, + cache_capacity_reservation: Some(reservation), + }; + + assert_eq!(*reserved_bytes.lock(), 120); + assert_eq!(cache.policy().max_capacity(), Some(680)); + + let handoff = inflight.cache_capacity_reservation.take(); + drop(inflight); + assert_eq!(*reserved_bytes.lock(), 120); + + drop(handoff); + assert_eq!(*reserved_bytes.lock(), 0); + assert_eq!(cache.policy().max_capacity(), Some(800)); + } + + #[test] + #[should_panic(expected = "cache capacity reservation underflow")] + fn reservation_release_fails_fast_on_counter_underflow() { + let reservation = NodeCacheCapacityReservation { + node_id: "node-a".to_string(), + bytes: 1, + base_capacity: 800, + reserved_bytes: Arc::new(Mutex::new(0)), + cache: new_cache(800), + }; + drop(reservation); + } } -// moved to crate::metrics::client #[cfg(test)] mod placement_metrics_tests { @@ -1663,3 +2047,84 @@ mod placement_metrics_tests { ); } } + +#[cfg(test)] +mod put_admission_tests { + use super::{MasterKvRouter, MasterKvRouterNewArg, OneKvNodesRoutes}; + use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError}; + use std::sync::Arc; + + #[::tokio::test] + async fn reject_if_exists_covers_inflight_and_committed_keys() { + let router = MasterKvRouter::construct(MasterKvRouterNewArg::default()) + .await + .expect("router construction must succeed"); + let key = "create-only-key"; + + router + .reserve_inflight_put_key(key, false, true) + .expect("first create-only reservation must succeed"); + let inflight_err = router + .reserve_inflight_put_key(key, false, true) + .expect_err("concurrent create-only reservation must be rejected"); + assert!(matches!( + inflight_err, + KvError::Api(ApiError::KeyAlreadyExists { key: ref rejected_key }) + if rejected_key == key + )); + let overwrite_err = router + .reserve_inflight_put_key(key, false, false) + .expect_err("create-only reservation must exclude ordinary overwrites"); + assert!(matches!( + overwrite_err, + KvError::Api(ApiError::KeyBeingWritten { key: ref rejected_key }) + if rejected_key == key + )); + + // Commit visibility must precede reservation release. A following create-only put + // then observes either state and cannot enter through a transition gap. + router.inner().kv_routes.insert( + key.to_string(), + Arc::new(OneKvNodesRoutes::new((1, 0), None)), + ); + router.release_inflight_put_key(key); + + let committed_err = router + .reserve_inflight_put_key(key, false, true) + .expect_err("committed key must be rejected"); + assert!(matches!( + committed_err, + KvError::Api(ApiError::KeyAlreadyExists { key: ref rejected_key }) + if rejected_key == key + )); + assert!( + !router.inner().inflight_put_key_counts.contains_key(key), + "a committed-key rejection must not leave a zero-count admission entry" + ); + } + + #[::tokio::test] + async fn existing_overwrite_and_inflight_only_contracts_remain_distinct() { + let router = MasterKvRouter::construct(MasterKvRouterNewArg::default()) + .await + .expect("router construction must succeed"); + let key = "overwrite-key"; + router.inner().kv_routes.insert( + key.to_string(), + Arc::new(OneKvNodesRoutes::new((1, 0), None)), + ); + + router + .reserve_inflight_put_key(key, false, false) + .expect("ordinary overwrite must remain allowed"); + let inflight_err = router + .reserve_inflight_put_key(key, true, false) + .expect_err("inflight-only admission must retain its typed error"); + assert!(matches!( + inflight_err, + KvError::Api(ApiError::KeyBeingWritten { key: ref rejected_key }) + if rejected_key == key + )); + router.release_inflight_put_key(key); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs index 9d5eb1d..7529de1 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs @@ -18,6 +18,13 @@ pub enum GetAllocationMode { DurableReplica = 2, } +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum GetSourceKind { + #[default] + Memory = 0, + Ssd = 1, +} + #[derive(Default, Debug, Clone, Encode, Decode)] pub struct GetStartReq { pub key: String, @@ -32,6 +39,7 @@ pub struct GetStartResp { pub get_id: u64, pub node_id: NodeIDString, pub put_id: PutIDForAKey, + pub source_kind: GetSourceKind, // absolute addresses because Mooncake transfer engine requires absolute addresses (not offsets) pub target_addr: u64, pub src_addr: u64, @@ -39,6 +47,8 @@ pub struct GetStartResp { pub target_base_addr: u64, pub src_base_addr: u64, pub len: u64, + /// SSD source staging bytes available at src_addr. Zero for memory sources. + pub ssd_stage_len: u64, pub error_code: ErrorCode, pub error_json: String, /// Server-side processing time in microseconds for this RPC handler @@ -56,6 +66,8 @@ impl RPCReq for GetStartReq { #[derive(Default, Debug, Clone, Encode, Decode)] pub struct GetRevokeReq { pub get_id: u64, + /// True only when an SSD stage failed and the source must be removed from routing. + pub drop_ssd_source: bool, } impl MsgPackSerializePart for GetRevokeReq { fn msg_id(&self) -> u32 { @@ -163,6 +175,7 @@ pub struct PutStartReq { pub key: String, pub len: u64, pub reject_if_inflight_same_key: bool, + pub reject_if_exists: bool, /// Prefer placing the target allocation on any kvclient within this sub_cluster. pub preferred_sub_cluster: Option, /// Optional source-node override for side-transfer workers that share an owner's mmap. @@ -250,6 +263,65 @@ impl RPCReq for PutDoneReq { type Resp = PutDoneResp; } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdReplicaCommitReq { + pub key: String, + pub put_id: PutIDForAKey, + pub node_id: NodeIDString, + pub len: u64, +} +impl MsgPackSerializePart for SsdReplicaCommitReq { + fn msg_id(&self) -> u32 { + MsgId::SsdReplicaCommitReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdReplicaCommitResp { + pub error_code: ErrorCode, + pub error_json: String, + /// Server-side processing time in microseconds for this RPC handler + pub server_process_us: i64, +} +impl MsgPackSerializePart for SsdReplicaCommitResp { + fn msg_id(&self) -> u32 { + MsgId::SsdReplicaCommitResp as u32 + } +} +impl RPCReq for SsdReplicaCommitReq { + type Resp = SsdReplicaCommitResp; +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash, Encode, Decode)] +pub struct SsdReplicaEviction { + pub key: String, + pub put_id: PutIDForAKey, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchSsdReplicaEvictReq { + pub evictions: Vec, +} +impl MsgPackSerializePart for BatchSsdReplicaEvictReq { + fn msg_id(&self) -> u32 { + MsgId::BatchSsdReplicaEvictReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchSsdReplicaEvictResp { + pub removed_count: u32, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for BatchSsdReplicaEvictResp { + fn msg_id(&self) -> u32 { + MsgId::BatchSsdReplicaEvictResp as u32 + } +} +impl RPCReq for BatchSsdReplicaEvictReq { + type Resp = BatchSsdReplicaEvictResp; +} + // --- RPC for MemHolder KeepAlive --- #[derive(Default, Debug, Clone, Encode, Decode)] diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs index d9f509f..9b5700c 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs @@ -1,7 +1,7 @@ -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; use crate::{ - cluster_manager::NodeID, + cluster_manager::{ClusterMember, NodeID, member_has_kv_ssd_storage}, master_seg_manager::one_seg_allocator::{Allocation, OneSegAllocator}, rpcresp_kvresult_convert::msg_and_error::KvError, }; @@ -13,15 +13,48 @@ use super::MasterKvRouterView; pub enum PutPlacementTarget { /// Place locally by reusing the requester's src allocation as the target. - Local { node_id: NodeID }, + Local { + node_id: NodeID, + persist_to_ssd: bool, + }, /// Place remotely with a pre-allocated target allocation. Remote { node_id: NodeID, allocator: Arc, allocation: Allocation, + persist_to_ssd: bool, }, } +#[derive(Debug)] +struct PutSsdPlacementScope { + capable_nodes: HashSet, +} + +impl PutSsdPlacementScope { + fn from_client_members(members: &[ClusterMember]) -> Self { + Self { + capable_nodes: members + .iter() + .filter(|member| member_has_kv_ssd_storage(member)) + .map(|member| member.id.clone().into()) + .collect(), + } + } + + fn requires_ssd(&self) -> bool { + !self.capable_nodes.is_empty() + } + + fn allows_node(&self, node_id: &str) -> bool { + !self.requires_ssd() || self.capable_nodes.contains(node_id) + } + + fn persists_to_ssd(&self, node_id: &str) -> bool { + self.capable_nodes.contains(node_id) + } +} + /// A trait for defining placement policies. #[async_trait] pub trait PlacementPolicy: Send + Sync { @@ -62,21 +95,26 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { len: u64, ) -> Result { let seg_manager = view.master_seg_manager(); + let client_members = view.cluster_manager().get_client_members(); + let ssd_scope = PutSsdPlacementScope::from_client_members(&client_members); let mut last_no_space_ctx: Option<(String, String, u64, u64)> = None; // (node, segment, total, free) if let Some(sc) = preferred_sub_cluster { - let mut preferred_nodes: Vec = view - .cluster_manager() - .get_client_members() - .into_iter() - .filter_map(|m| (m.sub_cluster.as_deref() == Some(sc)).then_some(m.id.into())) + let mut preferred_nodes: Vec = client_members + .iter() + .filter(|member| { + member.sub_cluster.as_deref() == Some(sc) + && ssd_scope.allows_node(member.id.as_str()) + }) + .map(|member| member.id.clone().into()) .collect(); if preferred_nodes.is_empty() { tracing::warn!( - "preferred_sub_cluster has no eligible kvclients: sub_cluster={:?}", - sc + "preferred_sub_cluster has no eligible kvclients: sub_cluster={:?} require_ssd={}", + sc, + ssd_scope.requires_ssd() ); } else { if preferred_nodes @@ -85,6 +123,7 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { { return Ok(PutPlacementTarget::Local { node_id: req_node_id.clone(), + persist_to_ssd: ssd_scope.persists_to_ssd(req_node_id.as_ref()), }); } @@ -115,6 +154,7 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { + persist_to_ssd: ssd_scope.persists_to_ssd(node_id.as_ref()), node_id, allocator, allocation, @@ -127,14 +167,19 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { // Local-first: prefer placing on the requesting node when possible. // This reduces cross-node transfers and enables src==target optimization. let local_allocators = seg_manager.get_node_allocators(req_node_id); - if !local_allocators.is_empty() { + if !local_allocators.is_empty() && ssd_scope.allows_node(req_node_id.as_ref()) { return Ok(PutPlacementTarget::Local { node_id: req_node_id.clone(), + persist_to_ssd: ssd_scope.persists_to_ssd(req_node_id.as_ref()), }); } for _attempt in 1..=3 { - let all_segs = seg_manager.get_all_segments_allocator(); + let all_segs = seg_manager + .get_all_segments_allocator() + .into_iter() + .filter(|(node_id, _)| ssd_scope.allows_node(node_id.as_ref())) + .collect::>(); if let Some((nodeid, allocator)) = all_segs.choose(&mut rand::thread_rng()).cloned() { let node_id: NodeID = nodeid.into(); let total = allocator.total_size_bytes(); @@ -148,6 +193,7 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { )); if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { + persist_to_ssd: ssd_scope.persists_to_ssd(node_id.as_ref()), node_id, allocator, allocation, @@ -198,21 +244,26 @@ impl PlacementPolicy for RandomPlacementPolicy { len: u64, ) -> Result { let seg_manager = view.master_seg_manager(); + let client_members = view.cluster_manager().get_client_members(); + let ssd_scope = PutSsdPlacementScope::from_client_members(&client_members); let mut last_no_space_ctx: Option<(String, String, u64, u64)> = None; // (node, segment, total, free) if let Some(sc) = preferred_sub_cluster { - let mut preferred_nodes: Vec = view - .cluster_manager() - .get_client_members() - .into_iter() - .filter_map(|m| (m.sub_cluster.as_deref() == Some(sc)).then_some(m.id.into())) + let mut preferred_nodes: Vec = client_members + .iter() + .filter(|member| { + member.sub_cluster.as_deref() == Some(sc) + && ssd_scope.allows_node(member.id.as_str()) + }) + .map(|member| member.id.clone().into()) .collect(); if preferred_nodes.is_empty() { tracing::warn!( - "preferred_sub_cluster has no eligible kvclients: sub_cluster={:?}", - sc + "preferred_sub_cluster has no eligible kvclients: sub_cluster={:?} require_ssd={}", + sc, + ssd_scope.requires_ssd() ); } else { if preferred_nodes @@ -223,6 +274,7 @@ impl PlacementPolicy for RandomPlacementPolicy { if !local_allocators.is_empty() { return Ok(PutPlacementTarget::Local { node_id: req_node_id.clone(), + persist_to_ssd: ssd_scope.persists_to_ssd(req_node_id.as_ref()), }); } } @@ -258,6 +310,7 @@ impl PlacementPolicy for RandomPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { + persist_to_ssd: ssd_scope.persists_to_ssd(node_id.as_ref()), node_id, allocator, allocation, @@ -268,7 +321,11 @@ impl PlacementPolicy for RandomPlacementPolicy { } for _attempt in 1..=3 { - let all_segs = seg_manager.get_all_segments_allocator(); + let all_segs = seg_manager + .get_all_segments_allocator() + .into_iter() + .filter(|(node_id, _)| ssd_scope.allows_node(node_id.as_ref())) + .collect::>(); if let Some((nodeid, allocator)) = all_segs.choose(&mut rand::thread_rng()).cloned() { let node_id: NodeID = nodeid.into(); if node_id.as_ref() == req_node_id.as_ref() { @@ -276,6 +333,7 @@ impl PlacementPolicy for RandomPlacementPolicy { if !local_allocators.is_empty() { return Ok(PutPlacementTarget::Local { node_id: req_node_id.clone(), + persist_to_ssd: ssd_scope.persists_to_ssd(req_node_id.as_ref()), }); } continue; @@ -292,6 +350,7 @@ impl PlacementPolicy for RandomPlacementPolicy { )); if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { + persist_to_ssd: ssd_scope.persists_to_ssd(node_id.as_ref()), node_id, allocator, allocation, @@ -322,3 +381,54 @@ impl PlacementPolicy for RandomPlacementPolicy { Err(err) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::cluster_manager::META_KEY_KV_SSD_STORAGE; + use std::collections::HashMap; + + fn client_member(id: &str, has_ssd: bool) -> ClusterMember { + let mut metadata = HashMap::from([("client".to_string(), "true".to_string())]); + if has_ssd { + metadata.insert(META_KEY_KV_SSD_STORAGE.to_string(), "true".to_string()); + } + ClusterMember { + id: id.to_string(), + addresses: Vec::new(), + port: None, + node_start_time: 1, + metadata, + sub_cluster: None, + network: None, + } + } + + #[test] + fn no_ssd_members_keep_memory_only_placement_enabled() { + let members = vec![ + client_member("owner-a", false), + client_member("owner-b", false), + ]; + let scope = PutSsdPlacementScope::from_client_members(&members); + + assert!(!scope.requires_ssd()); + assert!(scope.allows_node("owner-a")); + assert!(scope.allows_node("owner-b")); + assert!(!scope.persists_to_ssd("owner-a")); + } + + #[test] + fn ssd_members_restrict_targets_to_ssd_capable_owners() { + let members = vec![ + client_member("owner-memory", false), + client_member("owner-ssd", true), + ]; + let scope = PutSsdPlacementScope::from_client_members(&members); + + assert!(scope.requires_ssd()); + assert!(!scope.allows_node("owner-memory")); + assert!(scope.allows_node("owner-ssd")); + assert!(scope.persists_to_ssd("owner-ssd")); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs index 70d8858..3c35336 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs @@ -1,25 +1,26 @@ use super::NodeValueReplicaDesc; use super::{ - InflightPutAllocation, InflightPutInfo, KvRouteInfo, MasterKvRouterView, PutPlacementMode, - msg_pack::{PutDoneReq, PutDoneResp, PutRevokeReq, PutRevokeResp, PutStartReq, PutStartResp}, + InflightPutAllocation, InflightPutInfo, MasterKvRouterView, NodeCacheCapacityReservation, + PutPlacementMode, SsdReplicaCommitStatus, + msg_pack::{ + PutDoneReq, PutDoneResp, PutRevokeReq, PutRevokeResp, PutStartReq, PutStartResp, + SsdReplicaCommitReq, SsdReplicaCommitResp, + }, placement::PutPlacementTarget, }; +use crate::client_kv_api::msg_pack::SsdReplicaPersistReq; use crate::master_kv_router::OneKvNodesRoutes; use crate::master_kv_router::delete::DeleteKeyInfo; use crate::{ cluster_manager::{META_KEY_LOCAL_IPC_ROOT, NodeID}, master_seg_manager::one_seg_allocator::Allocation, - p2p::msg_pack::MsgPack, + p2p::msg_pack::{MsgPack, RPCCaller}, rpcresp_kvresult_convert::msg_and_error, }; use fluxon_commu::{META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME}; use parking_lot::Mutex; -use parking_lot::RwLock; use rand::seq::SliceRandom; -use std::{ - collections::HashMap, - sync::{Arc, atomic::AtomicU32}, -}; +use std::{sync::Arc, time::Duration}; pub type PutIDForAKey = (u64, u32); @@ -159,10 +160,11 @@ pub async fn handle_put_start( req_node_id: NodeID, ) -> (PutIDForAKey, MsgPack) { let key = req.serialize_part.key.clone(); - if let Err(err) = view - .master_kv_router() - .reserve_inflight_put_key(&key, req.serialize_part.reject_if_inflight_same_key) - { + if let Err(err) = view.master_kv_router().reserve_inflight_put_key( + &key, + req.serialize_part.reject_if_inflight_same_key, + req.serialize_part.reject_if_exists, + ) { let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); return ( (0, 0), @@ -267,6 +269,7 @@ pub async fn handle_put_start( let mut src_allocation = Some(src_allocation); let finalize = |node_id: NodeID, + persist_to_ssd: bool, inflight_alloc: InflightPutAllocation, src_addr: u64, target_addr: u64, @@ -278,6 +281,7 @@ pub async fn handle_put_start( key: key.clone(), len, req_node_id: req_node_id.clone(), + persist_to_ssd, src_target_allocation: Arc::new(Mutex::new(Some(inflight_alloc))), }; @@ -327,7 +331,10 @@ pub async fn handle_put_start( .await; match put_target { - Ok(PutPlacementTarget::Local { node_id }) => { + Ok(PutPlacementTarget::Local { + node_id, + persist_to_ssd, + }) => { if node_id != source_node_id { unreachable!( "Local placement must be the resolved source node; got node_id={} source_node_id={} requester_node_id={}", @@ -336,14 +343,15 @@ pub async fn handle_put_start( } tracing::debug!( - "put_start placement decided: local; put_id={:?} key={} requester_node_id={} source_node_id={} target_node_id={} preferred_sub_cluster={:?} len={}", + "put_start placement decided: local; put_id={:?} key={} requester_node_id={} source_node_id={} target_node_id={} preferred_sub_cluster={:?} len={} persist_to_ssd={}", put_id, key, req_node_id, source_node_id, node_id, req.serialize_part.preferred_sub_cluster, - req.serialize_part.len + req.serialize_part.len, + persist_to_ssd ); view.master_kv_router().record_put_placement_decision( req_node_id.as_ref(), @@ -364,6 +372,7 @@ pub async fn handle_put_start( .expect("src_allocation must exist when finalizing local put"); let fut = finalize( node_id, + persist_to_ssd, InflightPutAllocation::Local(src), abs, abs, @@ -378,6 +387,7 @@ pub async fn handle_put_start( Ok(PutPlacementTarget::Remote { node_id, allocation: target_allocation, + persist_to_ssd, .. }) => { let src_ref = src_allocation @@ -391,7 +401,7 @@ pub async fn handle_put_start( let allocation_size = target_allocation.size(); tracing::debug!( - "put_start placement decided: remote; put_id={:?} key={} requester_node_id={} source_node_id={} target_node_id={} preferred_sub_cluster={:?} len={} target_base_addr={} target_offset={} allocation_size={}", + "put_start placement decided: remote; put_id={:?} key={} requester_node_id={} source_node_id={} target_node_id={} preferred_sub_cluster={:?} len={} target_base_addr={} target_offset={} allocation_size={} persist_to_ssd={}", put_id, key, req_node_id, @@ -401,7 +411,8 @@ pub async fn handle_put_start( req.serialize_part.len, target_base, target_offset, - allocation_size + allocation_size, + persist_to_ssd ); view.master_kv_router().record_put_placement_decision( req_node_id.as_ref(), @@ -414,6 +425,7 @@ pub async fn handle_put_start( .expect("src_allocation must exist when finalizing remote put"); let fut = finalize( node_id, + persist_to_ssd, InflightPutAllocation::Remote { src, target: target_allocation, @@ -474,6 +486,229 @@ pub async fn handle_put_revoke( } } +fn insert_memory_replica_into_cache_if_current( + view: &MasterKvRouterView, + key: &str, + put_id: PutIDForAKey, + node_id: &NodeID, + weight_bytes: u32, +) { + let Some(route_ref) = view.master_kv_router().inner().kv_routes.get(key) else { + tracing::debug!( + "Skipping delayed cache insertion for missing key: key={} put_id=({},{}) node={}", + key, + put_id.0, + put_id.1, + node_id + ); + return; + }; + let route = route_ref.value().clone(); + drop(route_ref); + + if route.put_id != put_id || !route.has_memory_replica(node_id) { + tracing::debug!( + "Skipping delayed cache insertion for stale replica: key={} put_id=({},{}) node={}", + key, + put_id.0, + put_id.1, + node_id + ); + return; + } + + let Some(cache) = view + .master_kv_router() + .get_node_cache_controller(node_id.as_ref()) + else { + tracing::warn!( + "No cache controller found for node: {}, node is not ready", + node_id + ); + return; + }; + let desc = NodeValueReplicaDesc { + weight_bytes, + put_id, + }; + tracing::debug!("Inserting key: {:?} into cache", key); + cache.insert(key.to_string(), desc); + tracing::debug!( + "Inserted key: {:?} into cache, current cache size: {}", + key, + cache.weighted_size() + ); +} + +fn spawn_ssd_replica_persist_request( + view: &MasterKvRouterView, + key: String, + put_id: PutIDForAKey, + node_id: NodeID, + len: u64, + allocation: Arc, + cache_weight_bytes: Option, + pending_persist_reservation: Option, +) { + let target_addr = allocation.base_addr() + allocation.addr(); + let view = view.clone(); + let view_task = view.clone(); + let _ = view.spawn("post_put_ssd_replica_persist", async move { + let _allocation_guard = allocation; + let req = MsgPack { + serialize_part: SsdReplicaPersistReq { + key: key.clone(), + put_id, + target_addr, + len, + }, + raw_bytes: Vec::new(), + }; + let resp = RPCCaller::::new() + .call( + view_task.p2p_module(), + node_id.clone(), + req, + Some(Duration::from_secs(60)), + 2, + ) + .await; + match resp { + Ok(resp) => { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json, + ) { + tracing::warn!( + "SSD replica persist failed: key={} put_id=({},{}) node={} err={}", + key, + put_id.0, + put_id.1, + node_id, + err + ); + } else if resp.serialize_part.persisted { + tracing::debug!( + "SSD replica persist completed: key={} put_id=({},{}) node={}", + key, + put_id.0, + put_id.1, + node_id + ); + } else { + tracing::debug!( + "SSD replica persist skipped because owner has no SSD store: key={} put_id=({},{}) node={}", + key, + put_id.0, + put_id.1, + node_id + ); + } + } + Err(err) => { + tracing::warn!( + "SSD replica persist RPC failed: key={} put_id=({},{}) node={} err={:?}", + key, + put_id.0, + put_id.1, + node_id, + err + ); + } + } + + // The allocation becomes Moka-managed only after the persist attempt. + // Release its temporary reservation before inserting the same bytes. + drop(pending_persist_reservation); + if let Some(weight_bytes) = cache_weight_bytes { + insert_memory_replica_into_cache_if_current( + &view_task, + &key, + put_id, + &node_id, + weight_bytes, + ); + } + }); +} + +fn ok_ssd_replica_commit_resp() -> MsgPack { + MsgPack { + serialize_part: SsdReplicaCommitResp { + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_ssd_replica_commit( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + let req = req.serialize_part; + let node_id: NodeID = req.node_id.clone().into(); + let Some(route_ref) = view.master_kv_router().inner().kv_routes.get(&req.key) else { + tracing::debug!( + "Ignoring SSD replica commit for missing key: key={} put_id=({},{}) node={}", + req.key, + req.put_id.0, + req.put_id.1, + req.node_id + ); + return ok_ssd_replica_commit_resp(); + }; + let route = route_ref.value().clone(); + drop(route_ref); + + if route.put_id != req.put_id { + tracing::debug!( + "Ignoring stale SSD replica commit: key={} req_put_id=({},{}) current_put_id=({},{}) node={}", + req.key, + req.put_id.0, + req.put_id.1, + route.put_id.0, + route.put_id.1, + req.node_id + ); + return ok_ssd_replica_commit_resp(); + } + + match route.commit_ssd_replica(&node_id, req.len) { + SsdReplicaCommitStatus::MissingMemory => { + tracing::debug!( + "Ignoring SSD replica commit without matching memory replica: key={} put_id=({},{}) node={}", + req.key, + req.put_id.0, + req.put_id.1, + req.node_id + ); + return ok_ssd_replica_commit_resp(); + } + SsdReplicaCommitStatus::TombedNode => { + tracing::debug!( + "Ignoring SSD replica commit for tombed node: key={} put_id=({},{}) node={}", + req.key, + req.put_id.0, + req.put_id.1, + req.node_id + ); + return ok_ssd_replica_commit_resp(); + } + SsdReplicaCommitStatus::Committed => {} + } + tracing::debug!( + "Committed SSD replica route: key={} put_id=({},{}) node={} len={}", + req.key, + req.put_id.0, + req.put_id.1, + req.node_id, + req.len + ); + ok_ssd_replica_commit_resp() +} + pub async fn handle_put_done( view: MasterKvRouterView, req: MsgPack, @@ -488,6 +723,8 @@ pub async fn handle_put_done( if let Some(InflightPutInfo { node_id, key, + len, + persist_to_ssd, src_target_allocation, .. }) = view @@ -497,7 +734,8 @@ pub async fn handle_put_done( .remove(&full_put_id) .await { - view.master_kv_router().release_inflight_put_key(&key); + // Keep same-key admission reserved until the new route is visible. + let key_reservation = InflightPutKeyReservation::new(view.clone(), key.clone()); let Some(allocs) = src_target_allocation.lock().take() else { tracing::warn!( "Put operation with put_id {:?} not found for completion", @@ -574,10 +812,23 @@ pub async fn handle_put_done( } else { target_cap_bytes as u32 }; - // Note: moka cache insertion happens after commit in a spawned task - // using the same saturated weight; avoid unused local here. - // If lease is provided, attach first and fail fast on error + let mut pending_persist_reservation = None; if let Some(lease_id) = lease_id_opt { + let lease_cache_reservation = match view + .master_kv_router() + .reserve_node_cache_capacity(node_id.as_ref(), target_cap_bytes) + { + Ok(reservation) => reservation, + Err(e) => { + let kv_err: crate::rpcresp_kvresult_convert::msg_and_error::KvError = e.into(); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &kv_err, + ), + raw_bytes: Vec::new(), + }; + } + }; if let Err(e) = view .master_lease_manager() .attach_key(lease_id, key.clone(), put_id) @@ -589,50 +840,27 @@ pub async fn handle_put_done( raw_bytes: Vec::new(), }; } - // Reserve cache capacity on this node for the leased allocation now (fetch_sub semantics) - if let Err(e) = view + target_allocation.set_on_drop(move || drop(lease_cache_reservation)); + } else { + pending_persist_reservation = match view .master_kv_router() - .adjust_node_cache_capacity_for_lease(node_id.as_ref(), target_cap_bytes as i64) + .reserve_node_cache_capacity(node_id.as_ref(), target_cap_bytes) { - let kv_err: crate::rpcresp_kvresult_convert::msg_and_error::KvError = e.into(); - return MsgPack { - serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&kv_err), - raw_bytes: Vec::new(), - }; - } - // And attach an on-drop hook to restore it (fetch_add on Allocation drop) - let view_clone = view.clone(); - let node_id_string = node_id.as_ref().to_string(); - target_allocation.set_on_drop(move || { - match view_clone.try_adjust_node_cache_capacity_for_lease( - &node_id_string, - -(target_cap_bytes as i64), - ) { - Some(Ok(())) => {} - Some(Err(e)) => { - tracing::warn!( - "Failed to restore moka capacity on drop: node_id={}, bytes={}, err={}", - node_id_string, - target_cap_bytes, - e - ); - } - None => { - tracing::debug!( - "Skipped restoring moka capacity on drop because MasterKvRouterView is gone: node_id={}, bytes={}", - node_id_string, - target_cap_bytes - ); - } + Ok(reservation) => reservation, + Err(e) => { + let kv_err: crate::rpcresp_kvresult_convert::msg_and_error::KvError = e.into(); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &kv_err, + ), + raw_bytes: Vec::new(), + }; } - }); + }; } - let completed_info = KvRouteInfo { - node_id: node_id.clone(), - allocation: Arc::new(target_allocation), - tomb_tag, - }; + let target_allocation = Arc::new(target_allocation); + let target_allocation_for_ssd = persist_to_ssd.then(|| Arc::clone(&target_allocation)); // Insert into kv_routes with replica support let mut old_one_kv_routes: Option> = None; @@ -645,27 +873,48 @@ pub async fn handle_put_done( .entry(key.clone()) .or_insert_with(|| { inserted = true; - Arc::new(OneKvNodesRoutes { - put_id, - lease_id: lease_id_opt, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), - }) + Arc::new(OneKvNodesRoutes::new(put_id, lease_id_opt)) }); // we need to take out old one_kv_routes if it is not inserted if !inserted { old_one_kv_routes = Some(one_kv_routes.clone()); - *one_kv_routes = Arc::new(OneKvNodesRoutes { + *one_kv_routes = Arc::new(OneKvNodesRoutes::new(put_id, lease_id_opt)); + } + one_kv_routes.insert_memory_replica( + node_id.clone(), + target_allocation, + tomb_tag.clone(), + ); + } + // Publish the route before releasing admission so reject-if-exists observes either + // the inflight reservation or the committed route, with no visibility gap. + drop(key_reservation); + + let cache_weight_bytes = (lease_id_opt.is_none() + && view.master_kv_router().replica_cache_enabled()) + .then_some(saturated_weight_u32); + if let Some(target_allocation_for_ssd) = target_allocation_for_ssd { + spawn_ssd_replica_persist_request( + &view, + key.clone(), + put_id, + node_id.clone(), + len, + target_allocation_for_ssd, + cache_weight_bytes, + pending_persist_reservation, + ); + } else { + drop(pending_persist_reservation); + if let Some(weight_bytes) = cache_weight_bytes { + insert_memory_replica_into_cache_if_current( + &view, + &key, put_id, - lease_id: lease_id_opt, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), - }); + &node_id, + weight_bytes, + ); } - one_kv_routes - .nodes_replicas - .write() - .insert(node_id.clone(), completed_info); } if let Some(old) = old_one_kv_routes { @@ -684,54 +933,18 @@ pub async fn handle_put_done( } } - // Post-commit maintenance: update prefix-count index (for CountPrefix RPC) - // and, if applicable, update per-node cache controller. Run both in a - // spawned task to keep the PutDone RPC path lean and consistent with - // other async cache control operations. Deletion path already removes - // the index entry in delete.rs (do_delete_one_kv_all_replicas). + // Update the prefix-count index asynchronously to keep PutDone lean. + // SSD-backed targets delay cache insertion until the persist request completes. { let view_task = view.clone(); let key_for_spawn = key.clone(); - let node_for_spawn = node_id.clone(); let do_prefix_index_update = view.master_kv_router().prefix_index_enabled(); - let do_cache_insert = - lease_id_opt.is_none() && view.master_kv_router().replica_cache_enabled(); - // Reuse the saturated weight computed above for moka insertion - let cap_bytes_u32 = saturated_weight_u32; let _ = view.spawn("post_put_done_maintenance", async move { - // 1) Update the derived prefix-counting index asynchronously. - // This keeps PutDone lean, but also means CountPrefix visibility - // is not an immediate strong-consistency signal for this put. if do_prefix_index_update { let inner = view_task.master_kv_router().inner(); let mut tree = inner.prefix_index.write().await; tree.insert(&key_for_spawn); } - - // 2) Optionally update node cache controller (non-leased keys) - if do_cache_insert { - let cache = view_task - .master_kv_router() - .get_node_cache_controller(&node_for_spawn); - if let Some(cache) = cache { - let desc = NodeValueReplicaDesc { - weight_bytes: cap_bytes_u32, - put_id, - }; - tracing::debug!("Inserting key: {:?} into cache", key_for_spawn); - cache.insert(key_for_spawn.clone(), desc); - tracing::debug!( - "Inserted key: {:?} into cache, current cache size: {}", - key_for_spawn, - cache.weighted_size() - ); - } else { - tracing::warn!( - "No cache controller found for node: {}, node is not ready", - node_for_spawn - ); - } - } }); } diff --git a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs index 5c20cc1..1fcc539 100755 --- a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs +++ b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs @@ -22,7 +22,8 @@ async fn test1_lease_expire_removes_keys() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t1", "lease_client_t1").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t1", "lease_client_t1").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -82,11 +83,12 @@ async fn test2_rebind_to_new_lease_preserves_until_new_expire() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t2", "lease_client_t2").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t2", "lease_client_t2").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; - let l1 = grant_short_lease_for_test(master_fw.as_ref(), 2).await; + let l1 = grant_short_lease_for_test(master_fw.as_ref(), 300).await; let keys: Vec = (0..10).map(|i| format!("t2_key_{}", i)).collect(); let value_a = vec![3u8; 64]; for k in &keys { @@ -102,8 +104,8 @@ async fn test2_rebind_to_new_lease_preserves_until_new_expire() { keys.len() ); - // Rebind to a newer lease with longer TTL - let l2 = grant_short_lease_for_test(master_fw.as_ref(), 4).await; + // Rebind each key to a new lease and key version. + let l2 = grant_short_lease_for_test(master_fw.as_ref(), 300).await; let value_b = vec![9u8; 64]; for k in &keys { client_view @@ -117,8 +119,13 @@ async fn test2_rebind_to_new_lease_preserves_until_new_expire() { l2 ); - // After l1 expiry, keys should still exist due to l2 - sleep(Duration::from_secs(3)).await; + // Expiring the old lease must not delete the newer key versions bound to l2. + master_fw + .master_lease_manager_view() + .master_lease_manager() + .expire_lease_for_test(l1) + .await + .unwrap(); for k in &keys { assert!( client_view @@ -133,8 +140,13 @@ async fn test2_rebind_to_new_lease_preserves_until_new_expire() { "[test2-verify-still-present-after-first-expire] keys still exist after first lease expired, held by second lease" ); - // After l2 expiry + grace, keys should be gone - sleep(Duration::from_secs(3)).await; + // Expiring the current lease owns deletion of the current key versions. + master_fw + .master_lease_manager_view() + .master_lease_manager() + .expire_lease_for_test(l2) + .await + .unwrap(); for k in &keys { let exist = client_view .client_kv_api() @@ -161,7 +173,8 @@ async fn test3_keepalive() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t3", "lease_client_t3").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t3", "lease_client_t3").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -236,7 +249,8 @@ async fn test4_delete_under_lease_then_get_fails() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t4", "lease_client_t4").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t4", "lease_client_t4").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; diff --git a/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs b/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs index 86d6736..1d0068f 100644 --- a/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs +++ b/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs @@ -420,6 +420,11 @@ impl MasterLeaseManager { self.grant_lease_unchecked(ttl).await } + #[cfg(test)] + pub async fn expire_lease_for_test(&self, lease_id: LeaseID) -> Result<(), LeaseMgrError> { + Self::cleanup_single_lease_static(self.inner().view(), lease_id).await + } + async fn grant_lease_unchecked(&self, ttl: LeaseSecTTl) -> Result { let inner = self.inner(); diff --git a/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs b/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs index 6dcd136..1442f26 100755 --- a/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs +++ b/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs @@ -21,7 +21,7 @@ pub struct Allocation { allocator: Arc, /// Optional callback invoked when this allocation is dropped. /// Used by upper layers to perform side effects (e.g., capacity restoration). - on_drop: Option>, + on_drop: Option>, } // Custom Debug to avoid requiring Debug on the callback closure. @@ -76,7 +76,7 @@ impl Allocation { /// when this allocation is dropped. pub fn set_on_drop(&mut self, f: F) where - F: Fn() + Send + Sync + 'static, + F: FnOnce() + Send + Sync + 'static, { self.on_drop = Some(Box::new(f)); } diff --git a/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs b/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs index ad23b4d..6331282 100755 --- a/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs @@ -474,9 +474,9 @@ impl MemholderManagerTrait for MasterOwnerMemMgr { }; let mut targets = Vec::new(); - let nodes_replicas = nodes_kv_route_info.nodes_replicas.read(); - for (node_id, kv_route_info) in nodes_replicas.iter() { - if kv_route_info.tomb_tag.is_tomb() { + let node_replicas = nodes_kv_route_info.node_replicas.read(); + for (node_id, replicas) in node_replicas.iter() { + if replicas.tomb_tag.is_tomb() || replicas.memory.is_none() { continue; } let Some(member) = ctx @@ -517,11 +517,11 @@ impl MemholderManagerTrait for MasterOwnerMemMgr { continue; }; - let nodes_replicas = nodes_kv_route_info.nodes_replicas.read(); - let Some(kv_route_info) = nodes_replicas.get(target.node_id.as_str()) else { + let node_replicas = nodes_kv_route_info.node_replicas.read(); + let Some(replicas) = node_replicas.get(target.node_id.as_str()) else { continue; }; - if kv_route_info.tomb_tag.is_tomb() { + if replicas.tomb_tag.is_tomb() || replicas.memory.is_none() { continue; } diff --git a/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs b/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs index 692a9a0..78f5c6b 100644 --- a/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs @@ -101,6 +101,7 @@ fn new_client_config_with_size( large_file_paths: crate::config::LargeFilePaths { paths: vec![format!("/tmp/kvcache_large/{}", instance_key)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -134,6 +135,7 @@ fn new_zero_contribution_client_config( }, share_mem_path: format!("/tmp/kvcache_shared_memory/{}", owner_instance_key), large_file_paths: crate::config::LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -192,10 +194,12 @@ fn evict_all_replicas_for_key(master_view: &MasterKvRouterView, key: &str) -> us .get(key) .map(|route| { let node_ids = route - .nodes_replicas + .node_replicas .read() - .keys() - .cloned() + .iter() + .filter_map(|(node_id, replicas)| { + replicas.memory.is_some().then(|| node_id.clone()) + }) .collect::>(); (route.put_id, node_ids) }) diff --git a/fluxon_rs/fluxon_kv/src/memholder/mod.rs b/fluxon_rs/fluxon_kv/src/memholder/mod.rs index db0ada2..638a08c 100755 --- a/fluxon_rs/fluxon_kv/src/memholder/mod.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/mod.rs @@ -201,6 +201,7 @@ pub struct ExternalMemHolderInfo { pub offset: u64, pub len: u32, pub holder_id: u64, + pub source_kind: crate::master_kv_router::msg_pack::GetSourceKind, } impl ExternalMemHolder { /// Get the memory offset diff --git a/fluxon_rs/fluxon_kv/src/metric_reporter/mod.rs b/fluxon_rs/fluxon_kv/src/metric_reporter/mod.rs index 88fea03..d6ad694 100644 --- a/fluxon_rs/fluxon_kv/src/metric_reporter/mod.rs +++ b/fluxon_rs/fluxon_kv/src/metric_reporter/mod.rs @@ -1,3 +1,4 @@ +use crate::client_kv_api::{ClientKvApi, ClientKvApiAccessTrait}; use crate::cluster_manager::META_KEY_LOCAL_IPC_ROOT; use crate::cluster_manager::{ ClusterEvent, ClusterManager, ClusterManagerAccessTrait, ClusterManagerView, ClusterMember, @@ -304,6 +305,7 @@ define_module!( MetricReporter, (cluster_manager, ClusterManager), (p2p, P2pModule), + (client_kv_api, ClientKvApi), (master_seg_manager, MasterSegManager), (metric_reporter, MetricReporter) ); @@ -623,6 +625,7 @@ impl MetricReporter { pub async fn init2_after_prom_remote_write_wait(&self) -> Result<(), KvError> { self.spawn_master_only_collect_loop(); + self.spawn_owner_local_collect_loop(); Ok(()) } @@ -916,6 +919,32 @@ impl MetricReporter { }); } + fn spawn_owner_local_collect_loop(&self) { + let view = self.view().clone(); + let shutdown_poller = view.register_shutdown_poller(); + let mut shutdown_waiter = view.register_shutdown_waiter(); + let view_task = view.clone(); + let _ = view.spawn("metric_owner_local_collect_loop", async move { + let mut interval = + tokio::time::interval(Duration::from_secs(METRICS_FLUSH_INTERVAL_SECS)); + loop { + tokio::select! { + _ = shutdown_waiter.wait() => { + info!("Owner-local metric collector stopped by shutdown signal"); + break; + } + _ = interval.tick() => { + if !shutdown_poller.is_running() { + info!("Owner-local metric collector stopped by shutdown signal"); + break; + } + view_task.metric_reporter().owner_local_collect(); + } + } + } + }); + } + /// Collect per-segment capacity/usage from MasterSegManager (master only). fn collect_segment_metrics(&self) { let view = self.view().clone(); @@ -942,6 +971,22 @@ impl MetricReporter { self.collect_segment_metrics(); } + fn owner_local_collect(&self) { + let view = self.view().clone(); + let member = view.cluster_manager().get_self_info(); + if !matches!(member.node_role(), crate::cluster_manager::NodeRole::Client) { + return; + } + + let node = member.id.as_ref(); + for usage in view.client_kv_api().kv_ssd_storage_usage_snapshot() { + self.metrics + .set_kv_ssd_capacity_bytes(node, &usage.device, usage.capacity_bytes); + self.metrics + .set_kv_ssd_used_bytes(node, &usage.device, usage.used_bytes); + } + } + async fn wait_prom_remote_write_urls_best_effort(&self) -> Result<(), KvError> { if self.observability_disabled() { return Ok(()); diff --git a/fluxon_rs/fluxon_kv/src/metrics.rs b/fluxon_rs/fluxon_kv/src/metrics.rs index 88f1b9c..fb26969 100644 --- a/fluxon_rs/fluxon_kv/src/metrics.rs +++ b/fluxon_rs/fluxon_kv/src/metrics.rs @@ -252,6 +252,28 @@ impl MetricsHandle { }); } + pub fn set_kv_ssd_capacity_bytes(&self, node: &str, device: &str, bytes: u64) { + let Some(observe) = self.observe() else { + return; + }; + observe.try_submit(ObserveOp::SetKvSsdCapacityBytes { + node: node.to_string(), + device: device.to_string(), + bytes, + }); + } + + pub fn set_kv_ssd_used_bytes(&self, node: &str, device: &str, bytes: u64) { + let Some(observe) = self.observe() else { + return; + }; + observe.try_submit(ObserveOp::SetKvSsdUsedBytes { + node: node.to_string(), + device: device.to_string(), + bytes, + }); + } + pub fn set_fs_mount_fs_bytes( &self, mount_kind: FsMountKind, diff --git a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs index 42a9cbc..9ad9fe3 100644 --- a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs +++ b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs @@ -35,6 +35,10 @@ pub enum MsgId { DeleteAckResp = 3024, BatchDeleteAckReq = 3029, BatchDeleteAckResp = 3030, + SsdReplicaCommitReq = 3031, + SsdReplicaCommitResp = 3032, + BatchSsdReplicaEvictReq = 3033, + BatchSsdReplicaEvictResp = 3034, GetMetaReq = 3019, GetMetaResp = 3020, BatchDeleteClientKvMetaCacheReq = 3021, @@ -704,6 +708,9 @@ crate::define_err_group! { , (122, KeyBeingWritten { key: String }, msg: "Key is currently being written: key={key}") + , + (123, KeyAlreadyExists { key: String }, + msg: "Key already exists: key={key}") } } diff --git a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs index b6eb7d6..94d3fff 100755 --- a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs +++ b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs @@ -3,11 +3,13 @@ use super::msg_and_error::{ErrorCode, KvError, KvResult}; use crate::client_kv_api::msg_pack::{ ExternalDeleteAckResp, ExternalDeleteResp, ExternalGetResp, ExternalIsExistResp, ExternalPutCommitResp, ExternalPutRevokeResp, ExternalPutStartResp, ExternalPutTransferEndResp, + SsdReplicaPersistResp, SsdStageReadResp, }; use crate::master_kv_router::msg_pack::{ - BatchDeleteAckResp, BatchDeleteClientKvMetaCacheResp, DeleteAckResp, DeleteResp, GetDoneResp, - GetMasterOnlyMetricPartResp, GetMetaResp, GetRevokeResp, GetStartResp, MemHolderKeepAliveResp, - MemHolderReleaseResp, PutDoneResp, PutRevokeResp, PutStartResp, + BatchDeleteAckResp, BatchDeleteClientKvMetaCacheResp, BatchSsdReplicaEvictResp, DeleteAckResp, + DeleteResp, GetDoneResp, GetMasterOnlyMetricPartResp, GetMetaResp, GetRevokeResp, GetStartResp, + MemHolderKeepAliveResp, MemHolderReleaseResp, PutDoneResp, PutRevokeResp, PutStartResp, + SsdReplicaCommitResp, }; use crate::master_seg_manager::msg_pack::RequestSegmentRegistrationResp; use crate::memholder::ExternalMemHolderInfo; @@ -232,6 +234,26 @@ impl FromError for ExternalDeleteAckResp { } } } +impl FromError for SsdStageReadResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for SsdReplicaPersistResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} // ---- FromError for Master KV Router Resps ---- impl FromError for GetStartResp { @@ -294,6 +316,26 @@ impl FromError for PutDoneResp { } } } +impl FromError for SsdReplicaCommitResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for BatchSsdReplicaEvictResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} impl FromError for MemHolderKeepAliveResp { fn from_error(e: &KvError) -> Self { let code = e.code(); diff --git a/fluxon_rs/fluxon_mq/src/consumer.rs b/fluxon_rs/fluxon_mq/src/consumer.rs index c5e5fa4..6be6115 100644 --- a/fluxon_rs/fluxon_mq/src/consumer.rs +++ b/fluxon_rs/fluxon_mq/src/consumer.rs @@ -1364,8 +1364,19 @@ impl MpscConsumer { // 4) Bind consumer membership key under member lease, storing // ChanMemberMeta as JSON for future introspection. let consumer_key = keys::etcd_consumer_key(chan_id, &consumer_idx); - let compare = - etcd::Compare::create_revision(consumer_key.clone(), etcd::CompareOp::Equal, 0); + let compares = vec![ + etcd::Compare::create_revision( + keys::etcd_aborted_key(chan_id), + etcd::CompareOp::Equal, + 0, + ), + etcd::Compare::create_revision( + keys::etcd_meta_key(chan_id), + etcd::CompareOp::Greater, + 0, + ), + etcd::Compare::create_revision(consumer_key.clone(), etcd::CompareOp::Equal, 0), + ]; let member_meta = ChanMemberMeta { member_id: consumer_idx.clone(), @@ -1381,13 +1392,16 @@ impl MpscConsumer { meta_bytes, Some(etcd::PutOptions::new().with_lease(member_lease_id)), ); - let txn = etcd::Txn::new().when(vec![compare]).and_then(vec![put_op]); + let txn = etcd::Txn::new().when(compares).and_then(vec![put_op]); let txn_res = client .txn(txn) .await .with_context(|| format!("failed to bind consumer membership key {}", consumer_key))?; if !txn_res.succeeded() { - anyhow::bail!("consumer membership key {} already exists", consumer_key); + anyhow::bail!( + "consumer membership bind rejected for key {}: channel aborted, meta missing, or membership already exists", + consumer_key + ); } // 5) 创建预取 actor，并通过共享队列与之协作。 @@ -1535,19 +1549,30 @@ impl MpscConsumer { let member_lease_id = self.chan_mgr.member_lease_id(); let consumer_key = keys::etcd_consumer_key(self.chan_id, &self.consumer_idx); - let compare = - etcd::Compare::create_revision(consumer_key.clone(), etcd::CompareOp::Greater, 0); + let compares = vec![ + etcd::Compare::create_revision( + keys::etcd_aborted_key(self.chan_id), + etcd::CompareOp::Equal, + 0, + ), + etcd::Compare::create_revision( + keys::etcd_meta_key(self.chan_id), + etcd::CompareOp::Greater, + 0, + ), + etcd::Compare::create_revision(consumer_key.clone(), etcd::CompareOp::Greater, 0), + ]; let put_op = etcd::TxnOp::put( consumer_key.clone(), meta_bytes, Some(etcd::PutOptions::new().with_lease(member_lease_id)), ); - let txn = etcd::Txn::new().when(vec![compare]).and_then(vec![put_op]); + let txn = etcd::Txn::new().when(compares).and_then(vec![put_op]); let mut client = self.chan_mgr.etcd_client(); let txn_res = client.txn(txn).await?; if !txn_res.succeeded() { return Err(MpscError::Internal(format!( - "consumer membership key {} missing while syncing kvclient_sub_cluster", + "consumer membership sync rejected for key {}: channel aborted, meta missing, or membership missing", consumer_key ))); } diff --git a/fluxon_rs/fluxon_mq/src/create.rs b/fluxon_rs/fluxon_mq/src/create.rs index 4fbb753..ba2c854 100644 --- a/fluxon_rs/fluxon_mq/src/create.rs +++ b/fluxon_rs/fluxon_mq/src/create.rs @@ -1,7 +1,9 @@ use anyhow::Context; use etcd_client as etcd; +use std::collections::BTreeSet; +use std::time::Duration; -use fluxon_util::etcd::{get_cluster_lease_id, DistributeIdAllocator}; +use fluxon_util::etcd::{get_cluster_lease_id, DistributeIdAllocator, ManagedEtcdClient}; use fluxon_util::lease_manager::{ record_register_by as lm_record_register_by, LeaseBackendUid, LeaseManager, LeaseRegisterKind, }; @@ -10,6 +12,8 @@ use crate::error::MpscError; use crate::keys; use crate::manager::{get_chan_meta_with_version, ChanGlobalMeta, ChanManager}; +const ROLLBACK_FENCE_TXN_MAX_ATTEMPTS: usize = 5; + pub struct ChanCreateConfig { pub capacity: i64, pub ttl_seconds: i64, @@ -35,6 +39,211 @@ pub struct ChanCreateConfig { pub override_payload_lease_id: Option, } +/// Cleanup plan for a channel id allocated by `create_mpsc_channel` but not +/// successfully exposed to its caller. +/// +/// Only channel-owned etcd leases are recorded here. Caller-owned override +/// leases are protected and are never revoked by rollback. +#[derive(Debug, Clone)] +pub struct MpscCreateRollback { + chan_id: i64, + protected_etcd_lease_ids: BTreeSet, + channel_owned_cluster_lease_keys: BTreeSet, + channel_owned_etcd_lease_ids: BTreeSet, +} + +impl MpscCreateRollback { + fn new(chan_id: i64, protected_etcd_lease_ids: impl IntoIterator) -> Self { + Self { + chan_id, + protected_etcd_lease_ids: protected_etcd_lease_ids.into_iter().collect(), + channel_owned_cluster_lease_keys: BTreeSet::new(), + channel_owned_etcd_lease_ids: BTreeSet::new(), + } + } + + fn record_channel_owned_cluster_lease_key(&mut self, key: String) { + self.channel_owned_cluster_lease_keys.insert(key); + } + + fn record_channel_owned_etcd_lease(&mut self, lease_id: i64) { + if lease_id > 0 && !self.protected_etcd_lease_ids.contains(&lease_id) { + self.channel_owned_etcd_lease_ids.insert(lease_id); + } + } + + fn fence_cleanup_keys(&self) -> [String; 3] { + [ + keys::etcd_aborted_key(self.chan_id), + keys::etcd_meta_key(self.chan_id), + format!("/channels/{}/", self.chan_id), + ] + } + + fn exact_cleanup_keys(&self) -> [String; 3] { + [ + format!("dist_id_allocator/channels/{}", self.chan_id), + format!("cluster_lease/channels/{}", self.chan_id), + format!("cluster_lease/id_allocator/channels/{}", self.chan_id), + ] + } + + fn cleanup_prefixes(&self) -> [String; 1] { + [format!("dist_id_allocator/channels/{}/", self.chan_id)] + } + + /// Capture the rollback information before a newly-created manager is + /// consumed by producer/consumer binding. + pub fn from_created_manager( + chan_mgr: &ChanManager, + override_global_lease_id: Option, + override_member_lease_id: Option, + ) -> Self { + let mut rollback = Self::new( + chan_mgr.chan_id, + [override_global_lease_id, override_member_lease_id] + .into_iter() + .flatten(), + ); + if override_global_lease_id.is_none() { + rollback.record_channel_owned_cluster_lease_key(format!( + "cluster_lease/channels/{}", + chan_mgr.chan_id + )); + rollback.record_channel_owned_etcd_lease(chan_mgr.global_lease.id() as i64); + } + rollback.record_channel_owned_cluster_lease_key(format!( + "cluster_lease/id_allocator/channels/{}", + chan_mgr.chan_id + )); + rollback.record_channel_owned_etcd_lease(chan_mgr.global_long_lease.id() as i64); + if override_member_lease_id.is_none() { + rollback.record_channel_owned_etcd_lease(chan_mgr.member_lease.id() as i64); + } + rollback + } + + /// Delete all distributed metadata for the unpublished channel and revoke + /// only the etcd leases created exclusively for it. + pub async fn rollback(self, etcd_backend: &ManagedEtcdClient) -> anyhow::Result<()> { + let mut client = etcd_backend + .client() + .await + .context("connect etcd for new MPSC rollback")?; + self.rollback_with_client(&mut client).await + } + + async fn rollback_with_client(&self, client: &mut etcd::Client) -> anyhow::Result<()> { + self.install_abort_fence_and_delete_channel_state(client) + .await?; + + let mut errors = Vec::new(); + let chan_id = self.chan_id; + let mut channel_owned_etcd_lease_ids = self.channel_owned_etcd_lease_ids.clone(); + + // A failed cluster-lease helper RPC can be ambiguous: the key may have + // been committed even though the lease id was not returned. Recover the + // id from the authoritative key before deleting it. + for key in &self.channel_owned_cluster_lease_keys { + match client.get(key.clone(), None).await { + Ok(response) => { + if let Some(kv) = response.kvs().first() { + match std::str::from_utf8(kv.value()) + .ok() + .and_then(|value| value.parse::().ok()) + { + Some(lease_id) + if lease_id > 0 + && !self.protected_etcd_lease_ids.contains(&lease_id) => + { + channel_owned_etcd_lease_ids.insert(lease_id); + } + Some(_) => {} + None => errors.push(format!( + "decode channel-owned cluster lease id from key {key}" + )), + } + } + } + Err(error) => errors.push(format!("read cluster lease key {key}: {error}")), + } + } + + for key in self.exact_cleanup_keys() { + if let Err(error) = client.delete(key.clone(), None).await { + errors.push(format!("delete key {key}: {error}")); + } + } + + for prefix in self.cleanup_prefixes() { + let options = etcd::DeleteOptions::new().with_prefix(); + if let Err(error) = client.delete(prefix.clone(), Some(options)).await { + errors.push(format!("delete prefix {prefix}: {error}")); + } + } + + for lease_id in &channel_owned_etcd_lease_ids { + if let Err(error) = client.lease_revoke(*lease_id).await { + let message = error.to_string().to_ascii_lowercase(); + if !message.contains("lease not found") + && !message.contains("requested lease not found") + { + errors.push(format!("revoke lease {lease_id}: {error}")); + } + } + } + + if errors.is_empty() { + Ok(()) + } else { + anyhow::bail!( + "rollback for newly allocated MPSC chan_id={} had errors: {}", + chan_id, + errors.join("; ") + ) + } + } + + async fn install_abort_fence_and_delete_channel_state( + &self, + client: &mut etcd::Client, + ) -> anyhow::Result<()> { + let [aborted_key, meta_key, channel_prefix] = self.fence_cleanup_keys(); + let mut last_error = None; + + for attempt in 1..=ROLLBACK_FENCE_TXN_MAX_ATTEMPTS { + // This transaction is deliberately unconditional and idempotent. A + // lost response can be retried without opening a window in which a + // delayed create or bind write resurrects shared-lease metadata. + let txn = etcd::Txn::new().and_then(vec![ + etcd::TxnOp::put(aborted_key.clone(), "1", None), + etcd::TxnOp::delete(meta_key.clone(), None), + etcd::TxnOp::delete( + channel_prefix.clone(), + Some(etcd::DeleteOptions::new().with_prefix()), + ), + ]); + + match client.txn(txn).await { + Ok(_) => return Ok(()), + Err(error) => { + last_error = Some(error.to_string()); + if attempt < ROLLBACK_FENCE_TXN_MAX_ATTEMPTS { + tokio::time::sleep(Duration::from_millis(50 * attempt as u64)).await; + } + } + } + } + + anyhow::bail!( + "failed to install permanent abort fence for newly allocated MPSC chan_id={} after {} attempts: {}", + self.chan_id, + ROLLBACK_FENCE_TXN_MAX_ATTEMPTS, + last_error.unwrap_or_else(|| "unknown etcd transaction error".to_string()) + ) + } +} + /// Create a new MPSC channel id and write its metadata into etcd. /// /// Semantics mirror the Python ChanManager.create_chan design: @@ -55,19 +264,16 @@ pub struct ChanCreateConfig { /// `payload_lease_id` recorded in channel meta. pub async fn create_mpsc_channel( lease_manager: &LeaseManager, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, kv_backend_uid: LeaseBackendUid, cfg: ChanCreateConfig, rt_handle: tokio::runtime::Handle, ) -> Result { - // Build etcd client directly from provided endpoints; LeaseManager 仅负责 - // lease keepalive，不再对外暴露具体 client 实现。 - let mut client = etcd::Client::connect(etcd_endpoints.clone(), None) + let mut client = etcd_backend + .client() .await .map_err(|e| MpscError::Internal(format!("connect etcd failed: {}", e)))?; - let etcd_backend_uid = LeaseBackendUid::etcd_from(etcd_endpoints.clone()); - // 1) Allocate a new channel id using the global distributed // allocator under prefix "channels". // @@ -82,6 +288,56 @@ pub async fn create_mpsc_channel( tracing::debug!("allocated new mpsc chan_id={} ", chan_id); + let mut rollback = MpscCreateRollback::new( + chan_id, + [cfg.override_global_lease_id, cfg.override_member_lease_id] + .into_iter() + .flatten(), + ); + if cfg.override_global_lease_id.is_none() { + rollback + .record_channel_owned_cluster_lease_key(format!("cluster_lease/channels/{}", chan_id)); + } + rollback.record_channel_owned_cluster_lease_key(format!( + "cluster_lease/id_allocator/channels/{}", + chan_id + )); + let create_result = create_mpsc_channel_after_id( + lease_manager, + etcd_backend, + kv_backend_uid, + cfg, + rt_handle, + client.clone(), + chan_id, + &mut rollback, + ) + .await; + + match create_result { + Ok(chan_mgr) => Ok(chan_mgr), + Err(create_error) => match rollback.rollback_with_client(&mut client).await { + Ok(()) => Err(create_error), + Err(rollback_error) => Err(MpscError::Internal(format!( + "{}; rollback failed after allocating chan_id={}: {:#}", + create_error, chan_id, rollback_error + ))), + }, + } +} + +async fn create_mpsc_channel_after_id( + lease_manager: &LeaseManager, + etcd_backend: ManagedEtcdClient, + kv_backend_uid: LeaseBackendUid, + cfg: ChanCreateConfig, + rt_handle: tokio::runtime::Handle, + mut client: etcd::Client, + chan_id: i64, + rollback: &mut MpscCreateRollback, +) -> Result { + let etcd_backend_uid = LeaseBackendUid::etcd(etcd_backend.endpoints().clone()); + // 2) Acquire or reuse global lease for this channel's metadata; // this lease will be used to keep `/channels/meta/{chan_id}` and // `/channels/{chan_id}/next_producer_id` alive. @@ -95,6 +351,9 @@ pub async fn create_mpsc_channel( .await .map_err(|e| MpscError::Internal(format!("get_cluster_lease_id(meta) failed: {}", e)))?, }; + if cfg.override_global_lease_id.is_none() { + rollback.record_channel_owned_etcd_lease(global_lease_id); + } // 3) Acquire per-channel long-lived global lease id for id // allocator; this follows the design in `mpsc.md` where every @@ -109,6 +368,7 @@ pub async fn create_mpsc_channel( .map_err(|e| { MpscError::Internal(format!("get_cluster_lease_id(id_allocator) failed: {}", e)) })?; + rollback.record_channel_owned_etcd_lease(global_long_lease_id); // 3.1) If the caller overrides the per-channel member lease, validate that // lease before publishing any new channel metadata. Otherwise a dead @@ -190,15 +450,17 @@ pub async fn create_mpsc_channel( let meta_key = keys::etcd_meta_key(chan_id); - let compare = etcd::Compare::create_revision(meta_key.clone(), etcd::CompareOp::Equal, 0); + let aborted_key = keys::etcd_aborted_key(chan_id); + let compares = vec![ + etcd::Compare::create_revision(aborted_key, etcd::CompareOp::Equal, 0), + etcd::Compare::create_revision(meta_key.clone(), etcd::CompareOp::Equal, 0), + ]; let put_meta = etcd::TxnOp::put( meta_key.clone(), meta_bytes, Some(etcd::PutOptions::new().with_lease(global_lease_id)), ); - let txn = etcd::Txn::new() - .when(vec![compare]) - .and_then(vec![put_meta]); + let txn = etcd::Txn::new().when(compares).and_then(vec![put_meta]); let txn_res = client .txn(txn) .await @@ -207,7 +469,7 @@ pub async fn create_mpsc_channel( if !txn_res.succeeded() { return Err(MpscError::Internal(format!( - "meta key already exists for chan_id={}", + "meta create rejected because the key exists or channel creation was aborted, chan_id={}", chan_id ))); } @@ -275,6 +537,7 @@ pub async fn create_mpsc_channel( )) })?; let member_lease_id_u64 = member_resp.id() as u64; + rollback.record_channel_owned_etcd_lease(member_resp.id()); let member_lease_handle = lease_manager .register_lease_for_keepalive( etcd_backend_uid.clone(), @@ -311,6 +574,7 @@ pub async fn create_mpsc_channel( global_lease: global_lease_handle, global_long_lease: global_long_lease_handle, payload_lease: payload_lease_handle, + _etcd_backend: etcd_backend, etcd_client, }) } @@ -322,7 +586,7 @@ impl ChanManager { /// id-allocator). pub async fn new_with_chan_id( lease_manager: LeaseManager, - etcd_endpoints: Vec, + etcd_backend: ManagedEtcdClient, kv_backend_uid: LeaseBackendUid, chan_id: i64, rt_handle: tokio::runtime::Handle, @@ -337,8 +601,9 @@ impl ChanManager { // 1) 加载全局元数据，获取 TTL 以及（如有）记录下来的 // global_lease_id / global_long_lease_id，同时记录下等 // 待后续事务校验使用的 etcd 版本号。 - let etcd_backend_uid = LeaseBackendUid::etcd_from(etcd_endpoints.clone()); - let client = etcd::Client::connect(etcd_endpoints, None) + let etcd_backend_uid = LeaseBackendUid::etcd(etcd_backend.endpoints().clone()); + let client = etcd_backend + .client() .await .map_err(|e| anyhow::anyhow!("connect etcd failed: {}", e))?; let mut meta_client = client.clone(); @@ -534,7 +799,59 @@ impl ChanManager { global_lease, global_long_lease, payload_lease, + _etcd_backend: etcd_backend, etcd_client: client, }) } } + +#[cfg(test)] +mod tests { + use super::MpscCreateRollback; + use std::collections::BTreeSet; + + #[test] + fn rollback_plan_is_scoped_to_the_allocated_channel() { + let rollback = MpscCreateRollback::new(7, []); + + assert_eq!( + rollback.fence_cleanup_keys(), + [ + "/channels/aborted/7".to_string(), + "/channels/meta/7".to_string(), + "/channels/7/".to_string(), + ] + ); + assert_eq!( + rollback.exact_cleanup_keys(), + [ + "dist_id_allocator/channels/7".to_string(), + "cluster_lease/channels/7".to_string(), + "cluster_lease/id_allocator/channels/7".to_string(), + ] + ); + assert_eq!( + rollback.cleanup_prefixes(), + ["dist_id_allocator/channels/7/".to_string()] + ); + } + + #[test] + fn permanent_abort_marker_survives_channel_prefix_cleanup() { + let rollback = MpscCreateRollback::new(7, []); + let [aborted_key, _, channel_prefix] = rollback.fence_cleanup_keys(); + + assert!(!aborted_key.starts_with(&channel_prefix)); + } + + #[test] + fn rollback_never_records_override_etcd_leases_for_revoke() { + let mut rollback = MpscCreateRollback::new(7, [101, 202]); + rollback.record_channel_owned_etcd_lease(101); + rollback.record_channel_owned_etcd_lease(202); + rollback.record_channel_owned_etcd_lease(303); + rollback.record_channel_owned_etcd_lease(303); + + assert_eq!(rollback.channel_owned_etcd_lease_ids, BTreeSet::from([303])); + } +} diff --git a/fluxon_rs/fluxon_mq/src/keys.rs b/fluxon_rs/fluxon_mq/src/keys.rs index 1d55754..a1621dc 100644 --- a/fluxon_rs/fluxon_mq/src/keys.rs +++ b/fluxon_rs/fluxon_mq/src/keys.rs @@ -24,6 +24,14 @@ pub fn etcd_meta_key(chan_id: i64) -> String { s } +/// Permanent fence for a channel id whose creation was rolled back. +pub fn etcd_aborted_key(chan_id: i64) -> String { + let mut s = String::with_capacity(35); + s.push_str("/channels/aborted/"); + let _ = write!(&mut s, "{}", chan_id); + s +} + /// Etcd key for a producer registration under a channel. pub fn etcd_producer_key(chan_id: i64, producer_idx: &str) -> String { let mut s = String::with_capacity(64); @@ -217,3 +225,13 @@ pub fn etcd_register_consumer_idx_key(chan_id: i64, idx: i64) -> String { let _ = write!(&mut s, "{}", idx); s } + +#[cfg(test)] +mod tests { + use super::etcd_aborted_key; + + #[test] + fn aborted_key_uses_the_permanent_channel_fence_namespace() { + assert_eq!(etcd_aborted_key(7), "/channels/aborted/7"); + } +} diff --git a/fluxon_rs/fluxon_mq/src/lease_manager.rs b/fluxon_rs/fluxon_mq/src/lease_manager.rs index 05677f5..64b33ef 100755 --- a/fluxon_rs/fluxon_mq/src/lease_manager.rs +++ b/fluxon_rs/fluxon_mq/src/lease_manager.rs @@ -7,6 +7,7 @@ // explicitly to avoid implicit singletons. use anyhow::Result; +use fluxon_util::etcd::EtcdEndpointSet; use fluxon_util::lease_manager; use fluxon_util::run_async_from_sync::SyncAsyncBridge; use tokio::runtime::Runtime; @@ -28,7 +29,7 @@ pub use fluxon_util::lease_manager::{ /// config issues early. pub fn register_etcd_lease( rt: &Runtime, - endpoints: Vec, + endpoints: EtcdEndpointSet, ttl_seconds: i64, lease_id: u64, revoke_on_drop: bool, @@ -36,7 +37,7 @@ pub fn register_etcd_lease( let rth = rt.handle().clone(); let outer = rt .run_async_from_sync(async move { - let uid = LeaseBackendUid::etcd_from(endpoints); + let uid = LeaseBackendUid::etcd(endpoints); let mgr = &lease_manager::GLOBAL_LM; let gl = mgr .register_lease_for_keepalive( diff --git a/fluxon_rs/fluxon_mq/src/manager.rs b/fluxon_rs/fluxon_mq/src/manager.rs index b6d581a..2d3903a 100644 --- a/fluxon_rs/fluxon_mq/src/manager.rs +++ b/fluxon_rs/fluxon_mq/src/manager.rs @@ -1,12 +1,9 @@ use crate::keys; -use serde::{Deserialize, Serialize}; -use thiserror::Error; -use tracing::debug; - -use anyhow::Context; use etcd_client as etcd; -use fluxon_util::etcd::DistributeIdAllocator; +use fluxon_util::etcd::ManagedEtcdClient; use fluxon_util::lease_manager::{GeneralLease, LeaseBackendUid, LeaseManager}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; /// Initial produce offset when no messages have been produced. pub const PRODUCE_OFFSET_BEGIN: i64 = -1; @@ -169,8 +166,8 @@ pub async fn get_chan_meta( Ok(meta) } -/// Channel manager that operates on etcd metadata and cooperates with -/// the shared endpoints-scoped `LeaseManager` for lease registration. +/// Channel manager that operates on etcd metadata and cooperates with the +/// managed etcd client registry and the global lease manager. /// /// Backed by `LeaseManager` and its backend uid, this struct /// aggregates the channel id and the three etcd lease handles @@ -206,6 +203,9 @@ pub struct ChanManager { /// 决定好 payload lease id，并通过 LeaseManager 注册 /// 对应的 kvclient keepalive；此处始终持有一个有效句柄。 pub payload_lease: GeneralLease, + /// Keeps the canonical endpoint-scoped client registry entry alive. + pub(crate) _etcd_backend: ManagedEtcdClient, + /// Connected clone used by synchronous accessors throughout MQ actors. pub(crate) etcd_client: etcd::Client, } diff --git a/fluxon_rs/fluxon_mq/src/producer.rs b/fluxon_rs/fluxon_mq/src/producer.rs index fb4e4ea..796d8e2 100644 --- a/fluxon_rs/fluxon_mq/src/producer.rs +++ b/fluxon_rs/fluxon_mq/src/producer.rs @@ -166,6 +166,8 @@ impl MpscProducer { } }; let key = keys::etcd_producer_key(chan_id, &producer_idx); + let meta_key = keys::etcd_meta_key(chan_id); + let aborted_key = keys::etcd_aborted_key(chan_id); let member_meta = ProducerMemberMeta { producer_idx: producer_idx.clone(), @@ -174,19 +176,26 @@ impl MpscProducer { let member_meta_bytes = serde_json::to_vec(&member_meta) .map_err(|e| anyhow::anyhow!("serialize ProducerMemberMeta failed: {}", e))?; - let compare = etcd::Compare::create_revision(key.clone(), etcd::CompareOp::Equal, 0); + let compares = vec![ + etcd::Compare::create_revision(aborted_key.clone(), etcd::CompareOp::Equal, 0), + etcd::Compare::create_revision(meta_key.clone(), etcd::CompareOp::Greater, 0), + etcd::Compare::create_revision(key.clone(), etcd::CompareOp::Equal, 0), + ]; let put_op = etcd::TxnOp::put( key.clone(), member_meta_bytes, Some(etcd::PutOptions::new().with_lease(member_lease_id)), ); - let txn = etcd::Txn::new().when(vec![compare]).and_then(vec![put_op]); + let txn = etcd::Txn::new().when(compares).and_then(vec![put_op]); let txn_res = client .txn(txn) .await .with_context(|| format!("failed to bind producer membership key {}", key))?; if !txn_res.succeeded() { - anyhow::bail!("producer membership key {} already exists", key); + anyhow::bail!( + "producer membership bind rejected for key {}: channel aborted, meta missing, or membership already exists", + key + ); } // 4) Optionally write producer weight (default 1)，挂在 @@ -194,14 +203,29 @@ impl MpscProducer { let weight = weight.unwrap_or(1); let weight_key = keys::etcd_producer_weight_key(chan_id, &producer_idx); let global_lease_id = chan_mgr.global_lease.id() as i64; - client - .put( - weight_key.clone(), - weight.to_string(), - Some(etcd::PutOptions::new().with_lease(global_lease_id)), - ) + let weight_compares = vec![ + etcd::Compare::create_revision(aborted_key, etcd::CompareOp::Equal, 0), + etcd::Compare::create_revision(meta_key, etcd::CompareOp::Greater, 0), + etcd::Compare::create_revision(key.clone(), etcd::CompareOp::Greater, 0), + ]; + let put_weight = etcd::TxnOp::put( + weight_key.clone(), + weight.to_string(), + Some(etcd::PutOptions::new().with_lease(global_lease_id)), + ); + let weight_txn = etcd::Txn::new() + .when(weight_compares) + .and_then(vec![put_weight]); + let weight_txn_res = client + .txn(weight_txn) .await .with_context(|| format!("failed to write producer weight key {}", weight_key))?; + if !weight_txn_res.succeeded() { + anyhow::bail!( + "producer weight write rejected for key {}: channel aborted, meta missing, or membership missing", + weight_key.clone(), + ); + } let (consumer_bind_state_tx, consumer_bind_state_rx) = watch::channel(ConsumerBindState::NoneBound); diff --git a/fluxon_rs/fluxon_observability/src/keys.rs b/fluxon_rs/fluxon_observability/src/keys.rs index 88b5db0..c7bcd34 100644 --- a/fluxon_rs/fluxon_observability/src/keys.rs +++ b/fluxon_rs/fluxon_observability/src/keys.rs @@ -51,6 +51,13 @@ pub const PROM_METRIC_PROCESS_CPU_USAGE_PERCENT: &str = "process_cpu_usage_perce pub const PROM_METRIC_CONTAINER_MEMORY_USAGE_BYTES: &str = "container_memory_usage_bytes"; pub const PROM_METRIC_CONTAINER_MEMORY_LIMIT_BYTES: &str = "container_memory_limit_bytes"; +// ---------------- KV owner storage observe schema ---------------- +// +// Memory segment metrics are emitted by the master from MasterSegManager. KV SSD +// metrics are emitted by each owner because SSD shard/ring state is owner-local. +pub const PROM_METRIC_KV_SSD_CAPACITY_BYTES: &str = "kv_ssd_capacity_bytes"; +pub const PROM_METRIC_KV_SSD_USED_BYTES: &str = "kv_ssd_used_bytes"; + // ---------------- Tokio runtime observe schema (stable runtime metrics) ---------------- // // These gauges intentionally use only stable Tokio runtime metrics APIs. diff --git a/fluxon_rs/fluxon_observability/src/kv_metrics_actor.rs b/fluxon_rs/fluxon_observability/src/kv_metrics_actor.rs index 5690367..f0fe553 100644 --- a/fluxon_rs/fluxon_observability/src/kv_metrics_actor.rs +++ b/fluxon_rs/fluxon_observability/src/kv_metrics_actor.rs @@ -866,6 +866,16 @@ pub enum ObserveOp { device: String, bytes: u64, }, + SetKvSsdCapacityBytes { + node: String, + device: String, + bytes: u64, + }, + SetKvSsdUsedBytes { + node: String, + device: String, + bytes: u64, + }, SetFsMountFsBytes { mount_kind: FsMountKind, target_dir_abs: String, @@ -1170,6 +1180,8 @@ pub struct KvMetricsActorOwned { node_uptime_counter: CounterVec, segment_capacity_gauge: GaugeVec, segment_used_gauge: GaugeVec, + kv_ssd_capacity_gauge: GaugeVec, + kv_ssd_used_gauge: GaugeVec, node_network_transmit_bytes_counter: CounterVec, node_network_receive_bytes_counter: CounterVec, rdma_probe_port_count_gauge: GaugeVec, @@ -1624,6 +1636,24 @@ impl KvMetricsActorOwned { ) .expect("segment used gauge"); + let kv_ssd_capacity_gauge = GaugeVec::new( + Opts::new( + crate::keys::PROM_METRIC_KV_SSD_CAPACITY_BYTES, + "Total capacity in bytes for each owner-local KV SSD device", + ), + &["node", "device"], + ) + .expect("kv ssd capacity gauge"); + + let kv_ssd_used_gauge = GaugeVec::new( + Opts::new( + crate::keys::PROM_METRIC_KV_SSD_USED_BYTES, + "Used bytes for each owner-local KV SSD device", + ), + &["node", "device"], + ) + .expect("kv ssd used gauge"); + let node_network_transmit_bytes_counter = CounterVec::new( Opts::new( "node_network_transmit_bytes_total", @@ -1837,6 +1867,8 @@ impl KvMetricsActorOwned { register_collector(®istry, Box::new(node_uptime_counter.clone())); register_collector(®istry, Box::new(segment_capacity_gauge.clone())); register_collector(®istry, Box::new(segment_used_gauge.clone())); + register_collector(®istry, Box::new(kv_ssd_capacity_gauge.clone())); + register_collector(®istry, Box::new(kv_ssd_used_gauge.clone())); register_collector( ®istry, Box::new(node_network_transmit_bytes_counter.clone()), @@ -1911,6 +1943,8 @@ impl KvMetricsActorOwned { node_uptime_counter, segment_capacity_gauge, segment_used_gauge, + kv_ssd_capacity_gauge, + kv_ssd_used_gauge, node_network_transmit_bytes_counter, node_network_receive_bytes_counter, rdma_probe_port_count_gauge, @@ -2174,6 +2208,24 @@ impl KvMetricsActorOwned { .with_label_values(&[&node, &device]) .set(bytes as f64); } + ObserveOp::SetKvSsdCapacityBytes { + node, + device, + bytes, + } => { + self.kv_ssd_capacity_gauge + .with_label_values(&[&node, &device]) + .set(bytes as f64); + } + ObserveOp::SetKvSsdUsedBytes { + node, + device, + bytes, + } => { + self.kv_ssd_used_gauge + .with_label_values(&[&node, &device]) + .set(bytes as f64); + } ObserveOp::SetFsMountFsBytes { mount_kind, target_dir_abs, diff --git a/fluxon_rs/fluxon_ops/Cargo.toml b/fluxon_rs/fluxon_ops/Cargo.toml index f4f772a..96e73e8 100644 --- a/fluxon_rs/fluxon_ops/Cargo.toml +++ b/fluxon_rs/fluxon_ops/Cargo.toml @@ -16,7 +16,6 @@ uuid = { workspace = true } hyper = { workspace = true } libc = { workspace = true } askama = { workspace = true } -etcd-client = { workspace = true } clap = { workspace = true } futures = { workspace = true } hex = "0.4" diff --git a/fluxon_rs/fluxon_ops/build.rs b/fluxon_rs/fluxon_ops/build.rs index 585fbfc..51e95c4 100644 --- a/fluxon_rs/fluxon_ops/build.rs +++ b/fluxon_rs/fluxon_ops/build.rs @@ -59,9 +59,17 @@ print( } fn render_log_shard_helper(repo_root: &Path) -> String { - let helper_path = repo_root.join("deployment").join("utils").join("log_shard.py"); - fs::read_to_string(&helper_path) - .unwrap_or_else(|e| panic!("read log shard helper failed: {} ({})", helper_path.display(), e)) + let helper_path = repo_root + .join("deployment") + .join("utils") + .join("log_shard.py"); + fs::read_to_string(&helper_path).unwrap_or_else(|e| { + panic!( + "read log shard helper failed: {} ({})", + helper_path.display(), + e + ) + }) } fn main() { @@ -87,6 +95,10 @@ fn main() { ); println!( "cargo:rerun-if-changed={}", - repo_root.join("deployment").join("utils").join("log_shard.py").display() + repo_root + .join("deployment") + .join("utils") + .join("log_shard.py") + .display() ); } diff --git a/fluxon_rs/fluxon_ops/src/bin/fluxon_ops.rs b/fluxon_rs/fluxon_ops/src/bin/fluxon_ops.rs index 7317961..d10bd7c 100644 --- a/fluxon_rs/fluxon_ops/src/bin/fluxon_ops.rs +++ b/fluxon_rs/fluxon_ops/src/bin/fluxon_ops.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::time::Duration; use fluxon_proxy; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use fluxon_util::{ FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2, fluxon_cli_proxy_desc_etcd_key_v2, }; @@ -178,15 +179,15 @@ async fn run_serve(config: &Path, workdir: &Path) -> anyhow::Result<()> { let etcd_key = ops_panel_proxy_desc_etcd_key(fluxon_ops::OPS_SERVICE_NAME, &cli_cfg.cluster_name); - let mut etcd = etcd_client::Client::connect(cli_cfg.etcd_endpoints.clone(), None) - .await - .map_err(|e| { - anyhow::anyhow!( - "etcd connect failed while waiting for ops panel: key={} err={}", - etcd_key, - e - ) - })?; + let etcd_backend = + ManagedEtcdClient::acquire(EtcdEndpointSet::new(cli_cfg.etcd_endpoints.clone())?); + let mut etcd = etcd_backend.client().await.map_err(|e| { + anyhow::anyhow!( + "etcd connect failed while waiting for ops panel: key={} err={}", + etcd_key, + e + ) + })?; let ready_deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(30); loop { diff --git a/fluxon_rs/fluxon_ops/src/lib.rs b/fluxon_rs/fluxon_ops/src/lib.rs index 29d9434..d363030 100644 --- a/fluxon_rs/fluxon_ops/src/lib.rs +++ b/fluxon_rs/fluxon_ops/src/lib.rs @@ -14,7 +14,6 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; -use etcd_client as etcd; use futures::future::join_all; use hyper::body::HttpBody as _; use hyper::{Body, Method, Request, Response, StatusCode}; @@ -27,6 +26,7 @@ use fluxon_kv::user_rpc; use fluxon_kv::{ConfigArg, Framework, run_client}; use fluxon_proxy::{HeaderKv, PanelProxyMethod, PanelProxyResp}; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use fluxon_util::{ FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2, display_runtime_log_path, fluxon_cli_proxy_desc_etcd_key_v2, resolve_readable_log_path, @@ -80,7 +80,8 @@ const DELETE_APPLY_NO_WAIT_DELAY_SECONDS: u64 = 30; const EMBEDDED_SELECTION_SUPERVISOR_SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/selection_supervisor.py")); -const EMBEDDED_LOG_SHARD_HELPER_SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/log_shard.py")); +const EMBEDDED_LOG_SHARD_HELPER_SOURCE: &str = + include_str!(concat!(env!("OUT_DIR"), "/log_shard.py")); // Ops controller uses Fluxon user-RPC to talk to ops agents. // Keep the timeout as a fixed constant to avoid config surface area. @@ -351,7 +352,10 @@ fn workload_log_latest_shard_identity(logical_path: &Path) -> anyhow::Result anyhow::Result Ok(resolved) } -fn ensure_embedded_selection_supervisor_runtime(workdir: &Path) -> anyhow::Result<(PathBuf, PathBuf)> { +fn ensure_embedded_selection_supervisor_runtime( + workdir: &Path, +) -> anyhow::Result<(PathBuf, PathBuf)> { let runtime_dir = workdir.join(OPS_SELECTION_SUPERVISOR_DIR_NAME); std::fs::create_dir_all(&runtime_dir).with_context(|| { format!( @@ -1657,10 +1663,11 @@ fn selection_owner_supervisor( scope_key: Option<&str>, exclude_pid: Option, ) -> anyhow::Result> { - let owners: Vec = live_selection_supervisors(snapshot, Some(label), scope_key)? - .into_iter() - .filter(|supervisor| exclude_pid != Some(supervisor.pid())) - .collect(); + let owners: Vec = + live_selection_supervisors(snapshot, Some(label), scope_key)? + .into_iter() + .filter(|supervisor| exclude_pid != Some(supervisor.pid())) + .collect(); if owners.is_empty() { return Ok(None); } @@ -2068,7 +2075,16 @@ fn wait_for_selection_attached( argv: &[String], cwd: Option<&str>, ) -> anyhow::Result { - wait_for_selection_attached_for_scope(kind, name, authority, None, apply_id, owner_ts_ms, argv, cwd) + wait_for_selection_attached_for_scope( + kind, + name, + authority, + None, + apply_id, + owner_ts_ms, + argv, + cwd, + ) } fn wait_for_selection_attached_without_present_for_scope( @@ -2803,10 +2819,9 @@ impl SupervisorBackedWorkloads { fn list_workloads(&self) -> anyhow::Result> { let mut out: Vec = Vec::new(); let snapshot = selection_supervisor_proc_snapshot()?; - for status in observe_all_selection_statuses_for_snapshot( - &snapshot, - Some(self.scope_key.as_str()), - )? { + for status in + observe_all_selection_statuses_for_snapshot(&snapshot, Some(self.scope_key.as_str()))? + { let kind = status.kind.with_context(|| { format!( "selection supervisor list item missing kind: label={}", @@ -3159,7 +3174,10 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { None => { let Some(path) = resolve_readable_log_path(&logical_path) else { let resp = make_err_resp( - format!("log file is not available yet: logical_path={}", logical_path.display()), + format!( + "log file is not available yet: logical_path={}", + logical_path.display() + ), None, ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3186,138 +3204,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { }; let file_size = meta.len(); - let (start, end, start_cursor, end_cursor, effective_path, effective_file_size) = - match req.direction { - LogReadDirection::Forward => { - if let Some(cursor) = req.cursor.as_ref() { - if cursor.offset > file_size { - let resp = make_err_resp( - format!( - "cursor out of range: shard={} cursor={} file_size={}", - cursor.shard, cursor.offset, file_size - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - let mut effective_path = path.clone(); - let mut effective_shard = shard.clone(); - let mut effective_file_size = file_size; - let mut start = cursor.offset; - if cursor.offset == file_size { - if let Ok(Some(next_shard)) = - workload_log_next_shard(&logical_path, &cursor.shard) - { - let next_path = match workload_log_path_for_shard(&logical_path, &next_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&next_path) { - Ok(next_meta) => { - effective_file_size = next_meta.len(); - effective_path = next_path; - effective_shard = next_shard; - start = 0; - } - Err(e) => { - let resp = make_err_resp( - format!( - "stat next log shard failed: path={} err={}", - next_path.display(), - e - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - } - } else if let Ok(Some(latest_shard)) = - workload_log_latest_shard_identity(&logical_path) - { - if latest_shard != cursor.shard { - let latest_path = - match workload_log_path_for_shard(&logical_path, &latest_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&latest_path) { - Ok(latest_meta) => { - effective_file_size = latest_meta.len(); - effective_path = latest_path; - effective_shard = latest_shard; - start = 0; - } - Err(e) => { - let resp = make_err_resp( - format!( - "stat latest log shard failed: path={} err={}", - latest_path.display(), - e - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - } - } - } - } - let end = match max_bytes { - Some(max_bytes) => { - std::cmp::min(effective_file_size, start.saturating_add(max_bytes)) - } - None => effective_file_size, - }; - ( - start, - end, - Some(WorkloadLogCursor { - shard: effective_shard.clone(), - offset: start, - }), - Some(WorkloadLogCursor { - shard: effective_shard.clone(), - offset: end, - }), - effective_path, - effective_file_size, - ) - } else { - let end = file_size; - let start = match max_bytes { - Some(max_bytes) => end.saturating_sub(max_bytes), - None => 0, - }; - ( - start, - end, - Some(WorkloadLogCursor { - shard: shard.clone(), - offset: start, - }), - Some(WorkloadLogCursor { - shard: shard.clone(), - offset: end, - }), - path.clone(), - file_size, - ) - } - } - LogReadDirection::Backward => { - let Some(cursor) = req.cursor.as_ref() else { - let resp = make_err_resp( - "cursor is required for Backward reads".to_string(), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - }; + let (start, end, start_cursor, end_cursor, effective_path, effective_file_size) = match req + .direction + { + LogReadDirection::Forward => { + if let Some(cursor) = req.cursor.as_ref() { if cursor.offset > file_size { let resp = make_err_resp( format!( @@ -3331,30 +3222,31 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { let mut effective_path = path.clone(); let mut effective_shard = shard.clone(); let mut effective_file_size = file_size; - let mut end = cursor.offset; - if cursor.offset == 0 { - if let Ok(Some(prev_shard)) = - workload_log_previous_shard(&logical_path, &cursor.shard) + let mut start = cursor.offset; + if cursor.offset == file_size { + if let Ok(Some(next_shard)) = + workload_log_next_shard(&logical_path, &cursor.shard) { - let prev_path = match workload_log_path_for_shard(&logical_path, &prev_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&prev_path) { - Ok(prev_meta) => { - effective_file_size = prev_meta.len(); - effective_path = prev_path; - effective_shard = prev_shard; - end = effective_file_size; + let next_path = + match workload_log_path_for_shard(&logical_path, &next_shard) { + Ok(v) => v, + Err(e) => { + let resp = make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&next_path) { + Ok(next_meta) => { + effective_file_size = next_meta.len(); + effective_path = next_path; + effective_shard = next_shard; + start = 0; } Err(e) => { let resp = make_err_resp( format!( - "stat previous log shard failed: path={} err={}", - prev_path.display(), + "stat next log shard failed: path={} err={}", + next_path.display(), e ), Some(file_size), @@ -3362,11 +3254,47 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { return Ok(serde_json::to_vec(&resp).unwrap()); } } + } else if let Ok(Some(latest_shard)) = + workload_log_latest_shard_identity(&logical_path) + { + if latest_shard != cursor.shard { + let latest_path = + match workload_log_path_for_shard(&logical_path, &latest_shard) + { + Ok(v) => v, + Err(e) => { + let resp = + make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&latest_path) { + Ok(latest_meta) => { + effective_file_size = latest_meta.len(); + effective_path = latest_path; + effective_shard = latest_shard; + start = 0; + } + Err(e) => { + let resp = make_err_resp( + format!( + "stat latest log shard failed: path={} err={}", + latest_path.display(), + e + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + } + } } } - let start = match max_bytes { - Some(max_bytes) => end.saturating_sub(max_bytes), - None => 0, + let end = match max_bytes { + Some(max_bytes) => { + std::cmp::min(effective_file_size, start.saturating_add(max_bytes)) + } + None => effective_file_size, }; ( start, @@ -3382,8 +3310,103 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { effective_path, effective_file_size, ) + } else { + let end = file_size; + let start = match max_bytes { + Some(max_bytes) => end.saturating_sub(max_bytes), + None => 0, + }; + ( + start, + end, + Some(WorkloadLogCursor { + shard: shard.clone(), + offset: start, + }), + Some(WorkloadLogCursor { + shard: shard.clone(), + offset: end, + }), + path.clone(), + file_size, + ) } - }; + } + LogReadDirection::Backward => { + let Some(cursor) = req.cursor.as_ref() else { + let resp = make_err_resp( + "cursor is required for Backward reads".to_string(), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + }; + if cursor.offset > file_size { + let resp = make_err_resp( + format!( + "cursor out of range: shard={} cursor={} file_size={}", + cursor.shard, cursor.offset, file_size + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + let mut effective_path = path.clone(); + let mut effective_shard = shard.clone(); + let mut effective_file_size = file_size; + let mut end = cursor.offset; + if cursor.offset == 0 { + if let Ok(Some(prev_shard)) = + workload_log_previous_shard(&logical_path, &cursor.shard) + { + let prev_path = + match workload_log_path_for_shard(&logical_path, &prev_shard) { + Ok(v) => v, + Err(e) => { + let resp = make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&prev_path) { + Ok(prev_meta) => { + effective_file_size = prev_meta.len(); + effective_path = prev_path; + effective_shard = prev_shard; + end = effective_file_size; + } + Err(e) => { + let resp = make_err_resp( + format!( + "stat previous log shard failed: path={} err={}", + prev_path.display(), + e + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + } + } + } + let start = match max_bytes { + Some(max_bytes) => end.saturating_sub(max_bytes), + None => 0, + }; + ( + start, + end, + Some(WorkloadLogCursor { + shard: effective_shard.clone(), + offset: start, + }), + Some(WorkloadLogCursor { + shard: effective_shard.clone(), + offset: end, + }), + effective_path, + effective_file_size, + ) + } + }; if end < start { let resp = make_err_resp( @@ -3416,7 +3439,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { Ok(v) => v, Err(e) => { let resp = make_err_resp( - format!("open log failed: path={} err={}", effective_path.display(), e), + format!( + "open log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3425,7 +3452,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { if let Err(e) = std::io::Seek::seek(&mut f, std::io::SeekFrom::Start(start)) { let resp = make_err_resp( - format!("seek log failed: path={} err={}", effective_path.display(), e), + format!( + "seek log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3434,7 +3465,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { let mut buf: Vec = vec![0; len]; if let Err(e) = std::io::Read::read_exact(&mut f, &mut buf) { let resp = make_err_resp( - format!("read log failed: path={} err={}", effective_path.display(), e), + format!( + "read log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -4067,8 +4102,7 @@ fn desired_workload_matches_running( &desired.name, &desired.authority, Some(workloads.scope_key.as_str()), - ) - else { + ) else { return false; }; desired_workload_status_matches_goal(&status, desired) @@ -7334,7 +7368,7 @@ async fn handle_fluxon_embedded_proxy(
Query/stop a managed workload on a single node.

- +

@@ -12990,7 +13024,9 @@ pub async fn run_controller_blocking( let etcd_endpoints = fluxon_kv::config::normalize_etcd_addresses(&client_cfg.etcd_addresses_raw) .map_err(|e| anyhow::anyhow!("normalize_etcd_addresses failed: {}", e))?; - let mut client = etcd::Client::connect(etcd_endpoints, None) + let etcd_backend = ManagedEtcdClient::acquire(EtcdEndpointSet::new(etcd_endpoints)?); + let mut client = etcd_backend + .client() .await .with_context(|| { format!( @@ -14337,14 +14373,14 @@ mod tests { assert_eq!(scoped_b.len(), 1); assert_eq!(scoped_b[0].pid(), 22); - let listed_a = observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-a")) - .unwrap(); + let listed_a = + observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-a")).unwrap(); assert_eq!(listed_a.len(), 1); assert_eq!(listed_a[0].label, "DaemonSet/target"); assert_eq!(listed_a[0].pid, Some(11)); - let listed_b = observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-b")) - .unwrap(); + let listed_b = + observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-b")).unwrap(); assert_eq!(listed_b.len(), 1); assert_eq!(listed_b[0].label, "DaemonSet/target"); assert_eq!(listed_b[0].pid, Some(22)); @@ -14548,8 +14584,8 @@ mod tests { zombie_infos: Vec::new(), }; - let listed = observe_apply_runtime_statuses_for_snapshot("apply-1", &snapshot, None) - .unwrap(); + let listed = + observe_apply_runtime_statuses_for_snapshot("apply-1", &snapshot, None).unwrap(); assert_eq!(listed.len(), 1); assert_eq!(listed[0].name.as_deref(), Some("target-present")); assert!(listed[0].present); @@ -14774,12 +14810,8 @@ mod tests { None )); - let delete_old = workloads.delete_generation( - WorkloadKind::Deployment, - &name, - &name, - Some("apply-1"), - ); + let delete_old = + workloads.delete_generation(WorkloadKind::Deployment, &name, &name, Some("apply-1")); if !delete_old.ok { let err = delete_old.err.as_deref().unwrap_or_default(); assert!( @@ -14813,8 +14845,7 @@ mod tests { delete_current.ok, "unguarded delete should bind and retire the current visible generation: {delete_current:?}" ); - wait_for_selection_absent(WorkloadKind::Deployment, &name, &name, Some("apply-2")) - .unwrap(); + wait_for_selection_absent(WorkloadKind::Deployment, &name, &name, Some("apply-2")).unwrap(); } #[test] @@ -14826,9 +14857,12 @@ mod tests { python_exe.display() ); let workdir = tempfile::tempdir().unwrap(); - let runtime = - SelectionSupervisorRuntime::materialize(workdir.path(), workdir.path(), python_exe.as_path()) - .unwrap(); + let runtime = SelectionSupervisorRuntime::materialize( + workdir.path(), + workdir.path(), + python_exe.as_path(), + ) + .unwrap(); assert!(runtime.script_path.exists()); assert!( runtime @@ -14849,9 +14883,12 @@ mod tests { python_exe.display() ); let workdir = tempfile::tempdir().unwrap(); - let runtime = - SelectionSupervisorRuntime::materialize(workdir.path(), workdir.path(), python_exe.as_path()) - .unwrap(); + let runtime = SelectionSupervisorRuntime::materialize( + workdir.path(), + workdir.path(), + python_exe.as_path(), + ) + .unwrap(); let log_path = workdir.path().join("startup.log"); let command = vec![ python_exe.display().to_string(), @@ -14876,7 +14913,9 @@ mod tests { "--".to_string(), "/bin/true".to_string(), ]; - let pid = runtime.spawn_detached_command(&log_path, command.as_slice()).unwrap(); + let pid = runtime + .spawn_detached_command(&log_path, command.as_slice()) + .unwrap(); let deadline = Instant::now() + Duration::from_secs(10); let expected = "owner-ts-ms must be positive"; let mut saw_expected = false; @@ -15164,7 +15203,9 @@ mod tests { }), max_bytes: Some(65536), }; - let raw = handler.handle("n1".into(), &serde_json::to_vec(&req).unwrap()).unwrap(); + let raw = handler + .handle("n1".into(), &serde_json::to_vec(&req).unwrap()) + .unwrap(); let resp: ReadWorkloadLogResp = serde_json::from_slice(&raw).unwrap(); assert!(resp.ok, "{resp:?}"); assert_eq!(resp.text.as_deref(), Some("new\n")); @@ -15209,7 +15250,9 @@ mod tests { }), max_bytes: Some(65536), }; - let raw = handler.handle("n1".into(), &serde_json::to_vec(&req).unwrap()).unwrap(); + let raw = handler + .handle("n1".into(), &serde_json::to_vec(&req).unwrap()) + .unwrap(); let resp: ReadWorkloadLogResp = serde_json::from_slice(&raw).unwrap(); assert!(resp.ok, "{resp:?}"); assert_eq!(resp.text.as_deref(), Some("old\n")); diff --git a/fluxon_rs/fluxon_pyo3/Cargo.toml b/fluxon_rs/fluxon_pyo3/Cargo.toml index 74ec69e..b896d24 100644 --- a/fluxon_rs/fluxon_pyo3/Cargo.toml +++ b/fluxon_rs/fluxon_pyo3/Cargo.toml @@ -16,6 +16,7 @@ fastws_transport = ["fluxon_kv/fastws_transport"] sockudo_ws_transport = ["fluxon_kv/sockudo_ws_transport"] tcp_transport = ["fluxon_kv/tcp_transport"] tcp_thread_transport = ["fluxon_kv/tcp_thread_transport"] +fluxon_fs_video_ffmpeg = [] [dependencies] pyo3 = { version = "0.22", features = ["extension-module", "abi3-py38"] } @@ -47,3 +48,7 @@ fluxon_cli = { path = "../fluxon_cli" } fluxon_ops = { path = "../fluxon_ops" } fluxon_fs = { path = "../fluxon_fs" } fluxon_proxy = { path = "../fluxon_proxy" } + +[build-dependencies] +cc = { workspace = true } +pkg-config = "0.3" diff --git a/fluxon_rs/fluxon_pyo3/build.rs b/fluxon_rs/fluxon_pyo3/build.rs index 90c419b..8a745b1 100644 --- a/fluxon_rs/fluxon_pyo3/build.rs +++ b/fluxon_rs/fluxon_pyo3/build.rs @@ -60,6 +60,7 @@ print("\n".join(args)) fn main() { emit_python_test_embed_link_args(); + build_fluxon_fs_video_ffmpeg_shim(); let target_dir = get_target_dir(); let runtime_search_subdirs = load_runtime_search_subdirs(); @@ -102,10 +103,100 @@ fn main() { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=src/"); + println!("cargo:rerun-if-changed=src/video_reader_ffi.c"); println!("cargo:rerun-if-changed=../target/debug/"); println!("cargo:rerun-if-changed=../target/release/"); } +fn build_fluxon_fs_video_ffmpeg_shim() { + println!("cargo:rerun-if-env-changed=CARGO_FEATURE_FLUXON_FS_VIDEO_FFMPEG"); + if env::var_os("CARGO_FEATURE_FLUXON_FS_VIDEO_FFMPEG").is_none() { + return; + } + + let ffmpeg_libs = ["libavformat", "libavcodec", "libavutil", "libswscale"]; + let mut build = cc::Build::new(); + build + .file("src/video_reader_ffi.c") + .flag_if_supported("-Wno-deprecated-declarations"); + + let mut pkg_config_ok = true; + let mut link_paths = Vec::new(); + let mut link_libs = Vec::new(); + let mut framework_paths = Vec::new(); + let mut frameworks = Vec::new(); + for lib_name in ffmpeg_libs { + let mut config = pkg_config::Config::new(); + config.cargo_metadata(false); + match config.probe(lib_name) { + Ok(lib) => { + for include_path in lib.include_paths { + build.include(include_path); + } + link_paths.extend(lib.link_paths); + link_libs.extend(lib.libs); + framework_paths.extend(lib.framework_paths); + frameworks.extend(lib.frameworks); + } + Err(err) => { + pkg_config_ok = false; + println!( + "cargo:warning=pkg-config could not find {} for FluxonFS VideoReader: {}. Falling back to default compiler/linker paths.", + lib_name, err + ); + break; + } + } + } + + build.compile("fluxon_fs_video_reader_ffi"); + + if pkg_config_ok { + emit_unique_link_searches("native", link_paths); + emit_unique_link_searches("framework", framework_paths); + emit_unique_link_libs(link_libs); + emit_unique_frameworks(frameworks); + } else { + println!("cargo:rustc-link-lib=avformat"); + println!("cargo:rustc-link-lib=avcodec"); + println!("cargo:rustc-link-lib=avutil"); + println!("cargo:rustc-link-lib=swscale"); + } +} + +fn emit_unique_link_searches(kind: &str, paths: Vec) { + let mut seen = Vec::new(); + for path in paths { + if seen.iter().any(|existing| existing == &path) { + continue; + } + println!("cargo:rustc-link-search={kind}={}", path.display()); + seen.push(path); + } +} + +fn emit_unique_link_libs(libs: Vec) { + let mut seen = Vec::new(); + for lib in libs { + if seen.iter().any(|existing| existing == &lib) { + continue; + } + println!("cargo:rustc-link-lib={lib}"); + seen.push(lib); + } +} + +fn emit_unique_frameworks(frameworks: Vec) { + let mut seen = Vec::new(); + for framework in frameworks { + if seen.iter().any(|existing| existing == &framework) { + continue; + } + println!("cargo:rustc-link-lib=framework={framework}"); + seen.push(framework); + } +} + fn emit_python_test_embed_link_args() { println!("cargo:rerun-if-env-changed=PYTHON"); diff --git a/fluxon_rs/fluxon_pyo3/src/error.rs b/fluxon_rs/fluxon_pyo3/src/error.rs index 97ab680..57024b4 100644 --- a/fluxon_rs/fluxon_pyo3/src/error.rs +++ b/fluxon_rs/fluxon_pyo3/src/error.rs @@ -221,6 +221,23 @@ pub(crate) fn new_key_being_written_error( error_class.call((), Some(&kwargs)).unwrap().into() } +pub(crate) fn new_key_already_exists_error( + py: Python<'_>, + message: &str, + key: Option<&str>, +) -> PyObject { + let api_error_module = py.import_bound("fluxon_py.api_error").unwrap(); + let error_class = api_error_module.getattr("KeyAlreadyExistsError").unwrap(); + + let kwargs = PyDict::new_bound(py); + kwargs.set_item("message", message).unwrap(); + if let Some(k) = key { + kwargs.set_item("key", k).unwrap(); + } + + error_class.call((), Some(&kwargs)).unwrap().into() +} + pub(crate) fn new_storage_full_error( py: Python<'_>, message: &str, @@ -353,6 +370,7 @@ pub(crate) fn new_payload_lease_not_found_error( /// - Api::InvalidPutMasterState -> PutDoneFailedError /// - Api::KeyNotFound -> KeyNotFoundError（携带 key） /// - Api::KeyBeingWritten -> KeyBeingWrittenError（携带 key） +/// - Api::KeyAlreadyExists -> KeyAlreadyExistsError（携带 key） /// - Api::NoSpace -> StorageFullError（available_space=free_capacity） /// - 其他 -> NetworkError（携带格式化消息） pub(crate) fn py_error_from_kv_error( @@ -386,6 +404,11 @@ pub(crate) fn py_error_from_kv_error( &format!("{}: Key is currently being written: {}", prefix, key), Some(key), ), + KvError::Api(CoreApiError::KeyAlreadyExists { key }) => new_key_already_exists_error( + py, + &format!("{}: Key already exists: {}", prefix, key), + Some(key), + ), KvError::Api(CoreApiError::InvalidArgument { detail }) => { new_invalid_argument_error(py, &format!("{}: Invalid argument: {}", prefix, detail)) } diff --git a/fluxon_rs/fluxon_pyo3/src/etcd.rs b/fluxon_rs/fluxon_pyo3/src/etcd.rs index f3c27d7..cd7558c 100644 --- a/fluxon_rs/fluxon_pyo3/src/etcd.rs +++ b/fluxon_rs/fluxon_pyo3/src/etcd.rs @@ -1,12 +1,13 @@ use etcd_client as etcd; -use fluxon_util::auto_clean_map::{AutoCleanMap, AutoCleanMapEntry}; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; +use fluxon_util::lease_manager::{GLOBAL_LM, GeneralLease, LeaseBackendUid, LeaseRegisterKind}; use fluxon_util::run_async_from_sync::SyncAsyncBridge; use pyo3::prelude::*; use pyo3::pybacked::PyBackedBytes; use pyo3::types::{PyBytes, PyList, PyTuple}; use pyo3::{PyErr, PyObject}; use std::future::Future; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::runtime::Runtime; use tracing::debug; @@ -27,7 +28,10 @@ fn normalize_raw_endpoint(endpoint: &str) -> PyResult { Ok(format!("http://{}", endpoint)) } -fn normalize_raw_endpoints(endpoints: Vec, component: &str) -> PyResult> { +pub(crate) fn normalize_raw_endpoints( + endpoints: Vec, + component: &str, +) -> PyResult> { if endpoints.is_empty() { return Err(PyErr::new::(format!( "{} requires at least one endpoint", @@ -41,54 +45,12 @@ fn normalize_raw_endpoints(endpoints: Vec, component: &str) -> PyResult< Ok(normalized) } -struct EtcdKvBackend { +pub(crate) fn endpoint_set_from_raw( endpoints: Vec, - client: tokio::sync::RwLock>, -} - -impl EtcdKvBackend { - fn new(endpoints: Vec) -> Self { - Self { - endpoints, - client: tokio::sync::RwLock::new(None), - } - } - - async fn client(&self) -> anyhow::Result { - { - let guard = self.client.read().await; - if let Some(client) = guard.as_ref() { - return Ok(client.clone()); - } - } - - let mut guard = self.client.write().await; - if let Some(client) = guard.as_ref() { - return Ok(client.clone()); - } - - let client = etcd::Client::connect(self.endpoints.clone(), None) - .await - .map_err(|e| { - anyhow::anyhow!( - "failed to connect etcd endpoints={:?}: {:?}", - self.endpoints, - e - ) - })?; - *guard = Some(client.clone()); - Ok(client) - } - - async fn clear_client(&self) { - let mut guard = self.client.write().await; - *guard = None; - } -} - -fn etcd_kv_backend_map() -> &'static AutoCleanMap, EtcdKvBackend> { - static MAP: OnceLock, EtcdKvBackend>> = OnceLock::new(); - MAP.get_or_init(|| AutoCleanMap::new()) + component: &str, +) -> PyResult { + EtcdEndpointSet::new(normalize_raw_endpoints(endpoints, component)?) + .map_err(|e| PyErr::new::(e.to_string())) } fn is_reconnectable_etcd_error(err: &etcd::Error) -> bool { @@ -107,7 +69,7 @@ fn is_reconnectable_etcd_error_text(msg: &str) -> bool { } async fn run_etcd_op( - backend: AutoCleanMapEntry, EtcdKvBackend>, + backend: ManagedEtcdClient, context: String, mut op: F, ) -> anyhow::Result @@ -117,14 +79,14 @@ where { let mut last_err = None; for attempt in 1..=2 { - let client = backend.client().await?; - match op(client).await { + let snapshot = backend.snapshot().await?; + match op(snapshot.client()).await { Ok(value) => return Ok(value), Err(err) => { let should_retry = attempt == 1 && is_reconnectable_etcd_error(&err); last_err = Some(err); if should_retry { - backend.clear_client().await; + snapshot.invalidate().await; continue; } let err = last_err.take().expect("etcd error must be recorded"); @@ -140,17 +102,16 @@ where #[pyclass(name = "EtcdKvClient")] pub struct PyEtcdKvClient { rt: Arc, - endpoints: Vec, - backend: AutoCleanMapEntry, EtcdKvBackend>, + endpoints: EtcdEndpointSet, + backend: ManagedEtcdClient, } #[pymethods] impl PyEtcdKvClient { #[new] fn new(endpoints: Vec) -> PyResult { - let endpoints = normalize_raw_endpoints(endpoints, "EtcdKvClient")?; - let backend = etcd_kv_backend_map() - .get_or_init(endpoints.clone(), || EtcdKvBackend::new(endpoints.clone())); + let endpoints = endpoint_set_from_raw(endpoints, "EtcdKvClient")?; + let backend = ManagedEtcdClient::acquire(endpoints.clone()); Ok(Self { rt: crate::mpsc::get_global_runtime(), endpoints, @@ -408,19 +369,21 @@ impl PyEtcdKvClient { } fn __repr__(&self) -> String { - format!("", self.endpoints) + format!("", self.endpoints.as_slice()) } } #[pyclass(name = "EtcdLock")] pub struct PyEtcdLock { rt: Arc, - endpoints: Vec, + endpoints: EtcdEndpointSet, + backend: ManagedEtcdClient, name: String, ttl_seconds: i64, timeout_seconds: f64, lease_id: Option, lock_key: Option>, + lease_guard: Option, } #[pymethods] @@ -433,7 +396,7 @@ impl PyEtcdLock { ttl_seconds: i64, timeout_seconds: Option, ) -> PyResult { - let endpoints = normalize_raw_endpoints(endpoints, "EtcdLock")?; + let endpoints = endpoint_set_from_raw(endpoints, "EtcdLock")?; if ttl_seconds <= 0 { return Err(PyErr::new::(format!( "EtcdLock ttl_seconds must be > 0, got {}", @@ -450,12 +413,14 @@ impl PyEtcdLock { Ok(Self { rt: crate::mpsc::get_global_runtime(), + backend: ManagedEtcdClient::acquire(endpoints.clone()), endpoints, name, ttl_seconds, timeout_seconds, lease_id: None, lock_key: None, + lease_guard: None, }) } @@ -484,9 +449,11 @@ impl PyEtcdLock { } let endpoints = self.endpoints.clone(); + let backend = self.backend.clone(); let name = self.name.clone(); let ttl_seconds = self.ttl_seconds; let timeout_duration = Duration::from_secs_f64(timeout_seconds); + let rt_handle = self.rt.handle().clone(); let t0 = Instant::now(); debug!( @@ -500,7 +467,7 @@ impl PyEtcdLock { let outer = py .allow_threads(|| { self.rt.run_async_from_sync(async move { - let mut client = etcd::Client::connect(endpoints, None).await.map_err(|e| { + let mut client = backend.client().await.map_err(|e| { anyhow::anyhow!("failed to connect etcd for lock {}: {:?}", name, e) })?; @@ -508,6 +475,29 @@ impl PyEtcdLock { anyhow::anyhow!("failed to grant etcd lease for lock {}: {:?}", name, e) })?; let lease_id = lease_resp.id(); + let lease_guard = match GLOBAL_LM + .register_lease_for_keepalive( + LeaseBackendUid::etcd(endpoints), + ttl_seconds, + lease_id as u64, + LeaseRegisterKind::Etcd { + revoke_on_drop: false, + }, + rt_handle, + ) + .await + { + Ok(lease_guard) => lease_guard, + Err(err) => { + let _ = client.lease_revoke(lease_id).await; + return Err(anyhow::anyhow!( + "failed to start keepalive for etcd lock {} lease {}: {}", + name, + lease_id, + err + )); + } + }; match tokio::time::timeout( timeout_duration, @@ -518,8 +508,12 @@ impl PyEtcdLock { ) .await { - Ok(Ok(resp)) => Ok(Some((lease_id, resp.key().to_vec()))), + Ok(Ok(resp)) => { + let lock_key = resp.key().to_vec(); + Ok(Some((lease_id, lock_key, lease_guard))) + } Ok(Err(err)) => { + drop(lease_guard); let _ = client.lease_revoke(lease_id).await; Err(anyhow::anyhow!( "failed to acquire etcd lock {}: {:?}", @@ -528,6 +522,7 @@ impl PyEtcdLock { )) } Err(_) => { + drop(lease_guard); let _ = client.lease_revoke(lease_id).await; Ok(None) } @@ -541,9 +536,10 @@ impl PyEtcdLock { outer.map_err(|e| PyErr::new::(e.to_string()))?; match acquire_outcome { - Some((lease_id, lock_key)) => { + Some((lease_id, lock_key, lease_guard)) => { self.lease_id = Some(lease_id); self.lock_key = Some(lock_key); + self.lease_guard = Some(lease_guard); debug!( target: "fluxon_pyo3::etcd", "end etcd lock acquire: name={}, lease_id={}, elapsed_ms={}", @@ -567,15 +563,19 @@ impl PyEtcdLock { fn release(&mut self, py: Python<'_>) -> PyResult { let Some(lock_key) = self.lock_key.clone() else { + self.lease_id = None; + self.lease_guard = None; return Ok(false); }; let Some(lease_id) = self.lease_id else { self.lock_key = None; + self.lease_guard = None; return Ok(false); }; - let endpoints = self.endpoints.clone(); + let backend = self.backend.clone(); let name = self.name.clone(); + let lease_guard = self.lease_guard.take(); let t0 = Instant::now(); debug!( @@ -585,42 +585,48 @@ impl PyEtcdLock { lease_id ); - let outer = py - .allow_threads(|| { - self.rt.run_async_from_sync(async move { - let mut client = etcd::Client::connect(endpoints, None).await.map_err(|e| { - anyhow::anyhow!("failed to connect etcd for unlock {}: {:?}", name, e) - })?; + let bridge_result = py.allow_threads(|| { + self.rt.run_async_from_sync(async move { + let lease_guard = lease_guard; + let mut client = backend.client().await.map_err(|e| { + anyhow::anyhow!("failed to connect etcd for unlock {}: {:?}", name, e) + })?; - let unlock_result = client.unlock(lock_key).await.map(|_| true).map_err(|e| { + let unlock_result = + client.unlock(lock_key).await.map(|_| true).map_err(|e| { anyhow::anyhow!("failed to unlock etcd lock {}: {:?}", name, e) }); - let revoke_result = client.lease_revoke(lease_id).await; - match (unlock_result, revoke_result) { - (Ok(unlocked), Ok(_)) => Ok(unlocked), - (Ok(_), Err(err)) => Err(anyhow::anyhow!( - "failed to revoke etcd lease {} for lock {}: {:?}", - lease_id, - name, - err - )), - (Err(err), Ok(_)) => Err(err), - (Err(unlock_err), Err(revoke_err)) => Err(anyhow::anyhow!( - "failed to unlock etcd lock {}: {}; failed to revoke lease {}: {:?}", - name, - unlock_err, - lease_id, - revoke_err - )), - } - }) + // Keep the lease alive until the unlock request finishes, then + // unregister it before the final explicit revoke. + drop(lease_guard); + let revoke_result = client.lease_revoke(lease_id).await; + match (unlock_result, revoke_result) { + (Ok(unlocked), Ok(_)) => Ok(unlocked), + (Ok(_), Err(err)) => Err(anyhow::anyhow!( + "failed to revoke etcd lease {} for lock {}: {:?}", + lease_id, + name, + err + )), + (Err(err), Ok(_)) => Err(err), + (Err(unlock_err), Err(revoke_err)) => Err(anyhow::anyhow!( + "failed to unlock etcd lock {}: {}; failed to revoke lease {}: {:?}", + name, + unlock_err, + lease_id, + revoke_err + )), + } }) - .map_err(|e| anyhow::anyhow!("runtime bridge failed in EtcdLock.release: {}", e)) - .map_err(|e| PyErr::new::(e.to_string()))?; + }); self.lock_key = None; self.lease_id = None; + let outer = bridge_result + .map_err(|e| anyhow::anyhow!("runtime bridge failed in EtcdLock.release: {}", e)) + .map_err(|e| PyErr::new::(e.to_string()))?; + let released = outer.map_err(|e| PyErr::new::(e.to_string()))?; debug!( @@ -671,14 +677,6 @@ impl PyEtcdLock { #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - static TEST_KEY_SEQ: AtomicUsize = AtomicUsize::new(1); - - fn unique_test_endpoints() -> Vec { - let seq = TEST_KEY_SEQ.fetch_add(1, Ordering::Relaxed); - vec![format!("http://unit-test-etcd-backend-{}", seq)] - } #[test] fn normalize_raw_endpoint_accepts_raw_host_port() { @@ -715,7 +713,10 @@ mod tests { #[test] fn etcd_kv_client_constructor_normalizes_raw_endpoints() { let client = PyEtcdKvClient::new(vec!["127.0.0.1:2379".to_string()]).unwrap(); - assert_eq!(client.endpoints, vec!["http://127.0.0.1:2379"]); + assert_eq!( + client.endpoints.as_slice(), + ["http://127.0.0.1:2379".to_string()] + ); } #[test] @@ -727,7 +728,10 @@ mod tests { Some(1.0), ) .unwrap(); - assert_eq!(lock.endpoints, vec!["http://127.0.0.1:2379"]); + assert_eq!( + lock.endpoints.as_slice(), + ["http://127.0.0.1:2379".to_string()] + ); } #[test] @@ -743,30 +747,6 @@ mod tests { ); } - #[test] - fn etcd_kv_backend_map_reuses_and_auto_cleans_live_entries() { - let endpoints = unique_test_endpoints(); - let map = etcd_kv_backend_map(); - assert!(map.with_existing(&endpoints, |_| ()).is_none()); - - { - let entry_a = - map.get_or_init(endpoints.clone(), || EtcdKvBackend::new(endpoints.clone())); - assert!(map.with_existing(&endpoints, |_| ()).is_some()); - - { - let entry_b = map.get_or_init(endpoints.clone(), || { - panic!("live backend entry should be reused") - }); - assert!(std::ptr::eq(&*entry_a, &*entry_b)); - } - - assert!(map.with_existing(&endpoints, |_| ()).is_some()); - } - - assert!(map.with_existing(&endpoints, |_| ()).is_none()); - } - #[test] fn reconnectable_error_text_matches_transient_transport_failures() { assert!(is_reconnectable_etcd_error_text("StatusCode::UNAVAILABLE")); diff --git a/fluxon_rs/fluxon_pyo3/src/lease_manager.rs b/fluxon_rs/fluxon_pyo3/src/lease_manager.rs index 99b8a79..b682461 100755 --- a/fluxon_rs/fluxon_pyo3/src/lease_manager.rs +++ b/fluxon_rs/fluxon_pyo3/src/lease_manager.rs @@ -1,4 +1,5 @@ use anyhow::Result as AnyResult; +use fluxon_util::etcd::ManagedEtcdClient; use fluxon_util::lease_manager::GLOBAL_LM; use fluxon_util::lease_manager::snapshot_active_lease_debug as lm_snapshot_active_lease_debug; use fluxon_util::run_async_from_sync::SyncAsyncBridge; @@ -11,7 +12,6 @@ use tokio::runtime::Runtime; use tracing::debug; // Use the shared PyO3 helper that submits Python callables to a Python // ThreadPoolExecutor and waits without holding the GIL for the whole duration. -use etcd_client as etcd; use fluxon_mq::lease_manager::{LeaseBackendUid, LeaseRegisterKind}; use fluxon_util::pyo3::run_longtime_py_function; @@ -117,9 +117,12 @@ impl PyLeaseBackendUid { fn __repr__(&self) -> String { match self.inner.kind() { - fluxon_util::lease_manager::LeaseType::Etcd => match self.inner.endpoints() { - Some(v) if !v.is_empty() => { - format!("", v.join(",")) + fluxon_util::lease_manager::LeaseType::Etcd => match self.inner.etcd_endpoint_set() { + Some(v) => { + format!( + "", + v.as_slice().join(",") + ) } _ => "".to_string(), }, @@ -187,18 +190,20 @@ impl LeaseManagerHandle { "begin allocate_etcd_lease: endpoints={}, ttl_seconds={}, revoke_on_drop={}", endpoints.join(","), ttl_seconds, revoke ); + let endpoints = crate::etcd::endpoint_set_from_raw(endpoints, "LeaseManagerHandle")?; let rth = self.rt.handle().clone(); let outer = py .allow_threads(|| { self.rt.run_async_from_sync(async move { - let uid = LeaseBackendUid::etcd_from(endpoints.clone()); - let mut client = etcd::Client::connect(endpoints, None).await.map_err(|e| { + let uid = LeaseBackendUid::etcd(endpoints.clone()); + let backend = ManagedEtcdClient::acquire(endpoints); + let mut client = backend.client().await.map_err(|e| { anyhow::anyhow!("failed to connect etcd when allocating lease: {:?}", e) })?; let resp = client.lease_grant(ttl_seconds, None).await?; let id = resp.id() as u64; let rt = rth; - GLOBAL_LM + match GLOBAL_LM .register_lease_for_keepalive( uid, ttl_seconds, @@ -209,6 +214,13 @@ impl LeaseManagerHandle { rt, ) .await + { + Ok(lease) => Ok(lease), + Err(err) => { + let _ = client.lease_revoke(id as i64).await; + Err(err) + } + } }) }) .map_err(|e| anyhow::anyhow!("runtime bridge failed in allocate_etcd_lease: {}", e)) @@ -243,12 +255,13 @@ impl LeaseManagerHandle { "begin register_etcd_lease: endpoints={}, ttl_seconds={}, lease_id={}, revoke_on_drop={}, register_by={}", endpoints.join(","), ttl_seconds, lease_id, revoke, register_by ); + let endpoints = crate::etcd::endpoint_set_from_raw(endpoints, "LeaseManagerHandle")?; fluxon_mq::lease_manager::record_register_by(lease_id, register_by); let rth = self.rt.handle().clone(); let outer = py .allow_threads(|| { self.rt.run_async_from_sync(async move { - let uid = LeaseBackendUid::etcd_from(endpoints); + let uid = LeaseBackendUid::etcd(endpoints); let rt = rth; GLOBAL_LM .register_lease_for_keepalive( diff --git a/fluxon_rs/fluxon_pyo3/src/lib.rs b/fluxon_rs/fluxon_pyo3/src/lib.rs index 631f1ea..3a7e007 100644 --- a/fluxon_rs/fluxon_pyo3/src/lib.rs +++ b/fluxon_rs/fluxon_pyo3/src/lib.rs @@ -29,12 +29,13 @@ use fluxon_kv::user_api::FlatDict; use fluxon_kv::user_api::FlatValue; use fluxon_kv::user_api::FluxonUserApi; use fluxon_kv::{ - ConfigArg, Framework, KvClientTrait, KvGetResult, + ConfigArg, Framework, KvClientTrait, KvGetResultWithSource, config::{ClientConfig, MasterConfig}, run_client, run_master, }; use fluxon_ops; use fluxon_proxy; +use fluxon_util::etcd::{EtcdEndpointSet, ManagedEtcdClient}; use fluxon_util::run_async_from_sync::{SyncAsyncBridge, borrow_stable_owner}; use fluxon_util::{ FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2, fluxon_cli_proxy_desc_etcd_key_v2, @@ -67,6 +68,9 @@ pub use etcd::{PyEtcdKvClient, PyEtcdLock}; pub use mpsc::{MpscConsumerHandle, MpscContext, MpscProducerHandle}; mod lease_manager; pub use lease_manager::{LeaseManagerHandle, PyGeneralLease, PyLeaseBackendUid}; +mod video_reader; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +pub use video_reader::FluxonFsVideoReader; #[derive(Debug, Clone, PartialEq, Eq)] struct RdmavDriverEnvUpdate { @@ -1082,6 +1086,29 @@ impl FluxonFsAgent { res.map_err(pyerr_from_fs_agent_error) } + #[pyo3(signature = (export_name, relpath, height, width, num_threads, request_identity=None))] + fn open_video_reader( + &self, + export_name: String, + relpath: String, + height: i64, + width: i64, + num_threads: i64, + request_identity: Option<(String, String)>, + py: Python, + ) -> PyResult { + video_reader::open_video_reader_from_agent( + py, + self.inner.clone(), + export_name, + relpath, + height, + width, + num_threads, + request_identity, + ) + } + #[pyo3(signature = (file_abs, mode))] fn open_plan(&self, file_abs: String, mode: String, py: Python) -> PyResult { let panic_ctx = format!("open_plan panic: file_abs={} mode={}", file_abs, mode,); @@ -2126,9 +2153,12 @@ fn fluxon_ops_controller_blocking( // - If etcd connect/get stalls without returning, the controller would hang forever and never // bring up the HTTP endpoint needed by test_runner/start_test_bed. // - Therefore we hard-bound etcd operations so the supervisor can restart on persistent faults. + let etcd_backend = ManagedEtcdClient::acquire(EtcdEndpointSet::new( + cli_cfg.etcd_endpoints.clone(), + )?); let mut etcd = tokio::time::timeout( tokio::time::Duration::from_secs(5), - etcd_client::Client::connect(cli_cfg.etcd_endpoints.clone(), None), + etcd_backend.client(), ) .await .map_err(|_| anyhow::anyhow!( @@ -2938,13 +2968,14 @@ impl KvClient { /// on the caller to keep the pointed-to memory alive until the async call completes. /// /// The backend encoding/copy runs on the Rust runtime without holding the Python GIL. - #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, callback=None))] + #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, reject_if_exists=false, callback=None))] fn put( &self, key: &str, ptrs: Vec<(u8, u64, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, callback: Option, py: Python, ) -> PyObject { @@ -2954,6 +2985,7 @@ impl KvClient { ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, callback: Option, py: Python, ) -> ApiResult { @@ -2982,6 +3014,9 @@ impl KvClient { if reject_if_inflight_same_key { o.0.push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfInflightSameKey); } + if reject_if_exists { + o.0.push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfExists); + } o }; @@ -3041,6 +3076,7 @@ impl KvClient { ptrs_owned, lease_id, reject_if_inflight_same_key, + reject_if_exists, callback, py, ) @@ -3048,13 +3084,14 @@ impl KvClient { } /// Put a key-value pair and wait for completion before returning. - #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false))] + #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, reject_if_exists=false))] fn put_blocking( &self, key: &str, ptrs: Vec<(u8, u64, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, py: Python, ) -> PyObject { fn put_blocking_inner( @@ -3063,6 +3100,7 @@ impl KvClient { ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exists: bool, py: Python, ) -> ApiResult { if ptrs.len() > (u32::MAX as usize) { @@ -3094,6 +3132,11 @@ impl KvClient { .0 .push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfInflightSameKey); } + if reject_if_exists { + put_opts + .0 + .push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfExists); + } let result = match py.allow_threads(|| { runtime.run_async_from_sync(async { unsafe { framework.kv_put_ptrs(&key, ptrs, put_opts).await } @@ -3141,6 +3184,7 @@ impl KvClient { ptrs_owned, lease_id, reject_if_inflight_same_key, + reject_if_exists, py, ) .into_py_object(py) @@ -3162,53 +3206,20 @@ impl KvClient { let future = async move { tracing::debug!("KvClient.get future start: key={}", key); - let result = framework.kv_get(&key).await; - tracing::debug!("KvClient.get framework.kv_get returned: key={}", key); + let result = framework.kv_get_with_source(&key).await; + tracing::debug!( + "KvClient.get framework.kv_get_with_source returned: key={}", + key + ); match result { - Ok(KvGetResult::Owner(Some(rust_holder))) => Python::with_gil(|py| { - tracing::debug!( - "KvClient.get entering Python::with_gil owner path: key={}", - key - ); - let mem_holder = MemHolder::new(rust_holder); - let py_result = match mem_holder.into_py_mem_holder(py) { - ApiResult::Success(py_holder) => { - if let Some(cb) = callback { - let args = PyTuple::new_bound(py, &[py_holder.bind(py)]); - match cb.call1(py, args) { - Ok(result) => ApiResult::new_success(result), - Err(e) => ApiResult::new_error(new_general_error( - py, - &format!("Callback failed: {}", e), - )), - } - } else { - ApiResult::new_success(py_holder.into_any()) - } - } - err => err, - }; - tracing::debug!( - "KvClient.get leaving Python::with_gil owner path: key={}", - key - ); - py_result - }), - Ok(KvGetResult::Owner(None)) => Python::with_gil(|py| { - ApiResult::new_error(new_key_not_found_error( - py, - &format!("Key not found: {}", key), - Some(&key), - )) - }), - Ok(KvGetResult::External(Some(external_mem_holder))) => { + Ok(KvGetResultWithSource::Owner(Some((rust_holder, source_kind)))) => { Python::with_gil(|py| { tracing::debug!( - "KvClient.get entering Python::with_gil external path: key={}", + "KvClient.get entering Python::with_gil owner path: key={}", key ); - let pyo3_external = ExternalMemHolder::new(external_mem_holder); - let py_result = match pyo3_external.into_py_mem_holder(py) { + let mem_holder = MemHolder::new(rust_holder, source_kind); + let py_result = match mem_holder.into_py_mem_holder(py) { ApiResult::Success(py_holder) => { if let Some(cb) = callback { let args = PyTuple::new_bound(py, &[py_holder.bind(py)]); @@ -3226,13 +3237,53 @@ impl KvClient { err => err, }; tracing::debug!( - "KvClient.get leaving Python::with_gil external path: key={}", + "KvClient.get leaving Python::with_gil owner path: key={}", key ); py_result }) } - Ok(KvGetResult::External(None)) => Python::with_gil(|py| { + Ok(KvGetResultWithSource::Owner(None)) => Python::with_gil(|py| { + ApiResult::new_error(new_key_not_found_error( + py, + &format!("Key not found: {}", key), + Some(&key), + )) + }), + Ok(KvGetResultWithSource::External(Some(( + external_mem_holder, + source_kind, + )))) => Python::with_gil(|py| { + tracing::debug!( + "KvClient.get entering Python::with_gil external path: key={}", + key + ); + let pyo3_external = + ExternalMemHolder::new(external_mem_holder, source_kind); + let py_result = match pyo3_external.into_py_mem_holder(py) { + ApiResult::Success(py_holder) => { + if let Some(cb) = callback { + let args = PyTuple::new_bound(py, &[py_holder.bind(py)]); + match cb.call1(py, args) { + Ok(result) => ApiResult::new_success(result), + Err(e) => ApiResult::new_error(new_general_error( + py, + &format!("Callback failed: {}", e), + )), + } + } else { + ApiResult::new_success(py_holder.into_any()) + } + } + err => err, + }; + tracing::debug!( + "KvClient.get leaving Python::with_gil external path: key={}", + key + ); + py_result + }), + Ok(KvGetResultWithSource::External(None)) => Python::with_gil(|py| { ApiResult::new_error(new_key_not_found_error( py, &format!("Key not found: {}", key), @@ -3287,31 +3338,27 @@ impl KvClient { let framework = borrow_stable_owner(&framework); let result = match py.allow_threads(|| { - runtime.run_async_from_sync(async { framework.kv_get(&key).await }) + runtime.run_async_from_sync(async { framework.kv_get_with_source(&key).await }) }) { Ok(v) => v, Err(e) => Err(anyhow::anyhow!("runtime bridge failed: {}", e).into()), }; match result { - Ok(KvGetResult::Owner(Some(rust_holder))) => { - let mem_holder = MemHolder::new(rust_holder); + Ok(KvGetResultWithSource::Owner(Some((rust_holder, source_kind)))) => { + let mem_holder = MemHolder::new(rust_holder, source_kind); mem_holder.into_py_mem_holder(py) } - Ok(KvGetResult::Owner(None)) => ApiResult::new_error(new_key_not_found_error( - py, - &format!("Key not found: {}", key), - Some(&key), - )), - Ok(KvGetResult::External(Some(external_mem_holder))) => { - let pyo3_external = ExternalMemHolder::new(external_mem_holder); + Ok(KvGetResultWithSource::Owner(None)) => ApiResult::new_error( + new_key_not_found_error(py, &format!("Key not found: {}", key), Some(&key)), + ), + Ok(KvGetResultWithSource::External(Some((external_mem_holder, source_kind)))) => { + let pyo3_external = ExternalMemHolder::new(external_mem_holder, source_kind); pyo3_external.into_py_mem_holder(py) } - Ok(KvGetResult::External(None)) => ApiResult::new_error(new_key_not_found_error( - py, - &format!("Key not found: {}", key), - Some(&key), - )), + Ok(KvGetResultWithSource::External(None)) => ApiResult::new_error( + new_key_not_found_error(py, &format!("Key not found: {}", key), Some(&key)), + ), Err(e) => { let err_obj = crate::error::py_error_from_kv_error(py, &e, "Get failed"); ApiResult::new_error(err_obj) @@ -4154,6 +4201,8 @@ fn fluxon_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + #[cfg(feature = "fluxon_fs_video_ffmpeg")] + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/fluxon_rs/fluxon_pyo3/src/memholder.rs b/fluxon_rs/fluxon_pyo3/src/memholder.rs index b2750e6..bb7afba 100755 --- a/fluxon_rs/fluxon_pyo3/src/memholder.rs +++ b/fluxon_rs/fluxon_pyo3/src/memholder.rs @@ -1,9 +1,11 @@ use std::sync::Arc; +use fluxon_kv::master_kv_router::msg_pack::GetSourceKind; use fluxon_kv::memholder::{ ExternalMemHolder as RustExternalMemHolder, UserMemHolder as RustMemHolder, }; use parking_lot::RwLock; +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use super::{ApiResult, new_general_error}; @@ -18,11 +20,26 @@ enum MemHolderInner { #[pyclass] pub struct MemHolder { holder: MemHolderInner, + source_kind: Option, access_cache: RwLock>, } +fn benchmark_source_kind_name(source_kind: GetSourceKind) -> &'static str { + match source_kind { + GetSourceKind::Memory => "memory", + GetSourceKind::Ssd => "ssd", + } +} + #[pymethods] impl MemHolder { + /// Return the source selected for this GET operation. + fn _benchmark_source_kind(&self) -> PyResult<&'static str> { + self.source_kind + .map(benchmark_source_kind_name) + .ok_or_else(|| PyRuntimeError::new_err("GET source is unavailable for this holder")) + } + /// Decode the held flat dict as a Python dict. /// /// The binary parsing runs without the Python GIL; only Python object creation holds the GIL. @@ -55,9 +72,10 @@ impl MemHolder { } impl MemHolder { - pub(crate) fn new(holder: Arc) -> Self { + pub(crate) fn new(holder: Arc, source_kind: GetSourceKind) -> Self { Self { holder: MemHolderInner::Seg(holder), + source_kind: Some(source_kind), access_cache: RwLock::new(None), } } @@ -65,6 +83,7 @@ impl MemHolder { pub(crate) fn new_owned(bytes: Vec) -> Self { Self { holder: MemHolderInner::Owned(Arc::<[u8]>::from(bytes)), + source_kind: None, access_cache: RwLock::new(None), } } @@ -112,11 +131,17 @@ impl MemHolder { #[pyclass] pub struct ExternalMemHolder { pub(crate) holder: Arc, + source_kind: GetSourceKind, access_cache: RwLock>, } #[pymethods] impl ExternalMemHolder { + /// Return the source selected for this GET operation. + fn _benchmark_source_kind(&self) -> &'static str { + benchmark_source_kind_name(self.source_kind) + } + /// Decode the held flat dict as a Python dict. /// /// The binary parsing runs without the Python GIL; only Python object creation holds the GIL. @@ -144,9 +169,10 @@ impl ExternalMemHolder { } impl ExternalMemHolder { - pub(crate) fn new(holder: Arc) -> Self { + pub(crate) fn new(holder: Arc, source_kind: GetSourceKind) -> Self { Self { holder, + source_kind, access_cache: RwLock::new(None), } } diff --git a/fluxon_rs/fluxon_pyo3/src/mpsc.rs b/fluxon_rs/fluxon_pyo3/src/mpsc.rs index 22d0d07..8a44c2a 100644 --- a/fluxon_rs/fluxon_pyo3/src/mpsc.rs +++ b/fluxon_rs/fluxon_pyo3/src/mpsc.rs @@ -11,7 +11,7 @@ use fluxon_mq::consumer::{ use fluxon_mq::{ ChanManager, MpscConsumer as CoreMpscConsumer, MpscError as CoreMpscError, MpscProducer as CoreMpscProducer, ShutdownCtl, - create::{ChanCreateConfig, create_mpsc_channel}, + create::{ChanCreateConfig, MpscCreateRollback, create_mpsc_channel}, }; use pyo3::Py; use pyo3::PyErr; @@ -25,6 +25,7 @@ use crate::flatdict_zerocopy::{FlatDictDataOwner, decode_flat_dict_to_wrapped_py use crate::lease_manager::PyLeaseBackendUid; use fluxon_kv::{Framework as KvFramework, KvClientTrait}; use fluxon_mq::lease_manager::LeaseBackendUid; +use fluxon_util::etcd::ManagedEtcdClient; use fluxon_util::lease_manager::{GLOBAL_LM, LeaseManager}; use fluxon_util::run_async_from_sync::SyncAsyncBridge; use tracing::{debug, warn}; @@ -129,11 +130,10 @@ fn finalize_payload_result( // (LeaseManagerHandle and PyLease moved to lease_manager.rs) /// Shared MPSC context bound to a specific etcd endpoint set. -/// Holds only the endpoints; runtime and lease managers are -/// singletons under the hood. +/// Holds the auto-clean client registry handle for that endpoint set. #[pyclass] pub struct MpscContext { - endpoints: Vec, + etcd_backend: ManagedEtcdClient, kv_backend_uid: LeaseBackendUid, kv_framework: Arc, kv_runtime: Handle, @@ -144,9 +144,8 @@ pub struct MpscContext { impl MpscContext { /// Create a new MPSC context from etcd endpoints. /// - /// The runtime is created lazily on first use and shared - /// globally; `LeaseManager::for_endpoints` handles per-endpoint - /// singletons internally. + /// The endpoint-scoped etcd client is connected lazily on first use and + /// shared with other Rust/PyO3 components through `fluxon_util`. #[new] fn new( py: Python<'_>, @@ -157,8 +156,9 @@ impl MpscContext { let uid = kv_backend_uid.borrow(py).backend_uid().clone(); let kv_client_ref = kv_client.borrow(py); let mq_context = crate::new_fluxon_mq_context(&kv_client_ref)?; + let endpoints = crate::etcd::endpoint_set_from_raw(etcd_endpoints, "MpscContext")?; Ok(Self { - endpoints: etcd_endpoints, + etcd_backend: ManagedEtcdClient::acquire(endpoints), kv_backend_uid: uid, kv_framework: mq_context.kv_framework.clone(), kv_runtime: mq_context.runtime.clone(), @@ -205,7 +205,7 @@ impl MpscContext { } }; - let endpoints = self.endpoints.clone(); + let etcd_backend = self.etcd_backend.clone(); let kv_backend_uid = self.kv_backend_uid.clone(); let self_info = self .kv_framework @@ -231,13 +231,19 @@ impl MpscContext { .run_async_from_sync(async move { let lease_manager = GLOBAL_LM.clone(); + let category = match (parent_mpmc_id_opt, parent_mpmc_member_id_opt) { + (Some(pid), Some(_mid)) => fluxon_mq::keys::MqCategory::MpmcSub { parent_mpmc_id: pid }, + (None, None) => fluxon_mq::keys::MqCategory::Mpsc, + _ => return Err(anyhow::anyhow!("parent_mpmc_id_opt and parent_mpmc_member_id_opt must be both provided or both None")), + }; + // Construct channel manager either by creating a new // channel or by binding to an existing one. let chan_mgr: anyhow::Result = async { match chan_id { Some(id) => ChanManager::new_with_chan_id( lease_manager.clone(), - endpoints.clone(), + etcd_backend.clone(), kv_backend_uid.clone(), id, rth.clone(), @@ -260,7 +266,7 @@ impl MpscContext { }; create_mpsc_channel( &GLOBAL_LM, - endpoints.clone(), + etcd_backend.clone(), kv_backend_uid.clone(), cfg, rth.clone(), @@ -272,14 +278,15 @@ impl MpscContext { } .await; let chan_mgr = chan_mgr?; + let rollback = chan_id.is_none().then(|| { + MpscCreateRollback::from_created_manager( + &chan_mgr, + override_global_lease_id, + override_member_lease_id, + ) + }); - let category = match (parent_mpmc_id_opt, parent_mpmc_member_id_opt) { - (Some(pid), Some(_mid)) => fluxon_mq::keys::MqCategory::MpmcSub { parent_mpmc_id: pid }, - (None, None) => fluxon_mq::keys::MqCategory::Mpsc, - _ => return Err(anyhow::anyhow!("parent_mpmc_id_opt and parent_mpmc_member_id_opt must be both provided or both None")), - }; - - CoreMpscProducer::bind_mpsc( + let bind_result = CoreMpscProducer::bind_mpsc( chan_mgr, ttl_seconds, weight, @@ -292,7 +299,22 @@ impl MpscContext { observe_node_role, observe, ) - .await + .await; + match bind_result { + Ok(producer) => Ok(producer), + Err(bind_error) => { + if let Some(rollback) = rollback { + if let Err(rollback_error) = rollback.rollback(&etcd_backend).await { + return Err(anyhow::anyhow!( + "{}; rollback failed for newly created MPSC producer: {:#}", + bind_error, + rollback_error + )); + } + } + Err(bind_error) + } + } })) .map_err(|e| PyErr::new::(format!( "runtime bridge failed in new_producer: {}", @@ -347,7 +369,7 @@ impl MpscContext { } }; - let endpoints = self.endpoints.clone(); + let etcd_backend = self.etcd_backend.clone(); let kv_backend_uid = self.kv_backend_uid.clone(); let shutdown = ShutdownCtl::new(); let shutdown_for_core = shutdown.clone(); @@ -374,11 +396,17 @@ impl MpscContext { .run_async_from_sync(async move { let lease_manager = GLOBAL_LM.clone(); + let category = match (parent_mpmc_id_opt, parent_mpmc_member_id_opt) { + (Some(pid), Some(_mid)) => fluxon_mq::keys::MqCategory::MpmcSub { parent_mpmc_id: pid }, + (None, None) => fluxon_mq::keys::MqCategory::Mpsc, + _ => return Err(anyhow::anyhow!("parent_mpmc_id_opt and parent_mpmc_member_id_opt must be both provided or both None")), + }; + let chan_mgr: anyhow::Result = async { match chan_id { Some(id) => ChanManager::new_with_chan_id( lease_manager.clone(), - endpoints.clone(), + etcd_backend.clone(), kv_backend_uid.clone(), id, rth.clone(), @@ -403,7 +431,7 @@ impl MpscContext { }; create_mpsc_channel( &GLOBAL_LM, - endpoints.clone(), + etcd_backend.clone(), kv_backend_uid.clone(), cfg, rth.clone(), @@ -415,14 +443,15 @@ impl MpscContext { } .await; let chan_mgr = chan_mgr?; + let rollback = chan_id.is_none().then(|| { + MpscCreateRollback::from_created_manager( + &chan_mgr, + override_global_lease_id, + override_member_lease_id, + ) + }); - let category = match (parent_mpmc_id_opt, parent_mpmc_member_id_opt) { - (Some(pid), Some(_mid)) => fluxon_mq::keys::MqCategory::MpmcSub { parent_mpmc_id: pid }, - (None, None) => fluxon_mq::keys::MqCategory::Mpsc, - _ => return Err(anyhow::anyhow!("parent_mpmc_id_opt and parent_mpmc_member_id_opt must be both provided or both None")), - }; - - CoreMpscConsumer::bind_mpsc( + let bind_result = CoreMpscConsumer::bind_mpsc( chan_mgr, ttl_seconds, lifecycle, @@ -434,7 +463,22 @@ impl MpscContext { observe_node_role, observe, ) - .await + .await; + match bind_result { + Ok(consumer) => Ok(consumer), + Err(bind_error) => { + if let Some(rollback) = rollback { + if let Err(rollback_error) = rollback.rollback(&etcd_backend).await { + return Err(anyhow::anyhow!( + "{}; rollback failed for newly created MPSC consumer: {:#}", + bind_error, + rollback_error + )); + } + } + Err(bind_error) + } + } })) .map_err(|e| PyErr::new::(format!( "runtime bridge failed in new_consumer: {}", diff --git a/fluxon_rs/fluxon_pyo3/src/video_reader.rs b/fluxon_rs/fluxon_pyo3/src/video_reader.rs new file mode 100644 index 0000000..040d700 --- /dev/null +++ b/fluxon_rs/fluxon_pyo3/src/video_reader.rs @@ -0,0 +1,599 @@ +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use std::collections::{BTreeMap, VecDeque}; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use std::ffi::{c_char, c_int, c_void}; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use std::panic::{self, AssertUnwindSafe}; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use std::ptr; +use std::sync::Arc; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use fluxon_fs::config::FluxonFsRequestIdentity; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use parking_lot::Mutex; +use pyo3::exceptions::PyRuntimeError; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +use pyo3::types::{PyAny, PyBytes, PyDict, PyTuple}; + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +const VIDEO_PAGE_CACHE_MAX_PAGES: usize = 64; +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +const VIDEO_ERR_BUF_BYTES: usize = 4096; + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +#[repr(C)] +struct FluxonVideoIo { + user_data: *mut c_void, + read_at: Option< + unsafe extern "C" fn( + user_data: *mut c_void, + offset: i64, + buf: *mut u8, + buf_size: c_int, + out_len: *mut c_int, + err_buf: *mut c_char, + err_buf_len: usize, + ) -> c_int, + >, +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +unsafe extern "C" { + fn fluxon_fs_video_decode_frames( + io: *const FluxonVideoIo, + file_size: i64, + indices: *const i64, + indices_len: c_int, + out_width: c_int, + out_height: c_int, + num_threads: c_int, + out_data: *mut u8, + out_data_len: i64, + err_buf: *mut c_char, + err_buf_len: usize, + ) -> c_int; +} + +pub fn open_video_reader_from_agent( + py: Python<'_>, + agent: Arc, + export_name: String, + relpath: String, + height: i64, + width: i64, + num_threads: i64, + request_identity: Option<(String, String)>, +) -> PyResult { + #[cfg(not(feature = "fluxon_fs_video_ffmpeg"))] + { + let _ = ( + py, + agent, + export_name, + relpath, + height, + width, + num_threads, + request_identity, + ); + Err(PyRuntimeError::new_err( + "FluxonFS VideoReader requires building fluxon_pyo3 with --features fluxon_fs_video_ffmpeg", + )) + } + + #[cfg(feature = "fluxon_fs_video_ffmpeg")] + { + let request_identity = crate::py_request_identity_tuple_to_core(request_identity)?; + let reader = FluxonFsVideoReader::open( + agent, + export_name, + relpath, + height, + width, + num_threads, + request_identity, + py, + )?; + Ok(Py::new(py, reader)?.into_py(py)) + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +#[pyclass] +pub struct FluxonFsVideoReader { + agent: Arc, + export_name: String, + relpath: String, + path_for_err: String, + size: i64, + mtime_ns: i64, + height: i64, + width: i64, + num_threads: i64, + request_identity: Option, + page_bytes: i64, + page_cache: Mutex, + decode_calls: AtomicU64, + decode_frames_requested: AtomicU64, + decode_errors: AtomicU64, + read_at_calls: AtomicU64, + read_at_requested_bytes: AtomicU64, + read_at_returned_bytes: AtomicU64, + page_cache_hits: AtomicU64, + page_cache_misses: AtomicU64, + remote_read_calls: AtomicU64, + remote_read_bytes: AtomicU64, + closed: AtomicBool, +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +struct VideoPageCache { + pages: BTreeMap>, + lru: VecDeque, + max_pages: usize, +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +impl VideoPageCache { + fn new(max_pages: usize) -> Self { + Self { + pages: BTreeMap::new(), + lru: VecDeque::new(), + max_pages: max_pages.max(1), + } + } + + fn get(&mut self, page_idx: i64) -> Option> { + let data = self.pages.get(&page_idx)?.clone(); + self.touch(page_idx); + Some(data) + } + + fn insert(&mut self, page_idx: i64, data: Vec) { + self.pages.insert(page_idx, data); + self.touch(page_idx); + while self.pages.len() > self.max_pages { + let Some(evict_idx) = self.lru.pop_front() else { + break; + }; + if self.pages.remove(&evict_idx).is_some() { + continue; + } + } + } + + fn touch(&mut self, page_idx: i64) { + self.lru.retain(|idx| *idx != page_idx); + self.lru.push_back(page_idx); + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +impl FluxonFsVideoReader { + fn open( + agent: Arc, + export_name: String, + relpath: String, + height: i64, + width: i64, + num_threads: i64, + request_identity: Option, + py: Python<'_>, + ) -> PyResult { + validate_non_empty(&export_name, "export_name")?; + validate_non_empty(&relpath, "relpath")?; + validate_positive_c_int(height, "height")?; + validate_positive_c_int(width, "width")?; + validate_positive_c_int(num_threads, "num_threads")?; + + let path_for_err = format!("fluxonfs://{}/{}", export_name, relpath); + let stat = py + .allow_threads(|| { + agent.remote_stat_by_handle_with_identity( + &export_name, + &relpath, + &path_for_err, + request_identity.as_ref(), + ) + }) + .map_err(crate::pyerr_from_fs_agent_error)?; + + if !stat.exists { + return Err(PyRuntimeError::new_err(format!( + "FluxonFS video not found: {}", + path_for_err + ))); + } + if !stat.is_file || stat.is_dir { + return Err(PyRuntimeError::new_err(format!( + "FluxonFS video path is not a file: {}", + path_for_err + ))); + } + if stat.size < 0 || stat.mtime_ns < 0 { + return Err(PyRuntimeError::new_err(format!( + "FluxonFS video stat returned invalid size/mtime: {} size={} mtime_ns={}", + path_for_err, stat.size, stat.mtime_ns + ))); + } + + Ok(Self { + agent, + export_name, + relpath, + path_for_err, + size: stat.size, + mtime_ns: stat.mtime_ns, + height, + width, + num_threads, + request_identity, + page_bytes: fluxon_fs::agent::REMOTE_CHUNK_BYTES as i64, + page_cache: Mutex::new(VideoPageCache::new(VIDEO_PAGE_CACHE_MAX_PAGES)), + decode_calls: AtomicU64::new(0), + decode_frames_requested: AtomicU64::new(0), + decode_errors: AtomicU64::new(0), + read_at_calls: AtomicU64::new(0), + read_at_requested_bytes: AtomicU64::new(0), + read_at_returned_bytes: AtomicU64::new(0), + page_cache_hits: AtomicU64::new(0), + page_cache_misses: AtomicU64::new(0), + remote_read_calls: AtomicU64::new(0), + remote_read_bytes: AtomicU64::new(0), + closed: AtomicBool::new(false), + }) + } + + fn read_at_bytes(&self, offset: i64, n: usize) -> Result, String> { + self.read_at_calls.fetch_add(1, Ordering::Relaxed); + self.read_at_requested_bytes + .fetch_add(n as u64, Ordering::Relaxed); + if self.closed.load(Ordering::SeqCst) { + return Err("FluxonFS VideoReader is closed".to_string()); + } + if offset < 0 { + return Err(format!("read offset must be non-negative: {}", offset)); + } + if n == 0 || offset >= self.size { + return Ok(Vec::new()); + } + let available = self.size.saturating_sub(offset) as usize; + let want = n.min(available); + let mut out = Vec::with_capacity(want); + let mut pos = offset; + while out.len() < want { + let page_idx = pos / self.page_bytes; + let page_start = page_idx * self.page_bytes; + let page_off = (pos - page_start) as usize; + let page = self.read_page(page_idx, page_start)?; + if page_off >= page.len() { + break; + } + let take = (want - out.len()).min(page.len() - page_off); + out.extend_from_slice(&page[page_off..page_off + take]); + pos = pos.saturating_add(take as i64); + } + self.read_at_returned_bytes + .fetch_add(out.len() as u64, Ordering::Relaxed); + Ok(out) + } + + fn read_page(&self, page_idx: i64, page_start: i64) -> Result, String> { + { + let mut cache = self.page_cache.lock(); + if let Some(page) = cache.get(page_idx) { + self.page_cache_hits.fetch_add(1, Ordering::Relaxed); + return Ok(page); + } + } + self.page_cache_misses.fetch_add(1, Ordering::Relaxed); + + let n = self.page_bytes.min(self.size.saturating_sub(page_start)); + if n <= 0 { + return Ok(Vec::new()); + } + self.remote_read_calls.fetch_add(1, Ordering::Relaxed); + let page = self + .agent + .remote_read_chunk_by_handle_with_identity( + &self.export_name, + &self.relpath, + page_start, + n, + self.size, + self.mtime_ns, + true, + &self.path_for_err, + self.request_identity.as_ref(), + ) + .map_err(|err| err.to_string())?; + self.remote_read_bytes + .fetch_add(page.len() as u64, Ordering::Relaxed); + + let mut cache = self.page_cache.lock(); + cache.insert(page_idx, page.clone()); + Ok(page) + } + + fn ensure_open(&self) -> PyResult<()> { + if self.closed.load(Ordering::SeqCst) { + Err(PyRuntimeError::new_err("FluxonFS VideoReader is closed")) + } else { + Ok(()) + } + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +#[pymethods] +impl FluxonFsVideoReader { + fn read_frames_numpy(&self, indices: &Bound<'_, PyAny>, py: Python<'_>) -> PyResult { + self.ensure_open()?; + let frame_indices: Vec = indices.extract()?; + for idx in &frame_indices { + if *idx < 0 { + return Err(PyValueError::new_err( + "FluxonFS VideoReader frame indices must be non-negative", + )); + } + } + + let frame_count = frame_indices.len(); + if frame_count > c_int::MAX as usize { + return Err(PyValueError::new_err( + "FluxonFS VideoReader frame count is too large", + )); + } + self.decode_calls.fetch_add(1, Ordering::Relaxed); + self.decode_frames_requested + .fetch_add(frame_count as u64, Ordering::Relaxed); + let frame_bytes = checked_frame_bytes(self.height, self.width)?; + let out_len = frame_bytes + .checked_mul(frame_count) + .ok_or_else(|| PyValueError::new_err("FluxonFS VideoReader output is too large"))?; + let mut out = vec![0u8; out_len]; + + if frame_count > 0 { + let mut err_buf = vec![0i8; VIDEO_ERR_BUF_BYTES]; + let decode_res = py.allow_threads(|| { + let mut state = VideoReadCallbackState { reader: self }; + let io = FluxonVideoIo { + user_data: (&mut state as *mut VideoReadCallbackState<'_>).cast::(), + read_at: Some(fluxon_video_read_at), + }; + unsafe { + fluxon_fs_video_decode_frames( + &io as *const FluxonVideoIo, + self.size, + frame_indices.as_ptr(), + frame_count as c_int, + self.width as c_int, + self.height as c_int, + self.num_threads as c_int, + out.as_mut_ptr(), + out.len() as i64, + err_buf.as_mut_ptr(), + err_buf.len(), + ) + } + }); + if decode_res != 0 { + self.decode_errors.fetch_add(1, Ordering::Relaxed); + let msg = c_err_buf_to_string(&err_buf); + return Err(PyRuntimeError::new_err(if msg.is_empty() { + "FluxonFS VideoReader decode failed".to_string() + } else { + format!("FluxonFS VideoReader decode failed: {}", msg) + })); + } + } + + numpy_array_from_vec(py, out, frame_count, self.height, self.width) + } + + fn stats(&self, py: Python<'_>) -> PyResult { + let cache_pages = { + let cache = self.page_cache.lock(); + cache.pages.len() as u64 + }; + let stats = PyDict::new_bound(py); + stats.set_item("file_size", self.size)?; + stats.set_item("mtime_ns", self.mtime_ns)?; + stats.set_item("height", self.height)?; + stats.set_item("width", self.width)?; + stats.set_item("num_threads", self.num_threads)?; + stats.set_item("page_bytes", self.page_bytes)?; + stats.set_item("page_cache_pages", cache_pages)?; + stats.set_item("page_cache_max_pages", VIDEO_PAGE_CACHE_MAX_PAGES as u64)?; + stats.set_item("decode_calls", self.decode_calls.load(Ordering::Relaxed))?; + stats.set_item( + "decode_frames_requested", + self.decode_frames_requested.load(Ordering::Relaxed), + )?; + stats.set_item("decode_errors", self.decode_errors.load(Ordering::Relaxed))?; + stats.set_item("read_at_calls", self.read_at_calls.load(Ordering::Relaxed))?; + stats.set_item( + "read_at_requested_bytes", + self.read_at_requested_bytes.load(Ordering::Relaxed), + )?; + stats.set_item( + "read_at_returned_bytes", + self.read_at_returned_bytes.load(Ordering::Relaxed), + )?; + stats.set_item( + "page_cache_hits", + self.page_cache_hits.load(Ordering::Relaxed), + )?; + stats.set_item( + "page_cache_misses", + self.page_cache_misses.load(Ordering::Relaxed), + )?; + stats.set_item( + "remote_read_calls", + self.remote_read_calls.load(Ordering::Relaxed), + )?; + stats.set_item( + "remote_read_bytes", + self.remote_read_bytes.load(Ordering::Relaxed), + )?; + Ok(stats.into_py(py)) + } + + fn close(&self) { + self.closed.store(true, Ordering::SeqCst); + } + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + #[pyo3(signature = (_exc_type=None, _exc=None, _tb=None))] + fn __exit__( + &self, + _exc_type: Option<&Bound<'_, PyAny>>, + _exc: Option<&Bound<'_, PyAny>>, + _tb: Option<&Bound<'_, PyAny>>, + ) -> bool { + self.close(); + false + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +struct VideoReadCallbackState<'a> { + reader: &'a FluxonFsVideoReader, +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +unsafe extern "C" fn fluxon_video_read_at( + user_data: *mut c_void, + offset: i64, + buf: *mut u8, + buf_size: c_int, + out_len: *mut c_int, + err_buf: *mut c_char, + err_buf_len: usize, +) -> c_int { + let result = panic::catch_unwind(AssertUnwindSafe(|| { + if user_data.is_null() || buf.is_null() || out_len.is_null() || buf_size < 0 { + return Err("invalid FluxonFS video read callback argument".to_string()); + } + let state = unsafe { &mut *(user_data as *mut VideoReadCallbackState<'_>) }; + let data = state.reader.read_at_bytes(offset, buf_size as usize)?; + let data_len = data.len(); + if data_len > 0 { + unsafe { + ptr::copy_nonoverlapping(data.as_ptr(), buf, data_len); + } + } + unsafe { + *out_len = data_len as c_int; + } + Ok(()) + })); + + match result { + Ok(Ok(())) => 0, + Ok(Err(err)) => { + write_c_err(err_buf, err_buf_len, &err); + -1 + } + Err(_) => { + write_c_err( + err_buf, + err_buf_len, + "FluxonFS video read callback panicked", + ); + -1 + } + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn validate_non_empty(value: &str, name: &str) -> PyResult<()> { + if value.trim().is_empty() { + Err(PyValueError::new_err(format!("{} must be non-empty", name))) + } else { + Ok(()) + } +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn validate_positive_c_int(value: i64, name: &str) -> PyResult<()> { + if value <= 0 { + return Err(PyValueError::new_err(format!("{} must be positive", name))); + } + if value > c_int::MAX as i64 { + return Err(PyValueError::new_err(format!( + "{} must be <= {}", + name, + c_int::MAX + ))); + } + Ok(()) +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn checked_frame_bytes(height: i64, width: i64) -> PyResult { + let bytes = height + .checked_mul(width) + .and_then(|v| v.checked_mul(3)) + .ok_or_else(|| PyValueError::new_err("FluxonFS VideoReader frame shape is too large"))?; + usize::try_from(bytes) + .map_err(|_| PyValueError::new_err("FluxonFS VideoReader frame shape is too large")) +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn numpy_array_from_vec( + py: Python<'_>, + data: Vec, + frame_count: usize, + height: i64, + width: i64, +) -> PyResult { + let np = py.import_bound("numpy")?; + let bytes = PyBytes::new_bound(py, &data); + let arr = np.call_method1("frombuffer", (bytes, "uint8"))?; + let shape = PyTuple::new_bound( + py, + [ + (frame_count as i64).into_py(py), + height.into_py(py), + width.into_py(py), + 3i64.into_py(py), + ], + ); + let reshaped = arr.call_method1("reshape", (shape,))?; + let owned = reshaped.call_method0("copy")?; + Ok(owned.into_py(py)) +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn c_err_buf_to_string(buf: &[i8]) -> String { + let bytes: Vec = buf + .iter() + .copied() + .take_while(|b| *b != 0) + .map(|b| b as u8) + .collect(); + String::from_utf8_lossy(&bytes).trim().to_string() +} + +#[cfg(feature = "fluxon_fs_video_ffmpeg")] +fn write_c_err(err_buf: *mut c_char, err_buf_len: usize, msg: &str) { + if err_buf.is_null() || err_buf_len == 0 { + return; + } + let bytes = msg.as_bytes(); + let take = bytes.len().min(err_buf_len.saturating_sub(1)); + unsafe { + ptr::copy_nonoverlapping(bytes.as_ptr().cast::(), err_buf, take); + *err_buf.add(take) = 0; + } +} diff --git a/fluxon_rs/fluxon_pyo3/src/video_reader_ffi.c b/fluxon_rs/fluxon_pyo3/src/video_reader_ffi.c new file mode 100644 index 0000000..9f008db --- /dev/null +++ b/fluxon_rs/fluxon_pyo3/src/video_reader_ffi.c @@ -0,0 +1,450 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +typedef int (*FluxonVideoReadAtFn)( + void *user_data, + int64_t offset, + uint8_t *buf, + int buf_size, + int *out_len, + char *err_buf, + size_t err_buf_len); + +typedef struct FluxonVideoIo { + void *user_data; + FluxonVideoReadAtFn read_at; +} FluxonVideoIo; + +typedef struct FluxonAvioOpaque { + const FluxonVideoIo *io; + int64_t pos; + int64_t size; + char *err_buf; + size_t err_buf_len; +} FluxonAvioOpaque; + +static void fluxon_set_err(char *err_buf, size_t err_buf_len, const char *msg) { + if (err_buf == NULL || err_buf_len == 0) { + return; + } + snprintf(err_buf, err_buf_len, "%s", msg); +} + +static void fluxon_set_av_err(char *err_buf, size_t err_buf_len, const char *ctx, int err) { + char av_msg[AV_ERROR_MAX_STRING_SIZE] = {0}; + av_strerror(err, av_msg, sizeof(av_msg)); + if (err_buf == NULL || err_buf_len == 0) { + return; + } + snprintf(err_buf, err_buf_len, "%s: %s", ctx, av_msg); +} + +static int fluxon_avio_read(void *opaque, uint8_t *buf, int buf_size) { + FluxonAvioOpaque *state = (FluxonAvioOpaque *)opaque; + if (state == NULL || state->io == NULL || state->io->read_at == NULL) { + return AVERROR(EIO); + } + if (buf_size <= 0) { + return 0; + } + if (state->pos >= state->size) { + return AVERROR_EOF; + } + + int want = buf_size; + int64_t available = state->size - state->pos; + if (available < want) { + want = (int)available; + } + if (want <= 0) { + return AVERROR_EOF; + } + + int out_len = 0; + int ok = state->io->read_at( + state->io->user_data, + state->pos, + buf, + want, + &out_len, + state->err_buf, + state->err_buf_len); + if (ok != 0) { + return AVERROR(EIO); + } + if (out_len <= 0) { + return AVERROR_EOF; + } + state->pos += out_len; + return out_len; +} + +static int64_t fluxon_avio_seek(void *opaque, int64_t offset, int whence) { + FluxonAvioOpaque *state = (FluxonAvioOpaque *)opaque; + if (state == NULL) { + return AVERROR(EIO); + } + if (whence == AVSEEK_SIZE) { + return state->size; + } + + int whence_base = whence & ~AVSEEK_FORCE; + int64_t new_pos = 0; + if (whence_base == SEEK_SET) { + new_pos = offset; + } else if (whence_base == SEEK_CUR) { + new_pos = state->pos + offset; + } else if (whence_base == SEEK_END) { + new_pos = state->size + offset; + } else { + return AVERROR(EINVAL); + } + if (new_pos < 0) { + return AVERROR(EINVAL); + } + state->pos = new_pos; + return new_pos; +} + +static int fluxon_decode_receive_frames( + AVCodecContext *codec_ctx, + AVFrame *frame, + struct SwsContext **sws_ctx, + const int64_t *indices, + int indices_len, + uint8_t *filled, + int *filled_count, + int64_t *frame_index, + int out_width, + int out_height, + uint8_t *out_data, + int64_t frame_bytes, + char *err_buf, + size_t err_buf_len) { + while (1) { + int ret = avcodec_receive_frame(codec_ctx, frame); + if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) { + return 0; + } + if (ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avcodec_receive_frame failed", ret); + return -1; + } + + for (int i = 0; i < indices_len; i++) { + if (filled[i] || indices[i] != *frame_index) { + continue; + } + *sws_ctx = sws_getCachedContext( + *sws_ctx, + frame->width, + frame->height, + (enum AVPixelFormat)frame->format, + out_width, + out_height, + AV_PIX_FMT_RGB24, + SWS_BILINEAR, + NULL, + NULL, + NULL); + if (*sws_ctx == NULL) { + fluxon_set_err(err_buf, err_buf_len, "sws_getCachedContext failed"); + return -1; + } + + uint8_t *dst_data[4] = { + out_data + ((int64_t)i * frame_bytes), + NULL, + NULL, + NULL, + }; + int dst_linesize[4] = {out_width * 3, 0, 0, 0}; + int scaled = sws_scale( + *sws_ctx, + (const uint8_t *const *)frame->data, + frame->linesize, + 0, + frame->height, + dst_data, + dst_linesize); + if (scaled != out_height) { + fluxon_set_err(err_buf, err_buf_len, "sws_scale produced incomplete frame"); + return -1; + } + filled[i] = 1; + *filled_count += 1; + } + + *frame_index += 1; + av_frame_unref(frame); + if (*filled_count >= indices_len) { + return 1; + } + } +} + +int fluxon_fs_video_decode_frames( + const FluxonVideoIo *io, + int64_t file_size, + const int64_t *indices, + int indices_len, + int out_width, + int out_height, + int num_threads, + uint8_t *out_data, + int64_t out_data_len, + char *err_buf, + size_t err_buf_len) { + if (io == NULL || io->read_at == NULL) { + fluxon_set_err(err_buf, err_buf_len, "missing FluxonVideoIo read callback"); + return -1; + } + if (file_size < 0 || indices_len < 0 || out_width <= 0 || out_height <= 0 || num_threads <= 0) { + fluxon_set_err(err_buf, err_buf_len, "invalid video decode argument"); + return -1; + } + if (indices_len == 0) { + return 0; + } + for (int i = 0; i < indices_len; i++) { + if (indices[i] < 0) { + fluxon_set_err(err_buf, err_buf_len, "frame index must be non-negative"); + return -1; + } + } + + int64_t frame_bytes = (int64_t)out_width * (int64_t)out_height * 3; + if (frame_bytes <= 0 || out_data_len != frame_bytes * (int64_t)indices_len) { + fluxon_set_err(err_buf, err_buf_len, "output buffer size mismatch"); + return -1; + } + + int ret = 0; + const int avio_buffer_size = 64 * 1024; + uint8_t *avio_buffer = NULL; + AVIOContext *avio_ctx = NULL; + AVFormatContext *fmt_ctx = NULL; + AVCodecContext *codec_ctx = NULL; +#if LIBAVFORMAT_VERSION_MAJOR < 59 + AVCodec *decoder = NULL; +#else + const AVCodec *decoder = NULL; +#endif + AVPacket *packet = NULL; + AVFrame *frame = NULL; + struct SwsContext *sws_ctx = NULL; + uint8_t *filled = NULL; + + FluxonAvioOpaque opaque = { + .io = io, + .pos = 0, + .size = file_size, + .err_buf = err_buf, + .err_buf_len = err_buf_len, + }; + + avio_buffer = av_malloc(avio_buffer_size); + if (avio_buffer == NULL) { + fluxon_set_err(err_buf, err_buf_len, "av_malloc AVIO buffer failed"); + ret = -1; + goto cleanup; + } + avio_ctx = avio_alloc_context( + avio_buffer, + avio_buffer_size, + 0, + &opaque, + fluxon_avio_read, + NULL, + fluxon_avio_seek); + if (avio_ctx == NULL) { + fluxon_set_err(err_buf, err_buf_len, "avio_alloc_context failed"); + ret = -1; + goto cleanup; + } + avio_buffer = NULL; + + fmt_ctx = avformat_alloc_context(); + if (fmt_ctx == NULL) { + fluxon_set_err(err_buf, err_buf_len, "avformat_alloc_context failed"); + ret = -1; + goto cleanup; + } + fmt_ctx->pb = avio_ctx; + fmt_ctx->flags |= AVFMT_FLAG_CUSTOM_IO; + + int open_ret = avformat_open_input(&fmt_ctx, NULL, NULL, NULL); + if (open_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avformat_open_input failed", open_ret); + ret = -1; + goto cleanup; + } + + int info_ret = avformat_find_stream_info(fmt_ctx, NULL); + if (info_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avformat_find_stream_info failed", info_ret); + ret = -1; + goto cleanup; + } + + int stream_idx = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0); + if (stream_idx < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "av_find_best_stream failed", stream_idx); + ret = -1; + goto cleanup; + } + AVStream *stream = fmt_ctx->streams[stream_idx]; + if (decoder == NULL) { + decoder = avcodec_find_decoder(stream->codecpar->codec_id); + } + if (decoder == NULL) { + fluxon_set_err(err_buf, err_buf_len, "missing video decoder"); + ret = -1; + goto cleanup; + } + + codec_ctx = avcodec_alloc_context3(decoder); + if (codec_ctx == NULL) { + fluxon_set_err(err_buf, err_buf_len, "avcodec_alloc_context3 failed"); + ret = -1; + goto cleanup; + } + int params_ret = avcodec_parameters_to_context(codec_ctx, stream->codecpar); + if (params_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avcodec_parameters_to_context failed", params_ret); + ret = -1; + goto cleanup; + } + codec_ctx->thread_count = num_threads; + int codec_ret = avcodec_open2(codec_ctx, decoder, NULL); + if (codec_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avcodec_open2 failed", codec_ret); + ret = -1; + goto cleanup; + } + + packet = av_packet_alloc(); + frame = av_frame_alloc(); + filled = av_mallocz(indices_len); + if (packet == NULL || frame == NULL || filled == NULL) { + fluxon_set_err(err_buf, err_buf_len, "packet/frame allocation failed"); + ret = -1; + goto cleanup; + } + + int filled_count = 0; + int64_t frame_index = 0; + while (filled_count < indices_len) { + int read_ret = av_read_frame(fmt_ctx, packet); + if (read_ret == AVERROR_EOF) { + break; + } + if (read_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "av_read_frame failed", read_ret); + ret = -1; + goto cleanup; + } + + if (packet->stream_index == stream_idx) { + int send_ret = avcodec_send_packet(codec_ctx, packet); + if (send_ret < 0) { + av_packet_unref(packet); + fluxon_set_av_err(err_buf, err_buf_len, "avcodec_send_packet failed", send_ret); + ret = -1; + goto cleanup; + } + int recv = fluxon_decode_receive_frames( + codec_ctx, + frame, + &sws_ctx, + indices, + indices_len, + filled, + &filled_count, + &frame_index, + out_width, + out_height, + out_data, + frame_bytes, + err_buf, + err_buf_len); + if (recv < 0) { + av_packet_unref(packet); + ret = -1; + goto cleanup; + } + } + av_packet_unref(packet); + } + + if (filled_count < indices_len) { + int flush_ret = avcodec_send_packet(codec_ctx, NULL); + if (flush_ret < 0) { + fluxon_set_av_err(err_buf, err_buf_len, "avcodec_send_packet flush failed", flush_ret); + ret = -1; + goto cleanup; + } + int recv = fluxon_decode_receive_frames( + codec_ctx, + frame, + &sws_ctx, + indices, + indices_len, + filled, + &filled_count, + &frame_index, + out_width, + out_height, + out_data, + frame_bytes, + err_buf, + err_buf_len); + if (recv < 0) { + ret = -1; + goto cleanup; + } + } + + if (filled_count < indices_len) { + fluxon_set_err(err_buf, err_buf_len, "requested frame index is beyond decoded video"); + ret = -1; + goto cleanup; + } + +cleanup: + if (sws_ctx != NULL) { + sws_freeContext(sws_ctx); + } + if (frame != NULL) { + av_frame_free(&frame); + } + if (packet != NULL) { + av_packet_free(&packet); + } + if (codec_ctx != NULL) { + avcodec_free_context(&codec_ctx); + } + if (fmt_ctx != NULL) { + avformat_close_input(&fmt_ctx); + } + if (avio_ctx != NULL) { + av_freep(&avio_ctx->buffer); + avio_context_free(&avio_ctx); + } + if (avio_buffer != NULL) { + av_free(avio_buffer); + } + if (filled != NULL) { + av_free(filled); + } + return ret; +} diff --git a/fluxon_rs/fluxon_util/src/dev_config.rs b/fluxon_rs/fluxon_util/src/dev_config.rs index c860910..65d2ec3 100644 --- a/fluxon_rs/fluxon_util/src/dev_config.rs +++ b/fluxon_rs/fluxon_util/src/dev_config.rs @@ -1,4 +1,4 @@ -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde_yaml::Value; use std::fs; use std::path::{Path, PathBuf}; @@ -106,7 +106,7 @@ pub fn read_build_ext_config_value() -> Result { } /// Read etcd endpoint from build_config_ext.yml and ensure it has http/https scheme. -/// Accepts either a scalar string like "10.126.126.235:2579" or a YAML list of endpoints. +/// Accepts either a scalar string like "192.0.2.10:2379" or a YAML list of endpoints. pub fn read_etcd_endpoint_from_build_config() -> Result { let v = read_build_ext_config_value()?; // Try string first @@ -122,7 +122,7 @@ pub fn read_etcd_endpoint_from_build_config() -> Result { } } Err(anyhow!( - "Missing or invalid 'etcd' in build_config_ext.yml; please set e.g. etcd: 10.126.126.235:2579" + "Missing or invalid 'etcd' in build_config_ext.yml; please set e.g. etcd: 192.0.2.10:2379" )) } @@ -334,7 +334,7 @@ mod tests { .expect("write outer cargo toml"); fs::write( outer_root.join("build_config_ext.yml"), - "etcd: 10.0.0.1:2379\n", + "etcd: 192.0.2.10:2379\n", ) .expect("write outer build_config_ext"); diff --git a/fluxon_rs/fluxon_util/src/etcd/managed_client.rs b/fluxon_rs/fluxon_util/src/etcd/managed_client.rs new file mode 100644 index 0000000..a1dcfcb --- /dev/null +++ b/fluxon_rs/fluxon_util/src/etcd/managed_client.rs @@ -0,0 +1,305 @@ +use std::sync::OnceLock; + +use anyhow::{Context, Result, bail}; +use etcd_client::Client; +use tokio::runtime::{Builder, Runtime}; +use tokio::sync::RwLock; + +use crate::auto_clean_map::{AutoCleanMap, AutoCleanMapEntry}; + +/// Canonical identity of one etcd cluster connection target. +/// +/// Endpoints are internal URL strings with an explicit `http://` or `https://` +/// scheme. Construction sorts and deduplicates them because etcd balances over +/// the set and input order does not change cluster identity. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct EtcdEndpointSet(Vec); + +impl EtcdEndpointSet { + pub fn new(endpoints: Vec) -> Result { + if endpoints.is_empty() { + bail!("etcd endpoint set must not be empty"); + } + + let mut canonical = Vec::with_capacity(endpoints.len()); + for endpoint in endpoints { + let endpoint = endpoint.trim(); + if endpoint.is_empty() { + bail!("etcd endpoint must not be empty"); + } + if !(endpoint.starts_with("http://") || endpoint.starts_with("https://")) { + bail!( + "etcd endpoint must be an internal URL with http:// or https:// scheme, got: {}", + endpoint + ); + } + canonical.push(endpoint.to_string()); + } + canonical.sort(); + canonical.dedup(); + Ok(Self(canonical)) + } + + #[inline] + pub fn as_slice(&self) -> &[String] { + &self.0 + } +} + +struct EtcdClientSlot { + generation: u64, + client: Option, +} + +struct EtcdClientBackend { + endpoints: EtcdEndpointSet, + slot: RwLock, +} + +impl EtcdClientBackend { + fn new(endpoints: EtcdEndpointSet) -> Self { + Self { + endpoints, + slot: RwLock::new(EtcdClientSlot { + generation: 0, + client: None, + }), + } + } + + async fn snapshot(&self, owner: ManagedEtcdClient) -> Result { + { + let slot = self.slot.read().await; + if let Some(client) = slot.client.as_ref() { + return Ok(ManagedEtcdClientSnapshot { + owner, + generation: slot.generation, + client: client.clone(), + }); + } + } + + // Hold the write lock while connecting so concurrent cold-start callers + // share one connection attempt and one resulting client generation. + let mut slot = self.slot.write().await; + if let Some(client) = slot.client.as_ref() { + return Ok(ManagedEtcdClientSnapshot { + owner, + generation: slot.generation, + client: client.clone(), + }); + } + + let endpoints = self.endpoints.as_slice().to_vec(); + let error_endpoints = endpoints.clone(); + let client = managed_etcd_runtime() + .spawn(async move { Client::connect(endpoints, None).await }) + .await + .context("managed etcd client runtime task stopped")? + .with_context(|| { + format!("failed to connect etcd for endpoints {:?}", error_endpoints) + })?; + slot.generation = slot + .generation + .checked_add(1) + .expect("etcd client generation overflow"); + slot.client = Some(client.clone()); + Ok(ManagedEtcdClientSnapshot { + owner, + generation: slot.generation, + client, + }) + } + + async fn invalidate_generation(&self, generation: u64) -> bool { + let mut slot = self.slot.write().await; + if slot.generation != generation || slot.client.is_none() { + return false; + } + slot.client = None; + true + } +} + +/// Owns tonic channel workers independently of whichever runtime first asks +/// for an endpoint set. The runtime is intentionally bounded and process-wide. +fn managed_etcd_runtime() -> &'static Runtime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + Builder::new_multi_thread() + .worker_threads(1) + .thread_name("fluxon-etcd-client") + .enable_all() + .build() + .expect("failed to build managed etcd client runtime") + }) +} + +fn managed_etcd_client_map() -> &'static AutoCleanMap { + static MAP: OnceLock> = OnceLock::new(); + MAP.get_or_init(AutoCleanMap::new) +} + +/// Auto-clean handle for the process-wide etcd client registry. +/// +/// Every live owner for the same canonical endpoint set shares one lazily +/// connected `etcd_client::Client` generation. Dropping the last handle removes +/// the backend entry and releases its cached client. +#[derive(Clone)] +pub struct ManagedEtcdClient { + entry: AutoCleanMapEntry, +} + +impl ManagedEtcdClient { + pub fn acquire(endpoints: EtcdEndpointSet) -> Self { + let entry = managed_etcd_client_map() + .get_or_init(endpoints.clone(), || EtcdClientBackend::new(endpoints)); + Self { entry } + } + + #[inline] + pub fn endpoints(&self) -> &EtcdEndpointSet { + &self.entry.endpoints + } + + pub async fn snapshot(&self) -> Result { + self.entry.snapshot(self.clone()).await + } + + pub async fn client(&self) -> Result { + Ok(self.snapshot().await?.client) + } +} + +/// One connected generation borrowed from a `ManagedEtcdClient`. +/// +/// The snapshot keeps the registry guard alive. Invalidating it only clears the +/// same generation, so a delayed failure cannot discard a newer replacement. +#[derive(Clone)] +pub struct ManagedEtcdClientSnapshot { + owner: ManagedEtcdClient, + generation: u64, + client: Client, +} + +impl ManagedEtcdClientSnapshot { + #[inline] + pub fn client(&self) -> Client { + self.client.clone() + } + + pub async fn invalidate(&self) -> bool { + self.owner + .entry + .invalidate_generation(self.generation) + .await + } +} + +#[cfg(test)] +mod tests { + use std::net::TcpListener; + use std::sync::mpsc; + use std::thread; + use std::time::Duration; + + use super::*; + + fn unique_endpoint_set() -> EtcdEndpointSet { + EtcdEndpointSet::new(vec![format!( + "http://managed-etcd-client-test-{}:2379", + uuid::Uuid::new_v4() + )]) + .unwrap() + } + + #[test] + fn endpoint_set_sorts_and_deduplicates_urls() { + let endpoints = EtcdEndpointSet::new(vec![ + "http://b.example:2379".to_string(), + " http://a.example:2379 ".to_string(), + "http://b.example:2379".to_string(), + ]) + .unwrap(); + assert_eq!( + endpoints.as_slice(), + [ + "http://a.example:2379".to_string(), + "http://b.example:2379".to_string() + ] + ); + } + + #[test] + fn endpoint_set_rejects_raw_or_empty_endpoints() { + assert!(EtcdEndpointSet::new(Vec::new()).is_err()); + assert!(EtcdEndpointSet::new(vec!["".to_string()]).is_err()); + assert!(EtcdEndpointSet::new(vec!["127.0.0.1:2379".to_string()]).is_err()); + } + + #[test] + fn registry_reuses_and_auto_cleans_live_entries() { + let endpoints = unique_endpoint_set(); + let map = managed_etcd_client_map(); + assert!(map.with_existing(&endpoints, |_| ()).is_none()); + + { + let handle_a = ManagedEtcdClient::acquire(endpoints.clone()); + assert!(map.with_existing(&endpoints, |_| ()).is_some()); + + { + let handle_b = ManagedEtcdClient::acquire(endpoints.clone()); + assert!(std::ptr::eq(&*handle_a.entry, &*handle_b.entry)); + } + + assert!(map.with_existing(&endpoints, |_| ()).is_some()); + } + + assert!(map.with_existing(&endpoints, |_| ()).is_none()); + } + + #[tokio::test] + async fn stale_snapshot_cannot_invalidate_new_generation() { + let endpoints = EtcdEndpointSet::new(vec!["http://127.0.0.1:1".to_string()]).unwrap(); + let handle = ManagedEtcdClient::acquire(endpoints); + + let first = handle.snapshot().await.unwrap(); + assert!(first.invalidate().await); + + let second = handle.snapshot().await.unwrap(); + assert_ne!(first.generation, second.generation); + assert!(!first.invalidate().await); + + let current = handle.snapshot().await.unwrap(); + assert_eq!(second.generation, current.generation); + } + + #[test] + fn client_worker_outlives_the_runtime_that_requested_the_snapshot() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let (accepted_tx, accepted_rx) = mpsc::sync_channel(1); + let accept_thread = thread::spawn(move || { + let _connection = listener.accept().unwrap(); + accepted_tx.send(()).unwrap(); + }); + + let handle = ManagedEtcdClient::acquire( + EtcdEndpointSet::new(vec![format!("http://{address}")]).unwrap(), + ); + let caller_runtime = Builder::new_current_thread().enable_all().build().unwrap(); + let snapshot = caller_runtime.block_on(handle.snapshot()).unwrap(); + drop(caller_runtime); + + let request_runtime = Builder::new_current_thread().enable_all().build().unwrap(); + request_runtime.block_on(async move { + let mut client = snapshot.client(); + let _ = tokio::time::timeout(Duration::from_secs(1), client.status()).await; + }); + + accepted_rx + .recv_timeout(Duration::from_secs(1)) + .expect("managed etcd worker stopped with the caller runtime"); + accept_thread.join().unwrap(); + } +} diff --git a/fluxon_rs/fluxon_util/src/etcd/mod.rs b/fluxon_rs/fluxon_util/src/etcd/mod.rs index 1d99ed1..cc2d9ab 100755 --- a/fluxon_rs/fluxon_util/src/etcd/mod.rs +++ b/fluxon_rs/fluxon_util/src/etcd/mod.rs @@ -1,9 +1,11 @@ pub mod cluster_lease; pub mod id_allocator; +pub mod managed_client; pub mod prefix_watch_actor; pub use cluster_lease::get_cluster_lease_id; pub use id_allocator::DistributeIdAllocator; +pub use managed_client::{EtcdEndpointSet, ManagedEtcdClient, ManagedEtcdClientSnapshot}; pub use prefix_watch_actor::{ AsyncStopSignal, ETCD_PREFIX_WATCH_RESTART_SLEEP, EtcdPrefixWatchLoopControl, OwnedEtcdWatchEvent, OwnedEtcdWatchEventKind, run_prefix_watch_loop, diff --git a/fluxon_rs/fluxon_util/src/lease_manager/keepalive_actor.rs b/fluxon_rs/fluxon_util/src/lease_manager/keepalive_actor.rs index 5e492f1..272176c 100644 --- a/fluxon_rs/fluxon_util/src/lease_manager/keepalive_actor.rs +++ b/fluxon_rs/fluxon_util/src/lease_manager/keepalive_actor.rs @@ -3,7 +3,8 @@ use super::lease_backend_uid::LeaseBackendUid; use super::lease_handle::{LeaseEntry, LeaseEntryKind}; use super::lifecycle::OnKeepalive; use crate::auto_clean_map::{AutoCleanMap, AutoCleanMapEntry}; -use etcd_client::{Client, LeaseKeepAliveStream, LeaseKeeper}; +use crate::etcd::ManagedEtcdClient; +use etcd_client::{LeaseKeepAliveStream, LeaseKeeper}; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -70,7 +71,7 @@ where // ---------- OneTtlKeepAliveActor & registry ---------- pub(crate) struct EtcdState { - pub(crate) client: Client, + pub(crate) client: ManagedEtcdClient, pub(crate) lease_id: i64, pub(crate) keeper: Option, pub(crate) stream: Option, @@ -95,7 +96,14 @@ impl EtcdState { let mut last_err: Option = None; for attempts in 0..10 { self.last_stage = "ensure_stream.open_stream"; - match self.client.lease_keep_alive(lease_id).await { + let mut client = match self.client.client().await { + Ok(client) => client, + Err(e) => { + last_err = Some(e.to_string()); + continue; + } + }; + match client.lease_keep_alive(lease_id).await { Ok((keeper, stream)) => { debug!( "renewed keepalive stream for lease_id={} attempts={}", @@ -479,7 +487,7 @@ pub enum ActorRegisterInvocation { label: Option, }, Etcd { - client: Client, + client: ManagedEtcdClient, revoke_on_drop: bool, }, } diff --git a/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_handle.rs b/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_handle.rs index c9c2500..559499d 100644 --- a/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_handle.rs +++ b/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_handle.rs @@ -1,5 +1,4 @@ use anyhow::Result as AnyResult; -use etcd_client::Client; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -9,6 +8,7 @@ use super::lease_backend_uid::LeaseBackendUid; use super::lifecycle::debug_keepalive_log; use crate::auto_clean_map::AutoCleanMap; use crate::auto_clean_map::AutoCleanMapEntry; +use crate::etcd::ManagedEtcdClient; /// Backend resources actually used by keepalive actors. /// @@ -16,8 +16,7 @@ use crate::auto_clean_map::AutoCleanMapEntry; /// types when the handle also carries the map guard (AutoCleanMapEntry). pub enum LeaseBackendInner { Etcd { - _endpoints: Vec, - client: Client, + client: ManagedEtcdClient, /// Per-lease keepalive state keyed by lease_id. Auto-evicts when the last /// guard (AutoCleanMapEntry) for that lease is dropped. states: AutoCleanMap>>, @@ -56,7 +55,7 @@ impl LeaseBackendHandle { } #[inline] - pub fn etcd_client(&self) -> Option { + pub fn managed_etcd_client(&self) -> Option { match &*self.entry { LeaseBackendInner::Etcd { client, .. } => Some(client.clone()), _ => None, diff --git a/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_uid.rs b/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_uid.rs index ff562bf..aa8b854 100755 --- a/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_uid.rs +++ b/fluxon_rs/fluxon_util/src/lease_manager/lease_backend_uid.rs @@ -2,6 +2,8 @@ use anyhow::Result as AnyResult; use std::fmt; use std::sync::Arc; +use crate::etcd::EtcdEndpointSet; + /// Backend kind for leases supported by the unified lease manager. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum LeaseType { @@ -11,11 +13,10 @@ pub enum LeaseType { /// Unique identifier for a lease backend. /// -/// - Etcd: endpoints list (Vec) sorted lexicographically to make -/// identity stable regardless of input order. +/// - Etcd: canonical endpoint set. /// - KvClient: cluster name; carries allocate/keepalive Rust closures. pub enum LeaseBackendUid { - Etcd(Vec), + Etcd(EtcdEndpointSet), KvClientWithCallbacks { cluster: String, /// Allocate closure: input ttl_seconds -> lease_id @@ -27,11 +28,7 @@ pub enum LeaseBackendUid { } impl LeaseBackendUid { - /// Construct an etcd backend uid from endpoint list; endpoints are sorted - /// to ensure identical identity regardless of input order. - pub fn etcd_from(mut endpoints: Vec) -> Self { - // Sort in-place; caller must pass explicit endpoints, we don't add defaults. - endpoints.sort(); + pub fn etcd(endpoints: EtcdEndpointSet) -> Self { LeaseBackendUid::Etcd(endpoints) } @@ -55,9 +52,9 @@ impl LeaseBackendUid { } } - pub fn endpoints(&self) -> Option<&[String]> { + pub fn etcd_endpoint_set(&self) -> Option<&EtcdEndpointSet> { match self { - LeaseBackendUid::Etcd(v) => Some(v.as_slice()), + LeaseBackendUid::Etcd(v) => Some(v), _ => None, } } @@ -140,11 +137,8 @@ impl std::hash::Hash for LeaseBackendUid { fn hash(&self, state: &mut H) { match self { LeaseBackendUid::Etcd(endpoints) => { - // tag + endpoints (construction sorted; order is stable) 0u8.hash(state); - for e in endpoints { - e.hash(state); - } + endpoints.hash(state); } LeaseBackendUid::KvClientWithCallbacks { cluster, .. } => { 1u8.hash(state); @@ -157,7 +151,7 @@ impl std::hash::Hash for LeaseBackendUid { impl fmt::Debug for LeaseBackendUid { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - LeaseBackendUid::Etcd(v) => write!(f, "Etcd({:?})", v), + LeaseBackendUid::Etcd(v) => write!(f, "Etcd({:?})", v.as_slice()), LeaseBackendUid::KvClientWithCallbacks { cluster, .. } => { write!(f, "KvClientWithCallbacks(cluster={})", cluster) } diff --git a/fluxon_rs/fluxon_util/src/lease_manager/lease_handle.rs b/fluxon_rs/fluxon_util/src/lease_manager/lease_handle.rs index 1c24cf2..03c25a4 100755 --- a/fluxon_rs/fluxon_util/src/lease_manager/lease_handle.rs +++ b/fluxon_rs/fluxon_util/src/lease_manager/lease_handle.rs @@ -91,7 +91,10 @@ impl Drop for GeneralLease { } } -/// Endpoint-scoped lease manager, bound to a specific etcd address set. +/// Stateless facade over the process-wide lease registries. +/// +/// Etcd backend identity is carried explicitly by `LeaseBackendUid`; client +/// lifetime and reuse are owned by the managed etcd client registry. #[derive(Clone, Default)] pub struct LeaseManager; diff --git a/fluxon_rs/fluxon_util/src/lease_manager/lifecycle.rs b/fluxon_rs/fluxon_util/src/lease_manager/lifecycle.rs index c17ac33..ea96cb0 100644 --- a/fluxon_rs/fluxon_util/src/lease_manager/lifecycle.rs +++ b/fluxon_rs/fluxon_util/src/lease_manager/lifecycle.rs @@ -10,8 +10,7 @@ use std::sync::{Arc, OnceLock}; use std::time::Duration; use tokio::sync::Mutex; -use anyhow::{Context, Result as AnyResult}; -use etcd_client::Client; +use anyhow::Result as AnyResult; use super::keepalive_actor::{ self, ActorRegisterInvocation, EtcdState, LeaseKey, OneTtlKeepAliveInner, ensure_inner_running, @@ -20,6 +19,7 @@ use super::lease_backend_handle::{LeaseBackendHandle, LeaseBackendInner}; use super::lease_backend_uid::{LeaseBackendUid, LeaseRegisterKind, LeaseType}; use super::lease_handle::{GeneralLease, LeaseEntry, LeaseEntryKind}; use crate::auto_clean_map::{AutoCleanMap, AutoCleanMapEntry}; +use crate::etcd::ManagedEtcdClient; const INITIAL_ETCD_KEEPALIVE_PROBE_RETRIES: usize = 5; const INITIAL_ETCD_KEEPALIVE_PROBE_BUDGET_MS: u64 = 60_000; @@ -97,7 +97,7 @@ fn backend_map() -> &'static AutoCleanMap { pub fn acquire_backend_handle( uid: LeaseBackendUid, kv_cb: Option AnyResult<()> + Send + Sync + 'static>>, - etcd_client: Option, + etcd_client: Option, rt: tokio::runtime::Handle, ) -> LeaseBackendHandle { let entry: AutoCleanMapEntry = @@ -115,12 +115,7 @@ pub fn acquire_backend_handle( LeaseBackendUid::Etcd(_) => { let client = etcd_client.expect("etcd backend acquire requires client on first creation"); - let endpoints = uid - .endpoints() - .expect("etcd uid must carry endpoints") - .to_vec(); LeaseBackendInner::Etcd { - _endpoints: endpoints, client, states: AutoCleanMap::new(), rt: rt.clone(), @@ -294,11 +289,22 @@ impl Drop for LeaseEntry { .. } => { if *revoke_on_drop { - if let Some(mut cli) = handle.etcd_client() { + if let Some(client) = handle.managed_etcd_client() { let lid = lease_id as i64; // Use the runtime handle carried by the backend handle (LeaseBackendInner) let rt = handle.runtime_handle(); rt.spawn(async move { + let mut cli = match client.client().await { + Ok(client) => client, + Err(e) => { + tracing::warn!( + "failed to acquire etcd client to revoke lease_id={} on drop: {:?}", + lid, + e + ); + return; + } + }; if let Err(e) = cli.lease_revoke(lid).await { tracing::warn!( "failed to revoke lease_id={} on drop: {:?}", @@ -312,7 +318,7 @@ impl Drop for LeaseEntry { } else { tracing::warn!( lease_id, - "etcd revoke_on_drop: missing etcd client in backend handle" + "etcd revoke_on_drop: missing managed etcd client in backend handle" ); } } @@ -336,11 +342,10 @@ pub async fn register_lease_for_keepalive( LeaseType::Etcd => { record_register_by(lease_id, format!("{:?},ttl={}", &backend_uid, ttl_seconds)); let endpoints = backend_uid - .endpoints() - .expect("etcd backend must carry endpoints"); - let client = Client::connect(endpoints, None).await.with_context(|| { - format!("failed to connect etcd for endpoints {:?}", endpoints) - })?; + .etcd_endpoint_set() + .expect("etcd backend must carry endpoints") + .clone(); + let client = ManagedEtcdClient::acquire(endpoints); let backend_handle = acquire_backend_handle( backend_uid.clone(), diff --git a/fluxon_rs/fluxon_util/src/lib.rs b/fluxon_rs/fluxon_util/src/lib.rs index a85aed0..cde6dc7 100644 --- a/fluxon_rs/fluxon_util/src/lib.rs +++ b/fluxon_rs/fluxon_util/src/lib.rs @@ -37,10 +37,9 @@ pub mod limitrate; pub mod pyo3; // Re-export for stable public API: existing call sites can keep using `fluxon_util::init_log`. pub use log::{ - current_daily_sharded_log_path, current_log_file_path, daily_sharded_log_path, - display_runtime_log_path, init_log, init_log_test, init_log_with_extra_layer, - latest_existing_daily_sharded_log_path, resolve_readable_log_path, - DEFAULT_DAILY_LOG_RETENTION_DAYS, + DEFAULT_DAILY_LOG_RETENTION_DAYS, current_daily_sharded_log_path, current_log_file_path, + daily_sharded_log_path, display_runtime_log_path, init_log, init_log_test, + init_log_with_extra_layer, latest_existing_daily_sharded_log_path, resolve_readable_log_path, }; #[cfg(test)] mod test_util_test; @@ -251,7 +250,12 @@ mod tests { ); assert_logged_text( &active_log_path, - &["debug message", "info message", "warning message", "error message"], + &[ + "debug message", + "info message", + "warning message", + "error message", + ], ); } } diff --git a/fluxon_rs/fluxon_util/src/log.rs b/fluxon_rs/fluxon_util/src/log.rs index fc6066f..4db4ae4 100644 --- a/fluxon_rs/fluxon_util/src/log.rs +++ b/fluxon_rs/fluxon_util/src/log.rs @@ -61,9 +61,7 @@ fn read_test_log_shard_window_config() -> anyhow::Result 0"); @@ -85,7 +83,9 @@ fn read_test_log_shard_window_config() -> anyhow::Result) -> anyhow::Result { +fn resolve_shard_date_from_datetime( + now: chrono::DateTime, +) -> anyhow::Result { let Some(config) = read_test_log_shard_window_config()? else { return Ok(now.date_naive()); }; @@ -99,8 +99,8 @@ fn resolve_shard_date_from_datetime(now: chrono::DateTime) -> anyho ); } let bucket_index = delta_seconds / config.window_seconds; - let base_date = chrono::NaiveDate::from_ymd_opt(2026, 1, 1) - .expect("valid hard-coded synthetic base date"); + let base_date = + chrono::NaiveDate::from_ymd_opt(2026, 1, 1).expect("valid hard-coded synthetic base date"); Ok(base_date + chrono::Days::new(bucket_index as u64)) } @@ -108,10 +108,7 @@ fn current_shard_date() -> anyhow::Result { resolve_shard_date_from_datetime(chrono::Utc::now()) } -fn cleanup_old_daily_sharded_logs( - base_path: &Path, - retention_days: usize, -) -> anyhow::Result<()> { +fn cleanup_old_daily_sharded_logs(base_path: &Path, retention_days: usize) -> anyhow::Result<()> { let parent = match base_path.parent() { Some(parent) => parent, None => return Ok(()), @@ -124,7 +121,8 @@ fn cleanup_old_daily_sharded_logs( return Ok(()); }; fs::create_dir_all(parent)?; - let keep_since = current_shard_date()? - chrono::Days::new(retention_days.saturating_sub(1) as u64); + let keep_since = + current_shard_date()? - chrono::Days::new(retention_days.saturating_sub(1) as u64); let prefix = format!("{stem}."); for entry in std::fs::read_dir(parent)? { let entry = entry?; @@ -180,10 +178,7 @@ impl DailyShardedFileWriter { current_daily_sharded_log_path(&self.base_path) } - fn rotate_if_needed( - &self, - state: &mut DailyShardedFileWriterState, - ) -> io::Result<()> { + fn rotate_if_needed(&self, state: &mut DailyShardedFileWriterState) -> io::Result<()> { let next_path = self .current_path() .map_err(|err| io::Error::new(io::ErrorKind::Other, err.to_string()))?; @@ -314,20 +309,19 @@ pub fn daily_sharded_log_path( base_path: &Path, date: chrono::NaiveDate, ) -> anyhow::Result { - let file_name = base_path.file_name().and_then(|v| v.to_str()).ok_or_else(|| { - anyhow::anyhow!( - "log path must end with a valid utf-8 filename: {}", - base_path.display() - ) - })?; + let file_name = base_path + .file_name() + .and_then(|v| v.to_str()) + .ok_or_else(|| { + anyhow::anyhow!( + "log path must end with a valid utf-8 filename: {}", + base_path.display() + ) + })?; let stem = file_name .strip_suffix(".log") .ok_or_else(|| anyhow::anyhow!("log path must end with .log: {}", base_path.display()))?; - Ok(base_path.with_file_name(format!( - "{}.{}.log", - stem, - date.format("%Y-%m-%d") - ))) + Ok(base_path.with_file_name(format!("{}.{}.log", stem, date.format("%Y-%m-%d")))) } pub fn current_daily_sharded_log_path(base_path: &Path) -> anyhow::Result { diff --git a/fluxon_rs/fluxon_util/tests/log_mgmt.rs b/fluxon_rs/fluxon_util/tests/log_mgmt.rs index 431c5fc..f459337 100644 --- a/fluxon_rs/fluxon_util/tests/log_mgmt.rs +++ b/fluxon_rs/fluxon_util/tests/log_mgmt.rs @@ -59,7 +59,10 @@ fn kv_log_shards_roll_and_cleanup_with_test_window() { .expect("unix epoch") .as_secs() as i64; let _window_guard = EnvVarGuard::set(TEST_LOG_SHARD_WINDOW_SECONDS_ENV, "10"); - let _anchor_guard = EnvVarGuard::set(TEST_LOG_SHARD_ANCHOR_UNIX_SECONDS_ENV, (now - 2).to_string()); + let _anchor_guard = EnvVarGuard::set( + TEST_LOG_SHARD_ANCHOR_UNIX_SECONDS_ENV, + (now - 2).to_string(), + ); fluxon_util::init_log(log_path, instance_key); tracing::info!(target: "fluxon_util", "[kv-log-mgmt][phase=before] ts={}", now); @@ -115,7 +118,8 @@ fn resolve_readable_log_path_ignores_plain_base_log_when_daily_shards_exist() { let shard_path = temp_dir.path().join("startup.2026-06-21.log"); fs::write(&shard_path, "shard\n").expect("write shard log"); - let resolved = fluxon_util::resolve_readable_log_path(&base_path).expect("resolve readable log path"); + let resolved = + fluxon_util::resolve_readable_log_path(&base_path).expect("resolve readable log path"); assert_eq!(resolved, shard_path); } @@ -128,7 +132,7 @@ fn latest_existing_daily_sharded_log_path_skips_invalid_candidates() { fs::write(&invalid_shard_path, "invalid\n").expect("write invalid shard"); fs::write(&valid_shard_path, "valid\n").expect("write valid shard"); - let resolved = - fluxon_util::latest_existing_daily_sharded_log_path(&base_path).expect("resolve latest shard"); + let resolved = fluxon_util::latest_existing_daily_sharded_log_path(&base_path) + .expect("resolve latest shard"); assert_eq!(resolved, valid_shard_path); } diff --git a/fluxon_test_stack/benchmark_full_matrix.yaml b/fluxon_test_stack/benchmark_full_matrix.yaml index eaf0fbe..c225158 100644 --- a/fluxon_test_stack/benchmark_full_matrix.yaml +++ b/fluxon_test_stack/benchmark_full_matrix.yaml @@ -4,7 +4,7 @@ run: mode: full_once selectors: case_ids: ALL - profile_ids: [fluxon_fastws, fluxon_tquic, fluxon_sockudo_ws, fluxon_tcp, redis_sharded, alluxio_posix] + profile_ids: [fluxon_fastws, fluxon_tquic, fluxon_sockudo_ws, fluxon_tcp, fluxon_tcp_kv_ssd, mooncake_tcp, redis_sharded, alluxio_posix] command_ids: ALL test_ids: ALL @@ -52,6 +52,27 @@ scenes: scales: [n2_kvowner_dram_20gib, n5_kvowner_dram_20gib, n9_kvowner_dram_20gib] profiles: [fluxon_fastws, fluxon_tquic, fluxon_sockudo_ws, fluxon_tcp, redis_sharded] + kv_ssd_pressure_zipf: + test_stack: + mode: KVSTORE + read_ratio: 0.9 + write_ratio: 0.1 + request_distribution: zipfian + keyspace_size: 2048 + value_size_mode: RANDOM_WEIGHTED_SET + value_size_weighted_set: + - size_bytes: 1048576 + weight: 0.15 + - size_bytes: 2097152 + weight: 0.35 + - size_bytes: 4194304 + weight: 0.35 + - size_bytes: 8388608 + weight: 0.15 + select: + scales: [n1_kvowner_dram_1gib, n2_kvowner_dram_1gib] + profiles: [fluxon_tcp_kv_ssd, fluxon_tcp, mooncake_tcp] + kv_write_heavy_large_value: test_stack: mode: KVSTORE @@ -183,6 +204,31 @@ scales: owner_dram_bytes: 21474836480 benchmark: *bench_benchmark + n1_kvowner_dram_1gib: + duration_seconds: 60 + topology: 1 + owner: + owner_count: 1 + owner_dram_bytes: 1073741824 + benchmark: &ssd_pressure_benchmark + processes_per_target: 1 + threads_per_process: 4 + value_size: 1048576 + metric_warmup_seconds: 5 + start_idle_seconds: 10 + op_timeout_seconds: 60 + cluster_ready_timeout_seconds: 900 + value_size_list: [] + consumer_sim_handle_ms_range: [700, 1500] + + n2_kvowner_dram_1gib: + duration_seconds: 60 + topology: 2 + owner: + owner_count: 2 + owner_dram_bytes: 1073741824 + benchmark: *ssd_pressure_benchmark + artifact_sets: fluxon_fastws: release_source: &release_source_fastws @@ -296,19 +342,19 @@ profiles: - test_profile_adapter.py controller_url: __STACK_CONTROLLER_URL__ target_ip_map: &upstream_target_ip_map - node-1: 10.127.37.2 - node-2: 10.127.37.40 - node-3: 10.127.37.6 - node-4: 10.127.37.5 - node-5: 10.127.37.75 - node-6: 10.127.37.4 - node-7: 10.127.37.118 - node-8: 10.127.37.1 - node-9: 10.127.37.22 - node-10: 10.127.37.91 - primary-bastion: 10.233.75.142 + node-1: 192.0.2.11 + node-2: 192.0.2.12 + node-3: 192.0.2.13 + node-4: 192.0.2.14 + node-5: 192.0.2.15 + node-6: 192.0.2.16 + node-7: 192.0.2.17 + node-8: 192.0.2.18 + node-9: 192.0.2.19 + node-10: 192.0.2.20 + example-bastion: 198.51.100.10 bootstrap_ready_timeout_seconds: 300 - runtime_config: + runtime_config: &common_fluxon_runtime_config kv_base: &upstream_test_stack_kv_base instance_key: bench_base contribute_to_cluster_pool_size: {dram: 0, vram: {}} @@ -363,6 +409,89 @@ profiles: test_stack: <<: *common_test_stack_runtime + fluxon_tcp_kv_ssd: + artifact_set: fluxon_tcp + runtime: + test_stack: + <<: *common_test_stack_runtime + runtime_config: + <<: *common_fluxon_runtime_config + owner_kv_ssd: + large_limit_size: [17179869184] + + mooncake_tcp: + artifact_set: fluxon_tcp + runtime: + test_stack: + kind: MOONCAKE + coordinator_ready_timeout_seconds: 300 + port_alloc: + by_topology: + 1: + coordinator_port_base: 23779 + coordinator_port_stride: 10 + mooncake_rpc_port_base: 33000 + mooncake_rpc_port_stride: 10 + mooncake_metadata_port_base: 34000 + mooncake_metadata_port_stride: 10 + mooncake_metrics_port_base: 35000 + mooncake_metrics_port_stride: 10 + DEFAULT: + coordinator_port_base: 24779 + coordinator_port_stride: 10 + mooncake_rpc_port_base: 33100 + mooncake_rpc_port_stride: 10 + mooncake_metadata_port_base: 34100 + mooncake_metadata_port_stride: 10 + mooncake_metrics_port_base: 35100 + mooncake_metrics_port_stride: 10 + deploy: + adapter_cmd: + - test_profile_adapter.py + controller_url: __STACK_CONTROLLER_URL__ + target_ip_map: *upstream_target_ip_map + bootstrap_ready_timeout_seconds: 300 + runtime_config: + testbed_mooncake_storage: + mode: DEDICATED_OWNER + kv_base: + instance_key: bench_base + contribute_to_cluster_pool_size: {dram: 0, vram: {}} + protocol: + protocol_type: tcp + test_spec_config: + disable_observability: true + disable_master_replica_cache: true + disable_prefix_index: true + mooncake_spec: + local_buffer_size: 1073741824 + master: + wheel_name: mooncake_transfer_engine-0.3.0-cp310-cp310-linux_x86_64.whl + python_abi: cpython3.10 + cluster_id_prefix: fluxon-test-stack-mooncake + rpc_thread_num: 4 + deploy_templates: + coordinator: + lifecycle: service + k8s_ref: deployment/test_stack_coord + deployer: + command: [/bin/bash, -lc] + args: + - | + # Placeholder only. test_runner.py rewrites the actual TEST_STACK + # launch command and resolves the interpreter from target hostworkdir. + : + node: + lifecycle: job + k8s_ref: deployment/test_stack_node___INSTANCE_KEY__ + deployer: + command: [/bin/bash, -lc] + args: + - | + # Placeholder only. test_runner.py rewrites the actual TEST_STACK + # launch command and resolves the interpreter from target hostworkdir. + : "__INSTANCE_KEY__ __COORDINATOR__" + redis_sharded: artifact_set: fluxon_fastws runtime: diff --git a/fluxon_test_stack/benchmark_node_kv.py b/fluxon_test_stack/benchmark_node_kv.py index c9c559b..2767691 100644 --- a/fluxon_test_stack/benchmark_node_kv.py +++ b/fluxon_test_stack/benchmark_node_kv.py @@ -4,16 +4,22 @@ import bisect import copy +import ctypes +import ctypes.util import hashlib +import importlib.util import json import os import socket +import struct import threading import time +from collections import deque from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from enum import Enum, unique from functools import lru_cache +from pathlib import Path from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union from benchmark_role_names import ( @@ -26,6 +32,10 @@ from fluxon_py import FluxonKvClientConfig as KVCacheConfig from fluxon_py import new_store from fluxon_py.kvclient.kvclient_interface import KvClient, PutOptionalArgs +from fluxon_py.kvclient.nonzerocopy_encode import ( + DLPackBytesView, + _dlpack_cpu_tensor_info, +) TEST_MODE_MPMC = "MPMC" TEST_MODE_KVSTORE = "KVSTORE" @@ -42,6 +52,7 @@ KV_OPERATION_GET = "GET" BACKEND_KIND_FLUXON = "FLUXON" +BACKEND_KIND_MOONCAKE = "MOONCAKE" BACKEND_KIND_REDIS = "REDIS" BACKEND_KIND_ALLUXIO = "ALLUXIO" REDIS_BENCH_INFLIGHT_GUARD_PREFIX = "__fluxon_bench_inflight_guard__" @@ -56,11 +67,23 @@ BENCHMARK_KEY_KEYSPACE_SIZE = "keyspace_size" BENCHMARK_KEY_AFFINITY_LOCALITY_RATIO = "affinity_locality_ratio" BENCHMARK_KEY_AFFINITY_SLOT_COUNT = "affinity_slot_count" +BENCHMARK_KEY_KV_BOOTSTRAP_CONCURRENCY = "kv_bootstrap_concurrency" +BENCHMARK_KEY_KV_BOOTSTRAP_PUT_GAP_MS = "kv_bootstrap_put_gap_ms" +BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY = "kv_bootstrap_storage_full_policy" +BENCHMARK_KEY_KV_GET_OUTPUT = "kv_get_output" +BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX = "kv_cuda_device_index" +KV_BOOTSTRAP_STORAGE_FULL_POLICY_FAIL = "fail" +KV_BOOTSTRAP_STORAGE_FULL_POLICY_STOP = "stop" +KV_BOOTSTRAP_STORAGE_FULL_POLICIES = { + KV_BOOTSTRAP_STORAGE_FULL_POLICY_FAIL, + KV_BOOTSTRAP_STORAGE_FULL_POLICY_STOP, +} DEFAULT_KV_KEYSPACE_SIZE = 101 DEFAULT_ZIPFIAN_THETA = 0.99 STABLE_HASH_MODULUS = float(1 << 64) KV_SEED_BOOTSTRAP_MAX_CONCURRENCY = 16 +CUDA_H2D_PIPELINE_DEPTH = 2 KV_VERBOSE_PER_OP_LOG = str(os.environ.get("FLUXON_BENCH_KV_VERBOSE", "")).strip().lower() not in ("", "0", "false", "no") FLUXON_PHASE_LOG_INTERVAL_OPS = 128 FLUXON_PHASE_SLOW_OP_THRESHOLD_US = 50_000.0 @@ -254,6 +277,8 @@ FLUXON_PHASE_PATH_METRIC_OWNER_HANDLE_TO_RESP_SEND_US, ) _BENCHMARK_CLIENT_STRIP_TEST_SPEC_KEYS = ( + "kv_ssd_storage_backend", + "kv_ssd_uring_mode", "side_transfer_worker_count", "side_transfer_worker_p2p_port_base", "side_transfer_role", @@ -268,6 +293,11 @@ BENCHMARK_KEY_AFFINITY_LOCALITY_RATIO, BENCHMARK_KEY_AFFINITY_SLOT_COUNT, "kv_bootstrap_before_ready", + BENCHMARK_KEY_KV_BOOTSTRAP_CONCURRENCY, + BENCHMARK_KEY_KV_BOOTSTRAP_PUT_GAP_MS, + BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY, + BENCHMARK_KEY_KV_GET_OUTPUT, + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX, ) @@ -278,6 +308,50 @@ class KVGetResultKind(Enum): ERROR = "error" +@unique +class KVGetSourceKind(Enum): + MEMORY = "memory" + SSD = "ssd" + + +@unique +class KVGetOutput(Enum): + HOLDER = "holder" + BYTES = "bytes" + CUDA = "cuda" + + +def normalize_kv_get_output(raw: Any) -> KVGetOutput: + value = KVGetOutput.HOLDER.value if raw is None else str(raw).strip().lower() + try: + return KVGetOutput(value) + except ValueError as exc: + supported = ", ".join(output.value for output in KVGetOutput) + raise ValueError( + f"{BENCHMARK_KEY_KV_GET_OUTPUT} must be one of: {supported}; got {raw!r}" + ) from exc + + +def normalize_kv_get_source_kind(raw: Any) -> KVGetSourceKind: + value = str(raw).strip().lower() + try: + return KVGetSourceKind(value) + except ValueError as exc: + supported = ", ".join(source.value for source in KVGetSourceKind) + raise ValueError( + f"GET source must be one of: {supported}; got {raw!r}" + ) from exc + + +def normalize_kv_cuda_device_index(raw: Any) -> int: + value = 0 if raw is None else raw + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError( + f"{BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX} must be a non-negative integer; got {raw!r}" + ) + return int(value) + + def classify_kv_get_result(error_msg: Optional[str]) -> KVGetResultKind: if error_msg is None: return KVGetResultKind.CACHE_HIT @@ -304,6 +378,10 @@ def _is_key_being_written_error(error: Any) -> bool: return type(error).__name__ == "KeyBeingWrittenError" +def _is_key_already_exists_error(error: Any) -> bool: + return type(error).__name__ == "KeyAlreadyExistsError" + + def _is_mooncake_replica_not_ready_error(error: Any) -> bool: details = getattr(error, "details", None) if not isinstance(details, dict): @@ -312,7 +390,11 @@ def _is_mooncake_replica_not_ready_error(error: Any) -> bool: def _is_put_compat_success_error(error: Any) -> bool: - return _is_key_being_written_error(error) or _is_mooncake_replica_not_ready_error(error) + return ( + _is_key_being_written_error(error) + or _is_key_already_exists_error(error) + or _is_mooncake_replica_not_ready_error(error) + ) def normalize_kv_get_error(error_msg: Optional[str]) -> Optional[str]: @@ -1517,10 +1599,12 @@ def get_blocking( *, deadline_ts: float, ctx: str, + expected_payload_size: int, ) -> Optional[str]: _ = key _ = deadline_ts _ = ctx + _ = expected_payload_size return f"GET failed: backend {self.backend_kind} does not expose KV operations" def close(self) -> _SimpleResult: @@ -1803,6 +1887,7 @@ def get_blocking( *, deadline_ts: float, ctx: str, + expected_payload_size: int, ) -> Optional[str]: _ = deadline_ts _ = ctx @@ -1810,6 +1895,11 @@ def get_blocking( payload = self.get(key) if payload is None: return KV_GET_MISS_ERROR + if len(payload) != int(expected_payload_size): + return ( + "GET failed: payload length mismatch: " + f"expected={expected_payload_size} actual={len(payload)}" + ) return None except Exception as exc: return normalize_kv_get_error(f"GET failed: {exc}") @@ -1865,11 +1955,857 @@ def _sanitize_benchmark_client_kvcache_config(kvcache_config: dict[str, Any]) -> return sanitized -class FluxonBlockingStore: - def __init__(self, store: KvClient) -> None: - self.backend_kind = BACKEND_KIND_FLUXON +_FLAT_KV_TYPE_BYTES = 5 +_DLPACK_DTYPE_UINT = 1 +_CUDA_MEMCPY_HOST_TO_DEVICE = 1 +_CUDA_STREAM_NON_BLOCKING = 1 +_CUDA_ERROR_NOT_READY = 600 + + +def _flat_dict_payload_range( + data: memoryview, + expected_payload_size: int, +) -> tuple[memoryview, int, int]: + view = data if data.format == "B" and data.ndim == 1 else data.cast("B") + total_len = len(view) + if total_len < 4: + raise ValueError("flat dict payload is missing its entry-count header") + + (entry_count,) = struct.unpack_from(" total_len: + raise ValueError("flat dict payload has a truncated key length") + (key_len,) = struct.unpack_from(" total_len: + raise ValueError("flat dict payload has truncated key bytes") + key = bytes(view[pos : pos + key_len]).decode("utf-8") + pos += key_len + if pos + 5 > total_len: + raise ValueError("flat dict payload has a truncated value header") + type_id = int(view[pos]) + pos += 1 + (value_len,) = struct.unpack_from(" total_len: + raise ValueError("flat dict payload has truncated value bytes") + if key == "payload": + if type_id != _FLAT_KV_TYPE_BYTES: + raise TypeError(f"payload field must be bytes-compatible, got type id {type_id}") + if value_len != int(expected_payload_size): + raise ValueError( + "payload length mismatch: " + f"expected={expected_payload_size} actual={value_len}" + ) + return view, pos, value_len + pos += value_len + raise KeyError("flat dict payload field is missing") + + +def _flat_dict_payload_view(data: memoryview, expected_payload_size: int) -> memoryview: + view, offset, size = _flat_dict_payload_range(data, expected_payload_size) + return view[offset : offset + size] + + +def _mooncake_payload_view(buffer_handle: Any, expected_payload_size: int) -> memoryview: + return _flat_dict_payload_view(memoryview(buffer_handle), expected_payload_size) + + +def _cudart_library_candidates() -> list[str]: + candidates: list[str] = [] + discovered = ctypes.util.find_library("cudart") + if discovered: + candidates.append(discovered) + candidates.append("libcudart.so.12") + + try: + spec = importlib.util.find_spec("nvidia.cuda_runtime") + except (ImportError, ModuleNotFoundError): + spec = None + if spec is not None and spec.submodule_search_locations: + for package_dir in spec.submodule_search_locations: + candidates.append(str(Path(package_dir) / "lib" / "libcudart.so.12")) + + candidates.extend( + [ + "/usr/local/cuda/lib64/libcudart.so.12", + "/usr/local/cuda-12/lib64/libcudart.so.12", + ] + ) + return list(dict.fromkeys(candidates)) + + +class _CudaRuntime: + def __init__(self) -> None: + self._lib = self._load_library() + self._configure_functions() + + @staticmethod + def _load_library() -> Any: + failures: list[str] = [] + for candidate in _cudart_library_candidates(): + try: + return ctypes.CDLL(candidate) + except OSError as exc: + failures.append(f"{candidate}: {exc}") + raise RuntimeError("unable to load CUDA runtime: " + "; ".join(failures)) + + def _configure_functions(self) -> None: + self._lib.cudaSetDevice.argtypes = [ctypes.c_int] + self._lib.cudaSetDevice.restype = ctypes.c_int + self._lib.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] + self._lib.cudaMalloc.restype = ctypes.c_int + self._lib.cudaFree.argtypes = [ctypes.c_void_p] + self._lib.cudaFree.restype = ctypes.c_int + self._lib.cudaHostAlloc.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_size_t, + ctypes.c_uint, + ] + self._lib.cudaHostAlloc.restype = ctypes.c_int + self._lib.cudaFreeHost.argtypes = [ctypes.c_void_p] + self._lib.cudaFreeHost.restype = ctypes.c_int + self._lib.cudaStreamCreateWithFlags.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_uint, + ] + self._lib.cudaStreamCreateWithFlags.restype = ctypes.c_int + self._lib.cudaStreamSynchronize.argtypes = [ctypes.c_void_p] + self._lib.cudaStreamSynchronize.restype = ctypes.c_int + self._lib.cudaStreamDestroy.argtypes = [ctypes.c_void_p] + self._lib.cudaStreamDestroy.restype = ctypes.c_int + self._lib.cudaEventCreateWithFlags.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_uint, + ] + self._lib.cudaEventCreateWithFlags.restype = ctypes.c_int + self._lib.cudaEventRecord.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + self._lib.cudaEventRecord.restype = ctypes.c_int + self._lib.cudaEventQuery.argtypes = [ctypes.c_void_p] + self._lib.cudaEventQuery.restype = ctypes.c_int + self._lib.cudaEventSynchronize.argtypes = [ctypes.c_void_p] + self._lib.cudaEventSynchronize.restype = ctypes.c_int + self._lib.cudaEventElapsedTime.argtypes = [ + ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, + ctypes.c_void_p, + ] + self._lib.cudaEventElapsedTime.restype = ctypes.c_int + self._lib.cudaEventDestroy.argtypes = [ctypes.c_void_p] + self._lib.cudaEventDestroy.restype = ctypes.c_int + self._lib.cudaMemcpyAsync.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_void_p, + ] + self._lib.cudaMemcpyAsync.restype = ctypes.c_int + self._lib.cudaGetErrorString.argtypes = [ctypes.c_int] + self._lib.cudaGetErrorString.restype = ctypes.c_char_p + + def _check(self, code: int, operation: str) -> None: + if int(code) == 0: + return + raw_message = self._lib.cudaGetErrorString(int(code)) + message = raw_message.decode("utf-8", errors="replace") if raw_message else "unknown" + raise RuntimeError(f"{operation} failed: cuda_error={code} message={message}") + + def set_device(self, device_index: int) -> None: + self._check(self._lib.cudaSetDevice(int(device_index)), "cudaSetDevice") + + def malloc(self, size: int) -> int: + ptr = ctypes.c_void_p() + self._check( + self._lib.cudaMalloc(ctypes.byref(ptr), ctypes.c_size_t(size)), + "cudaMalloc", + ) + if not ptr.value: + raise RuntimeError("cudaMalloc returned a null device pointer") + return int(ptr.value) + + def free(self, ptr: int) -> None: + self._check(self._lib.cudaFree(ctypes.c_void_p(ptr)), "cudaFree") + + def host_alloc(self, size: int) -> int: + ptr = ctypes.c_void_p() + self._check( + self._lib.cudaHostAlloc( + ctypes.byref(ptr), + ctypes.c_size_t(size), + ctypes.c_uint(0), + ), + "cudaHostAlloc", + ) + if not ptr.value: + raise RuntimeError("cudaHostAlloc returned a null host pointer") + return int(ptr.value) + + def host_free(self, ptr: int) -> None: + self._check(self._lib.cudaFreeHost(ctypes.c_void_p(ptr)), "cudaFreeHost") + + def stream_create(self) -> int: + stream = ctypes.c_void_p() + self._check( + self._lib.cudaStreamCreateWithFlags( + ctypes.byref(stream), + ctypes.c_uint(_CUDA_STREAM_NON_BLOCKING), + ), + "cudaStreamCreateWithFlags", + ) + if not stream.value: + raise RuntimeError("cudaStreamCreateWithFlags returned a null stream") + return int(stream.value) + + def stream_synchronize(self, stream: int) -> None: + self._check( + self._lib.cudaStreamSynchronize(ctypes.c_void_p(stream)), + "cudaStreamSynchronize", + ) + + def stream_destroy(self, stream: int) -> None: + self._check( + self._lib.cudaStreamDestroy(ctypes.c_void_p(stream)), + "cudaStreamDestroy", + ) + + def event_create(self) -> int: + event = ctypes.c_void_p() + self._check( + self._lib.cudaEventCreateWithFlags( + ctypes.byref(event), + ctypes.c_uint(0), + ), + "cudaEventCreateWithFlags", + ) + if not event.value: + raise RuntimeError("cudaEventCreateWithFlags returned a null event") + return int(event.value) + + def event_record(self, event: int, stream: int) -> None: + self._check( + self._lib.cudaEventRecord( + ctypes.c_void_p(event), + ctypes.c_void_p(stream), + ), + "cudaEventRecord", + ) + + def event_query(self, event: int) -> bool: + code = int(self._lib.cudaEventQuery(ctypes.c_void_p(event))) + if code == 0: + return True + if code == _CUDA_ERROR_NOT_READY: + return False + self._check(code, "cudaEventQuery") + raise AssertionError("cudaEventQuery error handling must return or raise") + + def event_synchronize(self, event: int) -> None: + self._check( + self._lib.cudaEventSynchronize(ctypes.c_void_p(event)), + "cudaEventSynchronize", + ) + + def event_elapsed_us(self, start_event: int, done_event: int) -> float: + elapsed_ms = ctypes.c_float() + self._check( + self._lib.cudaEventElapsedTime( + ctypes.byref(elapsed_ms), + ctypes.c_void_p(start_event), + ctypes.c_void_p(done_event), + ), + "cudaEventElapsedTime", + ) + return max(0.0, float(elapsed_ms.value) * 1000.0) + + def event_destroy(self, event: int) -> None: + self._check( + self._lib.cudaEventDestroy(ctypes.c_void_p(event)), + "cudaEventDestroy", + ) + + def memcpy_host_to_device_async( + self, + *, + device_ptr: int, + host_ptr: int, + size: int, + stream: int, + ) -> None: + self._check( + self._lib.cudaMemcpyAsync( + ctypes.c_void_p(device_ptr), + ctypes.c_void_p(host_ptr), + ctypes.c_size_t(size), + _CUDA_MEMCPY_HOST_TO_DEVICE, + ctypes.c_void_p(stream), + ), + "cudaMemcpyAsync(host_to_device)", + ) + + +@dataclass +class _CudaPendingCopy: + token: Any + # Keep the native holder alive through the operation completion boundary. + keepalive: Any + enqueued_at: float + host_stage_us: float + submit_us: float + + +@dataclass(frozen=True) +class _CudaCopyCompletion: + token: Any + error_msg: Optional[str] + completed_at: float + completed_wall_ts: float + host_stage_us: float + submit_us: float + h2d_us: float + pipeline_residence_us: float + completion_wait_us: float + + +@dataclass +class _CudaPipelineSlot: + stream: int + start_event: int + done_event: int + host_ptr: int = 0 + host_capacity: int = 0 + device_ptr: int = 0 + device_capacity: int = 0 + pending: Optional[_CudaPendingCopy] = None + + +@dataclass +class _CudaThreadPipeline: + slots: list[_CudaPipelineSlot] + pending_slot_indexes: deque[int] = field(default_factory=deque) + lock: threading.Lock = field(default_factory=threading.Lock) + + +class _CudaHostToDevicePipeline: + def __init__( + self, + *, + device_index: int = 0, + depth: int = CUDA_H2D_PIPELINE_DEPTH, + runtime: Optional[Any] = None, + ) -> None: + if int(depth) <= 0: + raise ValueError(f"CUDA H2D pipeline depth must be > 0, got {depth}") + self._device_index = int(device_index) + self._depth = int(depth) + self._runtime = _CudaRuntime() if runtime is None else runtime + self._lock = threading.Lock() + self._states: dict[int, _CudaThreadPipeline] = {} + self._closed = False + + def _new_thread_state(self) -> _CudaThreadPipeline: + slots: list[_CudaPipelineSlot] = [] + try: + for _ in range(self._depth): + stream = self._runtime.stream_create() + start_event = 0 + done_event = 0 + try: + start_event = self._runtime.event_create() + done_event = self._runtime.event_create() + except Exception: + if start_event: + self._runtime.event_destroy(start_event) + self._runtime.stream_destroy(stream) + raise + slots.append( + _CudaPipelineSlot( + stream=stream, + start_event=start_event, + done_event=done_event, + ) + ) + except Exception: + for slot in reversed(slots): + self._runtime.event_destroy(slot.done_event) + self._runtime.event_destroy(slot.start_event) + self._runtime.stream_destroy(slot.stream) + raise + return _CudaThreadPipeline(slots=slots) + + def _state_for_current_thread(self) -> _CudaThreadPipeline: + thread_id = threading.get_ident() + with self._lock: + if self._closed: + raise RuntimeError("CUDA H2D pipeline is already closed") + state = self._states.get(thread_id) + if state is None: + self._runtime.set_device(self._device_index) + state = self._new_thread_state() + self._states[thread_id] = state + return state + + def _ensure_slot_capacity(self, slot: _CudaPipelineSlot, size: int) -> None: + if slot.host_capacity >= size and slot.device_capacity >= size: + return + new_host_ptr = self._runtime.host_alloc(size) + try: + new_device_ptr = self._runtime.malloc(size) + except Exception: + self._runtime.host_free(new_host_ptr) + raise + + old_host_ptr = slot.host_ptr + old_device_ptr = slot.device_ptr + slot.host_ptr = new_host_ptr + slot.host_capacity = int(size) + slot.device_ptr = new_device_ptr + slot.device_capacity = int(size) + if old_device_ptr: + self._runtime.free(old_device_ptr) + if old_host_ptr: + self._runtime.host_free(old_host_ptr) + + def _complete_front( + self, + state: _CudaThreadPipeline, + *, + block: bool, + ) -> Optional[_CudaCopyCompletion]: + if not state.pending_slot_indexes: + return None + slot_index = state.pending_slot_indexes[0] + slot = state.slots[slot_index] + pending = slot.pending + if pending is None: + raise RuntimeError("CUDA pipeline pending queue points to an empty slot") + + wait_started_at = time.perf_counter() + error_msg: Optional[str] = None + h2d_us = 0.0 + try: + if block: + self._runtime.event_synchronize(slot.done_event) + elif not self._runtime.event_query(slot.done_event): + return None + h2d_us = self._runtime.event_elapsed_us( + slot.start_event, + slot.done_event, + ) + except Exception as exc: + error_msg = str(exc) + try: + self._runtime.stream_synchronize(slot.stream) + except Exception as sync_exc: + error_msg = f"{error_msg}; stream cleanup failed: {sync_exc}" + wait_done_at = time.perf_counter() + completed_at = wait_done_at + completed_wall_ts = time.time() + + state.pending_slot_indexes.popleft() + slot.pending = None + return _CudaCopyCompletion( + token=pending.token, + error_msg=error_msg, + completed_at=completed_at, + completed_wall_ts=completed_wall_ts, + host_stage_us=pending.host_stage_us, + submit_us=pending.submit_us, + h2d_us=h2d_us, + pipeline_residence_us=max( + 0.0, + (completed_at - pending.enqueued_at) * 1_000_000.0, + ), + completion_wait_us=max( + 0.0, + (wait_done_at - wait_started_at) * 1_000_000.0, + ), + ) + + def _poll_locked(self, state: _CudaThreadPipeline) -> list[_CudaCopyCompletion]: + completions: list[_CudaCopyCompletion] = [] + while state.pending_slot_indexes: + completion = self._complete_front(state, block=False) + if completion is None: + break + completions.append(completion) + return completions + + def submit_from_host( + self, + *, + source_ptr: int, + size: int, + keepalive: Any, + token: Any, + ) -> list[_CudaCopyCompletion]: + if source_ptr <= 0: + raise ValueError(f"CUDA copy source pointer must be positive, got {source_ptr}") + if size <= 0: + raise ValueError(f"CUDA copy size must be positive, got {size}") + + state = self._state_for_current_thread() + with state.lock: + completions: list[_CudaCopyCompletion] = [] + if len(state.pending_slot_indexes) >= self._depth: + completion = self._complete_front(state, block=True) + if completion is None: + raise RuntimeError("full CUDA H2D pipeline has no pending completion") + completions.append(completion) + + slot_index = next( + ( + index + for index, slot in enumerate(state.slots) + if slot.pending is None + ), + None, + ) + if slot_index is None: + raise RuntimeError("CUDA H2D pipeline has no reusable slot") + slot = state.slots[slot_index] + enqueued_at = time.perf_counter() + try: + self._ensure_slot_capacity(slot, int(size)) + # Stage pageable backend memory into this slot's pinned buffer. + stage_started_at = time.perf_counter() + ctypes.memmove( + ctypes.c_void_p(slot.host_ptr), + ctypes.c_void_p(source_ptr), + ctypes.c_size_t(size), + ) + stage_done_at = time.perf_counter() + self._runtime.event_record(slot.start_event, slot.stream) + submit_started_at = time.perf_counter() + self._runtime.memcpy_host_to_device_async( + device_ptr=slot.device_ptr, + host_ptr=slot.host_ptr, + size=int(size), + stream=slot.stream, + ) + submit_done_at = time.perf_counter() + # Completion is collected by event polling or the final drain. + self._runtime.event_record(slot.done_event, slot.stream) + except Exception: + try: + self._runtime.stream_synchronize(slot.stream) + except Exception: + pass + raise + slot.pending = _CudaPendingCopy( + token=token, + keepalive=keepalive, + enqueued_at=enqueued_at, + host_stage_us=max( + 0.0, + (stage_done_at - stage_started_at) * 1_000_000.0, + ), + submit_us=max( + 0.0, + (submit_done_at - submit_started_at) * 1_000_000.0, + ), + ) + state.pending_slot_indexes.append(slot_index) + return completions + + def poll_current_thread(self) -> list[_CudaCopyCompletion]: + state = self._state_for_current_thread() + with state.lock: + return self._poll_locked(state) + + def drain_current_thread(self) -> list[_CudaCopyCompletion]: + state = self._state_for_current_thread() + with state.lock: + completions: list[_CudaCopyCompletion] = [] + while state.pending_slot_indexes: + completion = self._complete_front(state, block=True) + if completion is None: + raise RuntimeError("CUDA pipeline drain lost a pending completion") + completions.append(completion) + return completions + + def copy_from_host(self, source_ptr: int, size: int, *, keepalive: Any) -> None: + token = object() + completions = self.submit_from_host( + source_ptr=source_ptr, + size=size, + keepalive=keepalive, + token=token, + ) + completions.extend(self.drain_current_thread()) + matching = [completion for completion in completions if completion.token is token] + if len(matching) != 1: + raise RuntimeError( + "synchronous CUDA H2D copy did not produce exactly one completion" + ) + if matching[0].error_msg is not None: + raise RuntimeError(matching[0].error_msg) + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._closed = True + states = list(self._states.values()) + self._states.clear() + + errors: list[str] = [] + self._runtime.set_device(self._device_index) + for state in states: + with state.lock: + while state.pending_slot_indexes: + try: + self._complete_front(state, block=True) + except Exception as exc: + errors.append(str(exc)) + break + for slot in reversed(state.slots): + for operation, value in ( + (self._runtime.free, slot.device_ptr), + (self._runtime.host_free, slot.host_ptr), + (self._runtime.event_destroy, slot.done_event), + (self._runtime.event_destroy, slot.start_event), + (self._runtime.stream_destroy, slot.stream), + ): + if not value: + continue + try: + operation(value) + except Exception as exc: + errors.append(str(exc)) + if errors: + raise RuntimeError("CUDA H2D pipeline close failed: " + "; ".join(errors)) + + +@dataclass(frozen=True) +class _CudaHostPayload: + source_ptr: int + size: int + keepalive: Any + + +@dataclass(frozen=True) +class _PendingCudaGet: + token: Any + key: str + ctx: str + deadline_ts: float + expected_payload_size: int + started_at: float + host_ready_at: float + source_kind: Optional[KVGetSourceKind] + + +@dataclass(frozen=True) +class _PipelinedCudaGetCompletion: + token: Any + error_msg: Optional[str] + latency_us: float + finish_ts: float + expected_payload_size: int + source_kind: Optional[KVGetSourceKind] + + +@dataclass(frozen=True) +class KVBlockingGetCompletion: + error_msg: Optional[str] + source_kind: Optional[KVGetSourceKind] + + +@dataclass(frozen=True) +class _CudaWorkerGetToken: + op_idx: int + key: str + expected_data_size: int + inflight_at_start: int + + +class _MooncakeOffloadReadCounterWindow: + """Sample Mooncake's cumulative offload-RPC counter at window boundaries.""" + + def __init__(self, store: Any, *, start_ts: float, end_ts: float) -> None: + if end_ts < start_ts: + raise ValueError( + f"source counter window end precedes start: {end_ts} < {start_ts}" + ) self._store = store + self._start_ts = float(start_ts) + self._end_ts = float(end_ts) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._run, + name="mooncake-offload-read-counter-window", + daemon=True, + ) + self._lock = threading.Lock() + self._start_count: Optional[int] = None + self._end_count: Optional[int] = None + self._start_sample_ts: Optional[float] = None + self._end_sample_ts: Optional[float] = None + self._error: Optional[str] = None + + def start(self) -> None: + self._thread.start() + + def _wait_until(self, deadline_ts: float) -> bool: + delay = max(0.0, float(deadline_ts) - time.time()) + return self._stop.wait(timeout=delay) + + def _read_counter(self) -> int: + return self._store._benchmark_offload_rpc_read_count() + + def _run(self) -> None: + try: + if self._wait_until(self._start_ts): + return + start_count = self._read_counter() + start_sample_ts = time.time() + with self._lock: + self._start_count = start_count + self._start_sample_ts = start_sample_ts + + if self._wait_until(self._end_ts): + return + end_count = self._read_counter() + end_sample_ts = time.time() + with self._lock: + self._end_count = end_count + self._end_sample_ts = end_sample_ts + except Exception as exc: + with self._lock: + self._error = f"{type(exc).__name__}: {exc}" + + def finish(self) -> Dict[str, Any]: + self._thread.join(timeout=5.0) + if self._thread.is_alive(): + self._stop.set() + self._thread.join(timeout=1.0) + with self._lock: + if self._error is None: + self._error = "counter sampler did not reach the end boundary" + + with self._lock: + start_count = self._start_count + end_count = self._end_count + error = self._error + start_sample_ts = self._start_sample_ts + end_sample_ts = self._end_sample_ts + + delta: Optional[int] = None + if start_count is not None and end_count is not None: + delta = int(end_count) - int(start_count) + if delta < 0: + error = ( + "Mooncake offload counter decreased during the measurement window: " + f"start={start_count} end={end_count}" + ) + delta = None + return { + "observation_kind": "mooncake_offload_rpc_counter_window", + "counter_unit": "single_key_get_buffer_offload_rpc", + "window_start_ts": self._start_ts, + "window_end_ts": self._end_ts, + "start_sample_ts": start_sample_ts, + "end_sample_ts": end_sample_ts, + "start_count": start_count, + "end_count": end_count, + "ssd_read_count": delta, + "complete": delta is not None and error is None, + "error": error, + } + + def cancel(self) -> None: + self._stop.set() + self._thread.join(timeout=1.0) + + +class KVBenchmarkBlockingStore: + def __init__( + self, + store: KvClient, + *, + backend_kind: str, + get_output: KVGetOutput, + cuda_device_index: int = 0, + ) -> None: + self.backend_kind = str(backend_kind).strip().upper() + self._store = store + self._get_output = get_output self._phase_profiler = _FluxonPhaseProfiler() + self._put_dedupe_condition = threading.Condition() + self._put_confirmed_keys: set[str] = set() + self._put_inflight_keys: set[str] = set() + self._cuda_pipeline = ( + _CudaHostToDevicePipeline(device_index=cuda_device_index) + if self._get_output == KVGetOutput.CUDA + else None + ) + self._mooncake_store: Optional[Any] = None + self._source_counter_window: Optional[ + _MooncakeOffloadReadCounterWindow + ] = None + if self.backend_kind == BACKEND_KIND_MOONCAKE: + from fluxon_py.kvclient.mooncake import MooncakeStore + + if not isinstance(store, MooncakeStore): + raise TypeError( + "Mooncake benchmark backend must be backed by MooncakeStore; " + f"got {type(store)}" + ) + self._mooncake_store = store + + def begin_get_source_counter_window( + self, + *, + start_ts: float, + end_ts: float, + ) -> None: + if self._source_counter_window is not None: + raise RuntimeError("GET source counter window is already active") + if self.backend_kind != BACKEND_KIND_MOONCAKE: + return + assert self._mooncake_store is not None + sampler = _MooncakeOffloadReadCounterWindow( + self._mooncake_store, + start_ts=start_ts, + end_ts=end_ts, + ) + self._source_counter_window = sampler + sampler.start() + + def finish_get_source_counter_window(self) -> Dict[str, Any]: + sampler = self._source_counter_window + self._source_counter_window = None + if sampler is None: + return {} + return sampler.finish() + + def _begin_deduplicated_put( + self, + key: str, + *, + deadline_ts: float, + ) -> tuple[bool, Optional[str]]: + with self._put_dedupe_condition: + while key in self._put_inflight_keys: + remaining = float(deadline_ts) - time.time() + if remaining <= 0.0: + return False, f"PUT timed out waiting for same-key benchmark put: key={key!r}" + self._put_dedupe_condition.wait(timeout=remaining) + if key in self._put_confirmed_keys: + return False, None + self._put_inflight_keys.add(key) + return True, None + + def _finish_deduplicated_put(self, key: str, *, confirmed: bool) -> None: + with self._put_dedupe_condition: + if key not in self._put_inflight_keys: + raise RuntimeError(f"benchmark PUT key was not inflight: {key!r}") + self._put_inflight_keys.remove(key) + if confirmed: + self._put_confirmed_keys.add(key) + self._put_dedupe_condition.notify_all() def put_blocking( self, @@ -1882,17 +2818,61 @@ def put_blocking( try: _bench_kv_print(f"{ctx} PUT begin key={key!r} payload_size={len(payload)}", verbose_only=True) started_at = time.perf_counter() - result = self._store.put_blocking( + should_submit, dedupe_error = self._begin_deduplicated_put( key, - {"payload": payload}, - opts=PutOptionalArgs(reject_if_inflight_same_key=True), + deadline_ts=deadline_ts, ) + if not should_submit: + done_at = time.perf_counter() + phase_sample = _build_fluxon_sync_phase_sample( + started_at=started_at, + done_at=done_at, + deadline_ts=deadline_ts, + wall_done_ts=time.time(), + ) + self._phase_profiler.record( + op_name="PUT", + key=key, + sample=phase_sample, + error_msg=dedupe_error, + ) + if dedupe_error is not None: + return dedupe_error + _bench_kv_print( + f"{ctx} PUT deduplicated key={key!r}", + verbose_only=True, + ) + return None + + value: Any = payload + if self._get_output == KVGetOutput.CUDA: + value = DLPackBytesView( + payload, + dtype_code=_DLPACK_DTYPE_UINT, + bits=8, + lanes=1, + shape=(len(payload),), + ) + put_confirmed = False + put_error: Optional[Any] = None + try: + result = self._store.put_blocking( + key, + {"payload": value}, + opts=PutOptionalArgs(reject_if_exists=True), + ) + if result.is_ok(): + put_confirmed = True + else: + put_error = result.unwrap_error() + put_confirmed = _is_key_already_exists_error(put_error) + finally: + self._finish_deduplicated_put(key, confirmed=put_confirmed) done_at = time.perf_counter() wall_done_ts = time.time() err: Optional[str] = None compat_success = False - if not result.is_ok(): - put_error = result.unwrap_error() + if put_error is not None: if _is_put_compat_success_error(put_error): _bench_kv_print( f"{ctx} PUT compat-success key={key!r} reason={type(put_error).__name__}", @@ -1932,24 +2912,336 @@ def put_blocking( _bench_kv_print(f"{ctx} PUT exception key={key!r} err={exc}") return f"PUT exception: {exc}" + def _get_native_holder(self, key: str) -> Any: + if self.backend_kind == BACKEND_KIND_MOONCAKE: + assert self._mooncake_store is not None + return self._mooncake_store.get_buffer_blocking(key) + return self._store.get_blocking(key) + + def _get_source_kind(self, holder: Any) -> Optional[KVGetSourceKind]: + if self.backend_kind == BACKEND_KIND_MOONCAKE: + return None + return normalize_kv_get_source_kind(holder._benchmark_source_kind()) + + def _cuda_host_payload( + self, + holder: Any, + *, + expected_payload_size: int, + ) -> _CudaHostPayload: + if self.backend_kind == BACKEND_KIND_MOONCAKE: + raw_view, payload_offset, payload_size = _flat_dict_payload_range( + memoryview(holder), + expected_payload_size, + ) + if len(raw_view) != int(holder.size()): + raise ValueError( + "Mooncake buffer protocol size differs from BufferHandle.size(): " + f"view={len(raw_view)} handle={holder.size()}" + ) + return _CudaHostPayload( + source_ptr=int(holder.ptr()) + payload_offset, + size=payload_size, + keepalive=(holder, raw_view), + ) + + access_result = holder.access() + if not access_result.is_ok(): + raise RuntimeError( + f"MemHolder.access() failed: {access_result.unwrap_error()}" + ) + value = access_result.unwrap() + payload_view = value.get("payload") + dlpack_info = _dlpack_cpu_tensor_info(payload_view) + if not dlpack_info.is_ok(): + raise RuntimeError( + "Fluxon CUDA output requires a CPU DLPack payload: " + f"{dlpack_info.unwrap_error()}" + ) + ( + source_ptr, + payload_size, + dlpack_capsule, + _, + _, + _, + _, + ) = dlpack_info.unwrap() + if payload_size != int(expected_payload_size): + raise ValueError( + "DLPack payload length mismatch: " + f"expected={expected_payload_size} actual={payload_size}" + ) + return _CudaHostPayload( + source_ptr=int(source_ptr), + size=int(payload_size), + keepalive=(holder, value, payload_view, dlpack_capsule), + ) + + def _immediate_cuda_get_completion( + self, + *, + token: Any, + key: str, + deadline_ts: float, + expected_payload_size: int, + started_at: float, + error_msg: str, + ) -> _PipelinedCudaGetCompletion: + completed_at = time.perf_counter() + completed_wall_ts = time.time() + phase_sample = _build_fluxon_sync_phase_sample( + started_at=started_at, + done_at=completed_at, + deadline_ts=deadline_ts, + wall_done_ts=completed_wall_ts, + ) + self._phase_profiler.record( + op_name="GET", + key=key, + sample=phase_sample, + error_msg=error_msg, + ) + return _PipelinedCudaGetCompletion( + token=token, + error_msg=error_msg, + latency_us=max(0.0, (completed_at - started_at) * 1_000_000.0), + finish_ts=completed_wall_ts, + expected_payload_size=int(expected_payload_size), + source_kind=None, + ) + + def _finalize_cuda_copy_completion( + self, + completion: _CudaCopyCompletion, + ) -> _PipelinedCudaGetCompletion: + pending = completion.token + if not isinstance(pending, _PendingCudaGet): + raise TypeError( + "CUDA H2D pipeline returned an unknown completion token: " + f"{type(pending)}" + ) + + error_msg: Optional[str] = None + if completion.error_msg is not None: + error_msg = f"GET CUDA H2D failed: {completion.error_msg}" + elif completion.completed_wall_ts > pending.deadline_ts: + error_msg = ( + "GET timed out after CUDA H2D completion: " + f"deadline_ts={pending.deadline_ts:.3f} " + f"now_ts={completion.completed_wall_ts:.3f} " + f"now_ms={completion.completed_wall_ts * 1000.0:.1f}" + ) + + phase_sample = _build_fluxon_sync_phase_sample( + started_at=pending.started_at, + done_at=completion.completed_at, + deadline_ts=pending.deadline_ts, + wall_done_ts=completion.completed_wall_ts, + extra_payload={ + "cuda_backend_get_us": max( + 0.0, + (pending.host_ready_at - pending.started_at) * 1_000_000.0, + ), + "cuda_host_stage_us": completion.host_stage_us, + "cuda_submit_us": completion.submit_us, + "cuda_h2d_event_us": completion.h2d_us, + "cuda_pipeline_residence_us": completion.pipeline_residence_us, + "cuda_completion_wait_us": completion.completion_wait_us, + }, + ) + self._phase_profiler.record( + op_name="GET", + key=pending.key, + sample=phase_sample, + error_msg=error_msg, + ) + if error_msg is not None: + _bench_kv_print( + f"{pending.ctx} GET failed-after-CUDA key={pending.key!r} err={error_msg}" + ) + else: + _bench_kv_print( + f"{pending.ctx} GET CUDA done key={pending.key!r}", + verbose_only=True, + ) + return _PipelinedCudaGetCompletion( + token=pending.token, + error_msg=error_msg, + latency_us=max( + 0.0, + (completion.completed_at - pending.started_at) * 1_000_000.0, + ), + finish_ts=completion.completed_wall_ts, + expected_payload_size=pending.expected_payload_size, + source_kind=pending.source_kind, + ) + + def _finalize_cuda_copy_completions( + self, + completions: Sequence[_CudaCopyCompletion], + ) -> list[_PipelinedCudaGetCompletion]: + return [ + self._finalize_cuda_copy_completion(completion) + for completion in completions + ] + + def submit_cuda_get( + self, + key: str, + *, + deadline_ts: float, + ctx: str, + expected_payload_size: int, + token: Any, + ) -> list[_PipelinedCudaGetCompletion]: + if self._get_output != KVGetOutput.CUDA or self._cuda_pipeline is None: + raise RuntimeError("submit_cuda_get requires kv_get_output=cuda") + + _bench_kv_print(f"{ctx} GET begin key={key!r}", verbose_only=True) + started_at = time.perf_counter() + try: + result = self._get_native_holder(key) + if not result.is_ok(): + error_msg = normalize_kv_get_error( + f"GET failed: {result.unwrap_error()}" + ) + assert error_msg is not None + return [ + self._immediate_cuda_get_completion( + token=token, + key=key, + deadline_ts=deadline_ts, + expected_payload_size=expected_payload_size, + started_at=started_at, + error_msg=error_msg, + ) + ] + + holder = result.unwrap() + source_kind = self._get_source_kind(holder) + host_payload = self._cuda_host_payload( + holder, + expected_payload_size=expected_payload_size, + ) + host_ready_at = time.perf_counter() + pending = _PendingCudaGet( + token=token, + key=key, + ctx=ctx, + deadline_ts=float(deadline_ts), + expected_payload_size=int(expected_payload_size), + started_at=started_at, + host_ready_at=host_ready_at, + source_kind=source_kind, + ) + copy_completions = self._cuda_pipeline.submit_from_host( + source_ptr=host_payload.source_ptr, + size=host_payload.size, + keepalive=host_payload.keepalive, + token=pending, + ) + return self._finalize_cuda_copy_completions(copy_completions) + except Exception as exc: + error_msg = f"GET exception: {exc}" + _bench_kv_print(f"{ctx} GET exception key={key!r} err={exc}") + return [ + self._immediate_cuda_get_completion( + token=token, + key=key, + deadline_ts=deadline_ts, + expected_payload_size=expected_payload_size, + started_at=started_at, + error_msg=error_msg, + ) + ] + + def poll_cuda_gets(self) -> list[_PipelinedCudaGetCompletion]: + if self._get_output != KVGetOutput.CUDA or self._cuda_pipeline is None: + return [] + return self._finalize_cuda_copy_completions( + self._cuda_pipeline.poll_current_thread() + ) + + def drain_cuda_gets(self) -> list[_PipelinedCudaGetCompletion]: + if self._get_output != KVGetOutput.CUDA or self._cuda_pipeline is None: + return [] + return self._finalize_cuda_copy_completions( + self._cuda_pipeline.drain_current_thread() + ) + def get_blocking( self, key: str, *, deadline_ts: float, ctx: str, - ) -> Optional[str]: + expected_payload_size: int, + ) -> KVBlockingGetCompletion: + if self._get_output == KVGetOutput.CUDA: + token = object() + completions = self.submit_cuda_get( + key, + deadline_ts=deadline_ts, + ctx=ctx, + expected_payload_size=expected_payload_size, + token=token, + ) + completions.extend(self.drain_cuda_gets()) + matching = [completion for completion in completions if completion.token is token] + if len(matching) != 1: + return KVBlockingGetCompletion( + error_msg=( + "GET exception: synchronous CUDA GET did not produce exactly " + "one completion" + ), + source_kind=None, + ) + return KVBlockingGetCompletion( + error_msg=matching[0].error_msg, + source_kind=matching[0].source_kind, + ) + try: _bench_kv_print(f"{ctx} GET begin key={key!r}", verbose_only=True) started_at = time.perf_counter() - result = self._store.get_blocking(key) - done_at = time.perf_counter() - wall_done_ts = time.time() + result = self._get_native_holder(key) err: Optional[str] = None + source_kind: Optional[KVGetSourceKind] = None if not result.is_ok(): err = normalize_kv_get_error(f"GET failed: {result.unwrap_error()}") else: - _ = result.unwrap() + holder = result.unwrap() + source_kind = self._get_source_kind(holder) + if self._get_output == KVGetOutput.BYTES: + if self.backend_kind == BACKEND_KIND_MOONCAKE: + payload = bytes( + _mooncake_payload_view(holder, expected_payload_size) + ) + else: + access_result = holder.access() + if not access_result.is_ok(): + raise RuntimeError( + f"MemHolder.access() failed: {access_result.unwrap_error()}" + ) + value = access_result.unwrap() + payload = value.get("payload") + if not isinstance(payload, bytes): + raise TypeError( + "bytes output requires MemHolder payload to be bytes; " + f"got {type(payload)}" + ) + if len(payload) != int(expected_payload_size): + raise ValueError( + "materialized payload length mismatch: " + f"expected={expected_payload_size} actual={len(payload)}" + ) + if payload: + _ = payload[0] ^ payload[-1] + + done_at = time.perf_counter() + wall_done_ts = time.time() + if err is None: if wall_done_ts > deadline_ts: err = ( f"GET timed out after blocking wait: deadline_ts={deadline_ts:.3f} " @@ -1964,12 +3256,21 @@ def get_blocking( self._phase_profiler.record(op_name="GET", key=key, sample=phase_sample, error_msg=err) if err is not None: _bench_kv_print(f"{ctx} GET failed-after-block key={key!r} err={err}") - return err + return KVBlockingGetCompletion( + error_msg=err, + source_kind=source_kind, + ) _bench_kv_print(f"{ctx} GET done key={key!r}", verbose_only=True) - return None + return KVBlockingGetCompletion( + error_msg=None, + source_kind=source_kind, + ) except Exception as exc: _bench_kv_print(f"{ctx} GET exception key={key!r} err={exc}") - return f"GET exception: {exc}" + return KVBlockingGetCompletion( + error_msg=f"GET exception: {exc}", + source_kind=None, + ) def rpc_register(self, path: str, handler: Any) -> Any: return self._store.rpc_register(path, handler) @@ -1999,6 +3300,11 @@ def rpc_call_bytes( def close(self) -> _SimpleResult: try: + if self._source_counter_window is not None: + self._source_counter_window.cancel() + self._source_counter_window = None + if self._cuda_pipeline is not None: + self._cuda_pipeline.close() return self._store.close() except Exception as exc: return _SimpleResult.err(str(exc)) @@ -2016,8 +3322,21 @@ def flush_phase_summary(self) -> None: self._phase_profiler.flush_pending() -def init_kv_store(kvcache_config: dict[str, Any]) -> tuple[Optional[Any], Optional[str]]: +def init_kv_store( + kvcache_config: dict[str, Any], + *, + kv_get_output: Any = KVGetOutput.HOLDER.value, + kv_cuda_device_index: Any = 0, +) -> tuple[Optional[Any], Optional[str]]: backend_kind = str(kvcache_config.get("backend_kind", BACKEND_KIND_FLUXON)).strip().upper() + try: + get_output = normalize_kv_get_output(kv_get_output) + except ValueError as exc: + return None, str(exc) + try: + cuda_device_index = normalize_kv_cuda_device_index(kv_cuda_device_index) + except ValueError as exc: + return None, str(exc) if backend_kind == BACKEND_KIND_REDIS: try: redis_cfg = kvcache_config.get("redis") @@ -2070,7 +3389,15 @@ def init_kv_store(kvcache_config: dict[str, Any]) -> tuple[Optional[Any], Option store = result.unwrap() if store is None: return None, "Failed to create KVCache store: unwrap() returned None" - return FluxonBlockingStore(store), None + return ( + KVBenchmarkBlockingStore( + store, + backend_kind=backend_kind, + get_output=get_output, + cuda_device_index=cuda_device_index, + ), + None, + ) except Exception as exc: return None, f"Exception while creating KVCache store: {exc}" @@ -2094,14 +3421,39 @@ def kv_put_once( return store.put_blocking(key, bytes(payload), deadline_ts=deadline_ts, ctx=ctx) -def kv_get_once(store: Any, key: str, *, deadline_ts: float, ctx: str) -> Optional[str]: +def kv_get_once( + store: Any, + key: str, + *, + deadline_ts: float, + ctx: str, + expected_payload_size: int, +) -> KVBlockingGetCompletion: if store is None: - return "KV store is not initialized" + return KVBlockingGetCompletion( + error_msg="KV store is not initialized", + source_kind=None, + ) if not hasattr(store, "get_blocking"): backend_kind = getattr(store, "backend_kind", type(store).__name__) - return f"GET failed: backend {backend_kind} does not expose get_blocking" - return normalize_kv_get_error( - store.get_blocking(key, deadline_ts=deadline_ts, ctx=ctx) + return KVBlockingGetCompletion( + error_msg=f"GET failed: backend {backend_kind} does not expose get_blocking", + source_kind=None, + ) + completion = store.get_blocking( + key, + deadline_ts=deadline_ts, + ctx=ctx, + expected_payload_size=expected_payload_size, + ) + if not isinstance(completion, KVBlockingGetCompletion): + raise TypeError( + "benchmark get_blocking() must return KVBlockingGetCompletion; " + f"got {type(completion)}" + ) + return KVBlockingGetCompletion( + error_msg=normalize_kv_get_error(completion.error_msg), + source_kind=completion.source_kind, ) @@ -2356,6 +3708,48 @@ def _kv_bootstrap_before_ready_enabled(test_config: Mapping[str, Any]) -> bool: raise ValueError("kv_bootstrap_before_ready must be bool when present") +def _kv_bootstrap_concurrency(test_config: Mapping[str, Any], *, keyspace_size: int) -> int: + raw = test_config.get(BENCHMARK_KEY_KV_BOOTSTRAP_CONCURRENCY) + if raw is None: + return min(KV_SEED_BOOTSTRAP_MAX_CONCURRENCY, int(keyspace_size)) + concurrency = int(raw) + if concurrency <= 0: + raise ValueError(f"{BENCHMARK_KEY_KV_BOOTSTRAP_CONCURRENCY} must be > 0") + return min(concurrency, int(keyspace_size)) + + +def _kv_bootstrap_put_gap_seconds(test_config: Mapping[str, Any]) -> float: + raw = test_config.get(BENCHMARK_KEY_KV_BOOTSTRAP_PUT_GAP_MS, 0) + if not isinstance(raw, (int, float)): + raise ValueError(f"{BENCHMARK_KEY_KV_BOOTSTRAP_PUT_GAP_MS} must be number") + gap_ms = float(raw) + if gap_ms < 0.0: + raise ValueError(f"{BENCHMARK_KEY_KV_BOOTSTRAP_PUT_GAP_MS} must be >= 0") + return gap_ms / 1000.0 + + +def _kv_bootstrap_storage_full_policy(test_config: Mapping[str, Any]) -> str: + raw = test_config.get( + BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY, + KV_BOOTSTRAP_STORAGE_FULL_POLICY_FAIL, + ) + if not isinstance(raw, str): + raise ValueError(f"{BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY} must be string") + policy = raw.strip().lower() + if policy not in KV_BOOTSTRAP_STORAGE_FULL_POLICIES: + raise ValueError( + f"{BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY} must be one of " + f"{sorted(KV_BOOTSTRAP_STORAGE_FULL_POLICIES)}" + ) + return policy + + +def _is_kv_storage_full_error(error_msg: Optional[str]) -> bool: + if not error_msg: + return False + return "StorageFullError" in error_msg or "No space left" in error_msg + + def _build_operation_result( operation_result_cls: Any, *, @@ -2367,6 +3761,7 @@ def _build_operation_result( inflight_at_start: int, outcome_kind: Any, error_msg: Optional[str], + get_source_kind: Optional[str] = None, ) -> Any: return operation_result_cls( success=success, @@ -2377,6 +3772,7 @@ def _build_operation_result( inflight_at_start=inflight_at_start, outcome_kind=outcome_kind, error_msg=error_msg, + get_source_kind=get_source_kind, ) @@ -2431,15 +3827,34 @@ def prepare_kv_before_ready(benchmark_node: Any, *, logger: Any) -> bool: if fixed_value_size > 0: fixed_value = benchmark_node._generate_test_data(fixed_value_size) + bootstrap_concurrency = _kv_bootstrap_concurrency( + test_config, + keyspace_size=runtime_cfg.keyspace_size, + ) + bootstrap_put_gap_s = _kv_bootstrap_put_gap_seconds(test_config) + storage_full_policy = _kv_bootstrap_storage_full_policy(test_config) + if ( + storage_full_policy == KV_BOOTSTRAP_STORAGE_FULL_POLICY_STOP + and bootstrap_concurrency != 1 + ): + raise ValueError( + f"{BENCHMARK_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY}=stop requires " + f"{BENCHMARK_KEY_KV_BOOTSTRAP_CONCURRENCY}=1" + ) + logger.info( - "🧱 KV bootstrap before READY: mode=%s key_prefix=%s keyspace_size=%s distribution=%s", + "🧱 KV bootstrap before READY: mode=%s key_prefix=%s keyspace_size=%s distribution=%s " + "concurrency=%s put_gap_ms=%.3f storage_full_policy=%s", test_mode, runtime_cfg.key_prefix, runtime_cfg.keyspace_size, runtime_cfg.request_distribution, + bootstrap_concurrency, + bootstrap_put_gap_s * 1000.0, + storage_full_policy, ) - def _bootstrap_one_key(key_idx: int) -> None: + def _bootstrap_one_key(key_idx: int) -> bool: key = f"{runtime_cfg.key_prefix}_k{key_idx}" kv_value_size = _resolve_kv_value_size_for_key(benchmark_node, key_idx) value = ( @@ -2459,12 +3874,29 @@ def _bootstrap_one_key(key_idx: int) -> None: ), ) if not result.success: + if ( + storage_full_policy == KV_BOOTSTRAP_STORAGE_FULL_POLICY_STOP + and _is_kv_storage_full_error(result.error_msg) + ): + logger.warning( + "🧱 KV bootstrap stopped on storage full: key=%s key_idx=%s requested_keys=%s err=%s", + key, + key_idx, + runtime_cfg.keyspace_size, + result.error_msg, + ) + return False raise RuntimeError(f"KV bootstrap PUT failed: key={key} err={result.error_msg}") + if bootstrap_put_gap_s > 0.0: + time.sleep(bootstrap_put_gap_s) + return True - bootstrap_concurrency = min(KV_SEED_BOOTSTRAP_MAX_CONCURRENCY, runtime_cfg.keyspace_size) + completed_keys = 0 if bootstrap_concurrency <= 1: for key_idx in range(runtime_cfg.keyspace_size): - _bootstrap_one_key(key_idx) + if not _bootstrap_one_key(key_idx): + break + completed_keys += 1 else: logger.info( "🧱 KV bootstrap using parallel shards: concurrency=%s", @@ -2478,7 +3910,6 @@ def _bootstrap_shard(shard_idx: int) -> int: completed += 1 return completed - completed_keys = 0 with ThreadPoolExecutor(max_workers=bootstrap_concurrency) as executor: futures = [ executor.submit(_bootstrap_shard, shard_idx) @@ -2491,8 +3922,19 @@ def _bootstrap_shard(shard_idx: int) -> int: completed_keys, runtime_cfg.keyspace_size, ) + if completed_keys <= 0: + raise RuntimeError("KV bootstrap before READY did not insert any keys") + if completed_keys < runtime_cfg.keyspace_size: + test_config[BENCHMARK_KEY_KEYSPACE_SIZE] = int(completed_keys) + logger.warning( + "🧱 KV bootstrap using partial keyspace after storage pressure: requested_keys=%s " + "effective_keyspace_size=%s", + runtime_cfg.keyspace_size, + completed_keys, + ) logger.info( - "✅ KV bootstrap before READY completed: keys=%d", + "✅ KV bootstrap before READY completed: keys=%d requested_keys=%d", + completed_keys, runtime_cfg.keyspace_size, ) return True @@ -2525,15 +3967,102 @@ def run_kv_worker( key_prefix=key_prefix.strip(), ) + cuda_store: Optional[KVBenchmarkBlockingStore] = None + if normalize_kv_get_output( + test_config.get(BENCHMARK_KEY_KV_GET_OUTPUT) + ) == KVGetOutput.CUDA: + candidate_store = getattr(benchmark_node, "kv_store", None) + if not isinstance(candidate_store, KVBenchmarkBlockingStore): + raise TypeError( + "kv_get_output=cuda requires KVBenchmarkBlockingStore; " + f"got {type(candidate_store)}" + ) + cuda_store = candidate_store + results: list[Any] = [] op_idx = 0 + outstanding_cuda_tokens: set[_CudaWorkerGetToken] = set() fixed_value = None if benchmark_node.value_size_mode == VALUE_SIZE_MODE_FIXED: fixed_value_size = int(test_config.get("value_size", 0)) if fixed_value_size > 0: fixed_value = benchmark_node._generate_test_data(fixed_value_size) + def _record_result(result: Any, *, result_op_idx: int, finish_ts: float) -> None: + result.node_id = benchmark_node.node_id + result.worker_id = thread_id + result.finish_ts = float(finish_ts) + + if benchmark_node.start_time is not None: + warmup_deadline_ts = benchmark_node.start_time + metric_warmup_seconds + if result.finish_ts < warmup_deadline_ts: + benchmark_node._mark_progress( + thread_id=thread_id, + op_idx=result_op_idx, + finish_ts=result.finish_ts, + latency_us=result.latency_us, + ) + return + + benchmark_node._mark_progress( + thread_id=thread_id, + op_idx=result_op_idx, + finish_ts=result.finish_ts, + latency_us=result.latency_us, + ) + results.append(result) + + def _consume_cuda_completions( + completions: Sequence[_PipelinedCudaGetCompletion], + ) -> None: + for completion in completions: + token = completion.token + if not isinstance(token, _CudaWorkerGetToken): + raise TypeError( + "CUDA GET completion returned an unknown worker token: " + f"{type(token)}" + ) + if token not in outstanding_cuda_tokens: + raise RuntimeError( + "CUDA GET completion is duplicated or was never submitted: " + f"op_idx={token.op_idx} key={token.key!r}" + ) + outstanding_cuda_tokens.remove(token) + benchmark_node._inflight_end() + + error_msg = normalize_kv_get_error(completion.error_msg) + success = error_msg is None + if success: + outcome_kind = operation_outcome.CACHE_HIT + elif classify_kv_get_result(error_msg) == KVGetResultKind.CACHE_MISS: + outcome_kind = operation_outcome.CACHE_MISS + else: + outcome_kind = operation_outcome.ERROR + result = _build_operation_result( + operation_result_cls, + success=success, + latency_us=completion.latency_us, + operation_type=KV_OPERATION_GET, + key=token.key, + data_size=token.expected_data_size if success else 0, + inflight_at_start=token.inflight_at_start, + outcome_kind=outcome_kind, + error_msg=error_msg, + get_source_kind=( + completion.source_kind.value + if success and completion.source_kind is not None + else None + ), + ) + _record_result( + result, + result_op_idx=token.op_idx, + finish_ts=completion.finish_ts, + ) + while True: + if cuda_store is not None: + _consume_cuda_completions(cuda_store.poll_cuda_gets()) if benchmark_node._benchmark_stop.is_set(): debug_print( f"thread {thread_id} observed benchmark stop intent, " @@ -2545,6 +4074,7 @@ def run_kv_worker( break inflight_at_start = benchmark_node._inflight_begin() + defer_inflight_end = False try: op_timeout_s = float(test_config["op_timeout_seconds"]) op_deadline_ts = min(float(deadline_ts), time.time() + op_timeout_s) @@ -2567,6 +4097,14 @@ def run_kv_worker( thread_id=thread_id, op_idx=op_idx, ) + ctx = ( + f"node={benchmark_node.node_id} role={node_role} thread={thread_id} " + f"op={op_idx} key_idx={key_idx} op_kind={op_kind.lower()}" + ) + + if cuda_store is not None and op_kind != KV_OPERATION_GET: + _consume_cuda_completions(cuda_store.drain_cuda_gets()) + if op_kind == KV_OPERATION_PUT: value = ( fixed_value @@ -2578,22 +4116,41 @@ def run_kv_worker( value, inflight_at_start, deadline_ts=op_deadline_ts, - ctx=( - f"node={benchmark_node.node_id} role={node_role} thread={thread_id} " - f"op={op_idx} key_idx={key_idx} op_kind={op_kind.lower()}" - ), + ctx=ctx, ) elif op_kind == KV_OPERATION_GET: - result = benchmark_node._get_single_operation( - key, - inflight_at_start, - deadline_ts=op_deadline_ts, - expected_data_size=kv_value_size, - ctx=( - f"node={benchmark_node.node_id} role={node_role} thread={thread_id} " - f"op={op_idx} key_idx={key_idx} op_kind={op_kind.lower()}" - ), - ) + if cuda_store is not None: + cuda_token = _CudaWorkerGetToken( + op_idx=op_idx, + key=key, + expected_data_size=kv_value_size, + inflight_at_start=inflight_at_start, + ) + outstanding_cuda_tokens.add(cuda_token) + try: + completions = cuda_store.submit_cuda_get( + key, + deadline_ts=op_deadline_ts, + ctx=ctx, + expected_payload_size=kv_value_size, + token=cuda_token, + ) + except Exception: + outstanding_cuda_tokens.discard(cuda_token) + raise + # Completion owns the in-flight decrement for pipelined GETs. + defer_inflight_end = True + _consume_cuda_completions(completions) + op_idx += 1 + continue + else: + result = benchmark_node._get_single_operation( + key, + inflight_at_start, + deadline_ts=op_deadline_ts, + expected_data_size=kv_value_size, + ctx=ctx, + ) else: result = _build_operation_result( operation_result_cls, @@ -2619,34 +4176,40 @@ def run_kv_worker( error_msg=str(exc), ) finally: - benchmark_node._inflight_end() + if not defer_inflight_end: + benchmark_node._inflight_end() - result.node_id = benchmark_node.node_id - result.worker_id = thread_id - result.finish_ts = time.time() - op_finish_ts = result.finish_ts - - if benchmark_node.start_time is not None: - warmup_deadline_ts = benchmark_node.start_time + metric_warmup_seconds - if op_finish_ts < warmup_deadline_ts: - benchmark_node._mark_progress( - thread_id=thread_id, - op_idx=op_idx, - finish_ts=op_finish_ts, - latency_us=result.latency_us, - ) - op_idx += 1 - continue - - benchmark_node._mark_progress( - thread_id=thread_id, - op_idx=op_idx, - finish_ts=op_finish_ts, - latency_us=result.latency_us, + _record_result( + result, + result_op_idx=op_idx, + finish_ts=time.time(), ) - results.append(result) op_idx += 1 + if cuda_store is not None: + _consume_cuda_completions(cuda_store.drain_cuda_gets()) + if outstanding_cuda_tokens: + lost_tokens = sorted(outstanding_cuda_tokens, key=lambda token: token.op_idx) + outstanding_cuda_tokens.clear() + for token in lost_tokens: + benchmark_node._inflight_end() + result = _build_operation_result( + operation_result_cls, + success=False, + latency_us=0.0, + operation_type=KV_OPERATION_GET, + key=token.key, + data_size=0, + inflight_at_start=token.inflight_at_start, + outcome_kind=operation_outcome.ERROR, + error_msg="CUDA GET pipeline lost its completion", + ) + _record_result( + result, + result_op_idx=token.op_idx, + finish_ts=time.time(), + ) + debug_print( f"thread {thread_id} exit kv run loop, total_ops={len(results)}, last_op_idx={op_idx}" ) diff --git a/fluxon_test_stack/ci_2_virt_node.py b/fluxon_test_stack/ci_2_virt_node.py index f055426..bf30fe0 100644 --- a/fluxon_test_stack/ci_2_virt_node.py +++ b/fluxon_test_stack/ci_2_virt_node.py @@ -7,6 +7,7 @@ import json import os import re +import shutil import socket import subprocess import sys @@ -34,7 +35,7 @@ DEFAULT_TESTBED_BOOTSTRAP_MODE = "bare_then_apply" DEFAULT_TESTBED_UI_PORT = 18080 DEFAULT_TESTBED_CONTROLLER_PORT = 19080 -DEFAULT_TESTBED_HOSTWORKDIR = Path("/mnt/nvme0/store_team_dev/fluxon_deploy") +DEFAULT_TESTBED_HOSTWORKDIR = Path("/tmp/fluxon-example/deploy") LOCAL_PRIMARY_NODE_SUFFIX = "a" LOCAL_SECONDARY_NODE_SUFFIX = "b" TEST_STACK_START_TEST_BED_CONFIG_ENV = "FLUXON_TEST_STACK_START_TEST_BED_CONFIG" @@ -152,6 +153,19 @@ def _parse_args() -> argparse.Namespace: action="store_true", help="Print generated config paths before executing commands.", ) + parser.add_argument( + "--cleanup-pack-runtime-after-success", + action="store_true", + help=( + "Delete the temporary pack_release_runtime tree after packaging succeeds. " + "The published release artifacts remain under --release-dir." + ), + ) + parser.add_argument( + "--cleanup-successful-case-artifacts", + action="store_true", + help="Forward successful-case artifact cleanup to test_runner.", + ) return parser.parse_args() @@ -651,6 +665,44 @@ def _prepare_pack_release_runtime_dirs(*, project_data_root: Path) -> None: (root / relpath).mkdir(parents=True, exist_ok=True) +def _cleanup_pack_release_runtime_after_success(*, runtime_root: Path, workdir: Path) -> None: + workdir = workdir.resolve() + runtime_root = Path(os.path.abspath(os.fspath(runtime_root))) + expected_root = workdir / "pack_release_runtime" + if runtime_root != expected_root: + raise ValueError( + "refusing to clean an unexpected pack runtime path: " + f"runtime_root={runtime_root} expected={expected_root}" + ) + if not runtime_root.exists(): + print(f"[cleanup_pack_runtime] already absent: {runtime_root}", flush=True) + return + if runtime_root.is_symlink() or not runtime_root.is_dir(): + raise ValueError(f"pack runtime cleanup requires a regular directory: {runtime_root}") + + free_before = shutil.disk_usage(workdir).free + try: + shutil.rmtree(runtime_root) + except PermissionError: + # Containerized packaging can leave root-owned cache entries. The cleanup + # flag is explicit, and the path equality check above bounds sudo to the + # generated pack runtime owned by this CI invocation. + sudo = shutil.which("sudo") + if sudo is None: + raise + subprocess.check_call([sudo, "-n", "rm", "-rf", "--", str(runtime_root)]) + if runtime_root.exists(): + raise RuntimeError(f"pack runtime still exists after cleanup: {runtime_root}") + + free_after = shutil.disk_usage(workdir).free + reclaimed_bytes = max(0, int(free_after) - int(free_before)) + print( + "[cleanup_pack_runtime] " + f"removed={runtime_root} reclaimed_bytes={reclaimed_bytes} free_bytes={free_after}", + flush=True, + ) + + def _run(argv: list[str], *, env: dict[str, str] | None = None) -> None: print("RUN: " + " ".join(_shell_quote(part) for part in argv), flush=True) subprocess.check_call(argv, cwd=str(REPO_ROOT), env=env) @@ -892,6 +944,11 @@ def main() -> int: pack_env = os.environ.copy() pack_env["FLUXON_PACK_RELEASE_NIX_CONFIG"] = str(ci_nix_pack_config_path) _run(pack_cmd, env=pack_env) + if args.cleanup_pack_runtime_after_success: + _cleanup_pack_release_runtime_after_success( + runtime_root=pack_release_runtime_root, + workdir=workdir, + ) wheel_name = _find_single_wheel(release_dir, pattern="fluxon-*.whl", ctx="top-level release wheel") @@ -983,6 +1040,8 @@ def main() -> int: "-w", str(runner_workdir), ] + if args.cleanup_successful_case_artifacts: + runner_cmd.append("--cleanup-successful-case-artifacts") _run(runner_cmd, env=_runner_env(release_dir=release_dir, start_cfg_path=Path(metadata["start_test_bed_path"]))) return 0 diff --git a/fluxon_test_stack/ci_test_list.yaml b/fluxon_test_stack/ci_test_list.yaml index ebafe4a..6f5dea6 100644 --- a/fluxon_test_stack/ci_test_list.yaml +++ b/fluxon_test_stack/ci_test_list.yaml @@ -295,7 +295,7 @@ scales: duration_seconds: 60 topology: 1 targets: - hosts: [infra44-ThinkStation-PX] + hosts: [example-node-a] owner: owner_count: 1 owner_dram_bytes: 3221225472 @@ -313,7 +313,7 @@ scales: duration_seconds: 60 topology: 1 targets: - hosts: [infra44-ThinkStation-PX] + hosts: [example-node-a] owner: owner_count: 1 owner_dram_bytes: 21474836480 @@ -323,7 +323,7 @@ scales: duration_seconds: 60 topology: 2 targets: - hosts: [infra44-ThinkStation-PX, infra46-ThinkStation-PX] + hosts: [example-node-a, example-node-b] owner: owner_count: 2 owner_dram_bytes: 3221225472 @@ -333,7 +333,7 @@ scales: duration_seconds: 60 topology: 2 targets: - hosts: [infra44-ThinkStation-PX, infra46-ThinkStation-PX] + hosts: [example-node-a, example-node-b] owner: owner_count: 2 owner_dram_bytes: 21474836480 @@ -431,8 +431,8 @@ profiles: - test_profile_adapter.py controller_url: __STACK_CONTROLLER_URL__ target_ip_map: - infra44-ThinkStation-PX: 192.168.151.44 - infra46-ThinkStation-PX: 192.168.151.46 + example-node-a: 192.0.2.10 + example-node-b: 192.0.2.11 bootstrap_ready_timeout_seconds: 300 scene_configs: ci_top_attention_doc_page_build: @@ -451,12 +451,12 @@ profiles: cluster_kv_owner: &cluster_kv_owner_runtime base_runtime: etcd: - target: infra44-ThinkStation-PX + target: example-node-a endpoint: scheme: HTTP host_port: 32579 greptime: - target: infra44-ThinkStation-PX + target: example-node-a endpoint: scheme: HTTP host_port: 34030 @@ -465,7 +465,7 @@ profiles: lifecycle: service k8s_ref: deployment/master deployer: - target: infra44-ThinkStation-PX + target: example-node-a command: [/bin/bash, -lc] args: - | @@ -500,12 +500,12 @@ profiles: rust_self_managed: base_runtime: etcd: - target: infra44-ThinkStation-PX + target: example-node-a endpoint: scheme: HTTP host_port: 32579 greptime: - target: infra44-ThinkStation-PX + target: example-node-a endpoint: scheme: HTTP host_port: 34030 @@ -540,8 +540,8 @@ profiles: - test_profile_adapter.py controller_url: __STACK_CONTROLLER_URL__ target_ip_map: - infra44-ThinkStation-PX: 192.168.151.44 - infra46-ThinkStation-PX: 192.168.151.46 + example-node-a: 192.0.2.10 + example-node-b: 192.0.2.11 bootstrap_ready_timeout_seconds: 300 runtime_config: # This config is passed through into run_dir/benchmark_config.py. @@ -552,7 +552,7 @@ profiles: fluxonkv_spec: cluster_name: __STACK_CLUSTER_NAME__ # Self-host benchmark nodes are launched by the host-side deployer, not inside a - # container with /hostworkdir mounted. Use the real host path so both infra44/46 + # container with /hostworkdir mounted. Use the real host path so both example-node-a/example-node-b # benchmark nodes can open the same self-host shared-memory root directly. share_mem_path: __STACK_SHARE_MEM_PATH__ mq_base: @@ -564,7 +564,7 @@ profiles: lifecycle: service k8s_ref: deployment/test_stack_coord deployer: - target: infra44-ThinkStation-PX + target: example-node-a command: [/bin/bash, -lc] args: - | @@ -665,8 +665,8 @@ profiles: - test_profile_adapter.py controller_url: __STACK_CONTROLLER_URL__ target_ip_map: - infra44-ThinkStation-PX: 192.168.151.44 - infra46-ThinkStation-PX: 192.168.151.46 + example-node-a: 192.0.2.10 + example-node-b: 192.0.2.11 bootstrap_ready_timeout_seconds: 300 runtime_config: redis: @@ -681,7 +681,7 @@ profiles: lifecycle: service k8s_ref: deployment/test_stack_coord deployer: - target: infra44-ThinkStation-PX + target: example-node-a command: [/bin/bash, -lc] args: - | @@ -718,16 +718,16 @@ profiles: - test_profile_adapter.py controller_url: __STACK_CONTROLLER_URL__ target_ip_map: - infra44-ThinkStation-PX: 192.168.151.44 - infra46-ThinkStation-PX: 192.168.151.46 + example-node-a: 192.0.2.10 + example-node-b: 192.0.2.11 bootstrap_ready_timeout_seconds: 300 runtime_config: alluxio: bundle_test_rsc_relpath: baselines/alluxio/alluxio_bundle.tar.gz bundle_root_relpath: baselines/alluxio/alluxio_bundle mount_root_by_target: - infra44-ThinkStation-PX: /mnt/alluxio - infra46-ThinkStation-PX: /mnt/alluxio + example-node-a: /mnt/alluxio + example-node-b: /mnt/alluxio namespace_prefix: fluxon_test_stack/alluxio mount_command_template: | cd __ALLUXIO_BUNDLE_ROOT__ @@ -737,7 +737,7 @@ profiles: lifecycle: service k8s_ref: deployment/test_stack_coord deployer: - target: infra44-ThinkStation-PX + target: example-node-a command: [/bin/bash, -lc] args: - | diff --git a/fluxon_test_stack/deployconf_testbed.yml b/fluxon_test_stack/deployconf_testbed.yml index 4015beb..7cfbda1 100644 --- a/fluxon_test_stack/deployconf_testbed.yml +++ b/fluxon_test_stack/deployconf_testbed.yml @@ -6,7 +6,7 @@ image: fluxon_quick_start:0.2.1 # English note: # - fluxon-deployer reads manifests from hostworkdir (outside this repo). # - Mirroring generated DaemonSet YAMLs prevents "repo output updated but deployer still applies stale YAMLs". -gen_k8s_daemonset_mirror_outdir: /mnt/nvme0/store_team_dev/fluxon_deploy/gen_k8s_daemonset +gen_k8s_daemonset_mirror_outdir: /tmp/fluxon-example/deploy/gen_k8s_daemonset cluster_nodes: - hostname: example-node-a @@ -302,7 +302,7 @@ service: p2p_listen_port: 31100 network: subnet_whitelist: - - "192.168.151.0/24" + - "192.0.2.0/24" monitoring: prometheus_base_url: "${FLUXON_PROMETHEUS_BASE_URL}" prom_remote_write_url: diff --git a/fluxon_test_stack/distributed_benchmark_coordinator.py b/fluxon_test_stack/distributed_benchmark_coordinator.py index 0859560..f84ed7b 100644 --- a/fluxon_test_stack/distributed_benchmark_coordinator.py +++ b/fluxon_test_stack/distributed_benchmark_coordinator.py @@ -39,7 +39,11 @@ canonicalize_kv_node_role, ) from .benchmark_node_kv import ( + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX, + BENCHMARK_KEY_KV_GET_OUTPUT, merge_kv_benchmark_extras, + normalize_kv_cuda_device_index, + normalize_kv_get_output, ) from .benchmark_node_rpc import ( FLUXON_PHASE_PATH_BUCKET_FAST, @@ -57,7 +61,11 @@ canonicalize_kv_node_role, ) from benchmark_node_kv import ( + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX, + BENCHMARK_KEY_KV_GET_OUTPUT, merge_kv_benchmark_extras, + normalize_kv_cuda_device_index, + normalize_kv_get_output, ) from benchmark_node_rpc import ( FLUXON_PHASE_PATH_BUCKET_FAST, @@ -277,6 +285,20 @@ def get_op_timeout_seconds(config_path: str) -> float: return op_timeout_seconds +def get_kv_get_output(config_path: str) -> str: + benchmark_cfg = _load_benchmark_section(config_path) + return normalize_kv_get_output( + benchmark_cfg.get(BENCHMARK_KEY_KV_GET_OUTPUT) + ).value + + +def get_kv_cuda_device_index(config_path: str) -> int: + benchmark_cfg = _load_benchmark_section(config_path) + return normalize_kv_cuda_device_index( + benchmark_cfg.get(BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX) + ) + + def _parse_value_size_weighted_set( raw_val: Any, *, @@ -361,6 +383,8 @@ def _required_output_result_path(config_path: str) -> str: KVCACHE_CONFIG_PATH ) OP_TIMEOUT_SECONDS = get_op_timeout_seconds(KVCACHE_CONFIG_PATH) +KV_GET_OUTPUT = get_kv_get_output(KVCACHE_CONFIG_PATH) +KV_CUDA_DEVICE_INDEX = get_kv_cuda_device_index(KVCACHE_CONFIG_PATH) print( f"ℹ️ 从配置加载 THREADS_PER_PROCESS={THREADS_PER_PROCESS}, MAX_BENCHMARK_SECONDS={MAX_BENCHMARK_SECONDS}, " f"METRIC_WARMUP_SECONDS={METRIC_WARMUP_SECONDS}, START_IDLE_SECONDS={START_IDLE_SECONDS}" @@ -369,6 +393,8 @@ def _required_output_result_path(config_path: str) -> str: f"ℹ️ 从配置加载 value_size_list={VALUE_SIZE_SWEEP_LIST}" ) print(f"ℹ️ 从配置加载 OP_TIMEOUT_SECONDS: {OP_TIMEOUT_SECONDS}") +print(f"ℹ️ 从配置加载 KV_GET_OUTPUT: {KV_GET_OUTPUT}") +print(f"ℹ️ 从配置加载 KV_CUDA_DEVICE_INDEX: {KV_CUDA_DEVICE_INDEX}") def get_consumer_sim_handle_ms_range(config_path: str) -> Optional[Tuple[int, int]]: @@ -468,6 +494,8 @@ def _write_benchmark_result_file(all_summaries: List[Dict[str, Any]]) -> None: "value_size_weighted_set": VALUE_SIZE_WEIGHTED_SET, "max_benchmark_seconds": MAX_BENCHMARK_SECONDS, "metric_warmup_seconds": METRIC_WARMUP_SECONDS, + BENCHMARK_KEY_KV_GET_OUTPUT: KV_GET_OUTPUT, + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX: KV_CUDA_DEVICE_INDEX, "runs": all_summaries, } out_p = Path(out_path) @@ -671,6 +699,7 @@ class NodeMetrics: get_hit_operations: int get_miss_operations: int get_error_operations: int + get_source_summary: Dict[str, Any] avg_latency_us: float p50_latency_us: float p99_latency_us: float @@ -1420,6 +1449,85 @@ def _aggregate_p2p_rpc_completion_summary( ) +def _aggregate_get_source_summary(results: List[NodeMetrics]) -> Dict[str, Any]: + """Aggregate per-node GET source observations without inventing attribution.""" + memory_hits = 0 + ssd_hits = 0 + unknown_hits = 0 + observed_hits = 0 + complete_node_count = 0 + observation_kinds: set[str] = set() + + for result in results: + hit_count = max(0, int(result.get_hit_operations)) + raw_summary = result.get_source_summary + if not isinstance(raw_summary, dict): + unknown_hits += hit_count + observed_hits += hit_count + observation_kinds.add("unavailable") + continue + + raw_observation_kind = raw_summary.get("observation_kind") + observation_kind = ( + str(raw_observation_kind).strip() + if isinstance(raw_observation_kind, str) + else "" + ) + observation_kinds.add(observation_kind or "unavailable") + + count_keys = ( + "get_memory_hit_operations", + "get_ssd_hit_operations", + "get_unknown_source_operations", + "get_observed_hit_operations", + ) + raw_counts = [raw_summary.get(key) for key in count_keys] + counts_are_valid = all( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 + for value in raw_counts + ) + if not counts_are_valid: + unknown_hits += hit_count + observed_hits += hit_count + continue + + node_memory_hits, node_ssd_hits, node_unknown_hits, node_observed_hits = ( + int(value) for value in raw_counts + ) + counts_are_consistent = ( + node_observed_hits == hit_count + and node_memory_hits + node_ssd_hits + node_unknown_hits + == node_observed_hits + ) + if not counts_are_consistent: + unknown_hits += hit_count + observed_hits += hit_count + continue + + memory_hits += node_memory_hits + ssd_hits += node_ssd_hits + unknown_hits += node_unknown_hits + observed_hits += node_observed_hits + if bool(raw_summary.get("complete")) and node_unknown_hits == 0: + complete_node_count += 1 + + node_count = len(results) + return { + "observation_kinds": sorted(observation_kinds), + "get_memory_hit_operations": memory_hits, + "get_ssd_hit_operations": ssd_hits, + "get_unknown_source_operations": unknown_hits, + "get_observed_hit_operations": observed_hits, + "complete": ( + node_count > 0 + and complete_node_count == node_count + and unknown_hits == 0 + ), + "node_count": node_count, + "complete_node_count": complete_node_count, + } + + def _build_aggregated_bench_points(results: List[NodeMetrics]) -> Dict[str, Any]: total_ops = sum(r.total_operations for r in results) total_successful_ops = sum(r.successful_operations for r in results) @@ -1460,6 +1568,7 @@ def _build_aggregated_bench_points(results: List[NodeMetrics]) -> Dict[str, Any] [(float(r.inflight_avg), int(r.total_operations)) for r in results] ) inflight_max = max((int(r.inflight_max) for r in results), default=0) + get_source_summary = _aggregate_get_source_summary(results) return { "total_ops": total_ops, "total_successful_ops": total_successful_ops, @@ -1468,6 +1577,7 @@ def _build_aggregated_bench_points(results: List[NodeMetrics]) -> Dict[str, Any] "get_hit_operations": get_hit_operations, "get_miss_operations": get_miss_operations, "get_error_operations": get_error_operations, + "get_source_summary": get_source_summary, "total_duration_seconds": total_duration, "throughput_ops_per_sec": throughput_ops_per_sec, "total_throughput_ops_per_sec": total_throughput_ops_per_sec, @@ -2340,6 +2450,9 @@ def handle_report_results( get_hit_operations=results_data.get("get_hit_operations", 0), get_miss_operations=results_data.get("get_miss_operations", 0), get_error_operations=results_data.get("get_error_operations", 0), + get_source_summary=copy.deepcopy( + results_data.get("get_source_summary", {}) + ), avg_latency_us=results_data.get("avg_latency_us", 0.0), p50_latency_us=results_data.get("p50_latency_us", 0.0), p99_latency_us=results_data.get("p99_latency_us", 0.0), @@ -2745,6 +2858,7 @@ def wait_for_completion(self, *, timeout_s: float) -> bool: get_hit_operations=0, get_miss_operations=0, get_error_operations=0, + get_source_summary={}, avg_latency_us=0.0, p50_latency_us=0.0, p99_latency_us=0.0, @@ -2817,6 +2931,16 @@ def build_incomplete_run_summary( "get_hit_operations": 0, "get_miss_operations": 0, "get_error_operations": 0, + "get_source_summary": { + "observation_kinds": [], + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": 0, + "get_observed_hit_operations": 0, + "complete": False, + "node_count": 0, + "complete_node_count": 0, + }, "total_duration_seconds": 0.0, "total_bytes": 0, "overall_success_rate": 0.0, @@ -2854,6 +2978,7 @@ def build_incomplete_run_summary( "get_hit_operations": int(overall["get_hit_operations"]), "get_miss_operations": int(overall["get_miss_operations"]), "get_error_operations": int(overall["get_error_operations"]), + "get_source_summary": copy.deepcopy(overall["get_source_summary"]), "total_duration_seconds": float(overall["total_duration_seconds"]), "throughput_ops_per_sec": float(overall["overall_tps"]), "total_throughput_ops_per_sec": float(overall["overall_total_tps"]), @@ -3391,6 +3516,7 @@ def get_overall_summary(self) -> Optional[Dict[str, Any]]: "get_hit_operations": int(bench_points["get_hit_operations"]), "get_miss_operations": int(bench_points["get_miss_operations"]), "get_error_operations": int(bench_points["get_error_operations"]), + "get_source_summary": copy.deepcopy(bench_points["get_source_summary"]), "total_duration_seconds": total_duration, "total_bytes": total_bytes, "overall_success_rate": overall_success_rate, @@ -3439,6 +3565,7 @@ def get_node_summaries(self) -> Optional[List[Dict[str, Any]]]: "get_hit_operations": node.get_hit_operations, "get_miss_operations": node.get_miss_operations, "get_error_operations": node.get_error_operations, + "get_source_summary": copy.deepcopy(node.get_source_summary), "avg_latency_us": node.avg_latency_us, "p50_latency_us": node.p50_latency_us, "p95_latency_us": node.p95_latency_us, diff --git a/fluxon_test_stack/distributed_benchmark_node.py b/fluxon_test_stack/distributed_benchmark_node.py index 0065f5d..4aa9d54 100644 --- a/fluxon_test_stack/distributed_benchmark_node.py +++ b/fluxon_test_stack/distributed_benchmark_node.py @@ -55,8 +55,13 @@ from .benchmark_node_kv import ( KV_OPERATION_GET, KV_OPERATION_PUT, + BACKEND_KIND_FLUXON, + BACKEND_KIND_MOONCAKE, + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX, + BENCHMARK_KEY_KV_GET_OUTPUT, KV_NODE_ROLE_SEED, KV_NODE_ROLE_WORKER, + KVBenchmarkBlockingStore, KVGetResultKind, canonicalize_kv_node_role, classify_kv_get_result, @@ -95,8 +100,13 @@ from benchmark_node_kv import ( KV_OPERATION_GET, KV_OPERATION_PUT, + BACKEND_KIND_FLUXON, + BACKEND_KIND_MOONCAKE, + BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX, + BENCHMARK_KEY_KV_GET_OUTPUT, KV_NODE_ROLE_SEED, KV_NODE_ROLE_WORKER, + KVBenchmarkBlockingStore, KVGetResultKind, canonicalize_kv_node_role, classify_kv_get_result, @@ -1002,6 +1012,8 @@ class OperationResult: worker_id: Optional[int] = None # Operation completion time (wall clock), used for precise warmup filtering finish_ts: float = 0.0 + # Selected source for a successful KV GET: "memory", "ssd", or None. + get_source_kind: Optional[str] = None @dataclass @@ -1233,6 +1245,7 @@ def __init__(self): ) self._network_bandwidth_sampler: Optional[NetworkBandwidthSampler] = None self._network_bandwidth_summary: Dict[str, Any] = {} + self._kv_get_source_counter_window: Dict[str, Any] = {} self._benchmark_stop = threading.Event() self._forced_benchmark_result: Optional[Dict[str, Any]] = None self._prepared_mpmc_round: Optional[PreparedMPMCRound] = None @@ -2745,7 +2758,11 @@ def initialize_from_test_config(self) -> bool: logger.info("🔧 正在创建KVCache存储实例...") # KV store initialization is needed only once. A previous merge caused duplicate calls, # leading to repeated cluster member registration. - store, err = init_kv_store(kvcache_config) + store, err = init_kv_store( + kvcache_config, + kv_get_output=self.test_config.get(BENCHMARK_KEY_KV_GET_OUTPUT), + kv_cuda_device_index=self.test_config.get(BENCHMARK_KEY_KV_CUDA_DEVICE_INDEX), + ) if err is not None: logger.error(f"❌ KVCache存储实例创建失败: {err}") return False @@ -2872,6 +2889,80 @@ def _generate_test_data(self, size: int) -> bytes: return payload_pool[random.randrange(len(payload_pool))] return os.urandom(size_int) + def _build_kv_get_source_summary( + self, + get_hit_ops: List[OperationResult], + ) -> Dict[str, Any]: + hit_count = len(get_hit_ops) + store = self.kv_store + if not isinstance(store, KVBenchmarkBlockingStore): + return { + "observation_kind": "unavailable", + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": hit_count, + "get_observed_hit_operations": hit_count, + "complete": False, + "details": {}, + } + + if store.backend_kind == BACKEND_KIND_FLUXON: + memory_hits = sum( + 1 for op in get_hit_ops if op.get_source_kind == "memory" + ) + ssd_hits = sum(1 for op in get_hit_ops if op.get_source_kind == "ssd") + unknown_hits = max(0, hit_count - memory_hits - ssd_hits) + return { + "observation_kind": "fluxon_per_operation_route_source", + "get_memory_hit_operations": memory_hits, + "get_ssd_hit_operations": ssd_hits, + "get_unknown_source_operations": unknown_hits, + "get_observed_hit_operations": hit_count, + "complete": unknown_hits == 0, + "details": {}, + } + + if store.backend_kind == BACKEND_KIND_MOONCAKE: + details = copy.deepcopy(self._kv_get_source_counter_window) + raw_ssd_hits = details.get("ssd_read_count") + if ( + bool(details.get("complete")) + and isinstance(raw_ssd_hits, int) + and not isinstance(raw_ssd_hits, bool) + and 0 <= raw_ssd_hits <= hit_count + ): + ssd_hits = int(raw_ssd_hits) + return { + "observation_kind": "mooncake_offload_rpc_counter_window", + "get_memory_hit_operations": hit_count - ssd_hits, + "get_ssd_hit_operations": ssd_hits, + "get_unknown_source_operations": 0, + "get_observed_hit_operations": hit_count, + "complete": True, + "details": details, + } + if isinstance(raw_ssd_hits, int) and raw_ssd_hits > hit_count: + details["counter_excess_operations"] = raw_ssd_hits - hit_count + return { + "observation_kind": "mooncake_offload_rpc_counter_window", + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": hit_count, + "get_observed_hit_operations": hit_count, + "complete": False, + "details": details, + } + + return { + "observation_kind": "unavailable", + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": hit_count, + "get_observed_hit_operations": hit_count, + "complete": False, + "details": {}, + } + def _calculate_benchmark_results(self) -> Dict[str, Any]: """Compute benchmark results.""" def _empty_results() -> Dict[str, Any]: @@ -2887,6 +2978,15 @@ def _empty_results() -> Dict[str, Any]: "get_hit_operations": 0, "get_miss_operations": 0, "get_error_operations": 0, + "get_source_summary": { + "observation_kind": "unavailable", + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": 0, + "get_observed_hit_operations": 0, + "complete": False, + "details": {}, + }, "total_duration_seconds": 0, "avg_latency_us": 0, "p50_latency_us": 0, @@ -2951,6 +3051,7 @@ def _empty_results() -> Dict[str, Any]: get_error_ops = [ r for r in get_ops if r.outcome_kind == OperationOutcome.ERROR ] + get_source_summary = self._build_kv_get_source_summary(get_hit_ops) # Effective duration: exclude warmup; cut off at end_time effective_start = self.start_time + METRIC_WARMUP_SECONDS @@ -3088,6 +3189,7 @@ def _empty_results() -> Dict[str, Any]: "get_hit_operations": len(get_hit_ops), "get_miss_operations": len(get_miss_ops), "get_error_operations": len(get_error_ops), + "get_source_summary": get_source_summary, "total_duration_seconds": duration, "avg_latency_us": avg_latency, "p50_latency_us": p50_latency, @@ -3527,7 +3629,14 @@ def _get_single_operation( """Execute single GET operation and measure performance.""" op_start = time.perf_counter() try: - err = kv_get_once(self.kv_store, key, deadline_ts=deadline_ts, ctx=ctx) + completion = kv_get_once( + self.kv_store, + key, + deadline_ts=deadline_ts, + ctx=ctx, + expected_payload_size=expected_data_size, + ) + err = completion.error_msg op_end = time.perf_counter() if err is not None: get_outcome = classify_kv_get_result(err) @@ -3554,6 +3663,11 @@ def _get_single_operation( inflight_at_start=inflight_at_start, outcome_kind=OperationOutcome.CACHE_HIT, error_msg=None, + get_source_kind=( + completion.source_kind.value + if completion.source_kind is not None + else None + ), ) except Exception as e: op_end = time.perf_counter() @@ -3888,6 +4002,15 @@ def _set_forced_benchmark_result( "get_hit_operations": 0, "get_miss_operations": 0, "get_error_operations": 0, + "get_source_summary": { + "observation_kind": "unavailable", + "get_memory_hit_operations": 0, + "get_ssd_hit_operations": 0, + "get_unknown_source_operations": 0, + "get_observed_hit_operations": 0, + "complete": False, + "details": {}, + }, "total_duration_seconds": 0, "avg_latency_us": 0, "p50_latency_us": 0, @@ -4152,6 +4275,7 @@ def run_benchmark(self) -> Dict[str, Any]: self._kv_store_closed = False self._network_bandwidth_sampler = None self._network_bandwidth_summary = {} + self._kv_get_source_counter_window = {} try: if test_mode == TestMode.MPMC.value: @@ -4166,6 +4290,11 @@ def run_benchmark(self) -> Dict[str, Any]: # - KV mode: end_time is deadline_ts # - MPMC mode: overwritten to now when stop intent is triggered self.end_time = deadline_ts + if isinstance(self.kv_store, KVBenchmarkBlockingStore): + self.kv_store.begin_get_source_counter_window( + start_ts=self.start_time + METRIC_WARMUP_SECONDS, + end_ts=deadline_ts, + ) self._start_network_bandwidth_sampler() # Start heartbeat after time window is initialized, so it can report elapsed/remaining. self._start_heartbeat() @@ -4173,6 +4302,10 @@ def run_benchmark(self) -> Dict[str, Any]: finally: self._stop_heartbeat() self._stop_network_bandwidth_sampler() + if isinstance(self.kv_store, KVBenchmarkBlockingStore): + self._kv_get_source_counter_window = ( + self.kv_store.finish_get_source_counter_window() + ) # Keep previously set end_time (KV mode: deadline; MPMC mode: close time). if self.end_time is None: diff --git a/fluxon_test_stack/gitops/gitops.yaml b/fluxon_test_stack/gitops/gitops.yaml index 8488736..69f1454 100644 --- a/fluxon_test_stack/gitops/gitops.yaml +++ b/fluxon_test_stack/gitops/gitops.yaml @@ -6,7 +6,7 @@ retention: repos: - addr: git@github.com:Tele-AI/fluxon.git follow: - - branch: big_step2 + - branch: example-feature-branch run: name_prefix: fluxon_ci commands: @@ -39,20 +39,20 @@ repos: "fluxon_pyo3": fluxon_pyo3_wheels[0], }, "cluster": { - "etcd_endpoints": ["10.126.126.231:2379"], + "etcd_endpoints": ["192.0.2.10:2379"], "cluster_name": "fluxon_deployer_cluster", "monitoring": { - "prometheus_base_url": "http://10.126.126.231:4010/v1/prometheus", + "prometheus_base_url": "http://192.0.2.10:4010/v1/prometheus", }, }, "nodes": [ { - "name": "infra44-ThinkStation-PX", - "host": "10.126.126.231", + "name": "example-node-a", + "host": "192.0.2.10", "ssh": { "local": True, "port": 22, - "user": "zyc", + "user": "example_user", # Intentionally no password by default. # If your SSH requires a password, set ssh.password explicitly. # If you use an SSH key, set ssh.identity_file explicitly. diff --git a/fluxon_test_stack/single_node_ssd_bench/kv_ssd_pressure_chart.py b/fluxon_test_stack/single_node_ssd_bench/kv_ssd_pressure_chart.py new file mode 100644 index 0000000..87870b9 --- /dev/null +++ b/fluxon_test_stack/single_node_ssd_bench/kv_ssd_pressure_chart.py @@ -0,0 +1,924 @@ +#!/usr/bin/env python3 +"""Shared CPU/GPU renderer for the test-stack KV SSD-pressure matrix.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Mapping, Sequence + +from PIL import Image, ImageDraw, ImageFont + + +CONCURRENCIES = (2, 4, 8, 16) +BASE_IMPLEMENTATIONS = ( + "fluxon", + "mooncake_dedicated", + "mooncake_per_process", +) +IMPLEMENTATION_SHORT_LABELS = { + "fluxon": "F", + "mooncake_dedicated": "D", + "mooncake_per_process": "P", +} +OUTPUT_LABELS = { + "holder": "原生引用 / handle", + "bytes": "bytes", + "cuda": "cuda", +} +IMPLEMENTATION_COLORS = { + "fluxon": ("#A8C7FA", "#2563EB"), + "mooncake_dedicated": ("#FBC99D", "#E76F2E"), + "mooncake_per_process": ("#A7E3D5", "#168C7A"), +} +HIT_RATE_COLOR = "#D64A63" +SSD_BACKEND_IMPLEMENTATIONS = ("native", "foyer") +SSD_BACKEND_SHORT_LABELS = { + "native": "原生", + "foyer": "Foyer", +} +SSD_BACKEND_COLORS = { + "native": ("#BFDBFE", "#2563EB"), + "foyer": ("#DDD6FE", "#7C3AED"), +} + + +@dataclass(frozen=True) +class RowKey: + suite_group: str + output_mode: str + payload_mib: int + concurrency: int + implementation: str + ssd_enabled: bool + + +@dataclass(frozen=True) +class SsdBackendAblationKey: + output_mode: str + payload_mib: int + concurrency: int + ssd_backend_impl: str + + +def row_implementation(row: Mapping[str, object]) -> str: + backend = str(row["backend"]) + storage_mode = row.get("mooncake_storage_mode") + if backend == "fluxon": + if storage_mode is not None: + raise ValueError(f"Fluxon row has Mooncake storage mode: {row}") + return "fluxon" + if backend != "mooncake": + raise ValueError(f"unexpected backend: {backend!r}") + if storage_mode == "DEDICATED_OWNER": + return "mooncake_dedicated" + if storage_mode == "PER_BENCHMARK_PROCESS": + return "mooncake_per_process" + raise ValueError(f"unexpected Mooncake storage mode: {storage_mode!r}") + + +def row_key(row: Mapping[str, object]) -> RowKey: + return RowKey( + suite_group=str(row["suite_group"]), + output_mode=str(row["output_mode"]), + payload_mib=int(row["payload_mib"]), + concurrency=int(row["worker_concurrency"]), + implementation=row_implementation(row), + ssd_enabled=bool(row["ssd_enabled"]), + ) + + +def _hit_rate_percent(row: Mapping[str, object]) -> float: + hit = int(row["hit"]) + miss = int(row["miss"]) + completed = hit + miss + if completed == 0: + raise ValueError(f"row has no completed hit or miss operation: {row}") + return hit * 100.0 / completed + + +def _validate_row(row: Mapping[str, object], *, context: str) -> None: + if str(row.get("workload_kind")) != "pressure": + raise ValueError(f"non-pressure row in {context}: {row}") + if str(row.get("status")) != "SUCCESS" or not bool(row.get("valid")): + raise ValueError(f"invalid pressure row in {context}: {row}") + if int(row.get("error", 0)) != 0: + raise ValueError(f"pressure row contains errors in {context}: {row}") + if not bool(row.get("source_complete")): + raise ValueError(f"incomplete source counts in {context}: {row}") + if int(row.get("unknown_source_operations", 0)) != 0: + raise ValueError(f"unknown source operations in {context}: {row}") + + hit = int(row["hit"]) + memory_hits = int(row["memory_hit_operations"]) + ssd_hits = int(row["ssd_hit_operations"]) + if memory_hits + ssd_hits != hit: + raise ValueError(f"source counts do not match hits in {context}: {row}") + if not bool(row["ssd_enabled"]) and ssd_hits != 0: + raise ValueError(f"SSD-off row contains SSD hits in {context}: {row}") + + memory_gib_s = float(row["memory_logical_read_gib_per_second"]) + ssd_gib_s = float(row["ssd_logical_read_gib_per_second"]) + total_gib_s = float(row["gib_per_second"]) + if min(memory_gib_s, ssd_gib_s, total_gib_s) < 0: + raise ValueError(f"negative logical bandwidth in {context}: {row}") + if abs(memory_gib_s + ssd_gib_s - total_gib_s) > 0.005: + raise ValueError(f"source bandwidth does not match total in {context}: {row}") + _hit_rate_percent(row) + + +def load_pressure_matrix( + paths: Iterable[Path], + *, + suite_group: str, + output_modes: Sequence[str], + payloads: Sequence[int], + implementations: Sequence[str] = BASE_IMPLEMENTATIONS, +) -> tuple[dict[RowKey, dict], int]: + selected: dict[RowKey, dict] = {} + contexts: dict[RowKey, str] = {} + loaded = 0 + output_mode_set = set(output_modes) + payload_set = set(payloads) + for path in paths: + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), + start=1, + ): + if not raw_line.strip(): + continue + row = json.loads(raw_line) + loaded += 1 + if row.get("workload_kind") != "pressure": + continue + if row.get("suite_group") != suite_group: + continue + if row.get("output_mode") not in output_mode_set: + continue + if int(row.get("payload_mib", -1)) not in payload_set: + continue + if int(row.get("worker_concurrency", -1)) not in CONCURRENCIES: + continue + key = row_key(row) + selected[key] = row + contexts[key] = f"{path}:{line_number}" + + expected = { + RowKey( + suite_group=suite_group, + output_mode=output_mode, + payload_mib=payload_mib, + concurrency=concurrency, + implementation=implementation, + ssd_enabled=ssd_enabled, + ) + for output_mode in output_modes + for payload_mib in payloads + for concurrency in CONCURRENCIES + for implementation in implementations + for ssd_enabled in (False, True) + } + missing = sorted(expected - selected.keys(), key=repr) + unexpected = sorted(selected.keys() - expected, key=repr) + if missing or unexpected: + raise ValueError( + f"incomplete {suite_group} pressure matrix: " + f"rows={len(selected)}/{len(expected)} " + f"missing={missing} unexpected={unexpected}" + ) + for key, row in selected.items(): + _validate_row(row, context=contexts[key]) + return selected, loaded + + +class RasterCanvas: + def __init__(self, *, width: int, height: int) -> None: + self.width = width + self.height = height + self.image = Image.new("RGB", (width, height), "#F8FAFC") + self.draw = ImageDraw.Draw(self.image) + self._font_paths = { + "regular": self._find_font( + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ), + "bold": self._find_font( + "/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc", + "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", + ), + } + self._fonts: dict[tuple[str, int], ImageFont.FreeTypeFont] = {} + + @staticmethod + def _find_font(*candidates: str) -> Path: + for candidate in candidates: + path = Path(candidate) + if path.is_file(): + return path + raise RuntimeError(f"no usable chart font found in {candidates}") + + def _font(self, *, size: int, bold: bool = False) -> ImageFont.FreeTypeFont: + weight = "bold" if bold else "regular" + key = (weight, size) + font = self._fonts.get(key) + if font is None: + font = ImageFont.truetype(str(self._font_paths[weight]), size=size) + self._fonts[key] = font + return font + + def rect( + self, + x: float, + y: float, + width: float, + height: float, + *, + fill: str, + stroke: str, + stroke_width: float = 1.0, + rx: float = 0.0, + ) -> None: + self.draw.rounded_rectangle( + (round(x), round(y), round(x + width), round(y + height)), + radius=max(0, round(rx)), + fill=fill, + outline=stroke, + width=max(1, round(stroke_width)), + ) + + def line( + self, + x1: float, + y1: float, + x2: float, + y2: float, + *, + stroke: str, + stroke_width: float = 1.0, + dash: tuple[int, int] | None = None, + ) -> None: + width = max(1, round(stroke_width)) + if dash is None: + self.draw.line((x1, y1, x2, y2), fill=stroke, width=width) + return + dash_length, gap_length = dash + dx, dy = x2 - x1, y2 - y1 + length = math.hypot(dx, dy) + if length == 0: + return + ux, uy = dx / length, dy / length + cursor = 0.0 + while cursor < length: + end = min(length, cursor + dash_length) + self.draw.line( + ( + x1 + ux * cursor, + y1 + uy * cursor, + x1 + ux * end, + y1 + uy * end, + ), + fill=stroke, + width=width, + ) + cursor += dash_length + gap_length + + def circle( + self, + x: float, + y: float, + radius: float, + *, + fill: str, + stroke: str, + stroke_width: float = 1.0, + ) -> None: + self.draw.ellipse( + ( + round(x - radius), + round(y - radius), + round(x + radius), + round(y + radius), + ), + fill=fill, + outline=stroke, + width=max(1, round(stroke_width)), + ) + + def text( + self, + x: float, + y: float, + value: str, + *, + style: str, + anchor: str = "left", + fill: str | None = None, + ) -> None: + styles = { + "title": (28, True, "#0F172A"), + "subtitle": (15, False, "#475569"), + "panel-title": (18, True, "#0F172A"), + "axis": (12, False, "#334155"), + "axis-small": (10, False, "#64748B"), + "legend": (12, False, "#334155"), + "note": (13, False, "#475569"), + } + try: + size, bold, default_fill = styles[style] + except KeyError as exc: + raise ValueError(f"unknown text style: {style}") from exc + anchors = {"left": "ls", "middle": "ms", "end": "rs"} + self.draw.text( + (round(x), round(y)), + value, + font=self._font(size=size, bold=bold), + fill=fill or default_fill, + anchor=anchors[anchor], + ) + + def write(self, path: Path) -> None: + if path.suffix.lower() != ".png": + raise ValueError(f"raster chart output must use .png: {path}") + path.parent.mkdir(parents=True, exist_ok=True) + self.image.save(path, format="PNG", optimize=True) + + +def nice_axis_max(values: Iterable[float]) -> float: + peak = max(values, default=0.0) + if peak <= 0: + return 1.0 + raw_step = peak / 4.0 + magnitude = 10 ** math.floor(math.log10(raw_step)) + fraction = raw_step / magnitude + step_factor = next( + value for value in (1.0, 2.0, 2.5, 5.0, 10.0) if fraction <= value + ) + step = step_factor * magnitude + return step * math.ceil(peak / step) + + +def format_axis_value(value: float) -> str: + if value >= 10: + return f"{value:.0f}" + if value >= 1: + return f"{value:.1f}" + return f"{value:.2f}" + + +def _draw_axes( + canvas: RasterCanvas, + *, + x: float, + y: float, + width: float, + height: float, + maximum: float, +) -> None: + for tick_index in range(5): + throughput = maximum * tick_index / 4 + tick_y = y + height - height * tick_index / 4 + canvas.line(x, tick_y, x + width, tick_y, stroke="#E2E8F0") + canvas.text( + x - 7, + tick_y + 4, + format_axis_value(throughput), + style="axis-small", + anchor="end", + ) + for hit_rate in (0, 50, 100): + tick_y = y + height * (1.0 - hit_rate / 100.0) + canvas.text( + x + width + 7, + tick_y + 4, + f"{hit_rate}%", + style="axis-small", + ) + canvas.line(x, y, x, y + height, stroke="#94A3B8") + canvas.line(x + width, y, x + width, y + height, stroke="#F2A7B5") + canvas.line(x, y + height, x + width, y + height, stroke="#94A3B8") + + +def _draw_legend( + canvas: RasterCanvas, + *, + x: float, + y: float, + implementations: Sequence[str], +) -> None: + labels = { + "fluxon": ("Fluxon", 100), + "mooncake_dedicated": ("MC 独立 owner", 160), + "mooncake_per_process": ("MC 每进程", 130), + } + for implementation in implementations: + label, item_width = labels[implementation] + memory_color, ssd_color = IMPLEMENTATION_COLORS[implementation] + canvas.rect( + x, + y - 11, + 10, + 12, + fill=memory_color, + stroke=ssd_color, + rx=2, + ) + canvas.rect( + x + 10, + y - 11, + 10, + 12, + fill=ssd_color, + stroke=ssd_color, + rx=2, + ) + canvas.text(x + 27, y, label, style="legend") + x += item_width + canvas.circle( + x + 8, + y - 5, + 4, + fill=HIT_RATE_COLOR, + stroke="#ffffff", + stroke_width=1.5, + ) + canvas.text(x + 24, y, "命中率（右轴）", style="legend") + x += 135 + canvas.text( + x, + y, + "浅色=DRAM · 深色=SSD", + style="legend", + ) + + +def _draw_panel( + canvas: RasterCanvas, + *, + panel_x: float, + panel_y: float, + panel_width: float, + panel_height: float, + suite_group: str, + output_mode: str, + payload_mib: int, + maximum: float, + rows: Mapping[RowKey, Mapping[str, object]], + implementations: Sequence[str], +) -> None: + plot_x = panel_x + 50 + plot_y = panel_y + 58 + plot_width = panel_width - 88 + plot_height = panel_height - 142 + canvas.rect( + panel_x, + panel_y, + panel_width, + panel_height, + fill="#FFFFFF", + stroke="#E2E8F0", + rx=12, + ) + title = f"{payload_mib} MiB" + if output_mode != "cuda": + title = f"{OUTPUT_LABELS[output_mode]} · {title}" + canvas.text(panel_x + 18, panel_y + 29, title, style="panel-title") + canvas.text( + panel_x + panel_width - 18, + panel_y + 29, + f"0–{format_axis_value(maximum)} GiB/s", + style="axis", + anchor="end", + ) + _draw_axes( + canvas, + x=plot_x, + y=plot_y, + width=plot_width, + height=plot_height, + maximum=maximum, + ) + + group_width = plot_width / len(CONCURRENCIES) + bar_width = 11.0 + state_gap = 2.0 + implementation_gap = 4.0 + pair_width = 2 * bar_width + state_gap + series_width = ( + len(implementations) * pair_width + + (len(implementations) - 1) * implementation_gap + ) + baseline = plot_y + plot_height + for concurrency_index, concurrency in enumerate(CONCURRENCIES): + group_center = plot_x + (concurrency_index + 0.5) * group_width + cursor_x = group_center - series_width / 2 + for implementation_index, implementation in enumerate(implementations): + short = IMPLEMENTATION_SHORT_LABELS[implementation] + memory_color, ssd_color = IMPLEMENTATION_COLORS[implementation] + for state_index, ssd_enabled in enumerate((False, True)): + key = RowKey( + suite_group=suite_group, + output_mode=output_mode, + payload_mib=payload_mib, + concurrency=concurrency, + implementation=implementation, + ssd_enabled=ssd_enabled, + ) + row = rows[key] + bar_x = cursor_x + state_index * (bar_width + state_gap) + memory_gib_s = float(row["memory_logical_read_gib_per_second"]) + ssd_gib_s = float(row["ssd_logical_read_gib_per_second"]) + memory_height = memory_gib_s / maximum * plot_height + ssd_height = ssd_gib_s / maximum * plot_height + if memory_height > 0: + canvas.rect( + bar_x, + baseline - memory_height, + bar_width, + memory_height, + fill=memory_color, + stroke=ssd_color, + stroke_width=0.8, + ) + if ssd_height > 0: + canvas.rect( + bar_x, + baseline - memory_height - ssd_height, + bar_width, + ssd_height, + fill=ssd_color, + stroke=ssd_color, + stroke_width=0.8, + ) + hit_y = plot_y + plot_height * ( + 1.0 - _hit_rate_percent(row) / 100.0 + ) + canvas.circle( + bar_x + bar_width / 2, + hit_y, + 3.5, + fill=HIT_RATE_COLOR, + stroke="#ffffff", + stroke_width=1.2, + ) + state = "开" if ssd_enabled else "关" + canvas.text( + bar_x + bar_width / 2, + baseline + 14, + state, + style="axis-small", + anchor="middle", + ) + pair_center = cursor_x + pair_width / 2 + canvas.text( + pair_center, + baseline + 29, + short, + style="axis-small", + anchor="middle", + fill=ssd_color, + ) + cursor_x += pair_width + if implementation_index + 1 < len(implementations): + cursor_x += implementation_gap + canvas.text( + group_center, + baseline + 45, + f"c{concurrency}", + style="axis", + anchor="middle", + ) + if concurrency_index + 1 < len(CONCURRENCIES): + separator_x = plot_x + (concurrency_index + 1) * group_width + canvas.line( + separator_x, + baseline + 5, + separator_x, + baseline + 49, + stroke="#CBD5E1", + ) + + +def render_pressure_sweep( + *, + path: Path, + rows: Mapping[RowKey, Mapping[str, object]], + suite_group: str, + output_modes: Sequence[str], + payloads: Sequence[int], + title: str, + snapshot_label: str, + implementations: Sequence[str] = BASE_IMPLEMENTATIONS, +) -> None: + panel_width = 570 + panel_height = 475 + panel_gap = 18 + left = 48 + top = 132 + bottom = 62 + width = left * 2 + len(payloads) * panel_width + (len(payloads) - 1) * panel_gap + height = ( + top + + len(output_modes) * panel_height + + (len(output_modes) - 1) * panel_gap + + bottom + ) + canvas = RasterCanvas(width=width, height=height) + series_count = len(implementations) * 2 + series_count_label = {6: "六", 8: "八"}.get(series_count, str(series_count)) + canvas.text(left, 45, title, style="title") + canvas.text( + left, + 73, + f"每个并发度固定{series_count_label}柱：" + + "、".join( + { + "fluxon": "Fluxon", + "mooncake_dedicated": "MC 独立 owner", + "mooncake_per_process": "MC 每进程", + }[implementation] + for implementation in implementations + ) + + " × SSD 关/开；" + "颜色区分实现，浅/深区分 DRAM / SSD；各 payload 分面独立左轴", + style="subtitle", + ) + _draw_legend(canvas, x=left, y=106, implementations=implementations) + + for output_index, output_mode in enumerate(output_modes): + panel_y = top + output_index * (panel_height + panel_gap) + for payload_index, payload_mib in enumerate(payloads): + maximum = nice_axis_max( + float(row["memory_logical_read_gib_per_second"]) + + float(row["ssd_logical_read_gib_per_second"]) + for key, row in rows.items() + if key.output_mode == output_mode + and key.payload_mib == payload_mib + ) + panel_x = left + payload_index * (panel_width + panel_gap) + _draw_panel( + canvas, + panel_x=panel_x, + panel_y=panel_y, + panel_width=panel_width, + panel_height=panel_height, + suite_group=suite_group, + output_mode=output_mode, + payload_mib=payload_mib, + maximum=maximum, + rows=rows, + implementations=implementations, + ) + + canvas.text( + left, + height - 27, + f"快照：{snapshot_label}。柱高只计 hit payload；红点为 hit / (hit + miss)，" + "miss 是压力实验结果的一部分。", + style="note", + ) + canvas.write(path) + + +def _draw_ssd_backend_ablation_legend( + canvas: RasterCanvas, + *, + x: float, + y: float, +) -> None: + labels = { + "native": ("Fluxon 原生 SSD", 155), + "foyer": ("Fluxon Foyer SSD", 165), + } + for implementation in SSD_BACKEND_IMPLEMENTATIONS: + label, item_width = labels[implementation] + memory_color, ssd_color = SSD_BACKEND_COLORS[implementation] + canvas.rect( + x, + y - 11, + 10, + 12, + fill=memory_color, + stroke=ssd_color, + rx=2, + ) + canvas.rect( + x + 10, + y - 11, + 10, + 12, + fill=ssd_color, + stroke=ssd_color, + rx=2, + ) + canvas.text(x + 27, y, label, style="legend") + x += item_width + canvas.circle( + x + 8, + y - 5, + 4, + fill=HIT_RATE_COLOR, + stroke="#ffffff", + stroke_width=1.5, + ) + canvas.text(x + 24, y, "命中率（右轴）", style="legend") + x += 135 + canvas.text(x, y, "浅色=DRAM · 深色=SSD", style="legend") + + +def _draw_ssd_backend_ablation_panel( + canvas: RasterCanvas, + *, + panel_x: float, + panel_y: float, + panel_width: float, + panel_height: float, + output_mode: str, + payload_mib: int, + maximum: float, + rows: Mapping[SsdBackendAblationKey, Mapping[str, object]], +) -> None: + plot_x = panel_x + 50 + plot_y = panel_y + 58 + plot_width = panel_width - 88 + plot_height = panel_height - 142 + canvas.rect( + panel_x, + panel_y, + panel_width, + panel_height, + fill="#FFFFFF", + stroke="#E2E8F0", + rx=12, + ) + output_label = "MemHolder" if output_mode == "holder" else OUTPUT_LABELS[output_mode] + canvas.text( + panel_x + 18, + panel_y + 29, + f"{output_label} · {payload_mib} MiB", + style="panel-title", + ) + canvas.text( + panel_x + panel_width - 18, + panel_y + 29, + f"0–{format_axis_value(maximum)} GiB/s", + style="axis", + anchor="end", + ) + _draw_axes( + canvas, + x=plot_x, + y=plot_y, + width=plot_width, + height=plot_height, + maximum=maximum, + ) + + group_width = plot_width / len(CONCURRENCIES) + bar_width = 25.0 + backend_gap = 13.0 + series_width = len(SSD_BACKEND_IMPLEMENTATIONS) * bar_width + backend_gap + baseline = plot_y + plot_height + for concurrency_index, concurrency in enumerate(CONCURRENCIES): + group_center = plot_x + (concurrency_index + 0.5) * group_width + cursor_x = group_center - series_width / 2 + for backend_index, backend_impl in enumerate(SSD_BACKEND_IMPLEMENTATIONS): + key = SsdBackendAblationKey( + output_mode=output_mode, + payload_mib=payload_mib, + concurrency=concurrency, + ssd_backend_impl=backend_impl, + ) + row = rows[key] + memory_color, ssd_color = SSD_BACKEND_COLORS[backend_impl] + memory_gib_s = float(row["memory_logical_read_gib_per_second"]) + ssd_gib_s = float(row["ssd_logical_read_gib_per_second"]) + memory_height = memory_gib_s / maximum * plot_height + ssd_height = ssd_gib_s / maximum * plot_height + if memory_height > 0: + canvas.rect( + cursor_x, + baseline - memory_height, + bar_width, + memory_height, + fill=memory_color, + stroke=ssd_color, + stroke_width=0.8, + ) + if ssd_height > 0: + canvas.rect( + cursor_x, + baseline - memory_height - ssd_height, + bar_width, + ssd_height, + fill=ssd_color, + stroke=ssd_color, + stroke_width=0.8, + ) + hit_y = plot_y + plot_height * ( + 1.0 - _hit_rate_percent(row) / 100.0 + ) + canvas.circle( + cursor_x + bar_width / 2, + hit_y, + 4.0, + fill=HIT_RATE_COLOR, + stroke="#ffffff", + stroke_width=1.2, + ) + canvas.text( + cursor_x + bar_width / 2, + baseline + 18, + SSD_BACKEND_SHORT_LABELS[backend_impl], + style="axis-small", + anchor="middle", + fill=ssd_color, + ) + cursor_x += bar_width + if backend_index + 1 < len(SSD_BACKEND_IMPLEMENTATIONS): + cursor_x += backend_gap + canvas.text( + group_center, + baseline + 40, + f"c{concurrency}", + style="axis", + anchor="middle", + ) + if concurrency_index + 1 < len(CONCURRENCIES): + separator_x = plot_x + (concurrency_index + 1) * group_width + canvas.line( + separator_x, + baseline + 5, + separator_x, + baseline + 44, + stroke="#CBD5E1", + ) + + +def render_ssd_backend_ablation( + *, + path: Path, + rows: Mapping[SsdBackendAblationKey, Mapping[str, object]], + output_modes: Sequence[str], + payloads: Sequence[int], + title: str, + snapshot_label: str, +) -> None: + panel_width = 570 + panel_height = 475 + panel_gap = 18 + left = 48 + top = 132 + bottom = 62 + width = left * 2 + len(payloads) * panel_width + (len(payloads) - 1) * panel_gap + height = ( + top + + len(output_modes) * panel_height + + (len(output_modes) - 1) * panel_gap + + bottom + ) + canvas = RasterCanvas(width=width, height=height) + canvas.text(left, 45, title, style="title") + canvas.text( + left, + 73, + "每个并发度固定两柱：Fluxon 原生 SSD、Fluxon Foyer SSD；" + "颜色区分 SSD 实现，浅/深区分 DRAM / SSD；Fluxon DRAM 均为 2 GiB，" + "Foyer 内部 memory admission=0；两侧均使用 direct I/O", + style="subtitle", + ) + _draw_ssd_backend_ablation_legend(canvas, x=left, y=106) + + for output_index, output_mode in enumerate(output_modes): + panel_y = top + output_index * (panel_height + panel_gap) + for payload_index, payload_mib in enumerate(payloads): + maximum = nice_axis_max( + float(row["memory_logical_read_gib_per_second"]) + + float(row["ssd_logical_read_gib_per_second"]) + for key, row in rows.items() + if key.output_mode == output_mode and key.payload_mib == payload_mib + ) + panel_x = left + payload_index * (panel_width + panel_gap) + _draw_ssd_backend_ablation_panel( + canvas, + panel_x=panel_x, + panel_y=panel_y, + panel_width=panel_width, + panel_height=panel_height, + output_mode=output_mode, + payload_mib=payload_mib, + maximum=maximum, + rows=rows, + ) + + canvas.text( + left, + height - 27, + f"快照：{snapshot_label}。柱高只计 hit payload；红点为 hit / (hit + miss)，" + "两组都经过同一 Fluxon KV / FlatDict / route / transport 路径。", + style="note", + ) + canvas.write(path) diff --git a/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_cpu_sweep.py b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_cpu_sweep.py new file mode 100644 index 0000000..0188948 --- /dev/null +++ b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_cpu_sweep.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Render the test-stack CPU KV SSD-pressure matrix.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from kv_ssd_pressure_chart import ( + BASE_IMPLEMENTATIONS, + load_pressure_matrix, + render_pressure_sweep, +) + + +PAYLOADS = (1, 4, 8, 16) +OUTPUT_MODES = ("holder", "bytes") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--summary", + action="append", + type=Path, + required=True, + help="Input summary.jsonl; later files replace duplicate case keys.", + ) + parser.add_argument("--cpu-png-output", type=Path, required=True) + parser.add_argument("--snapshot-label", required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + rows, loaded = load_pressure_matrix( + args.summary, + suite_group="cpu", + output_modes=OUTPUT_MODES, + payloads=PAYLOADS, + implementations=BASE_IMPLEMENTATIONS, + ) + render_pressure_sweep( + path=args.cpu_png_output, + rows=rows, + suite_group="cpu", + output_modes=OUTPUT_MODES, + payloads=PAYLOADS, + title="CPU SSD-pressure：吞吐与命中率（原生引用 / handle、bytes）", + snapshot_label=args.snapshot_label, + implementations=BASE_IMPLEMENTATIONS, + ) + print( + "rendered_kv_ssd_cpu_sweep", + f"loaded_rows={loaded}", + f"matrix_rows={len(rows)}/192", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_foyer_ablation.py b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_foyer_ablation.py new file mode 100755 index 0000000..22fb12f --- /dev/null +++ b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_foyer_ablation.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Render the test-stack Fluxon native-vs-Foyer SSD backend ablation.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from kv_ssd_pressure_chart import ( + CONCURRENCIES, + SSD_BACKEND_IMPLEMENTATIONS, + SsdBackendAblationKey, + render_ssd_backend_ablation, +) + + +OUTPUT_MODES = ("holder", "bytes") +PAYLOADS = (1, 4, 8, 16) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--summary", + action="append", + type=Path, + required=True, + help="Input Foyer ablation summary.jsonl; later files replace duplicate keys.", + ) + parser.add_argument("--png-output", type=Path, required=True) + parser.add_argument("--snapshot-label", required=True) + return parser.parse_args() + + +def _key(row: dict) -> SsdBackendAblationKey: + return SsdBackendAblationKey( + output_mode=str(row["output_mode"]), + payload_mib=int(row["payload_mib"]), + concurrency=int(row["worker_concurrency"]), + ssd_backend_impl=str(row["ssd_backend_impl"]), + ) + + +def _validate_row(row: dict, *, context: str) -> None: + if row.get("workload_kind") != "pressure" or row.get("suite_group") != "cpu": + raise ValueError(f"unexpected workload in {context}: {row}") + expected_workload = { + "workload_id": "kv_ssd_pressure_zipf", + "request_distribution": "zipfian", + "read_ratio": 1.0, + "write_ratio": 0.0, + "keyspace_capacity_guard": False, + "kv_bootstrap_concurrency": 1, + } + mismatched_workload = { + field: (row.get(field), expected) + for field, expected in expected_workload.items() + if row.get(field) != expected + } + if mismatched_workload: + raise ValueError( + f"unexpected pressure controls in {context}: {mismatched_workload}" + ) + payload_mib = int(row["payload_mib"]) + concurrency = int(row["worker_concurrency"]) + mib = 1024 * 1024 + process_layout = { + 2: (1, 2), + 4: (1, 4), + 8: (2, 4), + 16: (4, 4), + } + processes_per_target, threads_per_process = process_layout[concurrency] + payload_bytes = payload_mib * mib - 128 + keyspace_size = 2560 // payload_mib + expected_controls = { + "payload_bytes": payload_bytes, + "keyspace_size": keyspace_size, + "dataset_payload_bytes": payload_bytes * keyspace_size, + "owner_dram_bytes": 2048 * mib, + "ssd_capacity_bytes": 16 * 1024 * mib, + "processes_per_target": processes_per_target, + "threads_per_process": threads_per_process, + "duration_seconds_configured": 35, + "metric_warmup_seconds": 5, + "kv_bootstrap_put_gap_ms": max(5, payload_mib * 1000 // 160), + } + mismatched_controls = { + field: (row.get(field), expected) + for field, expected in expected_controls.items() + if row.get(field) != expected + } + if mismatched_controls: + raise ValueError( + f"unexpected capacity or concurrency controls in {context}: " + f"{mismatched_controls}" + ) + if row.get("backend") != "fluxon": + raise ValueError(f"non-Fluxon row in {context}: {row}") + if row.get("status") != "SUCCESS" or not bool(row.get("valid")): + raise ValueError(f"invalid row in {context}: {row}") + if int(row.get("error", 0)) != 0: + raise ValueError(f"row contains errors in {context}: {row}") + if int(row.get("ssd_persist_failure_count", -1)) != 0: + raise ValueError(f"row contains SSD persist failures in {context}: {row}") + if not bool(row.get("source_complete")): + raise ValueError(f"incomplete source counts in {context}: {row}") + if int(row.get("unknown_source_operations", 0)) != 0: + raise ValueError(f"unknown source operations in {context}: {row}") + hit = int(row["hit"]) + memory_hits = int(row["memory_hit_operations"]) + ssd_hits = int(row["ssd_hit_operations"]) + if memory_hits + ssd_hits != hit: + raise ValueError(f"source counts do not match hits in {context}: {row}") + total = float(row["gib_per_second"]) + memory = float(row["memory_logical_read_gib_per_second"]) + ssd = float(row["ssd_logical_read_gib_per_second"]) + if min(total, memory, ssd) < 0 or abs(memory + ssd - total) > 0.005: + raise ValueError(f"source bandwidth does not match total in {context}: {row}") + + +def load_rows(paths: list[Path]) -> dict[SsdBackendAblationKey, dict]: + selected: dict[SsdBackendAblationKey, dict] = {} + contexts: dict[SsdBackendAblationKey, str] = {} + for path in paths: + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), start=1 + ): + if not raw_line.strip(): + continue + row = json.loads(raw_line) + if row.get("workload_kind") != "pressure" or row.get("suite_group") != "cpu": + continue + key = _key(row) + selected[key] = row + contexts[key] = f"{path}:{line_number}" + + expected = { + SsdBackendAblationKey( + output_mode=output_mode, + payload_mib=payload_mib, + concurrency=concurrency, + ssd_backend_impl=backend_impl, + ) + for output_mode in OUTPUT_MODES + for payload_mib in PAYLOADS + for concurrency in CONCURRENCIES + for backend_impl in SSD_BACKEND_IMPLEMENTATIONS + } + missing = sorted(expected - selected.keys(), key=repr) + unexpected = sorted(selected.keys() - expected, key=repr) + if missing or unexpected: + raise ValueError( + f"incomplete Fluxon SSD backend ablation: rows={len(selected)}/{len(expected)} " + f"missing={missing} unexpected={unexpected}" + ) + for key, row in selected.items(): + _validate_row(row, context=contexts[key]) + + control_fields = ( + "workload_id", + "request_distribution", + "read_ratio", + "write_ratio", + "keyspace_capacity_guard", + "kv_bootstrap_concurrency", + "kv_bootstrap_put_gap_ms", + "payload_bytes", + "keyspace_size", + "dataset_payload_bytes", + "owner_dram_bytes", + "ssd_capacity_bytes", + "worker_concurrency", + "processes_per_target", + "threads_per_process", + "duration_seconds_configured", + "metric_warmup_seconds", + ) + for output_mode in OUTPUT_MODES: + for payload_mib in PAYLOADS: + for concurrency in CONCURRENCIES: + native = selected[ + SsdBackendAblationKey(output_mode, payload_mib, concurrency, "native") + ] + foyer = selected[ + SsdBackendAblationKey(output_mode, payload_mib, concurrency, "foyer") + ] + differences = { + field: (native.get(field), foyer.get(field)) + for field in control_fields + if native.get(field) != foyer.get(field) + } + if differences: + raise ValueError( + "native/Foyer control mismatch for " + f"{output_mode} p{payload_mib} c{concurrency}: {differences}" + ) + return selected + + +def main() -> int: + args = parse_args() + rows = load_rows(args.summary) + render_ssd_backend_ablation( + path=args.png_output, + rows=rows, + output_modes=OUTPUT_MODES, + payloads=PAYLOADS, + title="Fluxon SSD backend 消融：原生实现 vs Foyer", + snapshot_label=args.snapshot_label, + ) + print(f"rendered_kv_ssd_foyer_ablation matrix_rows={len(rows)}/64") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_gpu_sweep.py b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_gpu_sweep.py new file mode 100644 index 0000000..a199aee --- /dev/null +++ b/fluxon_test_stack/single_node_ssd_bench/plot_kv_ssd_gpu_sweep.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +"""Render the test-stack GPU KV SSD-pressure matrix.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from kv_ssd_pressure_chart import load_pressure_matrix, render_pressure_sweep + + +PAYLOADS = (4, 8, 16) +OUTPUT_MODES = ("cuda",) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--summary", + action="append", + type=Path, + required=True, + help="Input summary.jsonl; later files replace duplicate case keys.", + ) + parser.add_argument("--gpu-png-output", type=Path, required=True) + parser.add_argument("--snapshot-label", required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + rows, loaded = load_pressure_matrix( + args.summary, + suite_group="gpu", + output_modes=OUTPUT_MODES, + payloads=PAYLOADS, + ) + render_pressure_sweep( + path=args.gpu_png_output, + rows=rows, + suite_group="gpu", + output_modes=OUTPUT_MODES, + payloads=PAYLOADS, + title="GPU SSD-pressure：吞吐与命中率", + snapshot_label=args.snapshot_label, + ) + print( + "rendered_kv_ssd_gpu_sweep", + f"loaded_rows={loaded}", + f"matrix_rows={len(rows)}/72", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/fluxon_test_stack/start_test_bed.yaml b/fluxon_test_stack/start_test_bed.yaml index 6a815ec..03b6eee 100644 --- a/fluxon_test_stack/start_test_bed.yaml +++ b/fluxon_test_stack/start_test_bed.yaml @@ -47,14 +47,14 @@ test_runner_ui: # deployconf.bootstrap_bare_services before launch-only bare start. bootstrap_phases: - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - etcd - greptime - tikv_pd - tikv - mode: fixed_bare - node: infra44-ThinkStation-PX + node: example-node-a services: - fluxon_core_controller - mode: coverage_bare diff --git a/fluxon_test_stack/test_runner.py b/fluxon_test_stack/test_runner.py index ac085ed..0c952c3 100644 --- a/fluxon_test_stack/test_runner.py +++ b/fluxon_test_stack/test_runner.py @@ -260,7 +260,7 @@ class _HttpGetJsonTransientError(RuntimeError): TEST_STACK_REDIS_INSTANCE_ID_PREFIX = "redis_node_" TEST_STACK_ALLUXIO_INSTANCE_ID_PREFIX = "alluxio_node_" TEST_STACK_MOONCAKE_MASTER_INSTANCE_ID = "mooncake_master" -TEST_STACK_BENCHMARK_FIXED_THREADS_PER_PROCESS = 4 +TEST_STACK_BENCHMARK_THREADS_PER_PROCESS_VALUES = (2, 4) # Owner-mode Fluxon KV configs in CI / TEST_STACK share the same canonical label. FLUXON_KV_OWNER_SUB_CLUSTER = "owner" @@ -274,6 +274,14 @@ class _HttpGetJsonTransientError(RuntimeError): TEST_STACK_BACKEND_ALLUXIO, TEST_STACK_BACKEND_MOONCAKE, } +TEST_STACK_MOONCAKE_STORAGE_MODE_DEDICATED_OWNER = "DEDICATED_OWNER" +TEST_STACK_MOONCAKE_STORAGE_MODE_PER_BENCHMARK_PROCESS = "PER_BENCHMARK_PROCESS" +TEST_STACK_MOONCAKE_STORAGE_MODES_ALLOWED = { + TEST_STACK_MOONCAKE_STORAGE_MODE_DEDICATED_OWNER, + TEST_STACK_MOONCAKE_STORAGE_MODE_PER_BENCHMARK_PROCESS, +} +TEST_STACK_MOONCAKE_SEGMENT_ALIGNMENT_BYTES = 16 * 1024 * 1024 +TEST_STACK_MOONCAKE_SSD_LIMIT_ENV = "MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES" TEST_STACK_COMPLETION_STATUS_SUCCESS = "SUCCESS" RUN_MODE_DEBUG_ONE_BY_ONE = "debug_one_by_one" @@ -381,14 +389,25 @@ def _test_stack_ops_namespace() -> str: TEST_STACK_REQUEST_DISTRIBUTION_ZIPFIAN, } TEST_STACK_SCENE_KEY_AFFINITY_LOCALITY_RATIO = "affinity_locality_ratio" +TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD = "keyspace_capacity_guard" +TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY = "kv_bootstrap_concurrency" +TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS = "kv_bootstrap_put_gap_ms" +TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY = "kv_bootstrap_storage_full_policy" +TEST_STACK_SCENE_KEY_KV_GET_OUTPUT = "kv_get_output" +TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX = "kv_cuda_device_index" +TEST_STACK_KV_BOOTSTRAP_STORAGE_FULL_POLICIES_ALLOWED = {"fail", "stop"} +TEST_STACK_KV_GET_OUTPUTS_ALLOWED = {"holder", "bytes", "cuda"} TEST_STACK_BENCHMARK_KEY_AFFINITY_SLOT_COUNT = "affinity_slot_count" +SCENE_ID_KV_DRAM_RESIDENT_ZIPF = "kv_dram_resident_zipf" SCENE_ID_KV_READ_HEAVY_AFFINITY = "kv_read_heavy_affinity" SCENE_ENUMS_ALLOWED = { "bench_mq", "fs_open_read_close_smallfiles", "fs_write_close_commit", + SCENE_ID_KV_DRAM_RESIDENT_ZIPF, "kv_read_heavy_zipf", SCENE_ID_KV_READ_HEAVY_AFFINITY, + "kv_ssd_pressure_zipf", "kv_write_heavy_large_value", "rpc_echo_small_payload", "rpc_echo_small_payload_zerorpc", @@ -646,6 +665,14 @@ def main() -> None: required=False, help="Optional GitOps config YAML. When set with --action ui, test_runner owns the GitOps poller/UI/API as part of the same service.", ) + parser.add_argument( + "--cleanup-successful-case-artifacts", + action="store_true", + help=( + "After a case and its teardown both succeed, keep only summary.yaml in its run directory. " + "Failed cases retain their complete diagnostic artifacts." + ), + ) parser.add_argument( "--top-attention-prefix", dest="top_attention_prefixes", @@ -871,6 +898,15 @@ def main() -> None: controller_url=_require_str(stack_identity.get("controller_url"), "stack_identity.controller_url"), case_id=case.case_id, ) + if args.cleanup_successful_case_artifacts: + previous_run_index = _require_int( + last_run.get("run_index"), + f"case_runs[{case.case_id}].last_run.run_index", + min_v=1, + ) + _cleanup_successful_run_artifacts( + results_root / case.case_id / f"run_{previous_run_index}" + ) continue _bench_case_preflight_cleanup( stack_identity=stack_identity, @@ -1093,6 +1129,21 @@ def main() -> None: print(f"ERROR: failed to write/update summary.yaml: {exc}") raise SystemExit(1) + if ( + args.cleanup_successful_case_artifacts + and outcome == RUN_OUTCOME_SUCCESS + and finalize_error is None + ): + try: + _cleanup_successful_run_artifacts(run_dir) + except Exception as exc: + print( + "ERROR: successful case artifact cleanup failed; stopping to avoid exhausting disk: " + f"case_id={case.case_id} run_dir={run_dir} err={type(exc).__name__}: {exc}", + flush=True, + ) + raise SystemExit(1) + if fatal_stop_after_finalize: raise SystemExit(1) @@ -1958,17 +2009,17 @@ def _remote_ssh_common_argv() -> List[str]: def _suite_cluster_name_for_workdir(workdir_root: Path) -> str: - """Return the fixed benchmark workload cluster namespace. + """Return a stable benchmark workload cluster namespace for this workdir. English note: - - Benchmark runs intentionally reuse one stable cluster namespace so shared-memory and - etcd state stay under a predictable top-level path. + - The cluster namespace is workdir-scoped so back-to-back benchmark profiles do not observe + stale KV membership from a previous suite run. - Isolation from the ops bed still comes from the explicit `cluster_name != ops_cluster_name` guard in `_load_stack_identity`. """ - _ = workdir_root - return "fluxon_benchmark" + digest = hashlib.sha256(str(workdir_root.resolve()).encode("utf-8")).hexdigest()[:12] + return f"fluxon_benchmark_{digest}" def _cluster_scoped_shared_dir(*, root_path: str, cluster_name: str) -> Path: @@ -2218,7 +2269,7 @@ def _is_runtime_token_placeholder(raw: str) -> bool: def _find_unresolved_runtime_tokens(raw: str) -> List[str]: - return sorted(set(re.findall(r"__[A-Z0-9_]+__", raw))) + return sorted(set(re.findall(r"__[A-Z][A-Z0-9_]*__", raw))) def _resolved_case_runtime_token_mapping(resolved_case: Dict[str, Any]) -> Dict[str, str]: @@ -3538,19 +3589,114 @@ def _resolved_run_dir_path(resolved_case: Dict[str, Any]) -> Path: return Path(_require_str(runtime.get("run_dir"), "runtime.run_dir")).resolve() -def _ci_share_mem_path(resolved_case: Dict[str, Any], *, run_dir: Path) -> str: +def _ci_share_mem_path(resolved_case: Dict[str, Any], *, owner_index: int) -> str: runtime = _require_dict(resolved_case.get("runtime"), "resolved_case.runtime") stack_identity = _require_dict(runtime.get("stack_identity"), "resolved_case.runtime.stack_identity") share_mem_root = _require_str( stack_identity.get("share_mem_path"), "resolved_case.runtime.stack_identity.share_mem_path", ) + owner_index_i = _require_int(owner_index, "owner_index", min_v=0) # English note: # - iceoryx2 uses share_mem_path as a base for per-node paths (e.g. .../nodes//iox2_/.service_tag). # - The per-node suffix can be long, and some filesystems enforce a max path length of 255 bytes. # - Therefore share_mem_path must be short and must not embed run_dir (which can be deep under repo/workdir). - token = hashlib.sha256(str(run_dir.resolve()).encode("utf-8")).hexdigest()[:16] - return str((Path(share_mem_root) / "ci" / token).resolve()) + # - CI owners use a dedicated namespace below the testbed root. This both isolates them from the + # long-lived testbed owner and gives every owner index one stable slot that is reused across cases. + return os.path.abspath(os.fspath(Path(share_mem_root) / "ci" / f"owner-{owner_index_i}")) + + +def _mutate_ci_owner_share_mem_dir( + resolved_case: Dict[str, Any], + *, + owner_index: int, + share_mem_path: str, + recreate: bool, +) -> None: + """Remove one bounded CI owner SHM slot, optionally recreating it empty. + + The exact-path check is intentional: the testbed owner stores its bundle directly below the + shared root under its own cluster name, while CI is allowed to mutate only ci/owner-{index}. + """ + owner_index_i = _require_int(owner_index, "owner_index", min_v=0) + expected_path = Path(_ci_share_mem_path(resolved_case, owner_index=owner_index_i)) + actual_path = Path(os.path.abspath(os.fspath(share_mem_path))) + expected_namespace = expected_path.parent + if actual_path != expected_path: + raise ValueError( + "refusing to mutate an unexpected CI owner share_mem_path: " + f"path={actual_path} expected={expected_path}" + ) + if expected_namespace.name != "ci" or expected_path.name != f"owner-{owner_index_i}": + raise ValueError( + "refusing to mutate a share-memory path outside the CI owner namespace: " + f"path={expected_path}" + ) + + instance_id = f"owner_{owner_index_i}" + remote_access = _instance_remote_target_access(resolved_case, instance_id=instance_id) + if remote_access is None: + if expected_namespace.is_symlink(): + raise ValueError(f"CI share-memory namespace must not be a symlink: {expected_namespace}") + if expected_path.is_symlink() or expected_path.is_file(): + expected_path.unlink() + elif expected_path.exists(): + shutil.rmtree(expected_path) + if recreate: + expected_path.mkdir(parents=True, exist_ok=False) + else: + target_name, node_cfg, dispatch_mod = remote_access + namespace_q = dispatch_mod.sh_quote(str(expected_namespace)) + owner_path_q = dispatch_mod.sh_quote(str(expected_path)) + remote_lines = [ + "set -euo pipefail", + f"if [ -L {namespace_q} ]; then", + f" echo 'refusing symlinked CI share-memory namespace: ' {namespace_q} >&2", + " exit 1", + "fi", + f"rm -rf -- {owner_path_q}", + ] + if recreate: + remote_lines.append(f"mkdir -p -- {owner_path_q}") + _run_remote_bash( + target_name=target_name, + node_cfg=node_cfg, + remote_cmd="\n".join(remote_lines), + ) + + action = "reset" if recreate else "removed" + print( + f"[CI owner share memory] {action} owner_index={owner_index_i} path={expected_path}", + flush=True, + ) + + +def _reset_ci_owner_share_mem_dir( + resolved_case: Dict[str, Any], + *, + owner_index: int, + share_mem_path: str, +) -> None: + _mutate_ci_owner_share_mem_dir( + resolved_case, + owner_index=owner_index, + share_mem_path=share_mem_path, + recreate=True, + ) + + +def _cleanup_ci_owner_share_mem_dir( + resolved_case: Dict[str, Any], + *, + owner_index: int, + share_mem_path: str, +) -> None: + _mutate_ci_owner_share_mem_dir( + resolved_case, + owner_index=owner_index, + share_mem_path=share_mem_path, + recreate=False, + ) def _ci_owner_shared_bundle_paths(run_dir: Path, *, owner_config_path: Path) -> List[Path]: @@ -3984,6 +4130,16 @@ def _require_int(v: Any, ctx: str, *, min_v: int, max_v: int | None = None) -> i return v +def _require_test_stack_benchmark_threads_per_process(v: Any, ctx: str) -> int: + threads_per_process = _require_int(v, ctx, min_v=1) + if threads_per_process not in TEST_STACK_BENCHMARK_THREADS_PER_PROCESS_VALUES: + supported = ", ".join( + str(value) for value in TEST_STACK_BENCHMARK_THREADS_PER_PROCESS_VALUES + ) + raise ValueError(f"{ctx} must be one of: {supported}") + return threads_per_process + + def _require_bool(v: Any, ctx: str) -> bool: if not isinstance(v, bool): raise ValueError(f"{ctx} must be bool") @@ -4067,6 +4223,8 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: "side_transfer_worker_count", "side_transfer_worker_p2p_port_base", "side_transfer_role", + "kv_ssd_storage_backend", + "kv_ssd_uring_mode", }, ctx, ) @@ -4095,6 +4253,30 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: f"{ctx}.p2p_transport_impl must be one of {sorted(allowed_p2p_transport_impls)}, got {p2p_transport_impl!r}" ) out["p2p_transport_impl"] = p2p_transport_impl + if "kv_ssd_storage_backend" in cfg: + kv_ssd_storage_backend = _require_str( + cfg.get("kv_ssd_storage_backend"), f"{ctx}.kv_ssd_storage_backend" + ) + allowed_kv_ssd_storage_backends = {"native", "foyer"} + if kv_ssd_storage_backend not in allowed_kv_ssd_storage_backends: + raise ValueError( + f"{ctx}.kv_ssd_storage_backend must be one of {sorted(allowed_kv_ssd_storage_backends)}, got {kv_ssd_storage_backend!r}" + ) + out["kv_ssd_storage_backend"] = kv_ssd_storage_backend + if "kv_ssd_uring_mode" in cfg: + kv_ssd_uring_mode = _require_str(cfg.get("kv_ssd_uring_mode"), f"{ctx}.kv_ssd_uring_mode") + allowed_kv_ssd_uring_modes = {"single_buffer", "iovec"} + if kv_ssd_uring_mode not in allowed_kv_ssd_uring_modes: + raise ValueError( + f"{ctx}.kv_ssd_uring_mode must be one of {sorted(allowed_kv_ssd_uring_modes)}, got {kv_ssd_uring_mode!r}" + ) + out["kv_ssd_uring_mode"] = kv_ssd_uring_mode + if out.get("kv_ssd_storage_backend") == "foyer" and out.get( + "kv_ssd_uring_mode", "single_buffer" + ) != "single_buffer": + raise ValueError( + f"{ctx}.kv_ssd_uring_mode is only valid with {ctx}.kv_ssd_storage_backend=native" + ) transport_mode_was_explicit = "transport_mode" in cfg side_transfer_role_raw = cfg.get("side_transfer_role") default_transport_mode = None if side_transfer_role_raw == "worker" else "transfer_with_rpc" @@ -4228,6 +4410,13 @@ def _resolve_test_stack_fluxon_protocol_cfg( return copy.deepcopy(_require_dict(raw_protocol, f"{ctx}.protocol")) +def _test_stack_protocol_requires_unlimited_memlock(protocol_cfg: Optional[Dict[str, Any]]) -> bool: + if protocol_cfg is None: + return False + protocol_type = _require_str(protocol_cfg.get("protocol_type"), "protocol.protocol_type").strip().lower() + return protocol_type == "rdma" + + def _rewrite_artifact_set_transport_impl( *, artifact_set_id: str, @@ -4618,6 +4807,65 @@ def _clean_workdir(workdir_root: Path) -> None: print(f"[clean_workdir] nothing to remove under: {workdir_root}", flush=True) +def _cleanup_successful_run_artifacts(run_dir: Path) -> None: + """Keep the terminal summary and remove bulky artifacts from a successful run.""" + run_dir = Path(run_dir).resolve() + if not run_dir.exists(): + print(f"[cleanup_successful_run_artifacts] run directory already absent: {run_dir}", flush=True) + return + if not run_dir.is_dir(): + raise ValueError(f"successful run path is not a directory: {run_dir}") + + summary_path = run_dir / "summary.yaml" + if not summary_path.is_file() or summary_path.is_symlink(): + raise ValueError(f"successful run must retain a regular summary.yaml before cleanup: {summary_path}") + summary = _require_dict(_load_yaml_file(summary_path), "successful run summary") + if summary.get("outcome") != RUN_OUTCOME_SUCCESS: + raise ValueError( + "refusing to clean a run whose summary is not SUCCESS: " + f"run_dir={run_dir} outcome={summary.get('outcome')!r}" + ) + test_stack_summary = summary.get("test_stack") + has_diagnostic_error = ( + summary.get("error") is not None + or summary.get("teardown_error") is not None + or ( + isinstance(test_stack_summary, dict) + and any( + test_stack_summary.get(field_name) is not None + for field_name in ("error", "collect_error", "teardown_error") + ) + ) + ) + if has_diagnostic_error: + print( + "[cleanup_successful_run_artifacts] preserving run with error diagnostics: " + f"{run_dir}", + flush=True, + ) + return + + free_before = shutil.disk_usage(run_dir).free + removed_entries = 0 + for path in list(run_dir.iterdir()): + if path.name == "summary.yaml": + continue + if path.is_symlink() or not path.is_dir(): + path.unlink() + else: + shutil.rmtree(path) + removed_entries += 1 + + free_after = shutil.disk_usage(run_dir).free + reclaimed_bytes = max(0, int(free_after) - int(free_before)) + print( + "[cleanup_successful_run_artifacts] " + f"run_dir={run_dir} removed_entries={removed_entries} reclaimed_bytes={reclaimed_bytes} " + f"free_bytes={free_after}", + flush=True, + ) + + def _repair_reserved_last_runs(case_runs: Dict[str, Any], *, results_root: Path) -> List[str]: """Repair reserved-only `last_run` entries based on existing run_dir artifacts. @@ -4837,11 +5085,20 @@ def _test_stack_backend_uses_external_fluxon_kv(*, backend_kind: str, mode: str) ) -def _test_stack_backend_uses_dedicated_kv_owners(*, backend_kind: str, mode: str) -> bool: +def _test_stack_backend_uses_dedicated_kv_owners( + *, + backend_kind: str, + mode: str, + mooncake_storage_mode: Optional[str] = None, +) -> bool: if backend_kind == TEST_STACK_BACKEND_FLUXON: return _test_stack_backend_uses_external_fluxon_kv(backend_kind=backend_kind, mode=mode) if backend_kind == TEST_STACK_BACKEND_MOONCAKE: - return mode in ( + if mooncake_storage_mode is None: + raise ValueError("Mooncake TEST_STACK requires an explicit storage mode") + if mooncake_storage_mode not in TEST_STACK_MOONCAKE_STORAGE_MODES_ALLOWED: + raise ValueError(f"invalid Mooncake TEST_STACK storage mode: {mooncake_storage_mode!r}") + return mooncake_storage_mode == TEST_STACK_MOONCAKE_STORAGE_MODE_DEDICATED_OWNER and mode in ( TEST_STACK_MODE_KVSTORE, TEST_STACK_MODE_KVSTORE_WITH_LOCAL_CACHE, TEST_STACK_MODE_RPC, @@ -5392,6 +5649,12 @@ def _parse_scene(item: Dict[str, Any], ctx: str) -> Dict[str, Any]: "write_ratio", "request_distribution", "keyspace_size", + TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY, + TEST_STACK_SCENE_KEY_KV_GET_OUTPUT, + TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX, TEST_STACK_SCENE_KEY_AFFINITY_LOCALITY_RATIO, "value_size_mode", "value_size_weighted_set", @@ -5462,6 +5725,78 @@ def _parse_scene(item: Dict[str, Any], ctx: str) -> Dict[str, Any]: min_v=1, ) + keyspace_capacity_guard = ts.get(TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD) + if keyspace_capacity_guard is not None: + if not isinstance(keyspace_capacity_guard, bool): + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD} must be bool" + ) + out_ts[TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD] = bool(keyspace_capacity_guard) + + bootstrap_concurrency = ts.get(TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY) + if bootstrap_concurrency is not None: + out_ts[TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY] = _require_int( + bootstrap_concurrency, + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY}", + min_v=1, + ) + + bootstrap_put_gap_ms = ts.get(TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS) + if bootstrap_put_gap_ms is not None: + if not isinstance(bootstrap_put_gap_ms, (int, float)): + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS} must be number" + ) + if float(bootstrap_put_gap_ms) < 0.0: + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS} must be >= 0" + ) + out_ts[TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS] = float(bootstrap_put_gap_ms) + + bootstrap_storage_full_policy = ts.get(TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY) + if bootstrap_storage_full_policy is not None: + policy = _require_str( + bootstrap_storage_full_policy, + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY}", + ).strip().lower() + if policy not in TEST_STACK_KV_BOOTSTRAP_STORAGE_FULL_POLICIES_ALLOWED: + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY} invalid: {policy!r}" + ) + out_ts[TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY] = policy + + kv_get_output_raw = ts.get(TEST_STACK_SCENE_KEY_KV_GET_OUTPUT) + kv_get_output = None + if kv_get_output_raw is not None: + if mode not in (TEST_STACK_MODE_KVSTORE, TEST_STACK_MODE_KVSTORE_WITH_LOCAL_CACHE): + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_GET_OUTPUT} " + "is only allowed for KV modes" + ) + kv_get_output = _require_str( + kv_get_output_raw, + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_GET_OUTPUT}", + ).strip().lower() + if kv_get_output not in TEST_STACK_KV_GET_OUTPUTS_ALLOWED: + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_GET_OUTPUT} " + f"invalid: {kv_get_output!r}" + ) + out_ts[TEST_STACK_SCENE_KEY_KV_GET_OUTPUT] = kv_get_output + + kv_cuda_device_index = ts.get(TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX) + if kv_cuda_device_index is not None: + if kv_get_output != "cuda": + raise ValueError( + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX} " + f"requires {TEST_STACK_SCENE_KEY_KV_GET_OUTPUT}=cuda" + ) + out_ts[TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX] = _require_int( + kv_cuda_device_index, + f"{ctx}.test_stack.{TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX}", + min_v=0, + ) + affinity_locality_ratio = ts.get(TEST_STACK_SCENE_KEY_AFFINITY_LOCALITY_RATIO) if affinity_locality_ratio is not None: if not isinstance(affinity_locality_ratio, (int, float)): @@ -5607,6 +5942,11 @@ def _parse_scene(item: Dict[str, Any], ctx: str) -> Dict[str, Any]: "write_ratio", "request_distribution", "keyspace_size", + TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS, + TEST_STACK_SCENE_KEY_KV_GET_OUTPUT, + TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX, "value_size_mode", "value_size_weighted_set", *TEST_STACK_RPC_SCENE_KEYS, @@ -5713,16 +6053,10 @@ def _parse_scale(item: Dict[str, Any], ctx: str) -> Dict[str, Any]: f"{ctx}.benchmark.processes_per_target", min_v=1, ) - threads_per_process = _require_int( + _ = _require_test_stack_benchmark_threads_per_process( bench.get("threads_per_process"), f"{ctx}.benchmark.threads_per_process", - min_v=1, ) - if int(threads_per_process) != TEST_STACK_BENCHMARK_FIXED_THREADS_PER_PROCESS: - raise ValueError( - f"{ctx}.benchmark.threads_per_process must be fixed to " - f"{TEST_STACK_BENCHMARK_FIXED_THREADS_PER_PROCESS}" - ) owner_group_processes = bench.get("owner_group_processes") if owner_group_processes is not None: _ = _require_int(owner_group_processes, f"{ctx}.benchmark.owner_group_processes", min_v=1) @@ -6012,10 +6346,106 @@ def _validate_profile_test_stack_monitoring_config(raw: Any, ctx: str) -> None: _ = _require_str(table_name, f"{ctx}.otlp_log_api.table_name") +def _normalize_test_stack_owner_kv_ssd_config(raw: Any, ctx: str) -> Optional[Dict[str, Any]]: + if raw is None: + return None + cfg = _require_dict(raw, ctx) + _forbid_unknown_keys(cfg, {"large_limit_size"}, ctx) + large_limit_size_raw = _require_list(cfg.get("large_limit_size"), f"{ctx}.large_limit_size") + if not large_limit_size_raw: + raise ValueError(f"{ctx}.large_limit_size must be non-empty") + large_limit_size: List[int] = [] + for index, raw_limit in enumerate(large_limit_size_raw): + large_limit_size.append( + _require_int( + raw_limit, + f"{ctx}.large_limit_size[{index}]", + min_v=512, + ) + ) + return {"large_limit_size": large_limit_size} + + +def _normalize_testbed_mooncake_storage_config(raw: Any, ctx: str) -> Dict[str, Any]: + cfg = _require_dict(raw, ctx) + _forbid_unknown_keys( + cfg, + {"mode", "ssd_offload_root", "ssd_capacity_bytes"}, + ctx, + ) + mode = _require_str(cfg.get("mode"), f"{ctx}.mode").strip().upper() + if mode not in TEST_STACK_MOONCAKE_STORAGE_MODES_ALLOWED: + raise ValueError( + f"{ctx}.mode must be one of {sorted(TEST_STACK_MOONCAKE_STORAGE_MODES_ALLOWED)}, " + f"got {mode!r}" + ) + + raw_ssd_root = cfg.get("ssd_offload_root") + raw_ssd_capacity = cfg.get("ssd_capacity_bytes") + if raw_ssd_root is None and raw_ssd_capacity is None: + return {"mode": mode} + if raw_ssd_root is None or raw_ssd_capacity is None: + raise ValueError( + f"{ctx}.ssd_offload_root and {ctx}.ssd_capacity_bytes must be set together" + ) + ssd_offload_root = _require_str( + raw_ssd_root, + f"{ctx}.ssd_offload_root", + ).strip() + if not ssd_offload_root: + raise ValueError(f"{ctx}.ssd_offload_root must be non-empty") + if not Path(ssd_offload_root).is_absolute(): + raise ValueError(f"{ctx}.ssd_offload_root must be absolute") + ssd_capacity_bytes = _require_int( + raw_ssd_capacity, + f"{ctx}.ssd_capacity_bytes", + min_v=1, + ) + return { + "mode": mode, + "ssd_offload_root": ssd_offload_root, + "ssd_capacity_bytes": int(ssd_capacity_bytes), + } + + +def _split_testbed_mooncake_capacity( + *, + total_bytes: int, + instance_count: int, + alignment_bytes: int, + ctx: str, +) -> int: + if instance_count <= 0: + raise ValueError(f"{ctx} instance_count must be > 0") + if total_bytes <= 0: + raise ValueError(f"{ctx} total_bytes must be > 0") + if total_bytes % instance_count != 0: + raise ValueError( + f"{ctx} total capacity must divide evenly across instances: " + f"total_bytes={total_bytes} instance_count={instance_count}" + ) + per_instance_bytes = total_bytes // instance_count + if alignment_bytes > 1 and per_instance_bytes % alignment_bytes != 0: + raise ValueError( + f"{ctx} per-instance capacity must be {alignment_bytes}-byte aligned: " + f"total_bytes={total_bytes} instance_count={instance_count} " + f"per_instance_bytes={per_instance_bytes}" + ) + return int(per_instance_bytes) + + def _validate_profile_test_stack_fluxon_runtime_config(rc: Dict[str, Any], ctx: str) -> None: _forbid_unknown_keys( rc, - {"kv_base", "mq_base", "kv_node_patch_template", "runtime_env", "monitoring", "owner_cpu_core_by_target"}, + { + "kv_base", + "mq_base", + "kv_node_patch_template", + "runtime_env", + "monitoring", + "owner_cpu_core_by_target", + "owner_kv_ssd", + }, ctx, ) kv_base = _require_dict(rc.get("kv_base"), f"{ctx}.kv_base") @@ -6031,6 +6461,7 @@ def _validate_profile_test_stack_fluxon_runtime_config(rc: Dict[str, Any], ctx: rc.get("owner_cpu_core_by_target"), f"{ctx}.owner_cpu_core_by_target", ) + _ = _normalize_test_stack_owner_kv_ssd_config(rc.get("owner_kv_ssd"), f"{ctx}.owner_kv_ssd") def _validate_profile_test_stack_redis_runtime_config(rc: Dict[str, Any], ctx: str) -> None: @@ -6081,7 +6512,7 @@ def _validate_profile_test_stack_redis_runtime_config(rc: Dict[str, Any], ctx: s def _validate_profile_test_stack_mooncake_runtime_config(rc: Dict[str, Any], ctx: str) -> None: - _forbid_unknown_keys(rc, {"kv_base", "master"}, ctx) + _forbid_unknown_keys(rc, {"kv_base", "master", "testbed_mooncake_storage"}, ctx) kv_base = _require_dict(rc.get("kv_base"), f"{ctx}.kv_base") if "rdma_device_names" in kv_base: raise ValueError(f"{ctx}.kv_base.rdma_device_names has been removed from Fluxon KV config") @@ -6093,10 +6524,12 @@ def _validate_profile_test_stack_mooncake_runtime_config(rc: Dict[str, Any], ctx _forbid_unknown_keys( mooncake_spec, { + "local_hostname", "local_buffer_size", "metadata_server", "master_server_address", "etcd_addresses", + "skip_get_size_on_get", }, f"{ctx}.kv_base.mooncake_spec", ) @@ -6105,6 +6538,10 @@ def _validate_profile_test_stack_mooncake_runtime_config(rc: Dict[str, Any], ctx f"{ctx}.kv_base.mooncake_spec.local_buffer_size", min_v=1, ) + if "local_hostname" in mooncake_spec: + raise ValueError( + f"{ctx}.kv_base.mooncake_spec.local_hostname is generated per TEST_STACK node; do not set it" + ) if "metadata_server" in mooncake_spec: raise ValueError( f"{ctx}.kv_base.mooncake_spec.metadata_server is generated by TEST_STACK Mooncake master; do not set it" @@ -6113,6 +6550,11 @@ def _validate_profile_test_stack_mooncake_runtime_config(rc: Dict[str, Any], ctx raise ValueError( f"{ctx}.kv_base.mooncake_spec.master_server_address is generated by TEST_STACK Mooncake master; do not set it" ) + if "skip_get_size_on_get" in mooncake_spec: + _ = _require_bool( + mooncake_spec.get("skip_get_size_on_get"), + f"{ctx}.kv_base.mooncake_spec.skip_get_size_on_get", + ) raw_etcd_addresses = mooncake_spec.get("etcd_addresses") if raw_etcd_addresses is not None: etcd_addresses = _require_list( @@ -6136,6 +6578,10 @@ def _validate_profile_test_stack_mooncake_runtime_config(rc: Dict[str, Any], ctx kv_base.get("benchmark_fast_path"), f"{ctx}.kv_base.benchmark_fast_path", ) + _ = _normalize_testbed_mooncake_storage_config( + rc.get("testbed_mooncake_storage"), + f"{ctx}.testbed_mooncake_storage", + ) pprof_duration_seconds = kv_base.get("pprof_duration_seconds") if pprof_duration_seconds is not None: _ = _require_int( @@ -8128,6 +8574,11 @@ def _test_stack_effective_kv_keyspace_size( if raw_keyspace_size is None: raise ValueError("scene.test_stack.keyspace_size is required for KV capacity sizing") requested_keyspace_size = _require_int(raw_keyspace_size, "scene.test_stack.keyspace_size", min_v=1) + raw_guard = ts_scene.get(TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD, True) + if not isinstance(raw_guard, bool): + raise ValueError(f"scene.test_stack.{TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD} must be bool") + if not bool(raw_guard): + return int(requested_keyspace_size) owner_scale = _require_dict(scale.get("owner"), "resolved_case.scale.owner") owner_count = _require_int(owner_scale.get("owner_count"), "resolved_case.scale.owner.owner_count", min_v=1) owner_dram_bytes = _require_int( @@ -8426,6 +8877,7 @@ def _build_test_stack_external_kv_owner_instances( runtime_env: Dict[str, str], owner_group_processes: Optional[int], owner_cpu_core_by_target: Dict[str, int], + owner_kv_ssd: Optional[Dict[str, Any]], ) -> List[Dict[str, Any]]: required_python_abi = _test_stack_runtime_required_python_abi( resolved_case=resolved_case, @@ -8476,6 +8928,25 @@ def _build_test_stack_external_kv_owner_instances( ) owner_services_dir = run_dir / "services" / "kv_owner" / target_slug owner_large_file_paths = _fluxon_kv_owner_large_file_paths(owner_work_root=owner_services_dir) + owner_fluxonkv_spec = { + "etcd_addresses": list(etcd_endpoints), + "cluster_name": cluster_name, + "share_mem_path": owner_share_mem_path, + "large_file_paths": owner_large_file_paths, + "sub_cluster": FLUXON_KV_OWNER_SUB_CLUSTER, + "p2p_listen_port": int(owner_p2p_listen_port), + } + if owner_kv_ssd is not None: + owner_large_limit_size = _require_list( + owner_kv_ssd.get("large_limit_size"), + "profile.test_stack.runtime_config.owner_kv_ssd.large_limit_size", + ) + if len(owner_large_limit_size) != len(owner_large_file_paths): + raise ValueError( + "profile.test_stack.runtime_config.owner_kv_ssd.large_limit_size length must match " + "generated owner large_file_paths length" + ) + owner_fluxonkv_spec["large_limit_size"] = copy.deepcopy(owner_large_limit_size) owner_cfg = { "instance_key": _test_stack_kv_owner_runtime_instance_key( runtime_instance_prefix=runtime_instance_prefix, @@ -8483,14 +8954,7 @@ def _build_test_stack_external_kv_owner_instances( ctx="external kv owner", ), "contribute_to_cluster_pool_size": {"dram": int(owner_dram_bytes), "vram": {}}, - "fluxonkv_spec": { - "etcd_addresses": list(etcd_endpoints), - "cluster_name": cluster_name, - "share_mem_path": owner_share_mem_path, - "large_file_paths": owner_large_file_paths, - "sub_cluster": FLUXON_KV_OWNER_SUB_CLUSTER, - "p2p_listen_port": int(owner_p2p_listen_port), - }, + "fluxonkv_spec": owner_fluxonkv_spec, } owner_protocol = kv_base.get("protocol") if owner_protocol is not None: @@ -8563,6 +9027,7 @@ def _build_test_stack_mooncake_owner_instances( test_spec_config: Dict[str, Any], perf_config: Optional[Dict[str, Any]], runtime_env: Dict[str, str], + testbed_mooncake_storage: Dict[str, Any], ) -> List[Dict[str, Any]]: required_python_abi = _test_stack_runtime_required_python_abi( resolved_case=resolved_case, @@ -8575,6 +9040,15 @@ def _build_test_stack_mooncake_owner_instances( "TEST_STACK Mooncake requires owner_count to match owner target list length: " f"owner_count={owner_count} targets={len(owner_targets)}" ) + storage_mode = _require_str( + testbed_mooncake_storage.get("mode"), + "profile.test_stack.runtime_config.testbed_mooncake_storage.mode", + ) + if storage_mode != TEST_STACK_MOONCAKE_STORAGE_MODE_DEDICATED_OWNER: + raise ValueError( + "Mooncake owner instances require testbed_mooncake_storage.mode=" + f"{TEST_STACK_MOONCAKE_STORAGE_MODE_DEDICATED_OWNER}, got {storage_mode!r}" + ) owner_dram_bytes = _require_int( owner_scale.get("owner_dram_bytes"), "scale.owner.owner_dram_bytes", @@ -8609,9 +9083,36 @@ def _build_test_stack_mooncake_owner_instances( "profile.test_stack.runtime_config.kv_base.mooncake_spec.local_buffer_size", min_v=16777216, ) + require_unlimited_memlock = _test_stack_protocol_requires_unlimited_memlock( + _require_dict( + kv_base.get("protocol"), + "profile.test_stack.runtime_config.kv_base.protocol", + ) + ) owner_instances: List[Dict[str, Any]] = [] + local_ipv4_addrs = _local_ipv4_addresses() + ssd_offload_root_raw = testbed_mooncake_storage.get("ssd_offload_root") + ssd_capacity_raw = testbed_mooncake_storage.get("ssd_capacity_bytes") + per_owner_ssd_capacity: Optional[int] = None + if ssd_offload_root_raw is not None: + per_owner_ssd_capacity = _split_testbed_mooncake_capacity( + total_bytes=_require_int( + ssd_capacity_raw, + "profile.test_stack.runtime_config.testbed_mooncake_storage.ssd_capacity_bytes", + min_v=1, + ), + instance_count=owner_count, + alignment_bytes=1, + ctx="Mooncake dedicated-owner SSD", + ) for target in owner_targets: + target_node_cfg = _require_dict(cluster_nodes.get(target), f"cluster_nodes[{target}]") + target_endpoint_host = _test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name=target, + local_ipv4_addrs=local_ipv4_addrs, + ) target_slug, instance_id = _test_stack_kv_owner_instance_id( owner_target=target, ctx="mooncake owner", @@ -8627,6 +9128,7 @@ def _build_test_stack_mooncake_owner_instances( "mooncake_spec": { # Mooncake owner/server should contribute storage capacity only. # The benchmark node keeps the local transfer buffer; owner stays pure server mode. + "local_hostname": target_endpoint_host, "local_buffer_size": 0, "metadata_server": metadata_server, "master_server_address": master_server_address, @@ -8639,6 +9141,27 @@ def _build_test_stack_mooncake_owner_instances( ) ), } + owner_pre_exec_shell = "" + owner_runtime_env = copy.deepcopy(runtime_env) + if ssd_offload_root_raw is not None: + owner_ssd_path = str( + (Path(_require_str(ssd_offload_root_raw, "ssd_offload_root")) / target_slug).resolve() + ) + owner_cfg["mooncake_spec"].update( + { + "enable_ssd_offload": True, + "ssd_offload_path": owner_ssd_path, + } + ) + owner_pre_exec_shell = ( + "mkdir -p " + + _shell_quote(owner_ssd_path) + + "\n" + ) + assert per_owner_ssd_capacity is not None + owner_runtime_env[TEST_STACK_MOONCAKE_SSD_LIMIT_ENV] = str( + per_owner_ssd_capacity + ) pprof_duration_seconds = kv_base.get("pprof_duration_seconds") if pprof_duration_seconds is not None: owner_cfg["pprof_duration_seconds"] = int(pprof_duration_seconds) @@ -8671,7 +9194,9 @@ def _build_test_stack_mooncake_owner_instances( "-w", str(owner_services_dir.resolve()), ], - runtime_env=runtime_env, + runtime_env=owner_runtime_env, + pre_exec_shell=owner_pre_exec_shell, + require_unlimited_memlock=require_unlimited_memlock, ) if perf_config is not None and "owner" in _require_list(perf_config.get("targets"), "perf_config.targets"): owner_cmd = _test_stack_perf_wrapper_command( @@ -8872,7 +9397,11 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - else: coord_target_str = _require_str(raw_coord_target, "coordinator_template.deployer.target") coord_target = coordinator_target if coord_target_str == "__TARGET__" else coord_target_str - coord_ip = _require_str(target_ip_map.get(coord_target), "deploy.target_ip_map[coordinator_target]") + coord_ip = _test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name=coord_target, + local_ipv4_addrs=_local_ipv4_addresses(), + ) coordinator_addr = f"{coord_ip}:{coordinator_port}" bench = _require_dict(scale.get("benchmark"), "scale.benchmark") @@ -8881,16 +9410,10 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - "scale.benchmark.processes_per_target", min_v=1, ) - threads_per_process = _require_int( + threads_per_process = _require_test_stack_benchmark_threads_per_process( bench.get("threads_per_process"), "scale.benchmark.threads_per_process", - min_v=1, ) - if int(threads_per_process) != TEST_STACK_BENCHMARK_FIXED_THREADS_PER_PROCESS: - raise ValueError( - "scale.benchmark.threads_per_process must be fixed to " - f"{TEST_STACK_BENCHMARK_FIXED_THREADS_PER_PROCESS}" - ) value_size = _require_int(bench.get("value_size"), "scale.benchmark.value_size", min_v=0) warmup = bench.get("metric_warmup_seconds") if not isinstance(warmup, (int, float)): @@ -8963,11 +9486,24 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - target_ip_map=target_ip_map, ctx="resolved_case.scale.owner.targets", ) + rc = _require_dict(ts_profile.get("runtime_config"), "profile.test_stack.runtime_config") + testbed_mooncake_storage: Optional[Dict[str, Any]] = None + mooncake_storage_mode: Optional[str] = None + if backend_kind == TEST_STACK_BACKEND_MOONCAKE: + testbed_mooncake_storage = _normalize_testbed_mooncake_storage_config( + rc.get("testbed_mooncake_storage"), + "profile.test_stack.runtime_config.testbed_mooncake_storage", + ) + mooncake_storage_mode = _require_str( + testbed_mooncake_storage.get("mode"), + "profile.test_stack.runtime_config.testbed_mooncake_storage.mode", + ) needs_kv_master = _test_stack_backend_requires_fluxon_kv_master(backend_kind=backend_kind, mode=scene_mode) uses_external_fluxon_kv = _test_stack_backend_uses_external_fluxon_kv(backend_kind=backend_kind, mode=scene_mode) uses_dedicated_kv_owners = _test_stack_backend_uses_dedicated_kv_owners( backend_kind=backend_kind, mode=scene_mode, + mooncake_storage_mode=mooncake_storage_mode, ) _require_explicit_owner_group_processes_for_multi_owner_same_machine( owner_targets_by_machine=owner_targets_by_machine, @@ -9003,6 +9539,11 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - node_instances: List[Dict[str, Any]] = [] node_roles: List[str] = [] node_overrides: List[Dict[str, Any]] = [] + mooncake_local_ipv4_addrs = ( + _local_ipv4_addresses() + if backend_kind == TEST_STACK_BACKEND_MOONCAKE + else set() + ) stack_cluster_name: Optional[str] = None stack_share_mem_path: Optional[str] = None if backend_kind == TEST_STACK_BACKEND_FLUXON: @@ -9016,13 +9557,15 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - "runtime.stack_identity.share_mem_path", ) - rc = _require_dict(ts_profile.get("runtime_config"), "profile.test_stack.runtime_config") kv_base: Dict[str, Any] kv_node_patch_template: Dict[str, Any] = {} mq_base: Optional[Any] = None rc_redis: Optional[Dict[str, Any]] = None rc_alluxio: Optional[Dict[str, Any]] = None rc_mooncake_master: Optional[Dict[str, Any]] = None + owner_kv_ssd: Optional[Dict[str, Any]] = None + mooncake_per_benchmark_process_dram_bytes: Optional[int] = None + mooncake_per_benchmark_process_ssd_capacity_bytes: Optional[int] = None owner_cpu_core_by_target: Dict[str, int] = {} alluxio_mount_root_by_target: Dict[str, str] = {} redis_instances: List[Dict[str, Any]] = [] @@ -9031,6 +9574,10 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - mooncake_metadata_port: Optional[int] = None mooncake_metrics_port: Optional[int] = None if backend_kind == TEST_STACK_BACKEND_FLUXON: + owner_kv_ssd = _normalize_test_stack_owner_kv_ssd_config( + rc.get("owner_kv_ssd"), + "profile.test_stack.runtime_config.owner_kv_ssd", + ) owner_cpu_core_by_target = _normalize_owner_cpu_core_by_target( rc.get("owner_cpu_core_by_target"), "profile.test_stack.runtime_config.owner_cpu_core_by_target", @@ -9189,6 +9736,8 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - ) kv_base["instance_key"] = f"{runtime_instance_prefix}__bench_base" elif backend_kind == TEST_STACK_BACKEND_MOONCAKE: + assert testbed_mooncake_storage is not None + assert mooncake_storage_mode is not None kv_base = copy.deepcopy(_require_dict(rc.get("kv_base"), "profile.test_stack.runtime_config.kv_base")) if "fluxonkv_spec" in kv_base: raise ValueError( @@ -9212,6 +9761,35 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - ) kv_base["backend_kind"] = TEST_STACK_BACKEND_MOONCAKE kv_base["instance_key"] = f"{runtime_instance_prefix}__bench_base" + if mooncake_storage_mode == TEST_STACK_MOONCAKE_STORAGE_MODE_PER_BENCHMARK_PROCESS: + owner_scale = _require_dict(scale.get("owner"), "resolved_case.scale.owner") + owner_count = _require_int( + owner_scale.get("owner_count"), + "resolved_case.scale.owner.owner_count", + min_v=1, + ) + owner_dram_bytes = _require_int( + owner_scale.get("owner_dram_bytes"), + "resolved_case.scale.owner.owner_dram_bytes", + min_v=TEST_STACK_MOONCAKE_SEGMENT_ALIGNMENT_BYTES, + ) + mooncake_per_benchmark_process_dram_bytes = _split_testbed_mooncake_capacity( + total_bytes=int(owner_count) * int(owner_dram_bytes), + instance_count=node_total, + alignment_bytes=TEST_STACK_MOONCAKE_SEGMENT_ALIGNMENT_BYTES, + ctx="Mooncake per-benchmark-process DRAM", + ) + if testbed_mooncake_storage.get("ssd_offload_root") is not None: + mooncake_per_benchmark_process_ssd_capacity_bytes = _split_testbed_mooncake_capacity( + total_bytes=_require_int( + testbed_mooncake_storage.get("ssd_capacity_bytes"), + "profile.test_stack.runtime_config.testbed_mooncake_storage.ssd_capacity_bytes", + min_v=1, + ), + instance_count=node_total, + alignment_bytes=1, + ctx="Mooncake per-benchmark-process SSD", + ) mooncake_rpc_port_base = _require_int( port_alloc.get("mooncake_rpc_port_base"), "profile.test_stack.port_alloc.mooncake_rpc_port_base", @@ -9280,6 +9858,7 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - mooncake_spec["metadata_server"] = f"http://{mooncake_host}:{int(mooncake_metadata_port)}/metadata" mooncake_spec["master_server_address"] = f"{mooncake_host}:{int(mooncake_rpc_port)}" mooncake_spec["etcd_addresses"] = list(etcd_endpoints) + mooncake_spec.setdefault("skip_get_size_on_get", True) raw_protocol = kv_base.get("protocol") if raw_protocol is None: kv_base["protocol"] = {"protocol_type": "tcp"} @@ -9344,6 +9923,15 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - kv_base["test_spec_config"] = copy.deepcopy(runtime_test_spec_config) if runtime_protocol_cfg is not None: kv_base["protocol"] = runtime_protocol_cfg + mooncake_requires_unlimited_memlock = ( + backend_kind == TEST_STACK_BACKEND_MOONCAKE + and _test_stack_protocol_requires_unlimited_memlock( + _require_dict( + kv_base.get("protocol"), + "profile.test_stack.runtime_config.kv_base.protocol", + ) + ) + ) # mq_new_or_bind_unique_key is used to derive etcd keys for MPMC channel creation, # including a distributed lock key "_lock". # @@ -9503,6 +10091,7 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - module_name="mooncake.cli", module_args=mooncake_args, runtime_env=runtime_env, + require_unlimited_memlock=mooncake_requires_unlimited_memlock, ) if perf_config is not None and TEST_STACK_MOONCAKE_MASTER_INSTANCE_ID in _require_list( perf_config.get("targets"), @@ -9590,8 +10179,10 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - runtime_env=runtime_env, owner_group_processes=owner_group_processes, owner_cpu_core_by_target=owner_cpu_core_by_target, + owner_kv_ssd=owner_kv_ssd, ) elif backend_kind == TEST_STACK_BACKEND_MOONCAKE and uses_dedicated_kv_owners: + assert testbed_mooncake_storage is not None owner_instances = _build_test_stack_mooncake_owner_instances( resolved_case=resolved_case, scale=scale, @@ -9605,6 +10196,7 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - test_spec_config=runtime_test_spec_config, perf_config=perf_config, runtime_env=runtime_env, + testbed_mooncake_storage=testbed_mooncake_storage, ) rpc_server_instance_keys: List[str] = [] @@ -9655,6 +10247,32 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - "__COORDINATOR__": coordinator_addr, "__TARGET__": target, } + node_runtime_env = copy.deepcopy(runtime_env) + node_pre_exec_shell = "" + mooncake_node_ssd_path: Optional[str] = None + if ( + backend_kind == TEST_STACK_BACKEND_MOONCAKE + and mooncake_storage_mode + == TEST_STACK_MOONCAKE_STORAGE_MODE_PER_BENCHMARK_PROCESS + ): + assert testbed_mooncake_storage is not None + ssd_offload_root = testbed_mooncake_storage.get("ssd_offload_root") + if ssd_offload_root is not None: + mooncake_node_ssd_path = str( + ( + Path(_require_str(ssd_offload_root, "ssd_offload_root")) + / runtime_instance_key + ).resolve() + ) + node_pre_exec_shell = ( + "mkdir -p " + + _shell_quote(mooncake_node_ssd_path) + + "\n" + ) + assert mooncake_per_benchmark_process_ssd_capacity_bytes is not None + node_runtime_env[TEST_STACK_MOONCAKE_SSD_LIMIT_ENV] = str( + mooncake_per_benchmark_process_ssd_capacity_bytes + ) inst = _require_dict(_subst_obj_tokens(node_tpl, mapping, "node_template"), "node_template.compiled") inst["id"] = instance_key inst_deployer = _require_dict(inst.get("deployer"), f"node_template[{instance_key}].deployer") @@ -9680,7 +10298,9 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - "--coordinator", coordinator_addr, ], - runtime_env=runtime_env, + runtime_env=node_runtime_env, + pre_exec_shell=node_pre_exec_shell, + require_unlimited_memlock=mooncake_requires_unlimited_memlock, ) ] node_instances.append(inst) @@ -9744,6 +10364,31 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - f"profile.test_stack.runtime_config.alluxio.mount_root_by_target[{target!r}]", ) } + elif backend_kind == TEST_STACK_BACKEND_MOONCAKE: + mooncake_spec_override: Dict[str, Any] = { + "local_hostname": _test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name=target, + local_ipv4_addrs=mooncake_local_ipv4_addrs, + ) + } + if ( + mooncake_storage_mode + == TEST_STACK_MOONCAKE_STORAGE_MODE_PER_BENCHMARK_PROCESS + ): + assert mooncake_per_benchmark_process_dram_bytes is not None + kv["contribute_to_cluster_pool_size"] = { + "dram": mooncake_per_benchmark_process_dram_bytes, + "vram": {}, + } + if mooncake_node_ssd_path is not None: + mooncake_spec_override.update( + { + "enable_ssd_offload": True, + "ssd_offload_path": mooncake_node_ssd_path, + } + ) + kv["mooncake_spec"] = mooncake_spec_override rec: Dict[str, Any] = {"kv": kv} if scene_mode == TEST_STACK_MODE_MPMC: rec["mq_role"] = role @@ -9806,6 +10451,7 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - script_path=test_stack_runtime["coordinator_script"], script_args=[], runtime_env=runtime_env, + require_unlimited_memlock=mooncake_requires_unlimited_memlock, ) ] @@ -9857,6 +10503,11 @@ def _compile_test_stack_case(resolved_case: Dict[str, Any], *, run_index: int) - "read_ratio", "write_ratio", "request_distribution", + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS, + TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY, + TEST_STACK_SCENE_KEY_KV_GET_OUTPUT, + TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX, TEST_STACK_SCENE_KEY_AFFINITY_LOCALITY_RATIO, ): if optional_key in ts_scene: @@ -10148,6 +10799,22 @@ def _cluster_node_is_local_host( return False +def _test_stack_target_advertise_host( + *, + cluster_nodes: Dict[str, Dict[str, Any]], + target_name: str, + local_ipv4_addrs: set[str], +) -> str: + node_cfg = _require_dict(cluster_nodes.get(target_name), f"cluster_nodes[{target_name}]") + if _cluster_node_is_local_host( + node_cfg, + target_name=target_name, + local_ipv4_addrs=local_ipv4_addrs, + ): + return "127.0.0.1" + return _require_str(node_cfg.get("ip"), f"cluster_nodes[{target_name}].ip") + + def _ensure_path_symlink(*, link_path: Path, target_path: Path) -> None: link_path = link_path.resolve(strict=False) target_path = target_path.resolve() @@ -10862,6 +11529,13 @@ def _ci_cleanup_runtime( ctx=f"CI cleanup runtime current_deployments {instance_id_text}", ) _wait_ci_ports_free(cleanup_case, timeout_s=timeout_s) + if _ci_has_instance(cleanup_case, instance_id="owner_0"): + owner_share_mem_path = _ci_share_mem_path(cleanup_case, owner_index=0) + _cleanup_ci_owner_share_mem_dir( + cleanup_case, + owner_index=0, + share_mem_path=owner_share_mem_path, + ) def _http_json_allow_error_status(req: urllib.request.Request) -> Tuple[int, Dict[str, Any]]: @@ -12314,13 +12988,20 @@ def _test_stack_runtime_command( script_path: str, script_args: List[str], runtime_env: Dict[str, str], + pre_exec_shell: str = "", + require_unlimited_memlock: bool = False, exec_wrapper_argv: Optional[List[str]] = None, ) -> str: + mooncake_wheel = _test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) assert_env_cmd = _test_stack_runtime_env_assert_command( venv_python=venv_python, - require_mooncake=_test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) is not None, + require_mooncake=mooncake_wheel is not None, + ) + mooncake_ld_library_path_cmd = ( + _test_stack_mooncake_cuda_runtime_ld_library_path_command(venv_python=venv_python) + if mooncake_wheel is not None + else "" ) - require_unlimited_memlock = _test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) is not None vendor_site_packages = ( run_dir / _TEST_STACK_RUNTIME_DIRNAME / _TEST_STACK_RUNTIME_VENDOR_SITE_PACKAGES_DIRNAME ).resolve() @@ -12362,10 +13043,12 @@ def _test_stack_runtime_command( + "export PYTHONPATH=" + _shell_quote(str(vendor_site_packages)) + ":${PYTHONPATH:-}\n" + + mooncake_ld_library_path_cmd + _test_stack_runtime_env_prelude_command() + runtime_env_exports + assert_env_cmd + "\n" + + pre_exec_shell + " ".join(exec_cmd) + "\n" ) @@ -12379,13 +13062,19 @@ def _test_stack_runtime_module_command( module_args: List[str], runtime_env: Dict[str, str], pre_exec_shell: str = "", + require_unlimited_memlock: bool = False, exec_wrapper_argv: Optional[List[str]] = None, ) -> str: + mooncake_wheel = _test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) assert_env_cmd = _test_stack_runtime_env_assert_command( venv_python=venv_python, - require_mooncake=_test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) is not None, + require_mooncake=mooncake_wheel is not None, + ) + mooncake_ld_library_path_cmd = ( + _test_stack_mooncake_cuda_runtime_ld_library_path_command(venv_python=venv_python) + if mooncake_wheel is not None + else "" ) - require_unlimited_memlock = _test_stack_runtime_mooncake_wheel_path_opt(run_dir=run_dir) is not None vendor_site_packages = ( run_dir / _TEST_STACK_RUNTIME_DIRNAME / _TEST_STACK_RUNTIME_VENDOR_SITE_PACKAGES_DIRNAME ).resolve() @@ -12428,6 +13117,7 @@ def _test_stack_runtime_module_command( + "export PYTHONPATH=" + _shell_quote(str(vendor_site_packages)) + ":${PYTHONPATH:-}\n" + + mooncake_ld_library_path_cmd + _test_stack_runtime_env_prelude_command() + runtime_env_exports + assert_env_cmd @@ -12615,6 +13305,36 @@ def _test_stack_runtime_env_prelude_command() -> str: return "" +def _test_stack_mooncake_cuda_runtime_ld_library_path_command(*, venv_python: Path) -> str: + return ( + "CUDA_RUNTIME_LIB_DIR=\"$(" + + _shell_quote(str(venv_python)) + + " - <<'PY'\n" + + "from pathlib import Path\n" + + "import site\n" + + "roots = []\n" + + "for getter in (site.getsitepackages,):\n" + + " try:\n" + + " roots.extend(getter())\n" + + " except Exception:\n" + + " pass\n" + + "try:\n" + + " roots.append(site.getusersitepackages())\n" + + "except Exception:\n" + + " pass\n" + + "for root in roots:\n" + + " candidate = Path(root) / 'nvidia' / 'cuda_runtime' / 'lib'\n" + + " if (candidate / 'libcudart.so.12').exists():\n" + + " print(candidate)\n" + + " break\n" + + "PY\n" + + ")\"\n" + + "if [ -n \"${CUDA_RUNTIME_LIB_DIR}\" ]; then\n" + + " export LD_LIBRARY_PATH=\"${CUDA_RUNTIME_LIB_DIR}:${LD_LIBRARY_PATH:-}\"\n" + + "fi\n" + ) + + def _test_stack_runtime_env_import_probe(*, require_mooncake: bool) -> str: probe = _TEST_STACK_HOST_VENV_IMPORT_PROBE_BASE if require_mooncake: @@ -12692,6 +13412,11 @@ def _test_stack_runtime_env_prepare_command( required_python_abi=required_python_abi, require_mooncake=mooncake_wheel is not None, ) + mooncake_ld_library_path_cmd = ( + _test_stack_mooncake_cuda_runtime_ld_library_path_command(venv_python=venv_python) + if mooncake_wheel is not None + else "" + ) python_bin_name = "python" + required_python_abi.removeprefix("cpython") venv_root = venv_python.parent.parent expected_sha256_inputs = _shell_quote(str(wheel)) @@ -12796,6 +13521,7 @@ def _test_stack_runtime_env_prepare_command( + _shell_quote(str(marker_path)) + "\n" + "fi\n" + + mooncake_ld_library_path_cmd + assert_env_cmd ) @@ -14344,7 +15070,7 @@ def _run_remote_bash_capture( argv=argv, password=None, ctx=ctx, - timeout_seconds=None, + timeout_seconds=_SSH_TRANSPORT_TIMEOUT_SECONDS, emit_output=False, max_attempts=1, ) diff --git a/fluxon_test_stack/test_runner_runtime_backend.py b/fluxon_test_stack/test_runner_runtime_backend.py index 30d1191..c5cbdda 100644 --- a/fluxon_test_stack/test_runner_runtime_backend.py +++ b/fluxon_test_stack/test_runner_runtime_backend.py @@ -64,8 +64,15 @@ def _prepare_ci_case( services_root = (run_dir / "services").resolve() services_root.mkdir(parents=True, exist_ok=True) (services_root / "share_mem").mkdir(parents=True, exist_ok=True) - share_mem_path = ctx._ci_share_mem_path(resolved_case, run_dir=run_dir) - Path(share_mem_path).mkdir(parents=True, exist_ok=True) + share_mem_path = ctx._ci_share_mem_path(resolved_case, owner_index=0) + if ctx._ci_has_instance(resolved_case, instance_id="owner_0"): + # One fixed owner slot is reused across all CI cases. Reset only this CI-scoped + # directory after stale deployments are gone; never touch the testbed owner's root. + ctx._reset_ci_owner_share_mem_dir( + resolved_case, + owner_index=0, + share_mem_path=share_mem_path, + ) venv_python = ctx._create_ci_runtime_venv(run_dir=run_dir) @@ -190,10 +197,28 @@ def _prepare_test_stack_case( profile_test_stack.get("kind"), "resolved_case.profile.test_stack.kind", ) + mooncake_storage_mode: Optional[str] = None + if backend_kind == ctx.TEST_STACK_BACKEND_MOONCAKE: + runtime_config = ctx._require_dict( + profile_test_stack.get("runtime_config"), + "resolved_case.profile.test_stack.runtime_config", + ) + mooncake_storage = ctx._normalize_testbed_mooncake_storage_config( + runtime_config.get("testbed_mooncake_storage"), + "resolved_case.profile.test_stack.runtime_config.testbed_mooncake_storage", + ) + mooncake_storage_mode = ctx._require_str( + mooncake_storage.get("mode"), + "resolved_case.profile.test_stack.runtime_config.testbed_mooncake_storage.mode", + ) owner_instance_ids: List[str] = [] share_mem_path: Optional[str] = None stack_cluster_name: Optional[str] = None - if ctx._test_stack_backend_uses_dedicated_kv_owners(backend_kind=backend_kind, mode=mode): + if ctx._test_stack_backend_uses_dedicated_kv_owners( + backend_kind=backend_kind, + mode=mode, + mooncake_storage_mode=mooncake_storage_mode, + ): runtime = ctx._require_dict(resolved_case.get("runtime"), "resolved_case.runtime") owner_instance_ids = [ iid @@ -321,7 +346,11 @@ def _prepare_test_stack_case( timeout_s=int(cluster_ready_timeout_seconds), ctx="TEST_STACK owner shared bundle", ) - elif ctx._test_stack_backend_uses_dedicated_kv_owners(backend_kind=backend_kind, mode=mode): + elif ctx._test_stack_backend_uses_dedicated_kv_owners( + backend_kind=backend_kind, + mode=mode, + mooncake_storage_mode=mooncake_storage_mode, + ): for owner_id in owner_instance_ids: ctx._wait_instance_running(resolved_case, instance_id=owner_id, timeout_s=60) ctx._stage_runtime_phase_run_dir(resolved_case, run_dir=run_dir, phase=node_runtime_phase) diff --git a/fluxon_test_stack/tests/test_benchmark_node_kv_contract.py b/fluxon_test_stack/tests/test_benchmark_node_kv_contract.py new file mode 100644 index 0000000..55f8636 --- /dev/null +++ b/fluxon_test_stack/tests/test_benchmark_node_kv_contract.py @@ -0,0 +1,691 @@ +from __future__ import annotations + +import ctypes +import gc +import logging +import struct +import sys +import threading +import time +import unittest +import weakref +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +TEST_STACK_DIR = REPO_ROOT / "fluxon_test_stack" +sys.path[:0] = [str(TEST_STACK_DIR), str(REPO_ROOT)] +try: + from fluxon_test_stack import benchmark_node_kv as _KV +finally: + for import_path in (str(TEST_STACK_DIR), str(REPO_ROOT)): + if import_path in sys.path: + sys.path.remove(import_path) + + +class _PutResult: + def __init__(self, *, success: bool, error_msg: str | None = None) -> None: + self.success = success + self.error_msg = error_msg + + +class _FakeHolder: + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.access_count = 0 + + def access(self): + self.access_count += 1 + return _KV._SimpleResult.ok({"payload": self.payload}) + + def _benchmark_source_kind(self) -> str: + return "memory" + + +class _FakeKvStore: + def __init__(self, holder: _FakeHolder) -> None: + self.holder = holder + self.put_value = None + self.put_keys: list[str] = [] + self.put_opts: list[object] = [] + + def get_blocking(self, key: str): + del key + return _KV._SimpleResult.ok(self.holder) + + def put_blocking(self, key: str, value, *, opts): + self.put_keys.append(key) + self.put_opts.append(opts) + self.put_value = value + return _KV._SimpleResult.ok(None) + + def close(self): + return _KV._SimpleResult.ok(None) + + +class _FakeCudaRuntime: + def __init__(self) -> None: + self._next_handle = 1 + self._host_buffers: dict[int, ctypes.Array] = {} + self._device_buffers: dict[int, ctypes.Array] = {} + self._event_ready: dict[int, bool] = {} + self.fail_next_event_synchronize = False + self.event_synchronize_started = threading.Event() + self.event_synchronize_gate: threading.Event | None = None + self.copied_payloads: list[bytes] = [] + + def _handle(self) -> int: + handle = self._next_handle + self._next_handle += 1 + return handle + + def set_device(self, device_index: int) -> None: + if device_index < 0: + raise ValueError("device index must be non-negative") + + def malloc(self, size: int) -> int: + buffer = ctypes.create_string_buffer(size) + ptr = ctypes.addressof(buffer) + self._device_buffers[ptr] = buffer + return ptr + + def free(self, ptr: int) -> None: + self._device_buffers.pop(ptr) + + def host_alloc(self, size: int) -> int: + buffer = ctypes.create_string_buffer(size) + ptr = ctypes.addressof(buffer) + self._host_buffers[ptr] = buffer + return ptr + + def host_free(self, ptr: int) -> None: + self._host_buffers.pop(ptr) + + def stream_create(self) -> int: + return self._handle() + + def stream_synchronize(self, stream: int) -> None: + del stream + + def stream_destroy(self, stream: int) -> None: + del stream + + def event_create(self) -> int: + event = self._handle() + self._event_ready[event] = False + return event + + def event_record(self, event: int, stream: int) -> None: + del stream + self._event_ready[event] = False + + def event_query(self, event: int) -> bool: + return self._event_ready[event] + + def event_synchronize(self, event: int) -> None: + self.event_synchronize_started.set() + if self.event_synchronize_gate is not None: + if not self.event_synchronize_gate.wait(timeout=1.0): + raise RuntimeError("timed out waiting for fake CUDA event gate") + if self.fail_next_event_synchronize: + self.fail_next_event_synchronize = False + raise RuntimeError("injected CUDA event failure") + self._event_ready[event] = True + + def event_elapsed_us(self, start_event: int, done_event: int) -> float: + del start_event, done_event + return 25.0 + + def complete_all_events(self) -> None: + for event in self._event_ready: + self._event_ready[event] = True + + def event_destroy(self, event: int) -> None: + self._event_ready.pop(event) + + def memcpy_host_to_device_async( + self, + *, + device_ptr: int, + host_ptr: int, + size: int, + stream: int, + ) -> None: + del stream + ctypes.memmove(device_ptr, host_ptr, size) + self.copied_payloads.append(ctypes.string_at(device_ptr, size)) + + +class _KeepAlive: + pass + + +class _FakeMooncakeBuffer(bytearray): + def ptr(self) -> int: + return ctypes.addressof(ctypes.c_ubyte.from_buffer(self)) + + def size(self) -> int: + return len(self) + + +class _FakeMooncakeStore: + def __init__(self, encoded_payload: bytes, *, on_get=None) -> None: + self._encoded_payload = encoded_payload + self._on_get = on_get + self.get_count = 0 + + def get_buffer_blocking(self, key: str): + del key + self.get_count += 1 + if self._on_get is not None: + self._on_get(self.get_count) + return _KV._SimpleResult.ok(_FakeMooncakeBuffer(self._encoded_payload)) + + +class _FakeCounterStore: + def __init__(self, counts: list[int]) -> None: + self._counts = iter(counts) + + def _benchmark_offload_rpc_read_count(self) -> int: + return next(self._counts) + + +class _FakeOperationResult: + def __init__(self, **kwargs) -> None: + for key, value in kwargs.items(): + setattr(self, key, value) + self.node_id = None + self.worker_id = None + self.finish_ts = 0.0 + + +class _FakeOperationOutcome: + SUCCESS = "success" + ERROR = "error" + CACHE_HIT = "cache_hit" + CACHE_MISS = "cache_miss" + + +class _FakeCudaBenchmarkNode: + def __init__(self, adapter, *, stop_after_gets: int) -> None: + self.test_config = { + "test_mode": _KV.TEST_MODE_KVSTORE, + "node_role": "worker", + "key_prefix": "bench", + "read_ratio": 1.0, + "write_ratio": 0.0, + "request_distribution": _KV.REQUEST_DISTRIBUTION_UNIFORM, + "keyspace_size": 1, + "op_timeout_seconds": 1.0, + "value_size": 12, + _KV.BENCHMARK_KEY_KV_GET_OUTPUT: _KV.KVGetOutput.CUDA.value, + } + self.kv_store = adapter + self.key_prefix = "bench" + self.value_size_mode = _KV.VALUE_SIZE_MODE_FIXED + self.node_id = "node0" + self.instance_key = "instance0" + self.start_time = None + self._benchmark_stop = threading.Event() + self._inflight = 0 + self.progress: list[int] = [] + self.stop_after_gets = int(stop_after_gets) + + def _resolve_kv_value_size(self, thread_id: int, key_idx: int) -> int: + del thread_id, key_idx + return 12 + + def _generate_test_data(self, value_size: int) -> bytes: + return b"x" * value_size + + def _inflight_begin(self) -> int: + self._inflight += 1 + return self._inflight + + def _inflight_end(self) -> int: + self._inflight -= 1 + if self._inflight < 0: + raise AssertionError("negative in-flight count") + return self._inflight + + def _mark_progress( + self, + *, + thread_id: int, + op_idx: int, + finish_ts: float, + latency_us: float, + ) -> None: + del thread_id, finish_ts, latency_us + self.progress.append(op_idx) + + def _get_single_operation(self, *args, **kwargs): + del args, kwargs + raise AssertionError("CUDA worker must use the pipelined GET path") + + +class _FakeBenchmarkNode: + def __init__(self) -> None: + self.test_config = { + "test_mode": _KV.TEST_MODE_KVSTORE, + "node_role": "worker", + "key_prefix": "bench", + "read_ratio": 0.9, + "write_ratio": 0.1, + "request_distribution": _KV.REQUEST_DISTRIBUTION_UNIFORM, + "keyspace_size": 5, + "kv_bootstrap_before_ready": True, + "kv_bootstrap_concurrency": 1, + "kv_bootstrap_storage_full_policy": _KV.KV_BOOTSTRAP_STORAGE_FULL_POLICY_STOP, + "op_timeout_seconds": 1, + "value_size_mode": _KV.VALUE_SIZE_MODE_FIXED, + "value_size": 1, + } + self.kv_store = object() + self.key_prefix = "bench" + self.value_size_mode = _KV.VALUE_SIZE_MODE_FIXED + self.node_id = "node0" + self.put_keys: list[str] = [] + + def _resolve_kv_value_size(self, thread_id: int, key_idx: int) -> int: + del thread_id, key_idx + return 1 + + def _generate_test_data(self, value_size: int) -> bytes: + return b"x" * int(value_size) + + def _put_single_operation( + self, + key: str, + value: bytes, + inflight_at_start: int, + *, + deadline_ts: float, + ctx: str, + ) -> _PutResult: + del value, inflight_at_start, deadline_ts, ctx + self.put_keys.append(key) + if len(self.put_keys) <= 3: + return _PutResult(success=True) + return _PutResult(success=False, error_msg="StorageFullError: No space left") + + +class TestBenchmarkNodeKvContract(unittest.TestCase): + def test_benchmark_client_strips_owner_only_ssd_backend_fields(self) -> None: + sanitized = _KV._sanitize_benchmark_client_kvcache_config( + { + "test_spec_config": { + "kv_ssd_storage_backend": "foyer", + "kv_ssd_uring_mode": "single_buffer", + "disable_observability": True, + } + } + ) + self.assertEqual( + sanitized["test_spec_config"], + {"disable_observability": True}, + ) + + def test_get_output_is_a_closed_enum(self) -> None: + self.assertEqual( + _KV.normalize_kv_get_output(None), + _KV.KVGetOutput.HOLDER, + ) + self.assertEqual( + _KV.normalize_kv_get_output("CUDA"), + _KV.KVGetOutput.CUDA, + ) + with self.assertRaisesRegex(ValueError, "kv_get_output must be one of"): + _KV.normalize_kv_get_output("length_only") + + def test_cuda_device_index_is_a_non_negative_integer(self) -> None: + self.assertEqual(_KV.normalize_kv_cuda_device_index(None), 0) + self.assertEqual(_KV.normalize_kv_cuda_device_index(6), 6) + for invalid in (-1, True, "6"): + with self.subTest(invalid=invalid): + with self.assertRaisesRegex(ValueError, "kv_cuda_device_index"): + _KV.normalize_kv_cuda_device_index(invalid) + + def test_flat_dict_payload_view_excludes_encoding_metadata(self) -> None: + payload = b"payload-data" + key = b"payload" + encoded = ( + struct.pack(" None: + payload = b"x" * 32 + holder = _FakeHolder(payload) + adapter = _KV.KVBenchmarkBlockingStore( + _FakeKvStore(holder), + backend_kind=_KV.BACKEND_KIND_FLUXON, + get_output=_KV.KVGetOutput.BYTES, + ) + + completion = adapter.get_blocking( + "key", + deadline_ts=10**10, + ctx="test", + expected_payload_size=len(payload), + ) + + self.assertIsNone(completion.error_msg) + self.assertEqual(completion.source_kind, _KV.KVGetSourceKind.MEMORY) + self.assertEqual(holder.access_count, 1) + + def test_mooncake_counter_window_reports_ssd_read_delta(self) -> None: + boundary_ts = time.time() - 1.0 + sampler = _KV._MooncakeOffloadReadCounterWindow( + _FakeCounterStore([17, 23]), + start_ts=boundary_ts, + end_ts=boundary_ts, + ) + + sampler.start() + summary = sampler.finish() + + self.assertTrue(summary["complete"]) + self.assertEqual(summary["ssd_read_count"], 6) + + def test_cuda_output_put_uses_cpu_dlpack_payload(self) -> None: + payload = b"x" * 32 + store = _FakeKvStore(_FakeHolder(payload)) + adapter = object.__new__(_KV.KVBenchmarkBlockingStore) + adapter.backend_kind = _KV.BACKEND_KIND_FLUXON + adapter._store = store + adapter._get_output = _KV.KVGetOutput.CUDA + adapter._phase_profiler = _KV._FluxonPhaseProfiler() + adapter._put_dedupe_condition = threading.Condition() + adapter._put_confirmed_keys = set() + adapter._put_inflight_keys = set() + + err = adapter.put_blocking( + "key", + payload, + deadline_ts=10**10, + ctx="test", + ) + + self.assertIsNone(err) + self.assertIsInstance(store.put_value["payload"], _KV.DLPackBytesView) + + def test_put_rejects_existing_keys_and_deduplicates_locally(self) -> None: + payload = b"payload-data" + store = _FakeKvStore(_FakeHolder(payload)) + adapter = _KV.KVBenchmarkBlockingStore( + store, + backend_kind=_KV.BACKEND_KIND_FLUXON, + get_output=_KV.KVGetOutput.HOLDER, + ) + + first_error = adapter.put_blocking( + "same-key", + payload, + deadline_ts=10**10, + ctx="first", + ) + duplicate_error = adapter.put_blocking( + "same-key", + payload, + deadline_ts=10**10, + ctx="duplicate", + ) + + self.assertIsNone(first_error) + self.assertIsNone(duplicate_error) + self.assertEqual(store.put_keys, ["same-key"]) + self.assertEqual(len(store.put_opts), 1) + self.assertTrue(store.put_opts[0].reject_if_exists) + + def test_cuda_pipeline_retains_host_reference_until_event_completion(self) -> None: + runtime = _FakeCudaRuntime() + pipeline = _KV._CudaHostToDevicePipeline( + device_index=0, + depth=2, + runtime=runtime, + ) + source = ctypes.create_string_buffer(b"payload") + keepalive = _KeepAlive() + keepalive_ref = weakref.ref(keepalive) + + completions = pipeline.submit_from_host( + source_ptr=ctypes.addressof(source), + size=len(source.raw), + keepalive=keepalive, + token="copy-0", + ) + del keepalive + gc.collect() + + self.assertEqual(completions, []) + self.assertIsNotNone(keepalive_ref()) + drained = pipeline.drain_current_thread() + gc.collect() + self.assertEqual([completion.token for completion in drained], ["copy-0"]) + self.assertIsNone(keepalive_ref()) + pipeline.close() + + def test_cuda_pipeline_submit_does_not_wait_for_the_h2d_event(self) -> None: + runtime = _FakeCudaRuntime() + runtime.event_synchronize_gate = threading.Event() + pipeline = _KV._CudaHostToDevicePipeline( + device_index=0, + depth=2, + runtime=runtime, + ) + source = ctypes.create_string_buffer(b"payload") + + submitted = pipeline.submit_from_host( + source_ptr=ctypes.addressof(source), + size=len(source.raw), + keepalive=source, + token="copy-0", + ) + + self.assertEqual(submitted, []) + self.assertFalse(runtime.event_synchronize_started.is_set()) + self.assertEqual(pipeline.poll_current_thread(), []) + runtime.event_synchronize_gate.set() + self.assertEqual( + [completion.token for completion in pipeline.drain_current_thread()], + ["copy-0"], + ) + pipeline.close() + + def test_cuda_pipeline_poll_collects_a_ready_event_without_waiting(self) -> None: + runtime = _FakeCudaRuntime() + pipeline = _KV._CudaHostToDevicePipeline( + device_index=0, + depth=2, + runtime=runtime, + ) + source = ctypes.create_string_buffer(b"payload") + + pipeline.submit_from_host( + source_ptr=ctypes.addressof(source), + size=len(source.raw), + keepalive=source, + token="copy-0", + ) + runtime.complete_all_events() + + completions = pipeline.poll_current_thread() + + self.assertEqual([completion.token for completion in completions], ["copy-0"]) + self.assertFalse(runtime.event_synchronize_started.is_set()) + pipeline.close() + + def test_cuda_pipeline_is_bounded_and_drains_in_submission_order(self) -> None: + runtime = _FakeCudaRuntime() + pipeline = _KV._CudaHostToDevicePipeline( + device_index=0, + depth=2, + runtime=runtime, + ) + sources = [ctypes.create_string_buffer(value) for value in (b"a", b"b", b"c")] + + first = pipeline.submit_from_host( + source_ptr=ctypes.addressof(sources[0]), + size=1, + keepalive=sources[0], + token="copy-0", + ) + second = pipeline.submit_from_host( + source_ptr=ctypes.addressof(sources[1]), + size=1, + keepalive=sources[1], + token="copy-1", + ) + third = pipeline.submit_from_host( + source_ptr=ctypes.addressof(sources[2]), + size=1, + keepalive=sources[2], + token="copy-2", + ) + + self.assertEqual(first, []) + self.assertEqual(second, []) + self.assertEqual([completion.token for completion in third], ["copy-0"]) + self.assertEqual( + [completion.token for completion in pipeline.drain_current_thread()], + ["copy-1", "copy-2"], + ) + self.assertEqual(runtime.copied_payloads, [b"a", b"b", b"c"]) + pipeline.close() + + def test_cuda_pipeline_reports_async_completion_error(self) -> None: + runtime = _FakeCudaRuntime() + pipeline = _KV._CudaHostToDevicePipeline( + device_index=0, + depth=1, + runtime=runtime, + ) + source = ctypes.create_string_buffer(b"x") + runtime.fail_next_event_synchronize = True + pipeline.submit_from_host( + source_ptr=ctypes.addressof(source), + size=1, + keepalive=source, + token="copy-0", + ) + + completions = pipeline.drain_current_thread() + + self.assertEqual(len(completions), 1) + self.assertEqual(completions[0].token, "copy-0") + self.assertIn("injected CUDA event failure", completions[0].error_msg or "") + pipeline.close() + + def test_cuda_store_counts_get_only_after_h2d_event_completion(self) -> None: + payload = b"payload-data" + key = b"payload" + encoded = ( + struct.pack(" None: + payload = b"payload-data" + key = b"payload" + encoded = ( + struct.pack("= node.stop_after_gets else None + ) + + results = _KV.run_kv_worker( + node, + thread_id=0, + deadline_ts=time.time() + 1.0, + operation_result_cls=_FakeOperationResult, + operation_outcome=_FakeOperationOutcome, + metric_warmup_seconds=0.0, + debug_print=lambda message: None, + ) + + self.assertIsNotNone(results) + assert results is not None + self.assertEqual(store.get_count, 3) + self.assertEqual(len(results), 3) + self.assertTrue(all(result.success for result in results)) + self.assertEqual(node._inflight, 0) + self.assertEqual(sorted(node.progress), [0, 1, 2]) + adapter._cuda_pipeline.close() + + def test_bootstrap_stop_policy_shrinks_effective_keyspace_on_storage_full(self) -> None: + node = _FakeBenchmarkNode() + + ok = _KV.prepare_kv_before_ready(node, logger=logging.getLogger(__name__)) + + self.assertTrue(ok) + self.assertEqual(node.put_keys, ["bench_k0", "bench_k1", "bench_k2", "bench_k3"]) + self.assertEqual(node.test_config["keyspace_size"], 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/fluxon_test_stack/tests/test_ci_2_virt_node_contract.py b/fluxon_test_stack/tests/test_ci_2_virt_node_contract.py index 1974927..625f512 100644 --- a/fluxon_test_stack/tests/test_ci_2_virt_node_contract.py +++ b/fluxon_test_stack/tests/test_ci_2_virt_node_contract.py @@ -51,7 +51,7 @@ def test_generated_suite_is_public_dual_local_nodes_ci_only(self) -> None: scene_ids=[self._DOC_SCENE_ID, self._KVTEST_SCENE_ID, self._LOG_MGMT_SCENE_ID, self._MQ_SCENE_ID], primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -80,7 +80,7 @@ def test_generated_suite_is_public_dual_local_nodes_ci_only(self) -> None: ) self.assertEqual( generated["profiles"]["fluxon_tcp_thread"]["runtime"]["ci"]["deploy"]["target_ip_map"], - {"local-node-a": "10.1.1.119", "local-node-b": "10.1.1.119"}, + {"local-node-a": "192.0.2.119", "local-node-b": "192.0.2.119"}, ) self.assertEqual( generated["profiles"]["fluxon_tcp_thread"]["runtime"]["ci"]["runtime_contracts"]["cluster_kv_owner"][ @@ -163,7 +163,7 @@ def test_generated_suite_supports_mq_core_ci_scene(self) -> None: scene_ids=[self._MQ_SCENE_ID], primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -193,7 +193,7 @@ def test_generated_suite_preserves_source_scene_configs(self) -> None: scene_ids=[self._KVTEST_SCENE_ID], primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -212,7 +212,7 @@ def test_generated_suite_injects_public_transport_feature_for_cargo_kv_unit(self scene_ids=[self._CARGO_KV_UNIT_SCENE_ID], primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -244,7 +244,7 @@ def test_generated_suite_supports_additional_cargo_ci_scenes(self) -> None: scene_ids=scene_ids, primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -268,7 +268,7 @@ def test_generated_suite_supports_doc_page_ci_scene(self) -> None: scene_ids=[self._DOC_SCENE_ID], primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", controller_port=19080, ) @@ -302,7 +302,7 @@ def test_generated_deployconf_rewrites_to_dual_local_nodes(self) -> None: deployconf_cfg=deployconf_cfg, primary_node_name="local-node-a", secondary_node_name="local-node-b", - host_ip="10.1.1.119", + host_ip="192.0.2.119", primary_hostworkdir=Path("/tmp/fluxon_testbed/a"), secondary_hostworkdir=Path("/tmp/fluxon_testbed/b"), wheel_name="fluxon-0.2.1-cp38-abi3-manylinux_2_28_x86_64.whl", @@ -324,7 +324,7 @@ def test_generated_deployconf_rewrites_to_dual_local_nodes(self) -> None: ) self.assertEqual( [node["ip"] for node in generated["cluster_nodes"]], - ["10.1.1.119", "10.1.1.119"], + ["192.0.2.119", "192.0.2.119"], ) self.assertEqual(generated["global_envs"]["FLUXON_CLUSTER_NODE_IDS"], "local-node-a local-node-b") self.assertEqual( @@ -362,7 +362,7 @@ def test_generated_deployconf_rewrites_to_dual_local_nodes(self) -> None: ) self.assertIn("local-node-a", generated["service"]["ops_agent"]["entrypoint"]) self.assertIn("local-node-b", generated["service"]["ops_agent"]["entrypoint"]) - self.assertIn(' - "10.1.1.119/32"', generated["service"]["master"]["entrypoint"]) + self.assertIn(' - "192.0.2.119/32"', generated["service"]["master"]["entrypoint"]) def test_generated_start_test_bed_config_points_to_local_authorities(self) -> None: start_cfg = _ENTRY._load_yaml_mapping(_ENTRY.DEFAULT_START_TEST_BED_TEMPLATE, ctx="start_test_bed") @@ -370,14 +370,14 @@ def test_generated_start_test_bed_config_points_to_local_authorities(self) -> No start_cfg=start_cfg, generated_deployconf_path=Path("/tmp/deployconf.yaml"), primary_node_name="local-node-a", - controller_access_ip="10.1.1.119", + controller_access_ip="192.0.2.119", controller_port=19080, ui_port=18080, ui_workdir=Path("/tmp/ui"), ) self.assertEqual(generated["deployconf_path"], "/tmp/deployconf.yaml") - self.assertEqual(generated["controller_url"], "http://10.1.1.119:19080/r/ops/fluxon_testbed") + self.assertEqual(generated["controller_url"], "http://192.0.2.119:19080/r/ops/fluxon_testbed") self.assertEqual(generated["controller_basic_auth"]["username"], "ops_admin") self.assertEqual(generated["controller_basic_auth"]["password"], "ops_password") self.assertEqual(generated["test_runner_ui"]["workdir"], "/tmp/ui") @@ -390,7 +390,7 @@ def test_generated_apply_check_config_excludes_control_plane_reapply(self) -> No start_cfg=start_cfg, generated_deployconf_path=Path("/tmp/deployconf.yaml"), primary_node_name="local-node-a", - controller_access_ip="10.1.1.119", + controller_access_ip="192.0.2.119", controller_port=19080, ui_port=18080, ui_workdir=Path("/tmp/ui"), @@ -514,6 +514,33 @@ def test_prepare_pack_release_runtime_dirs_creates_expected_layout(self) -> None self.assertTrue((root / "manylinux-cache" / "cargo-registry").is_dir()) self.assertTrue((root / "manylinux-cache" / "cargo-git").is_dir()) + def test_cleanup_pack_release_runtime_removes_only_owned_runtime(self) -> None: + with tempfile.TemporaryDirectory() as td: + workdir = Path(td) / "ci_workdir" + runtime_root = workdir / "pack_release_runtime" + nested = runtime_root / "project-data" / "cache" + nested.mkdir(parents=True) + (nested / "artifact.bin").write_bytes(b"artifact") + sibling = workdir / "generated" / "suite.yaml" + sibling.parent.mkdir(parents=True) + sibling.write_text("schema_version: 1\n", encoding="utf-8") + + _ENTRY._cleanup_pack_release_runtime_after_success( + runtime_root=runtime_root, + workdir=workdir, + ) + + self.assertFalse(runtime_root.exists()) + self.assertEqual(sibling.read_text(encoding="utf-8"), "schema_version: 1\n") + + unexpected = workdir / "unexpected" + unexpected.mkdir() + with self.assertRaisesRegex(ValueError, "unexpected pack runtime path"): + _ENTRY._cleanup_pack_release_runtime_after_success( + runtime_root=unexpected, + workdir=workdir, + ) + def test_sync_rather_no_git_submodule_uses_canonical_entrypoint(self) -> None: calls: list[list[str]] = [] @@ -529,8 +556,8 @@ def fake_run(argv: list[str], *, env=None) -> None: self.assertEqual(calls[0][1], str(_ENTRY.DEFAULT_RATHER_NO_GIT_SUBMODULE_SCRIPT.resolve())) def test_same_host_local_testbed_host_ip_requires_non_loopback(self) -> None: - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): - self.assertEqual(_ENTRY._same_host_local_testbed_host_ip(), "10.1.1.119") + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): + self.assertEqual(_ENTRY._same_host_local_testbed_host_ip(), "192.0.2.119") with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="127.0.0.1"): with self.assertRaisesRegex(RuntimeError, "requires a non-loopback IPv4 address"): _ENTRY._same_host_local_testbed_host_ip() @@ -561,12 +588,13 @@ def fake_run(argv: list[str], *, env=None) -> None: "--skip-pack", "--skip-dispatch", "--skip-start-testbed", + "--cleanup-successful-case-artifacts", ] original_argv = sys.argv[:] try: with mock.patch.object(_ENTRY, "_run", side_effect=fake_run): with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): sys.argv = argv rc = _ENTRY.main() finally: @@ -578,6 +606,7 @@ def fake_run(argv: list[str], *, env=None) -> None: runner_argv, runner_env = calls[-1] self.assertIsNotNone(runner_env) self.assertEqual(runner_argv[1], str((REPO_ROOT / "fluxon_test_stack" / "test_runner.py").resolve())) + self.assertIn("--cleanup-successful-case-artifacts", runner_argv) self.assertEqual( runner_env[_ENTRY.TEST_STACK_START_TEST_BED_CONFIG_ENV], str((workdir / "generated" / "start_test_bed.local.yaml").resolve()), @@ -631,7 +660,7 @@ def test_main_supports_explicit_suite_path(self) -> None: original_argv = sys.argv[:] try: with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): sys.argv = argv rc = _ENTRY.main() finally: @@ -697,7 +726,7 @@ def test_main_same_host_generated_configs_use_non_loopback_host_ip(self) -> None original_argv = sys.argv[:] try: with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): sys.argv = argv rc = _ENTRY.main() finally: @@ -715,13 +744,13 @@ def test_main_same_host_generated_configs_use_non_loopback_host_ip(self) -> None ) self.assertEqual( [node["ip"] for node in generated_deployconf["cluster_nodes"]], - ["10.1.1.119", "10.1.1.119"], + ["192.0.2.119", "192.0.2.119"], ) self.assertEqual( generated_start["controller_url"], - "http://10.1.1.119:19080/r/ops/fluxon_testbed", + "http://192.0.2.119:19080/r/ops/fluxon_testbed", ) - self.assertIn(' - "10.1.1.119/32"', generated_deployconf["service"]["master"]["entrypoint"]) + self.assertIn(' - "192.0.2.119/32"', generated_deployconf["service"]["master"]["entrypoint"]) def test_main_syncs_rather_no_git_submodule_before_pack(self) -> None: with tempfile.TemporaryDirectory() as td: @@ -749,12 +778,13 @@ def fake_run(argv: list[str], *, env=None) -> None: "--skip-dispatch", "--skip-start-testbed", "--skip-runner", + "--cleanup-pack-runtime-after-success", ] original_argv = sys.argv[:] try: with mock.patch.object(_ENTRY, "_run", side_effect=fake_run): with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): with mock.patch.object(_ENTRY, "_ensure_ci_pack_release_env", return_value=Path("/tmp/env.yaml")): with mock.patch.object(_ENTRY, "_render_ci_nix_pack_config", return_value=Path("/tmp/cfg.yaml")): sys.argv = argv @@ -773,6 +803,7 @@ def fake_run(argv: list[str], *, env=None) -> None: calls[1][0][1], str((REPO_ROOT / "fluxon_test_stack" / "pack_test_stack_rsc.py").resolve()), ) + self.assertFalse((workdir / "pack_release_runtime").exists()) def test_main_passes_explicit_release_dir_to_pack_stage(self) -> None: with tempfile.TemporaryDirectory() as td: @@ -807,7 +838,7 @@ def fake_run(argv: list[str], *, env=None) -> None: try: with mock.patch.object(_ENTRY, "_run", side_effect=fake_run): with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): with mock.patch.object(_ENTRY, "_ensure_ci_pack_release_env", return_value=Path("/tmp/env.yaml")): with mock.patch.object(_ENTRY, "_render_ci_nix_pack_config", return_value=Path("/tmp/cfg.yaml")): sys.argv = argv @@ -859,7 +890,7 @@ def fake_run(argv: list[str], *, env=None) -> None: try: with mock.patch.object(_ENTRY, "_run", side_effect=fake_run): with mock.patch.object(_ENTRY, "_detect_local_hostname", return_value="runner-host"): - with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="10.1.1.119"): + with mock.patch.object(_ENTRY, "_detect_local_ipv4", return_value="192.0.2.119"): sys.argv = argv rc = _ENTRY.main() finally: diff --git a/fluxon_test_stack/tests/test_test_runner_testbed_contract.py b/fluxon_test_stack/tests/test_test_runner_testbed_contract.py index 0e1cee6..18005e2 100644 --- a/fluxon_test_stack/tests/test_test_runner_testbed_contract.py +++ b/fluxon_test_stack/tests/test_test_runner_testbed_contract.py @@ -70,6 +70,129 @@ def _suite_cfg_with_declared_ci_commands(command_map: dict[str, list[dict]]) -> class TestTestRunnerTestbedContract(unittest.TestCase): + def test_normalize_test_spec_config_accepts_foyer_ssd_backend(self) -> None: + self.assertEqual( + _RUNNER._normalize_test_spec_config( + {"kv_ssd_storage_backend": "foyer"}, + "test_spec_config", + )["kv_ssd_storage_backend"], + "foyer", + ) + with self.assertRaisesRegex(ValueError, "kv_ssd_uring_mode is only valid"): + _RUNNER._normalize_test_spec_config( + { + "kv_ssd_storage_backend": "foyer", + "kv_ssd_uring_mode": "iovec", + }, + "test_spec_config", + ) + + def test_unresolved_runtime_tokens_exclude_numeric_run_scopes(self) -> None: + self.assertEqual( + _RUNNER._find_unresolved_runtime_tokens( + "bench__750222457007__worker __STACK_CONTROLLER_URL__" + ), + ["__STACK_CONTROLLER_URL__"], + ) + + def test_benchmark_threads_per_process_uses_bounded_values(self) -> None: + for value in (2, 4): + self.assertEqual( + _RUNNER._require_test_stack_benchmark_threads_per_process( + value, + "scale.benchmark.threads_per_process", + ), + value, + ) + + for value in (1, 3, 8): + with self.assertRaisesRegex(ValueError, "must be one of: 2, 4"): + _RUNNER._require_test_stack_benchmark_threads_per_process( + value, + "scale.benchmark.threads_per_process", + ) + + def test_suite_cluster_name_is_workdir_scoped(self) -> None: + with tempfile.TemporaryDirectory() as td: + root = Path(td) + first = _RUNNER._suite_cluster_name_for_workdir(root / "run_a") + first_again = _RUNNER._suite_cluster_name_for_workdir(root / "run_a") + second = _RUNNER._suite_cluster_name_for_workdir(root / "run_b") + + self.assertEqual(first, first_again) + self.assertNotEqual(first, second) + self.assertRegex(first, r"^fluxon_benchmark_[0-9a-f]{12}$") + + def test_ci_owner_share_mem_path_is_global_per_owner_index(self) -> None: + with tempfile.TemporaryDirectory() as td: + share_mem_root = Path(td) / "shared-root" + first_case = { + "runtime": { + "run_dir": "/tmp/results/case-a/run_1", + "stack_identity": {"share_mem_path": str(share_mem_root)}, + } + } + second_case = { + "runtime": { + "run_dir": "/tmp/results/case-b/run_9", + "stack_identity": {"share_mem_path": str(share_mem_root)}, + } + } + + first_owner_0 = _RUNNER._ci_share_mem_path(first_case, owner_index=0) + second_owner_0 = _RUNNER._ci_share_mem_path(second_case, owner_index=0) + first_owner_1 = _RUNNER._ci_share_mem_path(first_case, owner_index=1) + + self.assertEqual(first_owner_0, second_owner_0) + self.assertEqual(Path(first_owner_0), share_mem_root / "ci" / "owner-0") + self.assertEqual(Path(first_owner_1), share_mem_root / "ci" / "owner-1") + self.assertNotEqual(first_owner_0, first_owner_1) + + def test_ci_owner_share_mem_cleanup_cannot_touch_testbed_owner(self) -> None: + with tempfile.TemporaryDirectory() as td: + share_mem_root = Path(td) / "shared-root" + testbed_bundle = share_mem_root / "testbed-cluster" + testbed_bundle.mkdir(parents=True) + testbed_mmap = testbed_bundle / "mmap.file" + testbed_mmap.write_text("testbed-owner", encoding="utf-8") + resolved_case = { + "runtime": { + "run_dir": "/tmp/results/case-a/run_1", + "stack_identity": {"share_mem_path": str(share_mem_root)}, + } + } + owner_path = Path(_RUNNER._ci_share_mem_path(resolved_case, owner_index=0)) + old_case_bundle = owner_path / "old-ci-cluster" + old_case_bundle.mkdir(parents=True) + (old_case_bundle / "mmap.file").write_text("old-ci-owner", encoding="utf-8") + + with mock.patch.object(_RUNNER, "_instance_remote_target_access", return_value=None): + _RUNNER._reset_ci_owner_share_mem_dir( + resolved_case, + owner_index=0, + share_mem_path=str(owner_path), + ) + + self.assertTrue(owner_path.is_dir()) + self.assertEqual(list(owner_path.iterdir()), []) + self.assertEqual(testbed_mmap.read_text(encoding="utf-8"), "testbed-owner") + + with self.assertRaisesRegex(ValueError, "unexpected CI owner share_mem_path"): + _RUNNER._cleanup_ci_owner_share_mem_dir( + resolved_case, + owner_index=0, + share_mem_path=str(testbed_bundle), + ) + + _RUNNER._cleanup_ci_owner_share_mem_dir( + resolved_case, + owner_index=0, + share_mem_path=str(owner_path), + ) + + self.assertFalse(owner_path.exists()) + self.assertEqual(testbed_mmap.read_text(encoding="utf-8"), "testbed-owner") + def test_write_ci_master_owner_configs_emits_owner_large_file_paths(self) -> None: with tempfile.TemporaryDirectory() as td: run_dir = Path(td) @@ -321,6 +444,33 @@ def test_finalize_ci_case_runtime_deletes_each_apply_id_once(self) -> None: ) cleanup_runtime.assert_called_once_with(resolved_case, timeout_s=120) + def test_ci_cleanup_runtime_reclaims_only_fixed_ci_owner_slot(self) -> None: + cleanup_case = { + "runtime": { + "stack_identity": {"share_mem_path": "/tmp/testbed-shm"}, + }, + "runtime_model": _RUNNER._build_runtime_model(_RUNNER.CASE_FAMILY_CI), + "deploy": { + "instances": [ + {"id": "owner_0", "deployer": {"target": "node-a"}}, + ] + }, + } + expected_path = "/tmp/testbed-shm/ci/owner-0" + + with mock.patch.object(_RUNNER, "_ci_runtime_cleanup_case", return_value=cleanup_case): + with mock.patch.object(_RUNNER, "_ci_runtime_current_apply_ids", return_value=[]): + with mock.patch.object(_RUNNER, "_wait_ci_ports_free") as wait_ports_free: + with mock.patch.object(_RUNNER, "_cleanup_ci_owner_share_mem_dir") as cleanup_shm: + _RUNNER._ci_cleanup_runtime(cleanup_case, timeout_s=120) + + wait_ports_free.assert_called_once_with(cleanup_case, timeout_s=120) + cleanup_shm.assert_called_once_with( + cleanup_case, + owner_index=0, + share_mem_path=expected_path, + ) + def test_finalize_ci_case_runtime_preserves_structured_instance_ids(self) -> None: with tempfile.TemporaryDirectory() as td: run_dir = Path(td) @@ -462,6 +612,72 @@ def test_finalize_error_preserves_success_for_ci_and_bench(self) -> None: ) ) + def test_cleanup_successful_run_artifacts_keeps_only_summary(self) -> None: + with tempfile.TemporaryDirectory() as td: + root = Path(td) + run_dir = root / "run_1" + run_dir.mkdir() + summary = { + "schema_version": 1, + "outcome": _RUNNER.RUN_OUTCOME_SUCCESS, + } + (run_dir / "summary.yaml").write_text( + yaml.safe_dump(summary, sort_keys=False), + encoding="utf-8", + ) + (run_dir / "src" / "nested").mkdir(parents=True) + (run_dir / "src" / "nested" / "source.rs").write_text("source", encoding="utf-8") + (run_dir / "logs").mkdir() + (run_dir / "logs" / "stdout.log").write_text("log", encoding="utf-8") + (run_dir / "resolved_case.yaml").write_text("case: {}\n", encoding="utf-8") + outside = root / "outside.txt" + outside.write_text("keep", encoding="utf-8") + (run_dir / "outside-link").symlink_to(outside) + + _RUNNER._cleanup_successful_run_artifacts(run_dir) + + self.assertEqual([path.name for path in run_dir.iterdir()], ["summary.yaml"]) + self.assertEqual( + yaml.safe_load((run_dir / "summary.yaml").read_text(encoding="utf-8")), + summary, + ) + self.assertEqual(outside.read_text(encoding="utf-8"), "keep") + + def test_cleanup_successful_run_artifacts_refuses_failed_summary(self) -> None: + with tempfile.TemporaryDirectory() as td: + run_dir = Path(td) + (run_dir / "summary.yaml").write_text( + yaml.safe_dump({"outcome": _RUNNER.RUN_OUTCOME_FAILED}), + encoding="utf-8", + ) + diagnostic = run_dir / "exception.txt" + diagnostic.write_text("traceback", encoding="utf-8") + + with self.assertRaisesRegex(ValueError, "summary is not SUCCESS"): + _RUNNER._cleanup_successful_run_artifacts(run_dir) + + self.assertEqual(diagnostic.read_text(encoding="utf-8"), "traceback") + + def test_cleanup_successful_run_artifacts_preserves_teardown_diagnostics(self) -> None: + with tempfile.TemporaryDirectory() as td: + run_dir = Path(td) + (run_dir / "summary.yaml").write_text( + yaml.safe_dump( + { + "outcome": _RUNNER.RUN_OUTCOME_SUCCESS, + "teardown_error": "RuntimeError: owner cleanup failed", + } + ), + encoding="utf-8", + ) + diagnostic = run_dir / "logs" / "owner_0" / "stderr.log" + diagnostic.parent.mkdir(parents=True) + diagnostic.write_text("owner cleanup failed", encoding="utf-8") + + _RUNNER._cleanup_successful_run_artifacts(run_dir) + + self.assertEqual(diagnostic.read_text(encoding="utf-8"), "owner cleanup failed") + def test_write_ci_scene_config_yaml_emits_structured_scene_config(self) -> None: with tempfile.TemporaryDirectory() as td: run_dir = Path(td) @@ -513,7 +729,12 @@ def test_generated_test_stack_owner_config_emits_large_file_paths(self) -> None: target_slug = "local-node-a" runtime_instance_prefix = "case1" coord_tpl = {"deployer": {"target": ""}} - cluster_nodes = {"local-node-a": {"python_abi": "cpython3.10"}} + cluster_nodes = { + "local-node-a": { + "python_abi": "cpython3.10", + "ip": "192.0.2.10", + } + } resolved_case = { "runtime": { "run_dir": str(run_dir), @@ -553,6 +774,7 @@ def test_generated_test_stack_owner_config_emits_large_file_paths(self) -> None: runtime_env={}, owner_group_processes=None, owner_cpu_core_by_target={}, + owner_kv_ssd=None, ) self.assertEqual(len(owner_instances), 1) @@ -563,6 +785,565 @@ def test_generated_test_stack_owner_config_emits_large_file_paths(self) -> None: [str((run_dir / "services" / "kv_owner" / target_slug / "large").resolve())], ) + def test_generated_test_stack_owner_config_emits_owner_kv_ssd_limit(self) -> None: + with tempfile.TemporaryDirectory() as td: + run_dir = Path(td) + cfg_dir = run_dir / "configs" + cfg_dir.mkdir(parents=True) + owner_target = "local-node-a" + target_slug = "local-node-a" + runtime_instance_prefix = "case1" + coord_tpl = {"deployer": {"target": ""}} + cluster_nodes = {"local-node-a": {"python_abi": "cpython3.10"}} + resolved_case = { + "runtime": { + "run_dir": str(run_dir), + "stack_identity": { + "cluster_name": "bench_cluster", + "share_mem_path": "/tmp/bench_shm", + }, + } + } + + with mock.patch.object(_RUNNER, "_test_stack_runtime_required_python_abi", return_value="cpython3.10"): + with mock.patch.object(_RUNNER, "_test_stack_etcd_addresses", return_value=["127.0.0.1:19180"]): + with mock.patch.object(_RUNNER, "_test_stack_target_host_venv_python", return_value="/tmp/venv/bin/python3"): + with mock.patch.object(_RUNNER, "_test_stack_runtime_module_command", return_value="owner-cmd"): + owner_instances = _RUNNER._build_test_stack_external_kv_owner_instances( + scene_mode="KVSTORE", + resolved_case=resolved_case, + scale={"owner": {"owner_count": 1, "owner_dram_bytes": 1073741824}}, + runtime=resolved_case["runtime"], + run_dir=run_dir, + cfg_dir=cfg_dir, + coord_tpl=coord_tpl, + test_stack_runtime={}, + cluster_nodes=cluster_nodes, + owner_targets=[owner_target], + needs_kv_master=True, + kv_p2p_port_base=31000, + kv_p2p_port_stride=100, + kv_p2p_slot_offset=0, + p2p_ports_per_slot=10, + node_total=1, + run_index=1, + runtime_instance_prefix=runtime_instance_prefix, + kv_base={}, + test_spec_config={}, + perf_config=None, + runtime_env={}, + owner_group_processes=None, + owner_cpu_core_by_target={}, + owner_kv_ssd={"large_limit_size": [17179869184]}, + ) + + self.assertEqual(len(owner_instances), 1) + owner_cfg_path = cfg_dir / f"test_stack_kv_owner__{target_slug}.yaml" + owner_cfg = yaml.safe_load(owner_cfg_path.read_text(encoding="utf-8")) + self.assertEqual(owner_cfg["fluxonkv_spec"]["large_limit_size"], [17179869184]) + self.assertEqual( + owner_cfg["fluxonkv_spec"]["large_file_paths"], + [str((run_dir / "services" / "kv_owner" / target_slug / "large").resolve())], + ) + + def test_generated_mooncake_owner_config_emits_owner_ssd_offload(self) -> None: + with tempfile.TemporaryDirectory() as td: + run_dir = Path(td) + cfg_dir = run_dir / "configs" + cfg_dir.mkdir(parents=True) + owner_target = "local-node-a" + target_slug = "local-node-a" + runtime_instance_prefix = "case1" + coord_tpl = {"deployer": {"target": ""}} + cluster_nodes = { + "local-node-a": { + "python_abi": "cpython3.10", + "ip": "192.0.2.10", + } + } + resolved_case = {"runtime": {"run_dir": str(run_dir)}} + offload_root = "/tmp/mooncake-offload-contract" + offload_path = f"{offload_root}/{target_slug}" + + with mock.patch.object(_RUNNER, "_test_stack_runtime_required_python_abi", return_value="cpython3.10"): + with mock.patch.object(_RUNNER, "_test_stack_target_host_venv_python", return_value="/tmp/venv/bin/python3"): + with mock.patch.object(_RUNNER, "_test_stack_runtime_module_command", return_value="owner-cmd") as module_cmd: + owner_instances = _RUNNER._build_test_stack_mooncake_owner_instances( + resolved_case=resolved_case, + scale={"owner": {"owner_count": 1, "owner_dram_bytes": 1073741824}}, + run_dir=run_dir, + cfg_dir=cfg_dir, + coord_tpl=coord_tpl, + cluster_nodes=cluster_nodes, + owner_targets=[owner_target], + runtime_instance_prefix=runtime_instance_prefix, + kv_base={ + "mooncake_spec": { + "local_buffer_size": 1073741824, + "metadata_server": "http://127.0.0.1:34000/metadata", + "master_server_address": "127.0.0.1:33000", + "etcd_addresses": ["127.0.0.1:2379"], + }, + "protocol": {"protocol_type": "tcp"}, + }, + test_spec_config={}, + perf_config=None, + runtime_env={}, + testbed_mooncake_storage={ + "mode": "DEDICATED_OWNER", + "ssd_offload_root": offload_root, + "ssd_capacity_bytes": 17179869184, + }, + ) + + self.assertEqual(len(owner_instances), 1) + owner_cfg_path = cfg_dir / f"test_stack_kv_owner__{target_slug}.yaml" + owner_cfg = yaml.safe_load(owner_cfg_path.read_text(encoding="utf-8")) + self.assertEqual(owner_cfg["mooncake_spec"]["enable_ssd_offload"], True) + self.assertEqual(owner_cfg["mooncake_spec"]["ssd_offload_path"], offload_path) + self.assertEqual(owner_cfg["mooncake_spec"]["local_hostname"], "192.0.2.10") + self.assertIn("mkdir -p", module_cmd.call_args.kwargs["pre_exec_shell"]) + self.assertIn(offload_path, module_cmd.call_args.kwargs["pre_exec_shell"]) + self.assertEqual( + module_cmd.call_args.kwargs["runtime_env"][ + _RUNNER.TEST_STACK_MOONCAKE_SSD_LIMIT_ENV + ], + "17179869184", + ) + self.assertEqual(module_cmd.call_args.kwargs["require_unlimited_memlock"], False) + + def test_mooncake_memlock_is_required_only_for_rdma_protocol(self) -> None: + self.assertEqual(_RUNNER._test_stack_protocol_requires_unlimited_memlock({"protocol_type": "tcp"}), False) + self.assertEqual(_RUNNER._test_stack_protocol_requires_unlimited_memlock({"protocol_type": "rdma"}), True) + cmd = _RUNNER._test_stack_runtime_command( + run_dir=Path("/tmp/run"), + venv_python=Path("/tmp/venv/bin/python"), + script_path="/tmp/script.py", + script_args=[], + runtime_env={}, + require_unlimited_memlock=False, + ) + self.assertNotIn("ulimit -l unlimited", cmd) + cmd = _RUNNER._test_stack_runtime_command( + run_dir=Path("/tmp/run"), + venv_python=Path("/tmp/venv/bin/python"), + script_path="/tmp/script.py", + script_args=[], + runtime_env={}, + require_unlimited_memlock=True, + ) + self.assertIn("ulimit -l unlimited", cmd) + + def test_mooncake_storage_modes_and_capacity_split_are_bounded(self) -> None: + dedicated = _RUNNER._normalize_testbed_mooncake_storage_config( + { + "mode": "DEDICATED_OWNER", + "ssd_offload_root": "/tmp/mooncake-dedicated", + "ssd_capacity_bytes": 17179869184, + }, + "storage", + ) + self.assertEqual(dedicated["mode"], "DEDICATED_OWNER") + + per_process = _RUNNER._normalize_testbed_mooncake_storage_config( + {"mode": "PER_BENCHMARK_PROCESS"}, + "storage", + ) + self.assertEqual(per_process, {"mode": "PER_BENCHMARK_PROCESS"}) + self.assertEqual( + _RUNNER._split_testbed_mooncake_capacity( + total_bytes=2147483648, + instance_count=4, + alignment_bytes=16777216, + ctx="dram", + ), + 536870912, + ) + with self.assertRaisesRegex(ValueError, "must be set together"): + _RUNNER._normalize_testbed_mooncake_storage_config( + { + "mode": "PER_BENCHMARK_PROCESS", + "ssd_offload_root": "/tmp/mooncake-per-process", + }, + "storage", + ) + with self.assertRaisesRegex(ValueError, "divide evenly"): + _RUNNER._split_testbed_mooncake_capacity( + total_bytes=1073741824, + instance_count=3, + alignment_bytes=16777216, + ctx="dram", + ) + + def test_benchmark_full_matrix_includes_kv_ssd_pressure_comparison(self) -> None: + suite_cfg = yaml.safe_load( + (_RUNNER.RUNNER_REPO_ROOT / "fluxon_test_stack" / "benchmark_full_matrix.yaml").read_text(encoding="utf-8") + ) + suite = _RUNNER._parse_suite_config(suite_cfg) + cases = _RUNNER._expand_cases(suite) + + scene = suite.scenes["kv_ssd_pressure_zipf"] + self.assertEqual(scene["select"]["scales"], ["n1_kvowner_dram_1gib", "n2_kvowner_dram_1gib"]) + self.assertEqual(scene["select"]["profiles"], ["fluxon_tcp_kv_ssd", "fluxon_tcp", "mooncake_tcp"]) + self.assertEqual( + suite.profiles["fluxon_tcp_kv_ssd"]["runtime"]["test_stack"]["runtime_config"]["owner_kv_ssd"], + {"large_limit_size": [17179869184]}, + ) + self.assertEqual( + suite.profiles["mooncake_tcp"]["runtime"]["test_stack"]["runtime_config"][ + "testbed_mooncake_storage" + ], + {"mode": "DEDICATED_OWNER"}, + ) + self.assertEqual( + suite.scales["n1_kvowner_dram_1gib"]["owner"]["owner_dram_bytes"], + 1073741824, + ) + self.assertEqual( + sorted(case.case_id for case in cases if case.scene_id == "kv_ssd_pressure_zipf"), + [ + "kv_ssd_pressure_zipf__n1_kvowner_dram_1gib__fluxon_tcp", + "kv_ssd_pressure_zipf__n1_kvowner_dram_1gib__fluxon_tcp_kv_ssd", + "kv_ssd_pressure_zipf__n1_kvowner_dram_1gib__mooncake_tcp", + "kv_ssd_pressure_zipf__n2_kvowner_dram_1gib__fluxon_tcp", + "kv_ssd_pressure_zipf__n2_kvowner_dram_1gib__fluxon_tcp_kv_ssd", + "kv_ssd_pressure_zipf__n2_kvowner_dram_1gib__mooncake_tcp", + ], + ) + + def test_mooncake_benchmark_defaults_skip_get_size_on_get(self) -> None: + suite_cfg = yaml.safe_load( + (_RUNNER.RUNNER_REPO_ROOT / "fluxon_test_stack" / "benchmark_full_matrix.yaml").read_text(encoding="utf-8") + ) + suite = _RUNNER._parse_suite_config(suite_cfg) + case = next( + item + for item in _RUNNER._expand_cases(suite) + if item.scene_id == "kv_ssd_pressure_zipf" + and item.scale_id == "n1_kvowner_dram_1gib" + and item.profile_id == "mooncake_tcp" + ) + with tempfile.TemporaryDirectory() as td: + workdir_root = Path(td) + run_dir = workdir_root / case.case_id / "run_1" + run_dir.mkdir(parents=True) + resolved_case = _RUNNER._build_resolved_case_yaml( + case, + suite, + config_root=str(_RUNNER.RUNNER_REPO_ROOT), + workdir_root=str(workdir_root), + run_dir=str(run_dir), + ci_commands=None, + ci_prepare_steps=None, + execution_label=case.case_id, + command_id=None, + test_id=None, + stack_identity={ + "cluster_name": "test-cluster", + "share_mem_path": "/tmp/test-share", + "controller_url": "http://127.0.0.1:18080", + }, + ) + target_ip_map = resolved_case["deploy"]["target_ip_map"] + cluster_nodes = { + target: { + "python_abi": "cpython3.10", + "ip": ip, + "hostworkdir": "/tmp/test-hostworkdir", + } + for target, ip in target_ip_map.items() + } + + with mock.patch.object( + _RUNNER, + "_prepare_test_stack_runtime", + return_value={"coordinator_script": "/tmp/coordinator.py", "node_script": "/tmp/node.py"}, + ): + with mock.patch.object( + _RUNNER, + "_load_test_stack_cluster_nodes_and_dispatch", + return_value=(cluster_nodes, None), + ): + with mock.patch.object(_RUNNER, "_test_stack_runtime_required_python_abi", return_value="cpython3.10"): + with mock.patch.object( + _RUNNER, + "_test_stack_target_host_venv_python", + return_value=Path("/tmp/venv/bin/python3"), + ): + with mock.patch.object(_RUNNER, "_test_stack_runtime_module_command", return_value="module-cmd"): + with mock.patch.object( + _RUNNER, + "_test_stack_etcd_addresses", + return_value=["127.0.0.1:2379"], + ): + _RUNNER._compile_test_stack_case(resolved_case, run_index=1) + + benchmark_cfg = _RUNNER._load_test_stack_benchmark_config(run_dir) + self.assertEqual( + benchmark_cfg["kv_base"]["mooncake_spec"]["skip_get_size_on_get"], + True, + ) + self.assertEqual( + benchmark_cfg["node_overrides"][0]["kv"]["mooncake_spec"]["local_hostname"], + _RUNNER._test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name=next(iter(target_ip_map)), + local_ipv4_addrs=_RUNNER._local_ipv4_addresses(), + ), + ) + + def test_mooncake_per_benchmark_process_compilation_splits_total_capacity(self) -> None: + suite_cfg = yaml.safe_load( + (_RUNNER.RUNNER_REPO_ROOT / "fluxon_test_stack" / "benchmark_full_matrix.yaml").read_text( + encoding="utf-8" + ) + ) + suite = _RUNNER._parse_suite_config(suite_cfg) + case = next( + item + for item in _RUNNER._expand_cases(suite) + if item.scene_id == "kv_ssd_pressure_zipf" + and item.scale_id == "n1_kvowner_dram_1gib" + and item.profile_id == "mooncake_tcp" + ) + with tempfile.TemporaryDirectory() as td: + workdir_root = Path(td) + run_dir = workdir_root / case.case_id / "run_1" + run_dir.mkdir(parents=True) + resolved_case = _RUNNER._build_resolved_case_yaml( + case, + suite, + config_root=str(_RUNNER.RUNNER_REPO_ROOT), + workdir_root=str(workdir_root), + run_dir=str(run_dir), + ci_commands=None, + ci_prepare_steps=None, + execution_label=case.case_id, + command_id=None, + test_id=None, + stack_identity={ + "cluster_name": "test-cluster", + "share_mem_path": "/tmp/test-share", + "controller_url": "http://127.0.0.1:18080", + }, + ) + resolved_case["scale"]["benchmark"]["processes_per_target"] = 4 + mooncake_runtime = resolved_case["profile"]["test_stack"]["runtime_config"] + mooncake_runtime["testbed_mooncake_storage"] = { + "mode": "PER_BENCHMARK_PROCESS", + "ssd_offload_root": str(workdir_root / "mooncake-ssd"), + "ssd_capacity_bytes": 17179869184, + } + target_ip_map = resolved_case["deploy"]["target_ip_map"] + cluster_nodes = { + target: { + "python_abi": "cpython3.10", + "ip": ip, + "hostworkdir": "/tmp/test-hostworkdir", + } + for target, ip in target_ip_map.items() + } + + with mock.patch.object( + _RUNNER, + "_prepare_test_stack_runtime", + return_value={"coordinator_script": "/tmp/coordinator.py", "node_script": "/tmp/node.py"}, + ): + with mock.patch.object( + _RUNNER, + "_load_test_stack_cluster_nodes_and_dispatch", + return_value=(cluster_nodes, None), + ): + with mock.patch.object( + _RUNNER, + "_test_stack_runtime_required_python_abi", + return_value="cpython3.10", + ): + with mock.patch.object( + _RUNNER, + "_test_stack_target_host_venv_python", + return_value=Path("/tmp/venv/bin/python3"), + ): + with mock.patch.object( + _RUNNER, + "_test_stack_runtime_module_command", + return_value="module-cmd", + ): + with mock.patch.object( + _RUNNER, + "_test_stack_etcd_addresses", + return_value=["127.0.0.1:2379"], + ): + _RUNNER._compile_test_stack_case(resolved_case, run_index=1) + + benchmark_cfg = _RUNNER._load_test_stack_benchmark_config(run_dir) + overrides = benchmark_cfg["node_overrides"] + self.assertEqual(len(overrides), 4) + self.assertEqual( + { + item["kv"]["contribute_to_cluster_pool_size"]["dram"] + for item in overrides + }, + {268435456}, + ) + ssd_paths = { + item["kv"]["mooncake_spec"]["ssd_offload_path"] + for item in overrides + } + self.assertEqual(len(ssd_paths), 4) + instance_ids = { + item["id"] for item in resolved_case["deploy"]["instances"] + } + self.assertFalse( + any( + instance_id.startswith(_RUNNER.TEST_STACK_KV_OWNER_INSTANCE_ID_PREFIX) + for instance_id in instance_ids + ) + ) + node_commands = [ + item["deployer"]["args"][0] + for item in resolved_case["deploy"]["instances"] + if item["id"].startswith("worker_") + ] + self.assertEqual(len(node_commands), 4) + self.assertTrue( + all( + f"export {_RUNNER.TEST_STACK_MOONCAKE_SSD_LIMIT_ENV}=4294967296" + in command + for command in node_commands + ) + ) + + def test_mooncake_same_host_endpoint_uses_loopback(self) -> None: + cluster_nodes = { + "local": {"ip": "192.0.2.10"}, + "remote": {"ip": "192.0.2.11"}, + } + local_ipv4_addrs = {"127.0.0.1", "192.0.2.10"} + + self.assertEqual( + _RUNNER._test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name="local", + local_ipv4_addrs=local_ipv4_addrs, + ), + "127.0.0.1", + ) + self.assertEqual( + _RUNNER._test_stack_target_advertise_host( + cluster_nodes=cluster_nodes, + target_name="remote", + local_ipv4_addrs=local_ipv4_addrs, + ), + "192.0.2.11", + ) + + def test_kv_keyspace_capacity_guard_can_be_disabled_for_ssd_pressure(self) -> None: + ts_scene = { + "keyspace_size": 512, + "value_size_mode": "FIXED", + } + scale = { + "owner": { + "owner_count": 1, + "owner_dram_bytes": 512 * 1024 * 1024, + } + } + + guarded = _RUNNER._test_stack_effective_kv_keyspace_size( + case_id="guarded", + ts_scene=ts_scene, + scale=scale, + bench_value_size=4 * 1024 * 1024, + ) + unguarded = _RUNNER._test_stack_effective_kv_keyspace_size( + case_id="unguarded", + ts_scene={ + **ts_scene, + _RUNNER.TEST_STACK_SCENE_KEY_KEYSPACE_CAPACITY_GUARD: False, + }, + scale=scale, + bench_value_size=4 * 1024 * 1024, + ) + + self.assertLess(guarded, ts_scene["keyspace_size"]) + self.assertEqual(unguarded, ts_scene["keyspace_size"]) + + def test_kv_bootstrap_pressure_controls_parse(self) -> None: + scene = _RUNNER._parse_scene( + { + "test_stack": { + "mode": "KVSTORE", + "read_ratio": 0.9, + "write_ratio": 0.1, + "request_distribution": "zipfian", + "keyspace_size": 512, + "kv_bootstrap_concurrency": 1, + "kv_bootstrap_put_gap_ms": 10.5, + "kv_bootstrap_storage_full_policy": "stop", + "kv_get_output": "CUDA", + "kv_cuda_device_index": 6, + "value_size_mode": "FIXED", + }, + "select": {"scales": ["n1"], "profiles": ["fluxon"]}, + }, + "scenes.kv_ssd_pressure_zipf", + ) + + self.assertEqual( + scene["test_stack"][_RUNNER.TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_CONCURRENCY], + 1, + ) + self.assertEqual( + scene["test_stack"][_RUNNER.TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_PUT_GAP_MS], + 10.5, + ) + self.assertEqual( + scene["test_stack"][_RUNNER.TEST_STACK_SCENE_KEY_KV_BOOTSTRAP_STORAGE_FULL_POLICY], + "stop", + ) + self.assertEqual( + scene["test_stack"][_RUNNER.TEST_STACK_SCENE_KEY_KV_GET_OUTPUT], + "cuda", + ) + self.assertEqual( + scene["test_stack"][_RUNNER.TEST_STACK_SCENE_KEY_KV_CUDA_DEVICE_INDEX], + 6, + ) + + def test_kv_get_output_rejects_unbounded_variants(self) -> None: + with self.assertRaisesRegex(ValueError, "kv_get_output invalid"): + _RUNNER._parse_scene( + { + "test_stack": { + "mode": "KVSTORE", + "read_ratio": 1.0, + "write_ratio": 0.0, + "kv_get_output": "custom_callback", + "value_size_mode": "FIXED", + }, + "select": {"scales": ["n1"], "profiles": ["fluxon"]}, + }, + "scenes.kv_ssd_pressure_zipf", + ) + + def test_cuda_device_index_requires_cuda_output(self) -> None: + with self.assertRaisesRegex(ValueError, "kv_cuda_device_index requires kv_get_output=cuda"): + _RUNNER._parse_scene( + { + "test_stack": { + "mode": "KVSTORE", + "read_ratio": 1.0, + "write_ratio": 0.0, + "kv_get_output": "holder", + "kv_cuda_device_index": 6, + "value_size_mode": "FIXED", + }, + "select": {"scales": ["n1"], "profiles": ["fluxon"]}, + }, + "scenes.kv_ssd_pressure_zipf", + ) + def test_ci_source_overlay_includes_fluxon_test_stack(self) -> None: self.assertIn("fluxon_test_stack", _RUNNER._CI_SOURCE_OVERLAY_ROOTS) self.assertNotIn("quartz_prewarm", _RUNNER._CI_SOURCE_OVERLAY_ROOTS) @@ -1506,7 +2287,7 @@ def test_ci_base_runtime_service_target_ip_uses_loopback_for_same_host_local_nod ) resolved_case = { "deploy": { - "target_ip_map": {"logic-a": "192.168.1.10", "logic-b": "192.168.1.10"}, + "target_ip_map": {"logic-a": "192.0.2.10", "logic-b": "192.0.2.10"}, }, "profile": { "ci": { diff --git a/fluxon_test_stack/tests/test_top_attention_bin_kvtest_contract.py b/fluxon_test_stack/tests/test_top_attention_bin_kvtest_contract.py index 4d9b39c..3dd7d8f 100644 --- a/fluxon_test_stack/tests/test_top_attention_bin_kvtest_contract.py +++ b/fluxon_test_stack/tests/test_top_attention_bin_kvtest_contract.py @@ -54,7 +54,7 @@ def test_main_writes_build_config_ext_and_calls_cargo(self) -> None: }, "scene_config": { "kv_transport_feature": "tcp_thread_transport", - "kv_test_rounds": "p2p_only", + "kv_test_rounds": "p2p_only,p2p_only_ssd", }, "scene_runtime": { "etcd": {"ip": "127.0.0.1", "port": 19180}, @@ -95,7 +95,7 @@ def test_main_writes_build_config_ext_and_calls_cargo(self) -> None: ) self.assertEqual( run_cargo.call_args.kwargs["env"]["FLUXON_KV_TEST_ROUNDS"], - "p2p_only", + "p2p_only,p2p_only_ssd", ) diff --git a/fluxon_test_stack/tests/test_top_attention_mq_channel_contract.py b/fluxon_test_stack/tests/test_top_attention_mq_channel_contract.py index 4230d5d..b8c192e 100644 --- a/fluxon_test_stack/tests/test_top_attention_mq_channel_contract.py +++ b/fluxon_test_stack/tests/test_top_attention_mq_channel_contract.py @@ -45,6 +45,7 @@ def test_mq_channel_wrappers_accept_case_config_and_run_script_processes(self) - "fluxon_py/tests/test_api_chan_mpmc/test_api_chan_mpmc_quick_and_weighted_consume.py", "fluxon_py/tests/test_api_chan_mpmc/test_rebind_client.py", "fluxon_py/tests/test_api_chan_mpmc/test_ready_channels_access.py", + "fluxon_py/tests/test_api_chan_mpmc/test_create_reservation_contract.py", ], ), } diff --git a/fluxon_test_stack/top_attention_test_index/_bin_kvtest.py b/fluxon_test_stack/top_attention_test_index/_bin_kvtest.py index 3b3d859..4cd9c05 100644 --- a/fluxon_test_stack/top_attention_test_index/_bin_kvtest.py +++ b/fluxon_test_stack/top_attention_test_index/_bin_kvtest.py @@ -15,7 +15,13 @@ TEST_REQUIREMENTS = ["cargo", "etcd", "ops", "submodules"] SCENE_ID = "ci_top_attention_bin_kvtest" -KV_TEST_ROUND_NAMES = ("p2p_only", "rdma_transfer_only", "rdma_transfer_with_rpc") +KV_TEST_ROUND_NAMES = ( + "p2p_only", + "p2p_only_ssd", + "p2p_only_memory_ssd", + "rdma_transfer_only", + "rdma_transfer_with_rpc", +) def _parse_kv_test_rounds(raw: object) -> str: diff --git a/fluxon_test_stack/top_attention_test_index/_mq_mpmc.py b/fluxon_test_stack/top_attention_test_index/_mq_mpmc.py index e041b4b..cfd48b2 100755 --- a/fluxon_test_stack/top_attention_test_index/_mq_mpmc.py +++ b/fluxon_test_stack/top_attention_test_index/_mq_mpmc.py @@ -16,6 +16,7 @@ "fluxon_py/tests/test_api_chan_mpmc/test_api_chan_mpmc_quick_and_weighted_consume.py", "fluxon_py/tests/test_api_chan_mpmc/test_rebind_client.py", "fluxon_py/tests/test_api_chan_mpmc/test_ready_channels_access.py", + "fluxon_py/tests/test_api_chan_mpmc/test_create_reservation_contract.py", ] diff --git a/pics/kv_ssd_cpu_sweep.png b/pics/kv_ssd_cpu_sweep.png new file mode 100644 index 0000000..57fd357 Binary files /dev/null and b/pics/kv_ssd_cpu_sweep.png differ diff --git a/pics/kv_ssd_foyer_ablation.png b/pics/kv_ssd_foyer_ablation.png new file mode 100644 index 0000000..c6bde9a Binary files /dev/null and b/pics/kv_ssd_foyer_ablation.png differ diff --git a/pics/kv_ssd_gpu_sweep.png b/pics/kv_ssd_gpu_sweep.png new file mode 100644 index 0000000..802a3c7 Binary files /dev/null and b/pics/kv_ssd_gpu_sweep.png differ diff --git a/scripts/_build_doc_site_in_container_inner.py b/scripts/_build_doc_site_in_container_inner.py index ef15c30..a341685 100644 --- a/scripts/_build_doc_site_in_container_inner.py +++ b/scripts/_build_doc_site_in_container_inner.py @@ -1636,9 +1636,6 @@ def rewrite_target_path(raw_target: str) -> str: def rewrite_variant_target_path(raw_target: str, *, language: str) -> str: - if language != "cn": - return rewrite_target_path(raw_target) - unescaped_target = decode_markdown_target(raw_target) if ( not unescaped_target @@ -1653,6 +1650,19 @@ def rewrite_variant_target_path(raw_target: str, *, language: str) -> str: split = urllib.parse.urlsplit(unescaped_target) path_part = urllib.parse.unquote(split.path) + if path_part.startswith("../../pics/"): + # Source docs live one directory below the repository root. Quartz resolves + # staged page assets from the page route, so remove that repository-only hop. + rebuilt = urllib.parse.quote(path_part[len("../") :], safe="/.") + if split.query: + rebuilt += f"?{split.query}" + if split.fragment: + rebuilt += f"#{split.fragment}" + return rebuilt + + if language != "cn": + return rewrite_target_path(raw_target) + if path_part.endswith("README_CN.md"): normalized_path = normalize_readme_target_path(path_part) if normalized_path is not None: diff --git a/setup_and_pack/build_pack_fluxonkv_pylib_img.py b/setup_and_pack/build_pack_fluxonkv_pylib_img.py index 788cfac..e72d86f 100755 --- a/setup_and_pack/build_pack_fluxonkv_pylib_img.py +++ b/setup_and_pack/build_pack_fluxonkv_pylib_img.py @@ -23,9 +23,11 @@ GENERATED_CONFIG = SCRIPT_DIR / "build_pack_fluxonkv_pylib_img" / "pypack_builder_manylinux_2_28.generated.yaml" PREPARE_ARTIFACT_COPY_STEP_NAME = "copy_prepare_build_artifacts" LEGACY_COPY_STEP_NAME = "copy_sources_and_pyo3_build" +FLUXON_FS_VIDEO_FFMPEG_STEP_NAME = "install_fluxon_fs_video_ffmpeg" RUST_TOOLCHAIN_PLACEHOLDER = "__FLUXON_RUST_TOOLCHAIN__" FLUXON_RS_TOOLCHAIN_TOML = Path("fluxon_rs") / "rust-toolchain.toml" PREPARE_BUILD_YAML = Path("setup_and_pack") / "pub_prepare_build.yaml" +RPM_FUSION_FREE_RELEASE_RPM = "https://download1.rpmfusion.org/free/el/rpmfusion-free-release-8.noarch.rpm" def _host_arch_name() -> str: @@ -227,9 +229,29 @@ def add_path_copy(src_path: Path, dst_path: str) -> None: script_installs[:] = [ step for step in script_installs - if step.get("name") not in {PREPARE_ARTIFACT_COPY_STEP_NAME, LEGACY_COPY_STEP_NAME} + if step.get("name") + not in { + PREPARE_ARTIFACT_COPY_STEP_NAME, + LEGACY_COPY_STEP_NAME, + FLUXON_FS_VIDEO_FFMPEG_STEP_NAME, + } ] + script_installs.append( + { + "name": FLUXON_FS_VIDEO_FFMPEG_STEP_NAME, + "commands": [ + "rpm -q rpmfusion-free-release >/dev/null 2>&1 || " + f"yum install -y {RPM_FUSION_FREE_RELEASE_RPM}", + "yum install -y --setopt=tsflags=nodocs --setopt=install_weak_deps=False ffmpeg-devel", + "pkg-config --modversion libavformat libavcodec libavutil libswscale", + "pkg-config --cflags libavcodec", + "yum clean all", + "rm -rf /var/cache/yum", + ], + } + ) + if not copies_field: return diff --git a/setup_and_pack/nix/pack_fluxonkv_pylib.py b/setup_and_pack/nix/pack_fluxonkv_pylib.py index e12f8fe..e643966 100644 --- a/setup_and_pack/nix/pack_fluxonkv_pylib.py +++ b/setup_and_pack/nix/pack_fluxonkv_pylib.py @@ -178,7 +178,7 @@ def _cleanup_temp_workspace_mount_dirs() -> None: # Public fluxon_pyo3 only accepts its own transport toggles. Closed SDK linkage is # provided through fluxon_commu_closed_sdk_consumer via FLUXON_COMMU_CLOSED_SDK_ROOT, # so the wheel build must not inject closed-only feature names into the open crate. -PYO3_BASE_FEATURES = ("p2p_transfer",) +PYO3_BASE_FEATURES = ("p2p_transfer", "fluxon_fs_video_ffmpeg") RDMA_BACKEND_FEATURES = { "closed_sdk": [], } @@ -227,7 +227,11 @@ def _wheel_variant_key(transport_backend: str, rdma_backend: str) -> str: def _transport_backend_feature_csv(transport_backend: str, rdma_backend: str) -> str: - features = TRANSPORT_BACKEND_FEATURES[transport_backend] + RDMA_BACKEND_FEATURES[rdma_backend] + features = ( + list(PYO3_BASE_FEATURES) + + TRANSPORT_BACKEND_FEATURES[transport_backend] + + RDMA_BACKEND_FEATURES[rdma_backend] + ) return ",".join(features) diff --git a/setup_and_pack/tests/test_doc_site_builder_image_workflow.py b/setup_and_pack/tests/test_doc_site_builder_image_workflow.py index e91244c..9e6eb71 100644 --- a/setup_and_pack/tests/test_doc_site_builder_image_workflow.py +++ b/setup_and_pack/tests/test_doc_site_builder_image_workflow.py @@ -44,10 +44,19 @@ def test_main_testbed_workflow_keeps_suite_generation_in_workflow(self) -> None: self.assertIn("fluxon_test_stack/ci_2_virt_node.py", workflow_text) self.assertIn("Write ci_2_virt_node suite", workflow_text) self.assertIn("ci_top_attention_bin_kvtest", workflow_text) + self.assertIn("ci_top_attention_cargo_kv_unit", workflow_text) + self.assertIn('"kv_test_rounds": "p2p_only,p2p_only_ssd"', workflow_text) self.assertIn("ci_top_attention_doc_page_build", workflow_text) self.assertIn("ci_top_attention_mq_core", workflow_text) self.assertIn("doc_site_base_url", workflow_text) self.assertIn("rather_no_git_submodule.py", workflow_text) + self.assertIn("Reclaim runner disk before CI", workflow_text) + self.assertIn("docker system prune --all --force --volumes", workflow_text) + self.assertIn("/usr/local/lib/android", workflow_text) + self.assertIn("/usr/share/dotnet", workflow_text) + self.assertIn("--cleanup-pack-runtime-after-success", workflow_text) + self.assertIn("--cleanup-successful-case-artifacts", workflow_text) + self.assertIn("if: ${{ failure() }}", workflow_text) def test_docs_pages_uses_container_entrypoint(self) -> None: workflow_text = DOCS_PAGES_WORKFLOW_PATH.read_text(encoding="utf-8") diff --git a/setup_and_pack/utils/kill_by_cmd_name.py b/setup_and_pack/utils/kill_by_cmd_name.py index 5be690e..919a5b3 100755 --- a/setup_and_pack/utils/kill_by_cmd_name.py +++ b/setup_and_pack/utils/kill_by_cmd_name.py @@ -20,6 +20,6 @@ def kill_processes_by_exact_cmdline(target_cmdline): # 示例：匹配精确的命令行 kill_processes_by_exact_cmdline( "rust-analyzer" - # "/home/dehuazhang12/.vscode-server/cli/servers/Stable-c306e94f98122556ca081f527b466015e1bc37b0/server/node" - # "/home/pa/.cursor-server/extensions/rust-lang.rust-analyzer-0.3.2474-linux-x64/server/rust-analyzer" - ) \ No newline at end of file + # "/home/example-user-a/.vscode-server/cli/servers/Stable-c306e94f98122556ca081f527b466015e1bc37b0/server/node" + # "/home/example-user-b/.cursor-server/extensions/rust-lang.rust-analyzer-0.3.2474-linux-x64/server/rust-analyzer" + )

owner_idresource device used cap{{ d.resource_kind }} {{ d.device }} {{ d.used }} {{ d.cap }}