From 46e058d15e6fe41b0eb16df98cc5e19c001b0b7a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 7 Oct 2026 16:21:11 -0400 Subject: [PATCH 1/2] ci: defer the main.yaml integration-tests OCI lane off per-PR Completes the no-per-PR-OCI-builds intent: image-off-pr moved image.yaml off PRs and #2877 moved the nix integration job off PRs, but main.yaml's integration-tests (docker-compose image build + run_integration_tests.sh) fell between them and still ran on every PR. Defer to nightly + merge-to-main, matching #2880's pattern. Fork-local until folded into the upstream OCI/nightly set. --- .github/workflows/main.yaml | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 43dacaf1ea0..232952a3a33 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -1,17 +1,23 @@ name: CI # Controls when the workflow will run. +# +# The only live job here (integration-tests) builds the trace_machina/nativelink +# OCI image from the docker-compose Dockerfile and runs run_integration_tests.sh. +# That is a per-PR OCI build — the same class of heavy work #image-off-pr moved +# off PRs for image.yaml and #2877 moved off PRs for the nix integration job. +# This lane fell in the crack between those two, so defer it the same way: +# nightly cross-check + on every merge to main, not on every PR. Per-PR +# build/test signal is carried by the Nix "Bazel Dev" lane. on: push: branches: [ main ] paths-ignore: - '.github/styles/**' - 'web/**' - pull_request: - branches: [ main ] - paths-ignore: - - '.github/styles/**' - - 'web/**' + schedule: + - cron: '0 7 * * *' # 07:00 UTC daily (off-peak) + workflow_dispatch: permissions: read-all From 032270363850c6f900e2e8c62663949d9d5ce022 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 10:41:45 -0400 Subject: [PATCH 2/2] Reproduce Azure retries after destination failure in Rust --- .../vocabularies/TraceMachina/accept.txt | 1 + Cargo.lock | 27 + Cargo.toml | 28 + MODULE.bazel.lock | 1 + examples/azure_fleet/BUILD.bazel | 66 + examples/azure_fleet/README.md | 97 ++ examples/azure_fleet/account.rs | 1296 +++++++++++++++++ examples/azure_fleet/checks.rs | 312 ++++ examples/azure_fleet/main.rs | 976 +++++++++++++ examples/azure_fleet/node.rs | 529 +++++++ examples/azure_fleet/transport.rs | 52 + examples/azure_fleet/witness.rs | 748 ++++++++++ examples/azure_fleet/workload.rs | 311 ++++ nativelink-store/BUILD.bazel | 7 + nativelink-store/src/azure_blob_store.rs | 89 +- .../tests/azure_blob_store_test.rs | 190 +++ .../tests/azure_download_egress_test.rs | 958 ++++++++++++ .../fixtures/azure_download_response.json | 184 +++ 18 files changed, 5868 insertions(+), 4 deletions(-) create mode 100644 examples/azure_fleet/BUILD.bazel create mode 100644 examples/azure_fleet/README.md create mode 100644 examples/azure_fleet/account.rs create mode 100644 examples/azure_fleet/checks.rs create mode 100644 examples/azure_fleet/main.rs create mode 100644 examples/azure_fleet/node.rs create mode 100644 examples/azure_fleet/transport.rs create mode 100644 examples/azure_fleet/witness.rs create mode 100644 examples/azure_fleet/workload.rs create mode 100644 nativelink-store/tests/azure_download_egress_test.rs create mode 100644 nativelink-store/tests/fixtures/azure_download_response.json diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index c0832d0ef15..0c701944a2b 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -70,6 +70,7 @@ Qwik [Uu]psert Verilator Verilog +Zipf Zstandard alex [Ee]xecution diff --git a/Cargo.lock b/Cargo.lock index 1f9474622e9..f1a8e7a6f6f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -699,10 +699,13 @@ checksum = "8a18ed336352031311f4e0b4dd2ff392d4fbb370777c9d18d7fc9d7359f73871" dependencies = [ "axum-core", "bytes", + "form_urlencoded", "futures-util", "http 1.4.1", "http-body 1.0.1", "http-body-util", + "hyper", + "hyper-util", "itoa", "matchit", "memchr", @@ -710,7 +713,10 @@ dependencies = [ "percent-encoding", "pin-project-lite", "serde_core", + "serde_path_to_error", + "serde_urlencoded", "sync_wrapper", + "tokio", "tower", "tower-layer", "tower-service", @@ -2935,13 +2941,17 @@ name = "nativelink" version = "1.8.0" dependencies = [ "async-lock", + "async-trait", "axum", + "azure_core", + "azure_storage_blob", "bytes", "clap", "futures", "hex", "hyper", "hyper-util", + "libc", "mimalloc", "nativelink-config", "nativelink-error", @@ -2951,14 +2961,20 @@ dependencies = [ "nativelink-store", "nativelink-util", "nativelink-worker", + "prost", + "prost-types", "rand 0.9.4", + "reqwest", "rustls-pki-types", + "serde", + "serde_json", "sha2 0.10.9", "tokio", "tokio-rustls", "tonic", "tower", "tracing", + "tracing-subscriber", ] [[package]] @@ -4499,6 +4515,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" diff --git a/Cargo.toml b/Cargo.toml index 0f20bfb61f3..4d278d146b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,6 +82,34 @@ tonic = { version = "0.14.0", features = [ tower = { version = "0.5.2", default-features = false } tracing = { version = "0.1.41", default-features = false } +[dev-dependencies] +async-trait = { version = "0.1.88", default-features = false } +axum = { version = "0.8.3", default-features = false, features = [ + "http1", + "query", + "tokio", +] } +azure_core = { version = "1", default-features = false, features = [ + "reqwest", + "tokio", +] } +azure_storage_blob = { version = "1", default-features = false, features = [ + "tokio", +] } +libc = { version = "0.2.183", default-features = false } +prost = { version = "0.14.4", default-features = false } +prost-types = { version = "0.14.4", default-features = false } +reqwest = { version = "0.13", default-features = false, features = ["stream"] } +serde = { version = "1.0.219", default-features = false, features = ["derive"] } +serde_json = { version = "1.0.140", default-features = false, features = [ + "std", +] } +tokio = { version = "1.52.2", default-features = false, features = ["io-std"] } +tracing-subscriber = { version = "0.3.19", default-features = false, features = [ + "env-filter", + "fmt", +] } + [workspace.cargo-features-manager.keep] async-lock = ["std"] aws-sdk-s3 = ["rt-tokio"] diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index fceb1b0e9e9..7dc71bf6d2e 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1144,6 +1144,7 @@ "serde_json_1.0.145": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"automod\",\"req\":\"^1.0.11\"},{\"name\":\"indexmap\",\"optional\":true,\"req\":\"^2.2.3\"},{\"kind\":\"dev\",\"name\":\"indoc\",\"req\":\"^2.0.2\"},{\"name\":\"itoa\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"ref-cast\",\"req\":\"^1.0.18\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0.13\"},{\"name\":\"ryu\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde\",\"req\":\"^1.0.220\",\"target\":\"cfg(any())\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.194\"},{\"kind\":\"dev\",\"name\":\"serde_bytes\",\"req\":\"^0.11.10\"},{\"default_features\":false,\"name\":\"serde_core\",\"req\":\"^1.0.220\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0.166\"},{\"kind\":\"dev\",\"name\":\"serde_stacker\",\"req\":\"^0.1.8\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.108\"}],\"features\":{\"alloc\":[\"serde_core/alloc\"],\"arbitrary_precision\":[],\"default\":[\"std\"],\"float_roundtrip\":[],\"preserve_order\":[\"indexmap\",\"std\"],\"raw_value\":[],\"std\":[\"memchr/std\",\"serde_core/std\"],\"unbounded_depth\":[]}}", "serde_json_1.0.149": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"automod\",\"req\":\"^1.0.11\"},{\"name\":\"indexmap\",\"optional\":true,\"req\":\"^2.2.3\"},{\"kind\":\"dev\",\"name\":\"indoc\",\"req\":\"^2.0.2\"},{\"name\":\"itoa\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"ref-cast\",\"req\":\"^1.0.18\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0.13\"},{\"default_features\":false,\"name\":\"serde\",\"req\":\"^1.0.220\",\"target\":\"cfg(any())\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.194\"},{\"kind\":\"dev\",\"name\":\"serde_bytes\",\"req\":\"^0.11.10\"},{\"default_features\":false,\"name\":\"serde_core\",\"req\":\"^1.0.220\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0.166\"},{\"kind\":\"dev\",\"name\":\"serde_stacker\",\"req\":\"^0.1.8\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.108\"},{\"name\":\"zmij\",\"req\":\"^1.0\"}],\"features\":{\"alloc\":[\"serde_core/alloc\"],\"arbitrary_precision\":[],\"default\":[\"std\"],\"float_roundtrip\":[],\"preserve_order\":[\"indexmap\",\"std\"],\"raw_value\":[],\"std\":[\"memchr/std\",\"serde_core/std\"],\"unbounded_depth\":[]}}", "serde_json_1.0.150": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"automod\",\"req\":\"^1.0.11\"},{\"name\":\"indexmap\",\"optional\":true,\"req\":\"^2.2.3\"},{\"kind\":\"dev\",\"name\":\"indoc\",\"req\":\"^2.0.2\"},{\"name\":\"itoa\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"ref-cast\",\"req\":\"^1.0.18\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0.13\"},{\"default_features\":false,\"name\":\"serde\",\"req\":\"^1.0.220\",\"target\":\"cfg(any())\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.194\"},{\"kind\":\"dev\",\"name\":\"serde_bytes\",\"req\":\"^0.11.10\"},{\"default_features\":false,\"name\":\"serde_core\",\"req\":\"^1.0.220\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0.166\"},{\"kind\":\"dev\",\"name\":\"serde_stacker\",\"req\":\"^0.1.8\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.108\"},{\"name\":\"zmij\",\"req\":\"^1.0\"}],\"features\":{\"alloc\":[\"serde_core/alloc\"],\"arbitrary_precision\":[],\"default\":[\"std\"],\"float_roundtrip\":[],\"preserve_order\":[\"indexmap\",\"std\"],\"raw_value\":[],\"std\":[\"memchr/std\",\"serde_core/std\"],\"unbounded_depth\":[]}}", + "serde_path_to_error_0.1.20": "{\"dependencies\":[{\"name\":\"itoa\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde\",\"req\":\"^1.0.220\",\"target\":\"cfg(any())\"},{\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.220\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"serde_core\",\"req\":\"^1.0.220\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0.220\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.100\"}],\"features\":{}}", "serde_urlencoded_0.7.1": "{\"dependencies\":[{\"name\":\"form_urlencoded\",\"req\":\"^1\"},{\"name\":\"itoa\",\"req\":\"^1\"},{\"name\":\"ryu\",\"req\":\"^1\"},{\"name\":\"serde\",\"req\":\"^1.0.69\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1\"}],\"features\":{}}", "serde_with_3.21.0": "{\"dependencies\":[{\"default_features\":false,\"name\":\"base64\",\"optional\":true,\"req\":\"^0.22.1\"},{\"default_features\":false,\"name\":\"bs58\",\"optional\":true,\"req\":\"^0.5.1\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"chrono_0_4\",\"optional\":true,\"package\":\"chrono\",\"req\":\"^0.4.20\"},{\"name\":\"document-features\",\"optional\":true,\"req\":\"^0.2.7\"},{\"kind\":\"dev\",\"name\":\"expect-test\",\"req\":\"^1.5.1\"},{\"kind\":\"dev\",\"name\":\"fnv\",\"req\":\"^1.0.6\"},{\"kind\":\"dev\",\"name\":\"glob\",\"req\":\"^0.3.3\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"hashbrown_0_14\",\"optional\":true,\"package\":\"hashbrown\",\"req\":\"^0.14.0\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"hashbrown_0_15\",\"optional\":true,\"package\":\"hashbrown\",\"req\":\"^0.15.0\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"hashbrown_0_16\",\"optional\":true,\"package\":\"hashbrown\",\"req\":\"^0.16.0\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"hashbrown_0_17\",\"optional\":true,\"package\":\"hashbrown\",\"req\":\"^0.17.0\"},{\"default_features\":false,\"name\":\"hex\",\"optional\":true,\"req\":\"^0.4.3\"},{\"default_features\":false,\"features\":[\"serde-1\"],\"name\":\"indexmap_1\",\"optional\":true,\"package\":\"indexmap\",\"req\":\"^1.8\"},{\"default_features\":false,\"features\":[\"serde\"],\"name\":\"indexmap_2\",\"optional\":true,\"package\":\"indexmap\",\"req\":\"^2.0\"},{\"default_features\":false,\"features\":[\"resolve-file\"],\"kind\":\"dev\",\"name\":\"jsonschema\",\"req\":\"^0.33.0\"},{\"kind\":\"dev\",\"name\":\"mime\",\"req\":\"^0.3.16\"},{\"kind\":\"dev\",\"name\":\"pretty_assertions\",\"req\":\"^1.4.0\"},{\"default_features\":false,\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"regex\",\"req\":\"^1.12.1\"},{\"kind\":\"dev\",\"name\":\"rmp-serde\",\"req\":\"^1.3.0\"},{\"kind\":\"dev\",\"name\":\"ron\",\"req\":\"^0.12\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0.22\"},{\"default_features\":false,\"name\":\"schemars_0_8\",\"optional\":true,\"package\":\"schemars\",\"req\":\"^0.8.16\"},{\"kind\":\"dev\",\"name\":\"schemars_0_8\",\"package\":\"schemars\",\"req\":\"^0.8.16\"},{\"default_features\":false,\"name\":\"schemars_0_9\",\"optional\":true,\"package\":\"schemars\",\"req\":\"^0.9.0\"},{\"kind\":\"dev\",\"name\":\"schemars_0_9\",\"package\":\"schemars\",\"req\":\"^0.9.0\"},{\"default_features\":false,\"name\":\"schemars_1\",\"optional\":true,\"package\":\"schemars\",\"req\":\"^1.0.2\"},{\"kind\":\"dev\",\"name\":\"schemars_1\",\"package\":\"schemars\",\"req\":\"^1.0.2\"},{\"default_features\":false,\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.152\"},{\"kind\":\"dev\",\"name\":\"serde-xml-rs\",\"req\":\"^0.8.1\"},{\"default_features\":false,\"features\":[\"result\"],\"name\":\"serde_core\",\"req\":\"^1.0.225\"},{\"default_features\":false,\"name\":\"serde_json\",\"optional\":true,\"req\":\"^1.0.145\"},{\"features\":[\"preserve_order\"],\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.25\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.124\"},{\"name\":\"serde_with_macros\",\"optional\":true,\"req\":\"=3.21.0\"},{\"default_features\":false,\"name\":\"smallvec_1\",\"optional\":true,\"package\":\"smallvec\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"time_0_3\",\"optional\":true,\"package\":\"time\",\"req\":\"~0.3.47\"},{\"kind\":\"dev\",\"name\":\"yaml_serde\",\"req\":\"^0.10.3\"}],\"features\":{\"alloc\":[\"serde_core/alloc\",\"base64?/alloc\",\"bs58?/alloc\",\"chrono_0_4?/alloc\",\"hex?/alloc\",\"serde_json?/alloc\",\"time_0_3?/alloc\"],\"base58\":[\"dep:bs58\",\"alloc\"],\"base64\":[\"dep:base64\",\"alloc\"],\"chrono\":[\"chrono_0_4\"],\"chrono_0_4\":[\"dep:chrono_0_4\"],\"default\":[\"std\",\"macros\"],\"guide\":[\"dep:document-features\",\"macros\",\"std\"],\"hashbrown_0_14\":[\"dep:hashbrown_0_14\",\"alloc\"],\"hashbrown_0_15\":[\"dep:hashbrown_0_15\",\"alloc\"],\"hashbrown_0_16\":[\"dep:hashbrown_0_16\",\"alloc\"],\"hashbrown_0_17\":[\"dep:hashbrown_0_17\",\"alloc\"],\"hex\":[\"dep:hex\",\"alloc\"],\"indexmap\":[\"indexmap_1\"],\"indexmap_1\":[\"dep:indexmap_1\",\"alloc\"],\"indexmap_2\":[\"dep:indexmap_2\",\"alloc\"],\"json\":[\"dep:serde_json\",\"alloc\"],\"macros\":[\"dep:serde_with_macros\"],\"schemars_0_8\":[\"dep:schemars_0_8\",\"std\",\"serde_with_macros?/schemars_0_8\",\"dep:serde_json\"],\"schemars_0_9\":[\"dep:schemars_0_9\",\"alloc\",\"serde_with_macros?/schemars_0_9\",\"dep:serde_json\"],\"schemars_1\":[\"dep:schemars_1\",\"alloc\",\"serde_with_macros?/schemars_1\",\"dep:serde_json\"],\"smallvec_1\":[\"dep:smallvec_1\"],\"std\":[\"alloc\",\"bs58?/std\",\"serde_core/std\",\"chrono_0_4?/clock\",\"chrono_0_4?/std\",\"indexmap_1?/std\",\"indexmap_2?/std\",\"time_0_3?/serde-well-known\",\"time_0_3?/std\",\"schemars_0_9?/std\",\"schemars_1?/std\"],\"time_0_3\":[\"dep:time_0_3\"]}}", "serde_with_macros_3.21.0": "{\"dependencies\":[{\"name\":\"darling\",\"req\":\"^0.23.0\"},{\"kind\":\"dev\",\"name\":\"expect-test\",\"req\":\"^1.5.1\"},{\"kind\":\"dev\",\"name\":\"glob\",\"req\":\"^0.3.3\"},{\"kind\":\"dev\",\"name\":\"pretty_assertions\",\"req\":\"^1.4.0\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0.1\"},{\"name\":\"quote\",\"req\":\"^1.0.0\"},{\"default_features\":false,\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"regex\",\"req\":\"^1.12.1\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0.22\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.152\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.25\"},{\"features\":[\"extra-traits\",\"full\",\"parsing\"],\"name\":\"syn\",\"req\":\"^2.0.0\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.111\"}],\"features\":{\"schemars_0_8\":[],\"schemars_0_9\":[],\"schemars_1\":[]}}", diff --git a/examples/azure_fleet/BUILD.bazel b/examples/azure_fleet/BUILD.bazel new file mode 100644 index 00000000000..38953383e0b --- /dev/null +++ b/examples/azure_fleet/BUILD.bazel @@ -0,0 +1,66 @@ +load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test") + +_SOURCES = [ + "account.rs", + "checks.rs", + "main.rs", + "node.rs", + "transport.rs", + "witness.rs", + "workload.rs", +] + +_COMPILE_DATA = ["//nativelink-store:tests/fixtures/azure_download_response.json"] + +_DEPS = [ + "//nativelink-config", + "//nativelink-proto", + "//nativelink-scheduler", + "//nativelink-service", + "//nativelink-store", + "//nativelink-util", + "//nativelink-worker", + "@crates//:async-trait", + "@crates//:axum", + "@crates//:azure_core", + "@crates//:azure_storage_blob", + "@crates//:bytes", + "@crates//:clap", + "@crates//:futures", + "@crates//:hex", + "@crates//:libc", + "@crates//:prost", + "@crates//:prost-types", + "@crates//:reqwest", + "@crates//:serde", + "@crates//:serde_json", + "@crates//:sha2", + "@crates//:tokio", + "@crates//:tokio-rustls", + "@crates//:tonic", + "@crates//:tracing", + "@crates//:tracing-subscriber", +] + +# The complete experiment runs explicitly: it needs unshare, mount, taskset, +# and a writable output directory. Ordinary tests exercise the account's real +# SDK transport and disk accounting without mounting filesystems or running a fleet. +rust_binary( + name = "azure_fleet", + srcs = _SOURCES, + compile_data = _COMPILE_DATA, + crate_root = "main.rs", + target_compatible_with = ["@platforms//os:linux"], + deps = _DEPS, +) + +rust_test( + name = "azure_fleet_test", + size = "small", + timeout = "short", + srcs = _SOURCES, + compile_data = _COMPILE_DATA, + crate_root = "main.rs", + target_compatible_with = ["@platforms//os:linux"], + deps = _DEPS, +) diff --git a/examples/azure_fleet/README.md b/examples/azure_fleet/README.md new file mode 100644 index 00000000000..05d63624fe7 --- /dev/null +++ b/examples/azure_fleet/README.md @@ -0,0 +1,97 @@ +# Azure failed-fill fleet regression + +Four separate NativeLink worker processes use the real scheduler over gRPC, +filesystem caches, `FastSlowStore`, `VerifyStore`, and `AzureBlobStore`. The Azure +Blob SDK talks to a shared loopback HTTP account. Two workers' private tmpfs +filesystems are filled to real `ENOSPC`, then repaired without restarting them. +The same Rust executable checks the fault, recovery, and retry behavior. + +## Run + +Linux with unprivileged user/mount namespaces, `unshare`, `mount`, `taskset`, +`sh`, `sleep`, and `wc` is required. Allow several GiB of free memory for the four +private tmpfs mounts. No Azure credentials or Python tools are used. + +```bash +cargo test --example azure_fleet +cargo build --example azure_fleet +# Each output directory must not already exist. +target/debug/examples/azure_fleet --seed 2026101700 --scenario control \ + --output /tmp/azure-fleet-control +target/debug/examples/azure_fleet --seed 2026101700 --scenario failed-fill \ + --output /tmp/azure-fleet-failed-fill +``` + +Bazel also builds the executable and its ordinary unit tests: + +```bash +bazel test //examples/azure_fleet:azure_fleet_test +bazel build //examples/azure_fleet:azure_fleet +``` + +Run the failed-fill command on the regression commit, then repeat it with the +same seed and a new output directory on the fix commit. The executable has no +runtime bug switch: the linked `AzureBlobStore` implementation determines the +behavior. The ordinary tests exercise the HTTP account and acceptance checks; +the explicit fleet run needs the Linux facilities listed above. + +## Pass/fail contract + +The failed-fill run must observe worker `ENOSPC` on both filled filesystems, +repair them, complete every offered action without a censored client attempt, +and complete a warm-cache sentinel in each original worker process. Scheduler +samples must respect the two-action limit and end with the original workers +holding zero reservations. All workers must shut down cleanly. + +It fails at the first new outer download after the same read has observed a +failed send with a broken writer. Each worker records ordered read, attempt, +failed-send and completion events. Immutable SDK context carries the same read +and attempt identity to every HTTP partition, including tasks that outlive the +consuming download. The report joins these identities, checks complete trace +coverage, and separates bytes from forbidden new attempts from bytes belonging +to older SDK work. Missing or inconsistent evidence fails the run. + +Both filled workers must actually reach the failed-send boundary. A final 503 +cannot hide an earlier forbidden attempt, and legitimate origin retries ending +at the first failed send are not counted as forbidden retries. The old +retry-exhaustion log count remains a diagnostic, not the acceptance condition. +The control must satisfy the recovery checks without encountering `ENOSPC`. + +The broken implementation should fail the failed-fill check even when all builds +succeed. The terminal-destination fix should pass the identical check. Byte +counts and elapsed times are observations, never pass/fail thresholds. +`run.json`, the workload manifest, node logs, per-worker `downloads.jsonl` +traces, and cleanup outcomes are retained +on both acceptance success and failure. Setup/session errors also follow the +report and cleanup path after the account has been constructed. + +## Workload and measurement + +The printed seed fixes the finite input catalogue, bounded log-normal sizes, +file counts and execution holds, Zipf input popularity, Poisson arrivals after +an initial burst, response residual draws, and retry jitter. A 128 MiB input +ensures the failed write can propagate before the whole download is buffered. +These distributions define the stress workload; their parameters are explicit +fixture inputs. +The account's response timing uses the existing Azure response-law fixture; +small metadata responses use a declared one-millisecond delay. + +A 256 MiB/s token bucket with a 256 MiB burst admits requested ranges or returns +503. The SDK controls its real partition downloads. Each worker records its +process-visible CPU count. Seeds do not make OS scheduling or HTTP interleaving +deterministic. + +The report retains HTTP ranges, response-body bytes, time buckets, all client +attempts, worker process IDs, filesystem allocation, and scheduler snapshots. +Byte ledgers are checked for conservation. Payload means bytes yielded by the +mock HTTP server to Hyper; it is not confirmed cache delivery or Azure billing. +Already-started SDK tasks may continue after the fix prevents a new outer retry. +Their requests retain the original attempt identity even if they begin after a +newer attempt. The most recent HEAD remains a separate observation. + +The reproduction adds observation metadata and debug events to the Azure read +path while leaving its retry decisions unchanged. Only the replay client's +HTTP policy exports that metadata as identifying headers. Trace capture adds +measurement overhead; paired revisions use the same instrumentation. Four +processes on one host exercise the distributed paths while sharing a kernel +and loopback network; this is not four Azure VMs. diff --git a/examples/azure_fleet/account.rs b/examples/azure_fleet/account.rs new file mode 100644 index 00000000000..9a6a7fc1ed9 --- /dev/null +++ b/examples/azure_fleet/account.rs @@ -0,0 +1,1296 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shared HTTP account. A declared admission budget causes every 503. Payload +//! counts are bytes yielded by the server body to Hyper, not socket receipts or +//! billed Azure egress. The fitted timing describes client/proxy complete-body +//! latency; it does not reproduce Azure's internal admission or network stack. + +use core::pin::Pin; +use core::task::{Context, Poll}; +use core::time::Duration; +use std::collections::{BTreeMap, HashMap}; +use std::sync::{Arc, Mutex}; + +use axum::Router; +use axum::body::{Body, to_bytes}; +use axum::extract::{Path, Query, State}; +use axum::http::{HeaderMap, Method, Response, StatusCode}; +use axum::routing::any; +use azure_storage_blob::models::BlockLookupList; +use bytes::Bytes; +use futures::Stream; +use serde::Serialize; +use serde_json::Value; +use tokio::time::Instant; + +use crate::workload::{Blob, NODES, uniform}; + +const CHUNK: usize = 64 * 1024; +const PARTITION: usize = 4 * 1024 * 1024; +const CAP: f64 = 256.0 * 1024.0 * 1024.0; +// Admit two initial 8-range SDK waves across four workers, so the failed-fill +// scenario reaches its actual disk failure before admission retries dominate. +const BURST_BYTES: usize = 256 * 1024 * 1024; +const BURST: f64 = BURST_BYTES as f64; + +#[derive(Default, Clone, Serialize)] +pub(crate) struct Bucket { + pub(crate) bytes: u64, + pub(crate) repeated_bytes: u64, + pub(crate) get: u64, + pub(crate) head: u64, + pub(crate) put: u64, + pub(crate) throttled: u64, + pub(crate) ingress: u64, +} + +#[derive(Serialize)] +pub(crate) struct HttpRead { + /// Unique within this account ledger, independent of SDK attempt identity. + pub(crate) request_id: u64, + pub(crate) node: usize, + pub(crate) blob: String, + /// Context explicitly carried by the SDK, never inferred from a HEAD or time. + pub(crate) read_id: Option, + pub(crate) attempt: Option, + pub(crate) sent: Option, + pub(crate) writer_broken: Option, + pub(crate) offset: u64, + pub(crate) length: u64, + pub(crate) start_s: f64, + pub(crate) end_s: Option, + pub(crate) bytes: u64, + pub(crate) repeated_bytes: u64, + pub(crate) status: u16, + pub(crate) reason: &'static str, + pub(crate) phase: String, + pub(crate) local_inflight: u32, + pub(crate) timing_outside_fit: bool, + /// Most recent HEAD for this node/key. An observed request boundary, not a + /// general assertion that every HEAD starts a fill or every GET belongs to it. + pub(crate) head_sequence: Option, +} + +#[derive(Serialize)] +pub(crate) struct Head { + pub(crate) node: usize, + pub(crate) blob: String, + pub(crate) sequence: u64, + pub(crate) time_s: f64, + pub(crate) status: u16, + pub(crate) length: Option, + pub(crate) phase: String, +} + +#[derive(Clone)] +struct StoredBlob { + data: Bytes, + etag: String, +} + +struct Timing { + intercept: f64, + size_power: f64, + concurrency_power: f64, + residuals: Vec, +} + +impl Timing { + fn measured() -> Self { + let fit: Value = serde_json::from_str(include_str!( + "../../nativelink-store/tests/fixtures/azure_download_response.json" + )) + .unwrap(); + let residuals: Vec = fit["log_residual_quantiles"] + .as_array() + .unwrap() + .iter() + .map(|x| x.as_f64().unwrap()) + .collect(); + assert_eq!(residuals.len(), 129); + assert!(residuals.windows(2).all(|p| p[0] <= p[1])); + Self { + intercept: fit["baseline_ms"]["intercept"].as_f64().unwrap(), + size_power: fit["baseline_ms"]["size_power"].as_f64().unwrap(), + concurrency_power: fit["baseline_ms"]["concurrency_power"].as_f64().unwrap(), + residuals, + } + } + + fn duration(&self, inflight: u32, quantile: f64, length: usize) -> Duration { + // The workload's payloads are partition-aligned. Short metadata bodies + // use a declared one-millisecond delay, independently of the fitted law. + if length < PARTITION { + return Duration::from_millis(1); + } + assert!((0.0..1.0).contains(&quantile)); + let position = quantile * 128.0; + // Locate the lower empirical quantile without a float-to-index cast. + let index = (1_u32..=128) + .take_while(|&boundary| f64::from(boundary) <= position) + .count(); + let residual = (self.residuals[index + 1] - self.residuals[index]).mul_add( + position - f64::from(u32::try_from(index).unwrap()), + self.residuals[index], + ); + let size_mib = length as f64 / 1_048_576.0; + let ms = self.intercept + * (1.0 + size_mib).powf(self.size_power) + * (1.0 + f64::from(inflight)).powf(self.concurrency_power) + * residual.exp(); + Duration::from_secs_f64(ms / 1000.0) + } +} + +pub(crate) struct Account { + pub(crate) seed: u64, + pub(crate) start: Instant, + pub(crate) state: Mutex, + timing: Timing, +} + +pub(crate) struct AccountState { + blobs: HashMap, + // Retain blocks after commit: a repeated commit must be idempotent. + blocks: HashMap<(String, Vec), Bytes>, + version: u64, + tokens: f64, + last: Instant, + last_request: Instant, + active: [u32; NODES], + peak_active: [u32; NODES], + ordinal: HashMap<(usize, String, u64), u64>, + head_sequences: HashMap<(usize, String), u64>, + // Each node has one cache generation. Replication to a different node is + // first-use there, never classified as a within-generation repeat. + seen: HashMap<(usize, String), Vec<(u64, u64)>>, + pub(crate) phase: String, + pub(crate) buckets: BTreeMap, + pub(crate) reads: Vec, + pub(crate) heads: Vec, + pub(crate) admitted: u64, + pub(crate) completed_actions: u64, +} + +impl AccountState { + fn insert_blob(&mut self, name: String, data: Bytes) { + self.version += 1; + self.blobs.insert( + name, + StoredBlob { + data, + etag: format!("\"uploaded-{}\"", self.version), + }, + ); + } + + fn admit(&mut self, now: Instant, length: usize) -> bool { + self.tokens = CAP + .mul_add(now.duration_since(self.last).as_secs_f64(), self.tokens) + .min(BURST); + self.last = now; + if self.tokens < length as f64 { + return false; + } + self.tokens -= length as f64; + self.admitted += u64::try_from(length).unwrap(); + true + } +} + +impl Account { + pub(crate) fn new(seed: u64, blobs: &[Blob]) -> Arc { + let now = Instant::now(); + Arc::new(Self { + seed, + start: now, + timing: Timing::measured(), + state: Mutex::new(AccountState { + blobs: blobs + .iter() + .map(|b| { + ( + b.name(), + StoredBlob { + data: b.data.clone(), + etag: format!("\"{}\"", b.name()), + }, + ) + }) + .collect(), + blocks: HashMap::new(), + version: 0, + tokens: BURST, + last: now, + last_request: now, + active: [0; NODES], + peak_active: [0; NODES], + ordinal: HashMap::new(), + head_sequences: HashMap::new(), + seen: HashMap::new(), + phase: "warmup".into(), + buckets: BTreeMap::new(), + reads: Vec::new(), + heads: Vec::new(), + admitted: 0, + completed_actions: 0, + }), + }) + } + + pub(crate) fn phase(&self, value: &str) { + self.state.lock().unwrap().phase = value.into(); + } + + pub(crate) fn elapsed(&self) -> f64 { + self.start.elapsed().as_secs_f64() + } + + fn bin(&self) -> u64 { + u64::try_from(self.start.elapsed().as_millis() / 100).unwrap() + } + + pub(crate) fn report(&self) -> Value { + let state = self.state.lock().unwrap(); + let bytes: u64 = state.buckets.values().map(|b| b.bytes).sum(); + let read_bytes: u64 = state.reads.iter().map(|r| r.bytes).sum(); + let repeated: u64 = state.buckets.values().map(|b| b.repeated_bytes).sum(); + let (attributed_bytes, unknown_bytes) = + state + .reads + .iter() + .fold((0_u64, 0_u64), |(attributed, unknown), read| { + if read.read_id.is_some() + && read.attempt.is_some() + && read.sent.is_some() + && read.writer_broken.is_some() + { + (attributed + read.bytes, unknown) + } else { + (attributed, unknown + read.bytes) + } + }); + + // The time buckets and per-request ledger independently account for + // the same yielded bytes. Neither substitutes requested range length. + assert_eq!(bytes, read_bytes); + assert_eq!(bytes, attributed_bytes + unknown_bytes); + assert!( + state + .reads + .iter() + .enumerate() + .all(|(index, read)| { read.request_id == u64::try_from(index).unwrap() }) + ); + assert_eq!( + repeated, + state.reads.iter().map(|r| r.repeated_bytes).sum::() + ); + assert!(state.reads.iter().all(|r| { + r.bytes <= r.length + && r.repeated_bytes <= r.bytes + && ((200..300).contains(&r.status) || r.bytes == 0) + })); + assert_eq!( + state.reads.len() as u64, + state.buckets.values().map(|b| b.get).sum::() + ); + assert_eq!( + state.heads.len() as u64, + state.buckets.values().map(|b| b.head).sum::() + ); + assert!(bytes <= state.admitted); + assert!(state.admitted as f64 <= CAP.mul_add(self.elapsed(), BURST) + 1.0); + + serde_json::json!({ + "cap_bytes_s": CAP, + "burst_bytes": BURST, + "bucket_width_s": 0.1, + "capacity_boundary": "requested payload reserved at HTTP admission; not a wire-rate shaper", + "measurement_boundary": "successful response payload yielded by server Body to Hyper; excludes headers and error bodies; not confirmed socket receipt or billed egress", + "timing_boundary": "conditional client/proxy complete-body fit; independent empirical residual draws; metadata delay is declared, not calibrated", + "timing_outside_fit": state.reads.iter().filter(|r| r.timing_outside_fit).count(), + "admitted_payload": state.admitted, + "payload_bytes": bytes, + "attributed_payload_bytes": attributed_bytes, + "unknown_attribution_payload_bytes": unknown_bytes, + "attribution_boundary": "explicit SDK context: node/process + read_id + attempt; sent and writer_broken sampled before that SDK download; missing or invalid headers remain unknown", + "repeated_payload_bytes": repeated, + "active": state.active, + "peak_active": state.peak_active, + "head_sequence_boundary": "latest observed HEAD per node/key; the fixture's Verify-wrapped cold fill performs HEAD, Azure read retries do not; not a universal fill identifier", + "buckets": state.buckets, + "reads": state.reads, + "heads": state.heads, + "completed_actions": state.completed_actions, + }) + } + + /// Establish HTTP quiet, not absence of a sleeping application-level retry. + /// Call only after all submitted actions have reached their observed outcome. + pub(crate) async fn quiesce(&self) -> crate::Result<()> { + let start = Instant::now(); + loop { + let quiet = { + let state = self.state.lock().unwrap(); + state.active.iter().all(|&n| n == 0) + && state.last_request.elapsed() >= Duration::from_secs(1) + && state + .reads + .iter() + .filter_map(|r| r.end_s) + .fold(0.0, f64::max) + + 1.0 + <= self.elapsed() + }; + if quiet { + return Ok(()); + } + if start.elapsed() >= Duration::from_secs(10) { + return Err("account did not quiesce".into()); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } +} + +fn insert_range(ranges: &mut Vec<(u64, u64)>, start: u64, end: u64) -> u64 { + assert!(start <= end); + if start == end { + return 0; + } + let overlap = ranges + .iter() + .map(|&(a, b)| b.min(end).saturating_sub(a.max(start))) + .sum(); + let mut lo = start; + let mut hi = end; + ranges.retain(|&(a, b)| { + if a <= hi && b >= lo { + lo = lo.min(a); + hi = hi.max(b); + false + } else { + true + } + }); + ranges.push((lo, hi)); + ranges.sort_unstable(); + overlap +} + +struct Transfer { + account: Arc, + id: usize, + data: Bytes, + position: usize, + start: Instant, + duration: Duration, + sleep: Option>>, +} + +impl Stream for Transfer { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + use core::future::Future; + let transfer = self.get_mut(); + if transfer.position == transfer.data.len() { + return Poll::Ready(None); + } + let end = (transfer.position + CHUNK).min(transfer.data.len()); + let when = transfer.start + + transfer + .duration + .mul_f64(end as f64 / transfer.data.len() as f64); + let sleep = transfer + .sleep + .get_or_insert_with(|| Box::pin(tokio::time::sleep_until(when))); + if sleep.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + transfer.sleep = None; + let part = transfer.data.slice(transfer.position..end); + + // Charge only this yielded chunk. A requested range whose body is + // dropped partway through must not be counted as a complete transfer. + let mut state = transfer.account.state.lock().unwrap(); + let read = &state.reads[transfer.id]; + let (node, name, offset) = (read.node, read.blob.clone(), read.offset); + let repeated = insert_range( + state.seen.entry((node, name)).or_default(), + offset + u64::try_from(transfer.position).unwrap(), + offset + u64::try_from(end).unwrap(), + ); + let len = u64::try_from(part.len()).unwrap(); + let read = &mut state.reads[transfer.id]; + read.bytes += len; + read.repeated_bytes += repeated; + let bucket = state.buckets.entry(transfer.account.bin()).or_default(); + bucket.bytes += len; + bucket.repeated_bytes += repeated; + transfer.position = end; + Poll::Ready(Some(Ok(part))) + } +} + +impl Drop for Transfer { + fn drop(&mut self) { + // Releasing a body ends its active reservation even when the client + // stopped early. end_s therefore records disposal, not successful EOF. + let mut state = self.account.state.lock().unwrap(); + state.reads[self.id].end_s = Some(self.account.elapsed()); + let node = state.reads[self.id].node; + state.active[node] -= 1; + } +} + +fn response(status: StatusCode, data: impl Into) -> Response { + Response::builder() + .status(status) + .header("last-modified", "Thu, 01 Jan 1970 00:00:00 GMT") + .header("x-ms-blob-type", "BlockBlob") + .body(data.into()) + .unwrap() +} + +fn failure(status: StatusCode, code: &'static str) -> Response { + let mut r = response( + status, + format!("{code}Fixture {code}"), + ); + r.headers_mut() + .insert("x-ms-error-code", code.parse().unwrap()); + r +} + +// Azure accepts a range ending beyond EOF, but an empty blob rejects any range. +fn byte_range(headers: &HeaderMap, length: usize) -> Result<(usize, usize, bool), StatusCode> { + let Some(range) = headers.get("x-ms-range").or_else(|| headers.get("range")) else { + return Ok((0, length, false)); + }; + let (lo, hi) = range + .to_str() + .ok() + .and_then(|v| v.strip_prefix("bytes=")) + .and_then(|v| v.split_once('-')) + .ok_or(StatusCode::BAD_REQUEST)?; + let start: usize = lo.parse().map_err(|_| StatusCode::BAD_REQUEST)?; + let end = if hi.is_empty() { + length + } else { + hi.parse::() + .map_err(|_| StatusCode::BAD_REQUEST)? + .saturating_add(1) + .min(length) + }; + if start >= length || end <= start { + return Err(StatusCode::RANGE_NOT_SATISFIABLE); + } + Ok((start, end, true)) +} + +fn stage_block( + state: &mut AccountState, + name: String, + query: &HashMap, + data: Bytes, +) -> Result<(), Box>> { + let Some(id) = query + .get("blockid") + .and_then(|value| azure_core::base64::decode(value).ok()) + else { + return Err(Box::new(failure(StatusCode::BAD_REQUEST, "InvalidBlockId"))); + }; + + state.blocks.insert((name, id), data); + Ok(()) +} + +fn commit_blocks( + state: &mut AccountState, + name: String, + data: &[u8], +) -> Result<(), Box>> { + let Ok(list): Result = azure_core::xml::from_xml(data) else { + return Err(Box::new(failure( + StatusCode::BAD_REQUEST, + "InvalidXmlDocument", + ))); + }; + + // NativeLink uses Latest. Reject unsupported variants rather than commit + // an empty or truncated object. Keep staged blocks for repeated commits. + if list.committed.is_some() || list.uncommitted.is_some() { + return Err(Box::new(failure( + StatusCode::BAD_REQUEST, + "UnsupportedBlockList", + ))); + } + + let mut output = Vec::new(); + for id in list.latest.unwrap_or_default() { + let Some(block) = state.blocks.get(&(name.clone(), id)) else { + return Err(Box::new(failure( + StatusCode::BAD_REQUEST, + "InvalidBlockList", + ))); + }; + output.extend_from_slice(block); + } + + state.insert_blob(name, Bytes::from(output)); + Ok(()) +} + +async fn put( + account: &Account, + name: String, + query: HashMap, + body: Body, +) -> Response { + let Ok(data) = to_bytes(body, 64 * 1024 * 1024).await else { + return failure(StatusCode::BAD_REQUEST, "InvalidRequestBody"); + }; + let mut state = account.state.lock().unwrap(); + let bucket = state.buckets.entry(account.bin()).or_default(); + // Even a rejected block-list operation has already delivered this body. + bucket.ingress += u64::try_from(data.len()).unwrap(); + + let result = match query.get("comp").map(String::as_str) { + Some("block") => stage_block(&mut state, name, &query, data), + Some("blocklist") => commit_blocks(&mut state, name, &data), + None => { + state.insert_blob(name, data); + Ok(()) + } + _ => Err(Box::new(failure( + StatusCode::BAD_REQUEST, + "UnsupportedQueryParameter", + ))), + }; + + match result { + Ok(()) => response(StatusCode::CREATED, Body::empty()), + Err(response) => *response, + } +} + +fn head_response( + account: &Account, + state: &mut AccountState, + node: usize, + name: String, + blob: Option, +) -> Response { + state.buckets.entry(account.bin()).or_default().head += 1; + let sequence = state + .head_sequences + .entry((node, name.clone())) + .or_default(); + *sequence += 1; + let sequence = *sequence; + let phase = state.phase.clone(); + + state.heads.push(Head { + node, + blob: name, + sequence, + time_s: account.elapsed(), + status: if blob.is_some() { 200 } else { 404 }, + length: blob.as_ref().map(|blob| blob.data.len() as u64), + phase, + }); + + let Some(blob) = blob else { + return failure(StatusCode::NOT_FOUND, "BlobNotFound"); + }; + let mut result = response(StatusCode::OK, Body::empty()); + result.headers_mut().insert( + "content-length", + blob.data.len().to_string().parse().unwrap(), + ); + result + .headers_mut() + .insert("etag", blob.etag.parse().unwrap()); + result +} + +fn number_header(headers: &HeaderMap, name: &'static str) -> Option { + headers.get(name)?.to_str().ok()?.parse().ok() +} + +fn record_get( + account: &Account, + state: &mut AccountState, + node: usize, + name: &str, + headers: &HeaderMap, +) -> usize { + state.buckets.entry(account.bin()).or_default().get += 1; + let id = state.reads.len(); + let phase = state.phase.clone(); + let local_inflight = state.active[node]; + let head_sequence = state.head_sequences.get(&(node, name.to_owned())).copied(); + + // Record every GET before validation, including misses and refusals. Only + // admission opens its lifetime; Transfer::drop later closes that lifetime. + state.reads.push(HttpRead { + request_id: u64::try_from(id).unwrap(), + node, + blob: name.to_owned(), + read_id: number_header(headers, "x-nl-read-id"), + attempt: number_header(headers, "x-nl-attempt"), + sent: number_header(headers, "x-nl-sent"), + writer_broken: headers + .get("x-nl-writer-broken") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()), + offset: 0, + length: 0, + start_s: account.elapsed(), + end_s: Some(account.elapsed()), + bytes: 0, + repeated_bytes: 0, + status: 404, + reason: "missing", + phase, + local_inflight, + timing_outside_fit: false, + head_sequence, + }); + id +} + +fn admit_read( + account: &Account, + state: &mut AccountState, + id: usize, + length: usize, + ranged: bool, + now: Instant, +) -> Result>> { + if !state.admit(now, length) { + state.reads[id].status = 503; + state.reads[id].reason = "account_byte_budget"; + state.buckets.entry(account.bin()).or_default().throttled += 1; + + let mut result = failure(StatusCode::SERVICE_UNAVAILABLE, "ServerBusy"); + result + .headers_mut() + .insert("retry-after", "1".parse().unwrap()); + return Err(Box::new(result)); + } + + let node = state.reads[id].node; + state.active[node] += 1; + state.peak_active[node] = state.peak_active[node].max(state.active[node]); + let local_inflight = state.active[node]; + let status = if ranged { + StatusCode::PARTIAL_CONTENT + } else { + StatusCode::OK + }; + + let read = &mut state.reads[id]; + read.status = status.as_u16(); + read.reason = "admitted"; + read.end_s = None; + read.local_inflight = local_inflight; + read.timing_outside_fit = length != PARTITION || local_inflight > 32; + Ok(status) +} + +fn read_duration( + account: &Account, + state: &mut AccountState, + id: usize, + length: usize, +) -> Duration { + let read = &state.reads[id]; + let ordinal = state + .ordinal + .entry((read.node, read.blob.clone(), read.offset)) + .or_default(); + + // Stable object/range occurrence, independent of other nodes' HTTP order. + // The key and arithmetic here are part of replay compatibility. + let hash = read.blob.as_bytes().iter().fold(0_u64, |v, b| { + v.wrapping_mul(131).wrapping_add(u64::from(*b)) + }); + let quantile = uniform( + account.seed, + hash.wrapping_add((read.node as u64) << 56) + .wrapping_add(read.offset << 8) + .wrapping_add(*ordinal), + ); + *ordinal += 1; + + account + .timing + .duration(read.local_inflight, quantile, length) +} + +struct AdmittedDownload { + id: usize, + blob: StoredBlob, + start: usize, + end: usize, + ranged: bool, + status: StatusCode, + requested: Instant, + duration: Duration, +} + +fn prepare_download( + account: &Account, + state: &mut AccountState, + node: usize, + name: &str, + headers: &HeaderMap, + blob: Option, +) -> Result>> { + let id = record_get(account, state, node, name, headers); + let Some(blob) = blob else { + return Err(Box::new(failure(StatusCode::NOT_FOUND, "BlobNotFound"))); + }; + + if headers + .get("if-match") + .is_some_and(|value| value != "*" && value != blob.etag.as_str()) + { + state.reads[id].status = 412; + state.reads[id].reason = "etag_changed"; + return Err(Box::new(failure( + StatusCode::PRECONDITION_FAILED, + "ConditionNotMet", + ))); + } + + let (start, end, ranged) = match byte_range(headers, blob.data.len()) { + Ok(range) => range, + Err(status) => { + state.reads[id].status = status.as_u16(); + state.reads[id].reason = "invalid_range"; + return Err(Box::new(failure(status, "InvalidRange"))); + } + }; + + let length = end - start; + state.reads[id].offset = u64::try_from(start).unwrap(); + state.reads[id].length = u64::try_from(length).unwrap(); + let requested = Instant::now(); + let status = admit_read(account, state, id, length, ranged, requested)?; + let duration = read_duration(account, state, id, length); + + Ok(AdmittedDownload { + id, + blob, + start, + end, + ranged, + status, + requested, + duration, + }) +} + +fn download_response(account: Arc, download: AdmittedDownload) -> Response { + let AdmittedDownload { + id, + blob, + start, + end, + ranged, + status, + requested, + duration, + } = download; + + let blob_length = blob.data.len(); + let length = end - start; + let transfer = Transfer { + account, + id, + data: blob.data.slice(start..end), + position: 0, + start: requested, + duration, + sleep: None, + }; + let mut result = response(status, Body::from_stream(transfer)); + result + .headers_mut() + .insert("content-length", length.to_string().parse().unwrap()); + result + .headers_mut() + .insert("etag", blob.etag.parse().unwrap()); + if ranged { + result.headers_mut().insert( + "content-range", + format!("bytes {start}-{}/{blob_length}", end - 1) + .parse() + .unwrap(), + ); + } + result +} + +async fn handle( + State(account): State>, + Path((node, name)): Path<(usize, String)>, + Query(query): Query>, + method: Method, + headers: HeaderMap, + body: Body, +) -> Response { + if node >= NODES { + return response(StatusCode::BAD_REQUEST, "unknown node"); + } + + { + let mut state = account.state.lock().unwrap(); + state.last_request = Instant::now(); + let bucket = state.buckets.entry(account.bin()).or_default(); + if method == Method::PUT { + bucket.put += 1; + } + } + + // Upload collection can await the client. No account lock spans that wait. + if method == Method::PUT { + return put(&account, name, query, body).await; + } + + let mut state = account.state.lock().unwrap(); + let blob = state.blobs.get(&name).cloned(); + if method == Method::HEAD { + return head_response(&account, &mut state, node, name, blob); + } + if method != Method::GET { + return response(StatusCode::METHOD_NOT_ALLOWED, Body::empty()); + } + + // Blob lookup, admission, request ownership, and the timing draw remain in + // one critical section. Transfer then owns the active-read reservation. + let download = match prepare_download(&account, &mut state, node, &name, &headers, blob) { + Ok(download) => download, + Err(response) => return *response, + }; + drop(state); + + // Construct the streamed response after releasing the lock: dropping its + // body also locks the account to release that reservation and record its end. + download_response(account, download) +} + +pub(crate) async fn serve( + account: Arc, +) -> std::io::Result<(String, tokio::task::JoinHandle<()>)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let router = Router::new() + .route("/{node}/cas/{name}", any(handle)) + .with_state(account); + let task = nativelink_util::background_spawn!("fleet_http", async move { + axum::serve(listener, router).await.unwrap(); + }); + Ok((format!("http://{address}"), task)) +} + +#[cfg(test)] +mod tests { + use async_trait::async_trait; + use azure_core::http::policies::{Policy, PolicyResult}; + use azure_core::http::{Context as AzureContext, Request}; + use azure_storage_blob::models::{BlobClientDownloadOptions, HttpRange}; + use nativelink_store::azure_blob_store::AzureDownloadAttempt; + use tokio::sync::Notify; + + use super::*; + + // Hold an SDK child task before it reaches the server. This gives the test + // an exact old-task/new-attempt interleaving without timing-based guesses. + #[derive(Debug, Default)] + struct DelayedPartition { + entered: Notify, + release: Notify, + forwarded: Notify, + } + + #[async_trait] + impl Policy for DelayedPartition { + async fn send( + &self, + context: &AzureContext, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + let delayed = context + .value::() + .is_some_and(|identity| identity.read_id == 17 && identity.attempt == 1) + && request.headers().get_optional_str(&"range".into()) + == Some("bytes=4194304-8388607"); + if delayed { + self.entered.notify_one(); + self.release.notified().await; + } + + let result = next[0].send(context, request, &next[1..]).await; + if delayed { + self.forwarded.notify_one(); + } + result + } + } + + fn download_options(identity: AzureDownloadAttempt) -> BlobClientDownloadOptions<'static> { + let mut options = BlobClientDownloadOptions { + parallel: std::num::NonZeroUsize::new(2), + partition_size: std::num::NonZeroUsize::new(PARTITION), + range: Some(HttpRange::new( + identity.sent, + 2 * PARTITION as u64 - identity.sent, + )), + ..Default::default() + }; + options.method_options.context.insert(identity); + options + } + + #[test] + fn resumed_and_overlapping_reads_count_only_actual_repeated_bytes() { + // Compare interval accounting against a simple independent byte bitmap, + // including out-of-order, duplicate, adjoining, and bridging ranges. + let intervals = [ + (20, 30), + (0, 10), + (5, 25), + (3, 27), + (30, 40), + (0, 40), + (10, 10), + ]; + let mut seen = [false; 40]; + let mut ranges = Vec::new(); + for (start, end) in intervals { + let expected = seen[start..end].iter().filter(|&&v| v).count(); + assert_eq!( + insert_range(&mut ranges, start as u64, end as u64), + expected as u64 + ); + seen[start..end].fill(true); + } + assert_eq!(ranges, vec![(0, 40)]); + } + + #[test] + fn admission_refuses_excess_and_recovers_without_random_failure_injection() { + let account = Account::new(17, &[]); + let mut state = account.state.lock().unwrap(); + let start = state.last; + let initial_partitions = u64::try_from(BURST_BYTES / PARTITION).unwrap(); + for _ in 0..initial_partitions { + assert!(state.admit(start, PARTITION)); + } + assert!(!state.admit(start, PARTITION)); + assert!(!state.admit(start + Duration::from_millis(1), PARTITION)); + assert!(state.admit(start + Duration::from_secs_f64(1.0 / 64.0), PARTITION)); + assert_eq!(state.admitted, (initial_partitions + 1) * PARTITION as u64); + } + + #[test] + fn incomplete_or_invalid_identity_remains_unknown() { + let account = Account::new(31, &[]); + let mut headers = HeaderMap::new(); + headers.insert("x-nl-read-id", "17".parse().unwrap()); + headers.insert("x-nl-attempt", "invalid".parse().unwrap()); + headers.insert("x-nl-writer-broken", "maybe".parse().unwrap()); + let mut state = account.state.lock().unwrap(); + record_get(&account, &mut state, 0, "unknown", &headers); + let read = &state.reads[0]; + assert_eq!(read.request_id, 0); + assert_eq!(read.read_id, Some(17)); + assert_eq!(read.attempt, None); + assert_eq!(read.sent, None); + assert_eq!(read.writer_broken, None); + } + + #[expect( + clippy::disallowed_methods, + reason = "tokio::test creates its runtime and calls Runtime::block_on" + )] + #[tokio::test] + async fn dropped_body_counts_only_yielded_bytes_and_separates_node_replication() { + use futures::StreamExt; + let account = Account::new(29, &[]); + account + .state + .lock() + .unwrap() + .insert_blob("partial".into(), Bytes::from(vec![9; 2 * CHUNK])); + for node in [0, 1, 0] { + let response = handle( + State(account.clone()), + Path((node, "partial".into())), + Query(HashMap::new()), + Method::GET, + HeaderMap::new(), + Body::empty(), + ) + .await; + let mut body = response.into_body().into_data_stream(); + assert_eq!(body.next().await.unwrap().unwrap().len(), CHUNK); + drop(body); + } + let report = account.report(); + assert_eq!(report["payload_bytes"], 3 * CHUNK as u64); + assert_eq!(report["repeated_payload_bytes"], CHUNK as u64); + assert_eq!(report["attributed_payload_bytes"], 0); + assert_eq!( + report["unknown_attribution_payload_bytes"], + 3 * CHUNK as u64 + ); + let state = account.state.lock().unwrap(); + assert!(state.active.iter().all(|&n| n == 0)); + assert!( + state + .reads + .iter() + .all(|r| r.bytes < r.length && r.end_s.is_some()) + ); + assert_eq!(state.reads[1].repeated_bytes, 0); + } + + #[expect( + clippy::disallowed_methods, + reason = "tokio::test creates its runtime and calls Runtime::block_on" + )] + #[tokio::test] + async fn sdk_detached_partition_keeps_old_attempt_identity() -> crate::Result<()> { + use azure_core::http::{RetryOptions, Transport, Url}; + use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions}; + use futures::TryStreamExt; + + drop(tokio_rustls::rustls::crypto::ring::default_provider().install_default()); + let account = Account::new(20_261_021, &[]); + account + .state + .lock() + .unwrap() + .insert_blob("shared".into(), Bytes::from(vec![7; 2 * PARTITION])); + let (url, server) = serve(account.clone()).await?; + let gate = Arc::new(DelayedPartition::default()); + let mut options = BlobContainerClientOptions::default(); + options.client_options.retry = RetryOptions::none(); + crate::transport::attach(&mut options.client_options); + options.client_options.per_call_policies.push(gate.clone()); + options.client_options.transport = Some(Transport::new(Arc::new( + reqwest::Client::builder().no_proxy().build()?, + ))); + let client = + BlobContainerClient::new(Url::parse(&format!("{url}/0/cas"))?, None, Some(options))?; + let blob = client.blob_client("shared"); + + // Context values below are test labels; this test proves SDK provenance, + // while the store test proves when a real writer becomes broken. + let old = download_options(AzureDownloadAttempt { + read_id: 17, + attempt: 1, + sent: 0, + writer_broken: false, + }); + let mut old_body = blob.download(Some(old)).await?.body; + assert_eq!(old_body.try_next().await?.unwrap().len(), PARTITION); + tokio::time::timeout(Duration::from_secs(5), gate.entered.notified()).await?; + drop(old_body); + + // The resumed call reaches the same range first. The old partition is + // still alive, but has not yet made its HTTP request. + let resumed = download_options(AzureDownloadAttempt { + read_id: 17, + attempt: 2, + sent: PARTITION as u64, + writer_broken: true, + }); + let chunks: Vec = blob + .download(Some(resumed)) + .await? + .body + .try_collect() + .await?; + assert_eq!(chunks.concat(), vec![7; PARTITION]); + gate.release.notify_one(); + tokio::time::timeout(Duration::from_secs(5), gate.forwarded.notified()).await?; + account.quiesce().await?; + + let report = account.report(); + assert_eq!(report["payload_bytes"], 3 * PARTITION as u64); + assert_eq!(report["attributed_payload_bytes"], 3 * PARTITION as u64); + assert_eq!(report["unknown_attribution_payload_bytes"], 0); + assert_eq!(report["repeated_payload_bytes"], PARTITION as u64); + let state = account.state.lock().unwrap(); + let identities: Vec<_> = state + .reads + .iter() + .map(|read| { + ( + read.request_id, + read.read_id, + read.attempt, + read.offset, + read.sent, + read.writer_broken, + ) + }) + .collect(); + assert_eq!( + identities, + [ + (0, Some(17), Some(1), 0, Some(0), Some(false)), + ( + 1, + Some(17), + Some(2), + PARTITION as u64, + Some(PARTITION as u64), + Some(true) + ), + (2, Some(17), Some(1), PARTITION as u64, Some(0), Some(false)), + ] + ); + assert!(state.heads.is_empty()); + drop(state); + server.abort(); + Ok(()) + } + + #[expect( + clippy::disallowed_methods, + reason = "tokio::test creates its runtime and calls Runtime::block_on" + )] + #[tokio::test] + async fn sdk_round_trip_range_empty_blob_and_idempotent_commit() -> crate::Result<()> { + use azure_core::http::{RequestContent, RetryOptions, Transport, Url}; + use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions}; + use futures::TryStreamExt; + drop(tokio_rustls::rustls::crypto::ring::default_provider().install_default()); + let account = Account::new(20_261_017, &[]); + let (url, server) = serve(account.clone()).await?; + let mut options = BlobContainerClientOptions::default(); + options.client_options.retry = RetryOptions::none(); + crate::transport::attach(&mut options.client_options); + options.client_options.transport = Some(Transport::new(Arc::new( + reqwest::Client::builder().no_proxy().build()?, + ))); + let client = + BlobContainerClient::new(Url::parse(&format!("{url}/0/cas"))?, None, Some(options))?; + let blob = client.blob_client("payload"); + let block = blob.block_blob_client(); + let payload = Bytes::from(vec![7_u8; 2 * PARTITION + 17]); + block + .upload(RequestContent::from(payload.to_vec()), None) + .await?; + blob.get_properties(None).await?; + let chunks: Vec = blob.download(None).await?.body.try_collect().await?; + // The SDK pads its initial buffer for small objects; this object crosses + // partitions and must round-trip exactly, including the final 17 bytes. + assert_eq!(chunks.concat(), payload); + blob.get_properties(None).await?; + let committed = client.blob_client("committed"); + let block = committed.block_blob_client(); + block + .stage_block(b"one", 3, RequestContent::from(b"abc".to_vec()), None) + .await?; + block + .stage_block(b"two", 3, RequestContent::from(b"def".to_vec()), None) + .await?; + for _ in 0..2 { + block + .commit_block_list( + RequestContent::try_from(BlockLookupList { + latest: Some(vec![b"one".to_vec(), b"two".to_vec()]), + ..Default::default() + })?, + None, + ) + .await?; + } + let chunks: Vec = committed.download(None).await?.body.try_collect().await?; + assert_eq!(chunks.concat(), b"abcdef"); + let empty = client.blob_client("empty"); + empty + .block_blob_client() + .upload(RequestContent::from(Vec::new()), None) + .await?; + empty.get_properties(None).await?; + let empty_options = download_options(AzureDownloadAttempt { + read_id: 18, + attempt: 1, + sent: 0, + writer_broken: false, + }); + let chunks: Vec = empty + .download(Some(empty_options)) + .await? + .body + .try_collect() + .await?; + assert!(chunks.concat().is_empty()); + account.quiesce().await?; + let report = account.report(); + assert_eq!(report["payload_bytes"], payload.len() as u64 + 6); + let state = account.state.lock().unwrap(); + let empty_statuses: Vec = state + .reads + .iter() + .filter(|r| r.blob == "empty") + .map(|r| r.status) + .collect(); + assert_eq!(empty_statuses, [416, 200]); + // The SDK's empty-blob fallback is a second HTTP request belonging to + // the same outer download, not an application-level retry. + assert!(state.reads.iter().filter(|r| r.blob == "empty").all(|r| { + r.read_id == Some(18) + && r.attempt == Some(1) + && r.sent == Some(0) + && r.writer_broken == Some(false) + })); + assert!( + state + .reads + .iter() + .filter(|r| r.blob == "payload") + .all(|r| r.head_sequence == Some(1)) + ); + assert!( + state + .reads + .iter() + .filter(|r| r.blob == "empty") + .all(|r| r.head_sequence == Some(1)) + ); + let payload_heads: Vec<_> = state + .heads + .iter() + .filter(|h| h.blob == "payload") + .map(|h| h.sequence) + .collect(); + assert_eq!(payload_heads, [1, 2]); + drop(state); + server.abort(); + Ok(()) + } +} diff --git a/examples/azure_fleet/checks.rs b/examples/azure_fleet/checks.rs new file mode 100644 index 00000000000..d86354b267c --- /dev/null +++ b/examples/azure_fleet/checks.rs @@ -0,0 +1,312 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Acceptance depends on lifecycle outcomes and correlated download attempts, +//! never on a fitted byte saving or elapsed-time threshold. + +use std::collections::BTreeSet; + +use serde_json::Value; + +use crate::workload::NODES; + +pub(crate) fn closed_destination_exhaustions(log: &str) -> usize { + // Retain the legacy diagnostic for comparison, never as the retry verdict. + // Require one error chain: unrelated throttling and producer failures on + // separate lines do not establish a retry against a closed destination. + log.lines() + .filter(|line| { + line.contains("Retries exhausted") + && line.contains("client=\"azure\"") + && line.contains("Failed to send data to writer") + && (line.contains("Tried to send while stream is closed") + || line.contains("receiver disconnected")) + }) + .count() +} + +fn require(failures: &mut Vec, passed: bool, message: &str) { + if !passed { + failures.push(message.to_owned()); + } +} + +fn check_jobs(report: &Value, failures: &mut Vec) { + let jobs: &[Value] = report["jobs"].as_array().map_or(&[], Vec::as_slice); + require( + failures, + !jobs.is_empty() && Some(jobs.len() as u64) == report["expected_jobs"].as_u64(), + "not every planned job has an observed outcome", + ); + require( + failures, + jobs.iter().all(|job| { + job["outcome"] == "success" + && job["attempts"].as_array().is_some_and(|attempts| { + !attempts.is_empty() && attempts.iter().all(|a| a["outcome"] != "censored") + }) + }), + "work failed or a client attempt was censored", + ); +} + +fn check_worker_recovery(nodes: &[Value], failures: &mut Vec) { + require( + failures, + nodes.len() == NODES, + "not all four workers became ready", + ); + require( + failures, + nodes.iter().all(|node| { + node["sentinel_same_process"] == true && node["sentinel_cache_reused"] == true + }), + "same-worker sentinel or warm-cache reuse failed", + ); +} + +fn check_terminal_retries(nodes: &[Value], failures: &mut Vec) { + // A final error can hide earlier bad attempts, or follow only legitimate + // retries. Require a complete event/HTTP join before interpreting a zero. + require( + failures, + nodes + .iter() + .all(|node| node["download_witness"]["complete"] == true), + "download attempt evidence is incomplete", + ); + require( + failures, + nodes + .iter() + .all(|node| node["download_witness"]["new_attempts_after_failed_send"] == 0), + "Azure started a new download after destination failure", + ); +} + +fn check_filesystem_fault(report: &Value, nodes: &[Value], failures: &mut Vec) { + // Fault and control need different evidence: recovery alone cannot prove + // that either deliberately filled filesystem actually rejected a worker write. + if report["scenario"] == "failed-fill" { + require( + failures, + [0, 1].iter().all(|index| { + nodes + .iter() + .any(|node| node["node"] == *index && node["worker_enospc_observed"] == true) + }), + "intended worker ENOSPC was not observed on both filled filesystems", + ); + require( + failures, + report["fault"]["repaired_s"].is_number(), + "filesystem repair was not observed", + ); + require( + failures, + [0, 1].iter().all(|index| { + nodes.iter().any(|node| { + node["node"] == *index + && node["download_witness"]["failed_destination_reads"] + .as_u64() + .is_some_and(|count| count > 0) + }) + }), + "Azure did not observe destination failure on both filled filesystems", + ); + } else { + require( + failures, + nodes + .iter() + .all(|node| node["worker_enospc_observed"] == false), + "control unexpectedly encountered ENOSPC", + ); + } +} + +fn check_cleanup(report: &Value, node_count: usize, failures: &mut Vec) { + let cleanup: &[Value] = report["cleanup"].as_array().map_or(&[], Vec::as_slice); + require( + failures, + cleanup.len() == node_count && cleanup.iter().all(|node| node["success"] == true), + "worker shutdown failed or required forced termination", + ); +} + +fn check_scheduler(report: &Value, nodes: &[Value], failures: &mut Vec) { + // Sampled reservations and the final release are independent checks. + // The final snapshot must also retain every original worker identity. + let snapshots: &[Value] = report["scheduler_snapshots"] + .as_array() + .map_or(&[], Vec::as_slice); + require( + failures, + !snapshots.is_empty() + && snapshots.iter().all(|sample| { + sample["workers"].as_array().is_some_and(|workers| { + workers.iter().all(|worker| { + worker["max_inflight_tasks"] == 2 + && worker["running_actions"] + .as_u64() + .is_some_and(|count| count <= 2) + }) + }) + }), + "scheduler reservation samples are missing or exceed the advertised limit", + ); + let originals: BTreeSet<_> = nodes + .iter() + .filter_map(|node| node["worker_id"].as_str()) + .collect(); + require( + failures, + originals.len() == NODES + && snapshots.last().is_some_and(|sample| { + sample["final"] == true + && sample["workers"].as_array().is_some_and(|workers| { + let ids: BTreeSet<_> = workers + .iter() + .filter_map(|worker| worker["id"].as_str()) + .collect(); + ids == originals + && workers.iter().all(|worker| worker["running_actions"] == 0) + }) + }), + "final scheduler state did not retain the original workers with zero reservations", + ); +} + +pub(crate) fn failures(report: &Value) -> Vec { + let mut failures = Vec::new(); + let nodes: &[Value] = report["nodes"].as_array().map_or(&[], Vec::as_slice); + + require( + &mut failures, + report["session_error"].is_null(), + "fleet session failed", + ); + check_jobs(report, &mut failures); + check_worker_recovery(nodes, &mut failures); + check_terminal_retries(nodes, &mut failures); + check_filesystem_fault(report, nodes, &mut failures); + check_cleanup(report, nodes.len(), &mut failures); + check_scheduler(report, nodes, &mut failures); + + failures +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn terminal_retry_witness_requires_the_same_azure_error_chain() { + let prefix = "Retries exhausted client=\"azure\" attempt=7"; + assert_eq!( + closed_destination_exhaustions(&format!( + "{prefix} Failed to send data to writer: Tried to send while stream is closed" + )), + 1 + ); + assert_eq!( + closed_destination_exhaustions(&format!( + "{prefix} Failed to send data to writer: receiver disconnected" + )), + 1 + ); + assert_eq!( + closed_destination_exhaustions(&format!( + "{prefix} ServerBusy\nFailed to send data to writer: receiver disconnected" + )), + 0 + ); + assert_eq!( + closed_destination_exhaustions( + "Retries exhausted client=\"s3\" Failed to send data to writer: receiver disconnected" + ), + 0 + ); + assert_eq!( + closed_destination_exhaustions("Sender dropped before sending EOF"), + 0 + ); + } + + #[test] + fn complete_control_passes_but_terminal_retry_fails_without_a_byte_threshold() { + let nodes: Vec<_> = (0..NODES) + .map(|node| { + serde_json::json!({ + "node": node, + "worker_id": format!("worker-{node}"), + "sentinel_same_process": true, + "sentinel_cache_reused": true, + "closed_destination_retry_exhaustions": 0, + "download_witness": { + "complete": true, + "new_attempts_after_failed_send": 0, + "failed_destination_reads": 0 + }, + "worker_enospc_observed": false + }) + }) + .collect(); + let workers: Vec<_> = (0..NODES) + .map(|node| { + serde_json::json!({ + "id": format!("worker-{node}"), + "max_inflight_tasks": 2, + "running_actions": 0 + }) + }) + .collect(); + let mut report = serde_json::json!({ + "session_error": null, + "scenario": "control", + "expected_jobs": 1, + "jobs": [{ + "outcome": "success", + "attempts": [{"outcome": "success"}] + }], + "nodes": nodes, + "cleanup": vec![serde_json::json!({"success": true}); NODES], + "scheduler_snapshots": [{"final": true, "workers": workers}], + "account": {"payload_bytes": u64::MAX}, + "duration_s": 100_000 + }); + assert!(failures(&report).is_empty()); + // An exhausted final failed send alone does not prove that a new + // download followed it. Conversely, a final 503 cannot hide a witness. + report["nodes"][0]["closed_destination_retry_exhaustions"] = 1.into(); + assert!(failures(&report).is_empty()); + report["nodes"][0]["closed_destination_retry_exhaustions"] = 0.into(); + report["nodes"][0]["download_witness"]["new_attempts_after_failed_send"] = 1.into(); + assert_eq!( + failures(&report), + ["Azure started a new download after destination failure"] + ); + report["nodes"][0]["download_witness"]["complete"] = false.into(); + assert!( + failures(&report) + .iter() + .any(|reason| reason == "download attempt evidence is incomplete") + ); + } + + #[test] + fn missing_outcomes_and_log_evidence_cannot_pass() { + assert!(!failures(&serde_json::json!({})).is_empty()); + } +} diff --git a/examples/azure_fleet/main.rs b/examples/azure_fleet/main.rs new file mode 100644 index 00000000000..cdcf4f1d116 --- /dev/null +++ b/examples/azure_fleet/main.rs @@ -0,0 +1,976 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Real scheduler, worker processes, filesystem failures, and SDK HTTP downloads. + +mod account; +mod checks; +mod node; +mod transport; +mod witness; +mod workload; + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use clap::Parser; +use futures::future::join_all; +use nativelink_config::cas_server::{ExecutionConfig, WithInstanceName, WorkerApiConfig}; +use nativelink_config::schedulers::SchedulerSpec; +use nativelink_config::stores::MemorySpec; +use nativelink_proto::build::bazel::remote::execution::v2::execution_client::ExecutionClient; +use nativelink_proto::build::bazel::remote::execution::v2::{ + ExecuteRequest, ExecuteResponse, ExecutedActionMetadata, +}; +use nativelink_proto::google::longrunning::operation; +use nativelink_scheduler::default_scheduler_factory::scheduler_factory; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_service::execution_server::ExecutionServer; +use nativelink_service::worker_api_server::WorkerApiServer; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::store_manager::StoreManager; +use nativelink_util::store_trait::{Store, StoreLike}; +use prost::Message; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines}; +use tokio::process::{Child, ChildStdin, ChildStdout, Command}; +use tokio::time::{Instant, timeout}; +use tracing_subscriber::prelude::*; +use workload::{Job, NODES, Workload}; + +type Result = std::result::Result>; + +#[expect( + clippy::print_stdout, + reason = "the replay CLI prints its seed and result location" +)] +fn announce(message: std::fmt::Arguments<'_>) { + println!("{message}"); +} + +#[derive(Parser)] +struct Args { + #[arg(long)] + node: Option, + #[arg(long, default_value_t = 2026101700)] + seed: u64, + #[arg(long, default_value = "failed-fill", value_parser = ["control", "failed-fill"])] + scenario: String, + #[arg(long)] + output: Option, +} + +struct Node { + child: Child, + input: ChildStdin, + output: Lines>, + evidence: Value, +} + +impl Node { + async fn spawn( + config: node::Config, + directory: &Path, + executable: &Path, + cpus: &str, + ) -> Result { + let config_path = directory.join(format!("node-{}.json", config.node)); + std::fs::write(&config_path, serde_json::to_vec_pretty(&config)?)?; + let log = std::fs::File::create(directory.join(format!("node-{}.log", config.node)))?; + + // Each child owns a private mount namespace. stdout is the control + // protocol; all worker diagnostics go to its retained stderr log. + let mut child = Command::new("unshare") + .args([ + "--user", + "--map-root-user", + "--mount", + "--propagation", + "private", + "taskset", + "--cpu-list", + cpus, + ]) + .arg(executable) + .arg("--node") + .arg(config_path) + .env("RUST_LOG", "warn,nativelink_worker::local_worker=debug,nativelink_worker::running_actions_manager=debug") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(log) + .kill_on_drop(true) + .spawn()?; + let input = child.stdin.take().ok_or("missing child stdin")?; + let output = BufReader::new(child.stdout.take().ok_or("missing child stdout")?).lines(); + let mut node = Self { + child, + input, + output, + evidence: Value::Null, + }; + + match node.receive().await { + Ok(evidence) if evidence["ready"] == true => { + node.evidence = evidence; + Ok(node) + } + received => { + // A startup failure must not leave a private mount/worker alive. + drop(node.child.kill().await); + drop(node.child.wait().await); + Err(received + .err() + .unwrap_or_else(|| "worker did not become ready".into())) + } + } + } + + async fn receive(&mut self) -> Result { + let line = timeout(Duration::from_secs(20), self.output.next_line()) + .await?? + .ok_or("worker control channel closed; inspect node log")?; + Ok(serde_json::from_str(&line)?) + } + + async fn command(&mut self, command: &str) -> Result { + self.input + .write_all(format!("{command}\n").as_bytes()) + .await?; + self.input.flush().await?; + self.receive().await + } + + async fn stop(&mut self) -> Value { + let command = self.command("stop").await; + let mut forced = false; + + let status = if let Ok(status) = timeout(Duration::from_secs(8), self.child.wait()).await { + status + } else { + forced = true; + // kill() waits for the child, so even a failed drain is reaped. + match self.child.kill().await { + Ok(()) => self.child.wait().await, + Err(error) => Err(error), + } + }; + + json!({ + "node": self.evidence["node"], + "pid": self.evidence["pid"], + "success": command.is_ok() && status.as_ref().is_ok_and(std::process::ExitStatus::success) && !forced, + "forced": forced, + "reply": command.as_ref().ok(), + "command_error": command.as_ref().err().map(ToString::to_string), + "exit_code": status.as_ref().ok().and_then(std::process::ExitStatus::code), + "wait_error": status.as_ref().err().map(ToString::to_string), + }) + } +} + +async fn scheduler( + workload: &Workload, +) -> Result<( + String, + Arc, + tokio::task::JoinHandle<()>, +)> { + let stores = StoreManager::new(); + let cas = Store::new(MemoryStore::new(&MemorySpec::default())); + + // The scheduler needs action metadata. Workers still fetch cold payloads + // through their own Azure-backed stores and the HTTP account. + for blob in &workload.blobs { + cas.update_oneshot(blob.digest, blob.data.clone()).await?; + } + stores.add_store("cas", cas)?; + + let spec: SchedulerSpec = serde_json::from_value(json!({ + "simple": { + "supported_platform_properties": {"fleet_node": "exact"}, + "max_job_retries": 3, + "worker_timeout_s": 30, + "client_action_timeout_s": 180, + } + }))?; + let (actions, workers) = scheduler_factory(&spec, &stores, None).await?; + let workers = workers.ok_or("missing worker scheduler")?; + let execution = ExecutionServer::new( + &[WithInstanceName { + instance_name: String::new(), + config: ExecutionConfig { + experimental_buck2_invocation_isolation: false, + cas_store: "cas".into(), + scheduler: "scheduler".into(), + }, + }], + &HashMap::from([( + "scheduler".into(), + actions.ok_or("missing action scheduler")?, + )]), + &stores, + )?; + let api = WorkerApiServer::new( + &WorkerApiConfig { + scheduler: "scheduler".into(), + disable_kill_revoked_operations: false, + kill_revoked_operations_interval_s: 5, + }, + &HashMap::from([("scheduler".into(), workers.clone())]), + )?; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let address = format!("http://{}", listener.local_addr()?); + let incoming = futures::stream::unfold(listener, |listener| async { + Some((listener.accept().await.map(|(stream, _)| stream), listener)) + }); + let task = nativelink_util::background_spawn!("fleet_scheduler", async move { + tonic::transport::Server::builder() + .add_service(execution.into_service()) + .add_service(api.into_service()) + .serve_with_incoming(incoming) + .await + .expect("scheduler gRPC server"); + }); + + Ok((address, workers, task)) +} + +async fn execute(address: &str, job: &Job, account: &account::Account) -> Result { + let start = account.elapsed(); + let channel = tonic::transport::Endpoint::from_shared(address.to_owned())? + .connect() + .await?; + let mut client = ExecutionClient::new(channel); + let mut stream = client + .execute(ExecuteRequest { + action_digest: Some(job.digest().into()), + skip_cache_lookup: true, + digest_function: 1, + ..Default::default() + }) + .await? + .into_inner(); + while let Some(operation) = stream.message().await? { + if !operation.done { + continue; + } + match operation + .result + .ok_or("terminal operation without result")? + { + operation::Result::Response(response) => { + let response = ExecuteResponse::decode(response.value.as_slice())?; + return Ok(execution_result(&operation.name, response, start, account)); + } + operation::Result::Error(status) => { + return Ok(json!({ + "operation": operation.name, + "started_s": start, + "finished_s": account.elapsed(), + "outcome": "failed", + "status_code": status.code, + "status": status.message, + })); + } + } + } + Err("Execute stream ended without terminal operation".into()) +} + +fn timestamp(stamp: &Option) -> Option { + stamp + .as_ref() + .map(|stamp| json!({"seconds": stamp.seconds, "nanos": stamp.nanos})) +} + +fn lifecycle(metadata: &ExecutedActionMetadata) -> Value { + json!({ + "worker_start": timestamp(&metadata.worker_start_timestamp), + "input_fetch_start": timestamp(&metadata.input_fetch_start_timestamp), + "input_fetch_completed": timestamp(&metadata.input_fetch_completed_timestamp), + "execution_start": timestamp(&metadata.execution_start_timestamp), + "execution_completed": timestamp(&metadata.execution_completed_timestamp), + "output_upload_start": timestamp(&metadata.output_upload_start_timestamp), + "output_upload_completed": timestamp(&metadata.output_upload_completed_timestamp), + "worker_completed": timestamp(&metadata.worker_completed_timestamp), + }) +} + +fn execution_result( + name: &str, + response: ExecuteResponse, + start: f64, + account: &account::Account, +) -> Value { + let result = response.result.as_ref(); + let status = response.status.as_ref().map_or(0, |status| status.code); + let exit = result.map(|result| result.exit_code); + let metadata = result.and_then(|result| result.execution_metadata.as_ref()); + let worker = metadata.map(|metadata| metadata.worker.clone()); + let lifecycle = metadata.map(lifecycle); + let outputs = result.map(|result| { + result + .output_files + .iter() + .map(|file| { + json!({ + "path": file.path, + "digest": file.digest.as_ref().map(|digest| json!({ + "hash": digest.hash, + "size_bytes": digest.size_bytes, + })), + }) + }) + .collect::>() + }); + + json!({ + "operation": name, + "started_s": start, + "finished_s": account.elapsed(), + "outcome": if status == 0 && exit == Some(0) { "success" } else { "failed" }, + "exit_code": exit, + "status_code": status, + "status": response.status.map(|status| status.message), + "worker_id": worker, + "lifecycle": lifecycle, + "outputs": outputs, + }) +} + +async fn job_result( + address: String, + job: Job, + account: Arc, + phase: &'static str, + origin: Instant, + repair_at: Option, +) -> Value { + tokio::time::sleep_until(origin + Duration::from_secs_f64(job.arrival_s)).await; + let mut attempts = Vec::new(); + // One explicit client resubmission after repair, identical in both arms. + // Keep the original failed attempt in the record; never drop offered work. + for attempt in 0..2 { + let start = account.elapsed(); + let result = timeout(Duration::from_secs(130), execute(&address, &job, &account)).await; + let value = match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + json!({ + "started_s": start, + "finished_s": account.elapsed(), + "outcome": "failed", + "status": error.to_string(), + }) + } + Err(_) => { + json!({ + "started_s": start, + "finished_s": account.elapsed(), + "outcome": "censored", + "status": "130 second client deadline", + }) + } + }; + let success = value["outcome"] == "success"; + attempts.push(value); + if success || attempt == 1 || repair_at.is_none() { + break; + } + tokio::time::sleep_until(repair_at.unwrap() + Duration::from_millis(500)).await; + } + + // The top-level outcome describes the last attempt. The attempts array + // preserves earlier failures and deadlines for the acceptance checks. + let mut value = attempts.last().unwrap().clone(); + value["id"] = json!(job.id); + value["node"] = json!(job.node); + value["phase"] = json!(phase); + value["arrival_s"] = json!(origin.duration_since(account.start).as_secs_f64() + job.arrival_s); + value["planned_arrival_s"] = json!(job.arrival_s); + value["started_s"] = attempts[0]["started_s"].clone(); + value["attempts"] = json!(attempts); + + if value["outcome"] == "success" { + account.state.lock().unwrap().completed_actions += 1; + } + + value +} + +fn load() -> Value { + json!({ + "loadavg": std::fs::read_to_string("/proc/loadavg").unwrap_or_default(), + "pressure_cpu": std::fs::read_to_string("/proc/pressure/cpu").unwrap_or_default(), + "pressure_io": std::fs::read_to_string("/proc/pressure/io").unwrap_or_default(), + }) +} + +fn cpus() -> Result { + // Read the inherited affinity; do not assume host CPUs are available to us. + let status = std::fs::read_to_string("/proc/self/status")?; + let allowed = status + .lines() + .find_map(|l| l.strip_prefix("Cpus_allowed_list:\t")) + .ok_or("missing affinity")?; + let mut cpus = Vec::new(); + for piece in allowed.split(',') { + let ends: Vec<_> = piece + .split('-') + .map(str::parse::) + .collect::>()?; + cpus.extend(ends[0]..=ends.get(1).copied().unwrap_or(ends[0])); + } + Ok(cpus + .into_iter() + .take(8) + .map(|n| n.to_string()) + .collect::>() + .join(",")) +} + +// Run identity is kept apart from session evidence so teardown can +// still publish a complete report after a setup or workload error. +struct RunContext { + seed: u64, + scenario: String, + output: PathBuf, + executable: PathBuf, + binary_sha256: String, + workload_sha256: String, + start_unix_s: f64, + load_before: Value, +} + +impl RunContext { + fn new(args: Args) -> Result { + let output = args.output.ok_or("--output is required")?; + std::fs::create_dir(&output)?; + let output = std::fs::canonicalize(output)?; + let executable = std::env::current_exe()?; + let binary_sha256 = hex::encode(Sha256::digest(std::fs::read(&executable)?)); + let start_unix_s = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs_f64(); + let load_before = load(); + + announce(format_args!( + "seed={} scenario={} output={}", + args.seed, + args.scenario, + output.display() + )); + + Ok(Self { + seed: args.seed, + scenario: args.scenario, + output, + executable, + binary_sha256, + workload_sha256: String::new(), + start_unix_s, + load_before, + }) + } + + fn write_manifest(&mut self, workload: &Workload) -> Result<()> { + let manifest = json!({ + "seed": self.seed, + "catalogue_bytes": workload.catalogue_bytes, + "warmup": workload.warmup, + "jobs": workload.jobs, + "sentinel": workload.sentinel, + }); + let bytes = serde_json::to_vec_pretty(&manifest)?; + self.workload_sha256 = hex::encode(Sha256::digest(&bytes)); + std::fs::write(self.output.join("workload.json"), bytes)?; + Ok(()) + } +} + +#[derive(Default)] +struct Session { + nodes: Vec, + services: Vec>, + jobs: Vec, + events: Vec, + fault: Value, + before_sentinel: Value, + affinity: String, + snapshots: Arc>>, +} + +impl Session { + async fn start( + &mut self, + context: &RunContext, + workload: &Workload, + account: Arc, + ) -> Result<(String, Arc)> { + // Register each resource immediately. If the next startup step fails, + // the common teardown path still owns everything already started. + let (azure, http_task) = account::serve(account).await?; + self.services.push(http_task); + let (address, scheduler, grpc_task) = scheduler(workload).await?; + self.services.push(grpc_task); + self.affinity = cpus()?; + + for node in 0..NODES { + let root = context.output.join(format!("disk-{node}")); + std::fs::create_dir(&root)?; + self.nodes.push( + Node::spawn( + node::Config { + seed: context.seed, + node, + root, + azure: azure.clone(), + scheduler: address.clone(), + metadata: workload.metadata.clone(), + }, + &context.output, + &context.executable, + &self.affinity, + ) + .await?, + ); + } + + self.register_workers(&scheduler).await?; + Ok((address, scheduler)) + } + + async fn register_workers(&mut self, scheduler: &Arc) -> Result<()> { + timeout(Duration::from_secs(15), async { + loop { + if scheduler.worker_snapshot().await.len() == NODES { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await?; + + let registered = scheduler.worker_snapshot().await; + if !registered + .iter() + .all(|worker| worker.max_inflight_tasks == 2) + { + return Err("worker registered an unexpected inflight limit".into()); + } + + // Match the scheduler identity through the placement property, rather + // than relying on registration order or process IDs. + for (index, node) in self.nodes.iter_mut().enumerate() { + let worker = registered + .iter() + .find(|worker| { + worker.platform_properties.get("fleet_node") == Some(&index.to_string()) + }) + .ok_or("missing placement property")?; + node.evidence["worker_id"] = json!(worker.id); + } + Ok(()) + } + + fn start_sampler( + &mut self, + scheduler: Arc, + account: Arc, + ) { + let snapshots = self.snapshots.clone(); + let task = nativelink_util::background_spawn!("fleet_sampler", async move { + loop { + let workers = scheduler.worker_snapshot().await; + snapshots.lock().unwrap().push(json!({ + "time_s": account.elapsed(), + "workers": workers, + "load": load(), + })); + tokio::time::sleep(Duration::from_millis(200)).await; + } + }); + self.services.push(task); + } + + async fn warmup( + &mut self, + address: &str, + context: &RunContext, + workload: &Workload, + account: &Arc, + ) -> Result<()> { + self.jobs = join_all(workload.warmup.iter().cloned().map(|job| { + job_result( + address.to_owned(), + job, + account.clone(), + "warmup", + Instant::now(), + None, + ) + })) + .await; + + if self.jobs.iter().any(|job| job["outcome"] != "success") { + std::fs::write( + context.output.join("warmup-failed.json"), + serde_json::to_vec_pretty(&self.jobs)?, + )?; + return Err("warmup failed; inspect logs".into()); + } + + account.quiesce().await?; + Ok(()) + } + + async fn disk_command(&mut self, command: &str, account: &account::Account) -> Result<()> { + // Only the first two workers receive the storage fault. Keep every + // command reply with its timestamp, including partial setup failures. + for node in self.nodes.iter_mut().take(2) { + self.events.push(json!({ + "time_s": account.elapsed(), + "node": node.evidence["node"], + "reply": node.command(command).await?, + })); + } + Ok(()) + } + + async fn work( + &mut self, + address: &str, + failed_fill: bool, + workload: &Workload, + account: &Arc, + ) -> Result<()> { + if failed_fill { + self.disk_command("fill", account).await?; + } + + account.phase("work"); + let origin = Instant::now(); + let repair_at = origin + Duration::from_secs(2); + self.fault = if failed_fill { + json!({"started_s": account.elapsed(), "repaired_s": null, "nodes": [0, 1]}) + } else { + Value::Null + }; + + let work = join_all(workload.jobs.iter().cloned().map(|job| { + job_result( + address.to_owned(), + job, + account.clone(), + "work", + origin, + failed_fill.then_some(repair_at), + ) + })); + let repair = async { + if failed_fill { + tokio::time::sleep_until(repair_at).await; + self.disk_command("repair", account).await?; + self.fault["repaired_s"] = json!(account.elapsed()); + } + Ok::<_, Box>(()) + }; + + // Repair proceeds while actions are still running. Wait for both sides + // and retain completed work even if the repair itself reports an error. + let (completed, repaired) = tokio::join!(work, repair); + self.jobs.extend(completed); + repaired?; + account.quiesce().await?; + Ok(()) + } + + async fn sentinels( + &mut self, + address: &str, + workload: &Workload, + account: &Arc, + ) -> Result<()> { + self.before_sentinel = account.report(); + account.phase("sentinel"); + let sentinel = join_all(workload.sentinel.iter().cloned().map(|job| { + job_result( + address.to_owned(), + job, + account.clone(), + "sentinel", + Instant::now(), + None, + ) + })) + .await; + + // Retain completed sentinels even if a subsequent snapshot fails. + self.jobs.extend(sentinel.iter().cloned()); + for (index, node) in self.nodes.iter_mut().enumerate() { + let evidence = node.command("snapshot").await?; + node.evidence["final"] = evidence.clone(); + node.evidence["sentinel_same_process"] = json!( + sentinel[index]["outcome"] == "success" + && evidence["pid"] == node.evidence["pid"] + && sentinel[index]["worker_id"] == node.evidence["worker_id"] + ); + } + + account.quiesce().await?; + Ok(()) + } + + async fn run( + &mut self, + context: &RunContext, + workload: &Workload, + account: Arc, + ) -> Result<()> { + let (address, scheduler) = self.start(context, workload, account.clone()).await?; + self.start_sampler(scheduler.clone(), account.clone()); + + self.warmup(&address, context, workload, &account).await?; + self.work( + &address, + context.scenario == "failed-fill", + workload, + &account, + ) + .await?; + self.sentinels(&address, workload, &account).await?; + + // Join the sampler before taking the final snapshot; shutdown observations + // must not be mistaken for lost workers during the experiment. + let sampler = self + .services + .pop() + .expect("sampler is the last started service"); + sampler.abort(); + drop(sampler.await); + let workers = scheduler.worker_snapshot().await; + self.snapshots.lock().unwrap().push(json!({ + "time_s": account.elapsed(), + "workers": workers, + "final": true, + })); + Ok(()) + } + + async fn stop(&mut self) -> Vec { + let mut cleanup = Vec::new(); + for node in &mut self.nodes { + cleanup.push(node.stop().await); + } + + for service in self.services.drain(..) { + service.abort(); + drop(service.await); + } + cleanup + } + + fn collect_evidence(&mut self, output: &Path, account_report: &Value) { + for node in &mut self.nodes { + node.evidence["sentinel_cache_reused"] = json!( + !account_report["reads"] + .as_array() + .unwrap() + .iter() + .any(|read| { + read["phase"] == "sentinel" && read["node"] == node.evidence["node"] + }) + ); + + let log_path = output.join(format!("node-{}.log", node.evidence["node"])); + match std::fs::read_to_string(log_path) { + Ok(log) => { + node.evidence["error_witnesses"] = json!( + log.lines() + .filter(|line| { + line.contains("os error 28") + || line.contains("Retries exhausted") + || line.contains("Tried to send while stream is closed") + || line.contains("receiver disconnected") + }) + .collect::>() + ); + node.evidence["worker_enospc_observed"] = json!(log.contains("os error 28")); + node.evidence["closed_destination_retry_exhaustions"] = + json!(checks::closed_destination_exhaustions(&log)); + } + Err(error) => node.evidence["log_error"] = json!(error.to_string()), + } + + // Request ownership comes from the SDK context, not HEAD order or + // matching clocks between worker and HTTP account processes. + let reads: Vec<_> = account_report["reads"] + .as_array() + .unwrap() + .iter() + .filter(|read| read["node"] == node.evidence["node"]) + .cloned() + .collect(); + let trace_path = output.join(format!("node-{}.downloads.jsonl", node.evidence["node"])); + node.evidence["download_witness"] = witness::analyze(&trace_path, &reads); + } + } + + fn report( + &self, + context: &RunContext, + workload: &Workload, + session_error: Option<&str>, + cleanup: &[Value], + account_report: &Value, + duration_s: f64, + ) -> Value { + json!({ + "schema_version": 2, + "scenario": context.scenario, + "seed": context.seed, + "session_error": session_error, + "cleanup": cleanup, + "expected_jobs": workload.warmup.len() + workload.jobs.len() + workload.sentinel.len(), + "binary_sha256": context.binary_sha256, + "workload_sha256": context.workload_sha256, + "start_unix_s": context.start_unix_s, + "duration_s": duration_s, + "load": {"before": context.load_before, "after": load()}, + "fault": self.fault, + "events": self.events, + "jobs": self.jobs, + "nodes": self.nodes.iter().map(|node| &node.evidence).collect::>(), + "account": account_report, + "before_sentinel": self.before_sentinel, + "scheduler_snapshots": *self.snapshots.lock().unwrap(), + "cpu_affinity": self.affinity, + "measurement": { + "topology": "four separate LocalWorker processes, private mounts, real scheduler over gRPC, shared HTTP origin on one host", + "retry": { + "max_retries": 6, + "delay": 0.3, + "jitter": 0.5, + "throttle_retry_delay_s": 2, + "throttle_retry_max_delay_s": 60, + "throttle_retry_budget_s": 0, + "sdk_retry": "disabled", + }, + "client_policy": "one resubmission after repair for failed work; every attempt retained", + "started_s": "client submission, not command start", + } + }) + } +} + +fn write_report(output: &Path, mut report: Value) -> Result<()> { + let failures = checks::failures(&report); + report["accepted"] = json!(failures.is_empty()); + report["failures"] = json!(failures); + let path = output.join("run.json"); + std::fs::write(&path, serde_json::to_vec_pretty(&report)?)?; + + announce(format_args!( + "payload_bytes={} completed={} results={}", + report["account"]["payload_bytes"], + report["account"]["completed_actions"], + path.display() + )); + for node in report["nodes"].as_array().into_iter().flatten() { + announce(format_args!( + "node={}: {}", + node["node"], + witness::explain(&node["download_witness"]) + )); + } + if !failures.is_empty() { + return Err(format!( + "{}; results retained at {}", + failures.join("; "), + path.display() + ) + .into()); + } + Ok(()) +} + +async fn run(args: Args) -> Result<()> { + let mut context = RunContext::new(args)?; + let workload = Workload::new(context.seed); + context.write_manifest(&workload)?; + let account = account::Account::new(context.seed, &workload.blobs); + let mut session = Session::default(); + let outcome = session.run(&context, &workload, account.clone()).await; + + // Acceptance errors and setup failures use the same teardown/report path. + // Keep all observed work, HTTP ranges and raw logs even when the run fails. + let cleanup = session.stop().await; + let account_report = account.report(); + session.collect_evidence(&context.output, &account_report); + let session_error = outcome.as_ref().err().map(ToString::to_string); + let report = session.report( + &context, + &workload, + session_error.as_deref(), + &cleanup, + &account_report, + account.elapsed(), + ); + write_report(&context.output, report) +} + +#[expect( + clippy::disallowed_methods, + reason = "the standalone executable owns its Tokio runtime" +)] +#[tokio::main(flavor = "multi_thread", worker_threads = 4)] +async fn main() -> Result<()> { + drop(tokio_rustls::rustls::crypto::ring::default_provider().install_default()); + let args = Args::parse(); + let capture = args + .node + .as_ref() + .map(|path| witness::capture(&Path::new(path).with_extension("downloads.jsonl"))) + .transpose()?; + let (capture_layer, capture_handle) = match capture { + Some((layer, handle)) => (Some(layer), Some(handle)), + None => (None, None), + }; + let logs = tracing_subscriber::fmt::layer() + .with_writer(std::io::stderr) + .with_ansi(false) + .with_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "warn".into()), + ); + tracing_subscriber::registry() + .with(logs) + .with(capture_layer.map(|layer| { + layer.with_filter(tracing_subscriber::filter::filter_fn(|metadata| { + metadata.target() == "nativelink_store::azure_blob_store::download" + })) + })) + .init(); + let result = if let Some(path) = args.node { + node::run(&path).await + } else { + run(args).await + }; + if let Some(handle) = capture_handle { + handle.finish()?; + } + result +} diff --git a/examples/azure_fleet/node.rs b/examples/azure_fleet/node.rs new file mode 100644 index 00000000000..6a0fd1c8f5f --- /dev/null +++ b/examples/azure_fleet/node.rs @@ -0,0 +1,529 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! One real `LocalWorker` per process and private mount namespace. + +use std::collections::BTreeMap; +use std::fmt::Debug; +use std::io::{BufRead, Write}; +use std::os::unix::fs::MetadataExt; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, SystemTime}; + +use azure_core::http::{RetryOptions, Transport, Url}; +use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions}; +use bytes::Bytes; +use nativelink_config::cas_server::LocalWorkerConfig; +use nativelink_config::stores::{ + CommonObjectSpec, EvictionPolicy, ExperimentalAzureSpec, ExperimentalCloudObjectSpec, + FastSlowSpec, FilesystemSpec, MemorySpec, Retry, StoreDirection, StoreSpec, VerifySpec, +}; +use nativelink_store::azure_blob_store::AzureBlobStore; +use nativelink_store::fast_slow_store::FastSlowStore; +use nativelink_store::filesystem_store::FilesystemStore; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::verify_store::VerifyStore; +use nativelink_util::common::DigestInfo; +use nativelink_util::shutdown_guard::{Priority, ShutdownGuard}; +use nativelink_util::store_trait::{Store, StoreLike}; +use nativelink_worker::local_worker::new_local_worker; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use tokio::io::AsyncWriteExt; +use tokio::sync::{broadcast, mpsc}; +use tokio::task::JoinHandle; + +use crate::Result; +use crate::workload::{MIB, Metadata, uniform}; + +type FleetAzureStore = AzureBlobStore SystemTime>; + +#[derive(Serialize, Deserialize)] +pub(crate) struct Config { + pub(crate) seed: u64, + pub(crate) node: usize, + pub(crate) root: PathBuf, + pub(crate) azure: String, + pub(crate) scheduler: String, + pub(crate) metadata: Vec, +} + +#[derive(Clone, Copy, Ord, PartialOrd, Eq, PartialEq)] +enum Allocation { + Filler, + Executable, + Content, + Work, + Temporary, + Other, +} + +fn allocation(root: &Path, path: &Path) -> Allocation { + let relative = path + .strip_prefix(root) + .expect("snapshot path is within root"); + match relative + .components() + .next() + .and_then(|p| p.as_os_str().to_str()) + { + Some("pressure") => Allocation::Filler, + Some("content.exec") => Allocation::Executable, + Some("content") => Allocation::Content, + Some("work") => Allocation::Work, + Some("temp") => Allocation::Temporary, + _ => Allocation::Other, + } +} + +fn snapshot(root: &Path) -> Result { + let mut stack = vec![root.to_path_buf()]; + let mut inodes = BTreeMap::<(u64, u64), (u64, Allocation)>::new(); + let mut vanished_entries = 0_u64; + + while let Some(path) = stack.pop() { + let entries = match std::fs::read_dir(path) { + Ok(entries) => entries, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + vanished_entries += 1; + continue; + } + Err(e) => return Err(e.into()), + }; + for e in entries { + let e = e?; + let m = match e.metadata() { + Ok(metadata) => metadata, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + vanished_entries += 1; + continue; + } + Err(e) => return Err(e.into()), + }; + if m.is_dir() { + stack.push(e.path()); + continue; + } + if !m.is_file() { + continue; + } + let bytes = m.blocks() * 512; + let category = allocation(root, &e.path()); + inodes + .entry((m.dev(), m.ino())) + .and_modify(|entry| { + // Charge shared input links to their cache, independent of traversal order. + // Concurrent writes can change allocation during this non-atomic snapshot. + entry.0 = entry.0.max(bytes); + entry.1 = entry.1.min(category); + }) + .or_insert((bytes, category)); + } + } + // Aggregate after inode deduplication: work-directory hardlinks must not + // count as extra copies of their corresponding cache entries. + let (mut content, mut exec, mut work, mut filler, mut temporary, mut other, mut total) = + (0_u64, 0, 0, 0, 0, 0, 0); + for (bytes, category) in inodes.values() { + total += bytes; + match category { + Allocation::Filler => filler += bytes, + Allocation::Executable => exec += bytes, + Allocation::Content => content += bytes, + Allocation::Work => work += bytes, + Allocation::Temporary => temporary += bytes, + Allocation::Other => other += bytes, + } + } + Ok(json!({ + "content_allocated": content, + "exec_allocated": exec, + "work_unique_allocated": work, + "filler_allocated": filler, + "temporary_allocated": temporary, + "other_allocated": other, + "total_allocated": total, + "unique_file_inodes": inodes.len(), + "vanished_entries": vanished_entries, + "atomic": false, + "boundary": "allocated file blocks; unique device/inode" + })) +} + +fn fill(root: &Path) -> Result { + let mut f = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(root.join("pressure"))?; + let block = vec![0_u8; 1024 * 1024]; + let mut bytes = 0_u64; + + loop { + match f.write(&block) { + Ok(0) => return Err(std::io::Error::other("zero write before ENOSPC").into()), + Ok(n) => bytes += u64::try_from(n)?, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(e) if e.raw_os_error() == Some(libc::ENOSPC) => break, + Err(e) => return Err(e.into()), + } + } + + let mut result = snapshot(root)?; + result["filler_written"] = json!(bytes); + result["errno"] = json!(libc::ENOSPC); + Ok(result) +} + +async fn reply(value: Value) -> Result<()> { + let mut out = tokio::io::stdout(); + out.write_all(format!("{value}\n").as_bytes()).await?; + out.flush().await?; + Ok(()) +} + +async fn disk_snapshot(root: PathBuf) -> Result { + nativelink_util::spawn_blocking!("fleet_disk_snapshot", move || snapshot(&root)).await? +} + +async fn mount_cache(root: &Path) -> Result<()> { + tokio::fs::create_dir_all(root).await?; + if tokio::fs::read_dir(root) + .await? + .next_entry() + .await? + .is_some() + { + return Err(std::io::Error::other("worker mount point must be empty").into()); + } + let status = tokio::process::Command::new("mount") + .args(["-t", "tmpfs", "-o", "size=1g,nosuid,nodev", "tmpfs"]) + .arg(root) + .status() + .await?; + if !status.success() { + return Err(std::io::Error::other("private tmpfs mount failed; start node through unshare --user --map-root-user --mount --propagation private").into()); + } + Ok(()) +} + +async fn create_fast_store(config: &Config) -> Result<(FilesystemSpec, Store)> { + let fs_spec = FilesystemSpec { + content_path: config.root.join("content").to_string_lossy().into_owned(), + temp_path: config.root.join("temp").to_string_lossy().into_owned(), + eviction_policy: Some(EvictionPolicy { + max_bytes: usize::try_from(512 * MIB)?, + ..Default::default() + }), + ..Default::default() + }; + let filesystem: Arc = FilesystemStore::new(&fs_spec).await?; + let fast = Store::new(filesystem); + + // Seed only metadata. Payloads always traverse HTTP on a cold miss. + for m in &config.metadata { + fast.update_oneshot( + DigestInfo::try_new(&m.hash, m.size)?, + Bytes::from(hex::decode(&m.hex)?), + ) + .await?; + } + Ok((fs_spec, fast)) +} + +fn create_azure_store(config: &Config) -> Result<(ExperimentalAzureSpec, Arc)> { + let retry = Retry { + max_retries: 6, + delay: 0.3, + jitter: 0.5, + ..Default::default() + }; + let spec = ExperimentalAzureSpec { + account_name: "fleet".into(), + container: "cas".into(), + common: CommonObjectSpec { + retry, + throttle_retry_budget_s: 0.0, + ..Default::default() + }, + ..Default::default() + }; + + // Disable SDK-level retries so the observed retry policy belongs to the store. + let mut options = BlobContainerClientOptions::default(); + options.client_options.retry = RetryOptions::none(); + crate::transport::attach(&mut options.client_options); + options.client_options.transport = Some(Transport::new(Arc::new( + reqwest::Client::builder().no_proxy().build()?, + ))); + let client = BlobContainerClient::new( + Url::parse(&format!("{}/{}/cas", config.azure, config.node))?, + None, + Some(options), + )?; + // Each worker keeps its own deterministic jitter draw sequence. + let count = AtomicU64::new(0); + let seed = config.seed; + let node = u64::try_from(config.node)?; + let azure = FleetAzureStore::new_with_client_and_jitter( + &spec, + client, + Arc::new(move |delay: Duration| { + let index = count.fetch_add(1, Ordering::Relaxed); + #[expect( + clippy::cast_possible_truncation, + reason = "The seeded draw is bounded to (0, 1); rounding to f32 preserves production jitter arithmetic" + )] + let q = uniform(seed, 100_000 + node * 10_000 + index) as f32; + delay.mul_f32(0.5_f32.mul_add(q - 0.5, 1.0)) + }), + SystemTime::now, + )?; + Ok((spec, azure)) +} + +async fn create_cas(config: &Config) -> Result { + let (fs_spec, fast) = create_fast_store(config).await?; + let (spec, azure) = create_azure_store(config)?; + + // Preserve the normal verified slow-store / filesystem fast-store path. + let verify_spec = VerifySpec { + backend: StoreSpec::ExperimentalCloudObjectStore(ExperimentalCloudObjectSpec::Azure(spec)), + verify_size: true, + verify_hash: false, + }; + let verify = VerifyStore::new(&verify_spec, Store::new(azure)); + Ok(Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fs_spec), + slow: StoreSpec::Verify(Box::new(verify_spec)), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast, + Store::new(verify), + ))) +} + +fn worker_config(config: &Config) -> Result { + Ok(serde_json::from_value(json!({ + "name": format!("fleet-{}", config.node), + "worker_api_endpoint": {"uri": config.scheduler}, + "max_inflight_tasks": 2, + "cas_fast_slow_store": "cas", + "work_directory": config.root.join("work"), + "platform_properties": { + "fleet_node": {"values": [config.node.to_string()]} + }, + "max_action_timeout_s": 120, + "max_download_timeout_s": 90, + "max_upload_timeout_s": 30, + "drain_on_shutdown": true, + "upload_action_result": {"upload_ac_results_strategy": "never"} + }))?) +} + +fn read_commands() -> Result>> { + // Tokio's stdin uses an uncancellable blocking read that can hold runtime + // shutdown open after a worker failure. A detached reader thread does not. + let (commands_tx, commands) = mpsc::channel(16); + #[expect( + clippy::disallowed_methods, + reason = "Uncancellable stdin must use a detached OS thread independent of Tokio runtime shutdown" + )] + let _control_reader = std::thread::Builder::new() + .name("fleet-control".into()) + .spawn(move || { + for line in std::io::stdin().lock().lines() { + if commands_tx.blocking_send(line).is_err() { + break; + } + } + })?; + Ok(commands) +} + +async fn control_loop( + root: &Path, + commands: &mut mpsc::Receiver>, + worker_task: &mut JoinHandle, + worker_finished: &mut bool, +) -> Result<()> { + loop { + let line = tokio::select! { + result = &mut *worker_task => { + *worker_finished = true; + return Err(std::io::Error::other(format!("worker exited before controller stop: {result:?}")).into()); + }, + line = commands.recv() => line.transpose()?, + }; + let Some(line) = line else { + return Err(std::io::Error::other("node controller disconnected").into()); + }; + if apply_command(root, &line).await? { + return Ok(()); + } + } +} + +async fn apply_command(root: &Path, command: &str) -> Result { + match command { + "fill" => { + let root = root.to_path_buf(); + let disk = + nativelink_util::spawn_blocking!("fleet_fill", move || fill(&root)).await??; + reply(json!({ + "event": "filled", + "pid": std::process::id(), + "disk": disk + })) + .await?; + } + "repair" => { + // Free the same filesystem; neither the worker nor its cache is replaced. + tokio::fs::remove_file(root.join("pressure")).await?; + let disk = disk_snapshot(root.to_path_buf()).await?; + reply(json!({ + "event": "repaired", + "pid": std::process::id(), + "disk": disk + })) + .await?; + } + "snapshot" => { + let disk = disk_snapshot(root.to_path_buf()).await?; + reply(json!({ + "event": "snapshot", + "pid": std::process::id(), + "disk": disk + })) + .await?; + } + "stop" => return Ok(true), + _ => return Err(std::io::Error::other("unknown node control command").into()), + } + Ok(false) +} + +async fn drain_worker( + tx: &broadcast::Sender, + worker_task: JoinHandle, +) -> bool { + // LocalWorker acknowledges drain by dropping its ShutdownGuard; its service + // loop stays alive until the process exits, just as in the server. + let mut guard = ShutdownGuard::default(); + let sent = tx.send(guard.clone()).is_ok(); + let drained = sent + && tokio::time::timeout(Duration::from_secs(5), guard.wait_for(Priority::P0)) + .await + .is_ok(); + + // Reap the worker task even when drain fails. The controller decides which + // error to report after this cleanup has finished. + worker_task.abort(); + drop(worker_task.await); + drained +} + +pub(crate) async fn run(path: &str) -> Result<()> { + let config: Config = serde_json::from_slice(&tokio::fs::read(path).await?)?; + mount_cache(&config.root).await?; + let cas = create_cas(&config).await?; + let worker = new_local_worker( + Arc::new(worker_config(&config)?), + cas, + None, + Store::new(MemoryStore::new(&MemorySpec::default())), + ) + .await?; + let (tx, rx) = broadcast::channel(1); + let mut worker_task = nativelink_util::background_spawn!("fleet_worker", worker.run(rx)); + + reply(json!({ + "ready": true, + "pid": std::process::id(), + "node": config.node, + "process_cpus": std::thread::available_parallelism()?.get(), + "disk": disk_snapshot(config.root.clone()).await?, + "physical_limit_bytes": 1024 * MIB, + "logical_cache_bytes": 512 * MIB + })) + .await?; + + let mut commands = read_commands()?; + let mut worker_finished = false; + let control_result = control_loop( + &config.root, + &mut commands, + &mut worker_task, + &mut worker_finished, + ) + .await; + + // Preserve a controller error after cleanup; a failed drain becomes the + // reported error only when the controller otherwise completed normally. + if !worker_finished { + let drained = drain_worker(&tx, worker_task).await; + if !drained && control_result.is_ok() { + return Err(std::io::Error::other("worker did not drain within five seconds").into()); + } + } + control_result?; + + reply(json!({ + "event": "stopped", + "pid": std::process::id() + })) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn snapshot_counts_shared_inputs_once_and_charges_them_to_the_cache() -> Result<()> { + struct Directory(PathBuf); + impl Drop for Directory { + fn drop(&mut self) { + drop(std::fs::remove_dir_all(&self.0)); + } + } + let root = Directory(std::env::temp_dir().join(format!( + "azure-fleet-snapshot-{}-{}", + std::process::id(), + SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_nanos() + ))); + for name in ["content", "content.exec", "work"] { + std::fs::create_dir_all(root.0.join(name))?; + } + std::fs::write(root.0.join("content/blob"), vec![1_u8; 8192])?; + std::fs::write(root.0.join("content.exec/blob"), vec![1_u8; 8192])?; + std::fs::write(root.0.join("work/output"), vec![2_u8; 4096])?; + std::fs::hard_link(root.0.join("content/blob"), root.0.join("work/input"))?; + std::fs::hard_link(root.0.join("content.exec/blob"), root.0.join("work/tool"))?; + let content = std::fs::metadata(root.0.join("content/blob"))?.blocks() * 512; + let exec = std::fs::metadata(root.0.join("content.exec/blob"))?.blocks() * 512; + let work = std::fs::metadata(root.0.join("work/output"))?.blocks() * 512; + let measured = snapshot(&root.0)?; + assert_eq!(measured["unique_file_inodes"], 3); + assert_eq!(measured["content_allocated"], content); + assert_eq!(measured["exec_allocated"], exec); + assert_eq!(measured["work_unique_allocated"], work); + assert_eq!(measured["total_allocated"], content + exec + work); + Ok(()) + } +} diff --git a/examples/azure_fleet/transport.rs b/examples/azure_fleet/transport.rs new file mode 100644 index 00000000000..f5d15d445f4 --- /dev/null +++ b/examples/azure_fleet/transport.rs @@ -0,0 +1,52 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Export SDK-owned download identity to the local account's HTTP ledger. + +use std::sync::Arc; + +use async_trait::async_trait; +use azure_core::http::policies::{Policy, PolicyResult}; +use azure_core::http::{ClientOptions, Context, Request}; +use nativelink_store::azure_blob_store::AzureDownloadAttempt; + +/// Only the replay client installs this policy. Ordinary Azure clients retain +/// the typed context locally without adding these fixture-specific headers. +pub(crate) fn attach(options: &mut ClientOptions) { + options.per_call_policies.push(Arc::new(DownloadIdentity)); +} + +#[derive(Debug)] +struct DownloadIdentity; + +#[async_trait] +impl Policy for DownloadIdentity { + async fn send( + &self, + context: &Context, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + // The SDK owns this context and clones it into every partition task. + // A detached old task therefore keeps its original attempt identity. + if let Some(attempt) = context.value::() { + request.insert_header("x-nl-read-id", attempt.read_id.to_string()); + request.insert_header("x-nl-attempt", attempt.attempt.to_string()); + request.insert_header("x-nl-sent", attempt.sent.to_string()); + request.insert_header("x-nl-writer-broken", attempt.writer_broken.to_string()); + } + + next[0].send(context, request, &next[1..]).await + } +} diff --git a/examples/azure_fleet/witness.rs b/examples/azure_fleet/witness.rs new file mode 100644 index 00000000000..a405272412f --- /dev/null +++ b/examples/azure_fleet/witness.rs @@ -0,0 +1,748 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Ordered worker events joined to HTTP requests by SDK-carried identity. +//! +//! The witness is a failed send to a known broken writer followed by another +//! outer download of that read. Neither final error text nor cross-process +//! clock comparisons establish that ordering. Requests belonging to an older +//! download remain in that attempt, even when SDK tasks outlive the read. +//! The sent counter measures bytes accepted by the writer channel, not bytes +//! received by its downstream consumer. + +use std::collections::btree_map::Entry; +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::{self, File, OpenOptions}; +use std::io::{self, BufWriter, Write}; +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use tracing::field::{Field, Visit}; +use tracing::{Event as TracingEvent, Subscriber}; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context; + +const TARGET: &str = "nativelink_store::azure_blob_store::download"; + +struct Capture { + writer: BufWriter, + started: Instant, + sequence: u64, + error: Option, + finished: bool, +} + +impl Capture { + fn write(&mut self, mut fields: Map) { + self.sequence += 1; + fields.insert("sequence".into(), self.sequence.into()); + fields.insert( + "elapsed_s".into(), + self.started.elapsed().as_secs_f64().into(), + ); + let result = serde_json::to_writer(&mut self.writer, &fields) + .map_err(io::Error::other) + .and_then(|()| self.writer.write_all(b"\n")); + if let Err(error) = result { + self.error.get_or_insert_with(|| error.to_string()); + } + } +} + +pub(crate) struct CaptureLayer(Arc>); +pub(crate) struct CaptureHandle(Arc>); + +/// A fresh file avoids confusing two worker process generations. +pub(crate) fn capture(path: &Path) -> io::Result<(CaptureLayer, CaptureHandle)> { + let writer = BufWriter::new(OpenOptions::new().write(true).create_new(true).open(path)?); + let shared = Arc::new(Mutex::new(Capture { + writer, + started: Instant::now(), + sequence: 0, + error: None, + finished: false, + })); + Ok((CaptureLayer(Arc::clone(&shared)), CaptureHandle(shared))) +} + +impl CaptureHandle { + /// Called after the worker drains. The marker closes event coverage, not + /// the SDK's HTTP tasks; the account ledger measures those separately. + pub(crate) fn finish(self) -> io::Result<()> { + let mut state = self.0.lock().map_err(|e| io::Error::other(e.to_string()))?; + let mut footer = Map::new(); + footer.insert("event".into(), "trace_finished".into()); + footer.insert("capture_error".into(), json!(state.error)); + state.write(footer); + state.finished = true; + if let Err(error) = state.writer.flush() { + state.error.get_or_insert_with(|| error.to_string()); + } + match &state.error { + Some(error) => Err(io::Error::other(error.clone())), + None => Ok(()), + } + } +} + +#[derive(Default)] +struct Fields(Map); + +impl Visit for Fields { + fn record_u64(&mut self, field: &Field, value: u64) { + self.0.insert(field.name().into(), value.into()); + } + + fn record_i64(&mut self, field: &Field, value: i64) { + self.0.insert(field.name().into(), value.into()); + } + + fn record_bool(&mut self, field: &Field, value: bool) { + self.0.insert(field.name().into(), value.into()); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.0.insert(field.name().into(), value.into()); + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + self.0 + .insert(field.name().into(), format!("{value:?}").into()); + } +} + +impl Layer for CaptureLayer { + fn on_event(&self, event: &TracingEvent<'_>, _context: Context<'_, S>) { + if event.metadata().target() != TARGET { + return; + } + let mut fields = Fields::default(); + event.record(&mut fields); + // A poisoned mutex also makes finish fail. Do not let a diagnostic + // layer panic inside the production operation being observed. + if let Ok(mut state) = self.0.lock() { + if state.finished { + state + .error + .get_or_insert_with(|| "event after trace_finished".into()); + } + state.write(fields.0); + if state.finished { + // Preserve late events so analysis rejects the early footer. + if let Err(error) = state.writer.flush() { + state.error.get_or_insert_with(|| error.to_string()); + } + } + } + } +} + +#[derive(Clone, Deserialize, Serialize)] +struct Event { + sequence: u64, + elapsed_s: f64, + #[serde(rename = "event")] + kind: String, + read_id: u64, + attempt: u64, + sent: u64, + writer_broken: bool, + #[serde(skip_serializing_if = "Option::is_none")] + blob: Option, + #[serde(skip_serializing_if = "Option::is_none")] + stage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + outcome: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error_code: Option, +} + +#[derive(Serialize)] +struct Attempt { + started: Event, + failed_sends: Vec, + after_failed_send_sequence: Option, + http_requests: Vec, + body_bytes: u64, +} + +#[derive(Serialize)] +struct Read { + started: Option, + finished: Option, + first_failed_send: Option, + attempts: BTreeMap, + last_sent: u64, +} + +impl Read { + const fn new() -> Self { + Self { + started: None, + finished: None, + first_failed_send: None, + attempts: BTreeMap::new(), + last_sent: 0, + } + } + + fn observe(&mut self, event: Event, errors: &mut Vec) { + let prefix = format!("read {} event {}", event.read_id, event.sequence); + if self.finished.is_some() { + errors.push(format!("{prefix}: event after read_finished")); + } + if event.sent < self.last_sent { + errors.push(format!("{prefix}: writer byte counter decreased")); + } + if self + .first_failed_send + .as_ref() + .is_some_and(|failed| event.sent != failed.sent) + { + errors.push(format!( + "{prefix}: writer bytes changed after terminal send failure" + )); + } + if self.first_failed_send.is_some() && !event.writer_broken { + errors.push(format!("{prefix}: broken destination became writable")); + } + self.last_sent = event.sent; + match event.kind.as_str() { + "read_started" => { + if self.started.is_some() + || !self.attempts.is_empty() + || event.attempt != 0 + || event.blob.is_none() + { + errors.push(format!("{prefix}: invalid or duplicate read_started")); + } + self.started.get_or_insert(event); + } + "attempt_started" => { + let expected = self + .attempts + .last_key_value() + .map_or(Some(1), |(id, _)| id.checked_add(1)); + if self.started.is_none() || Some(event.attempt) != expected { + errors.push(format!( + "{prefix}: expected attempt {expected:?} after read_started" + )); + } + if event.writer_broken && self.first_failed_send.is_none() { + errors.push(format!( + "{prefix}: destination broke without observed failed send" + )); + } + let attempt = Attempt { + after_failed_send_sequence: self.first_failed_send.as_ref().map(|e| e.sequence), + started: event, + failed_sends: Vec::new(), + http_requests: Vec::new(), + body_bytes: 0, + }; + match self.attempts.entry(attempt.started.attempt) { + Entry::Occupied(_) => { + errors.push(format!("{prefix}: duplicate attempt_started")); + } + Entry::Vacant(entry) => { + entry.insert(attempt); + } + } + } + "send_failed" => { + if !matches!(event.stage.as_deref(), Some("data" | "eof")) { + errors.push(format!("{prefix}: unknown failed-send stage")); + } + if self.attempts.last_key_value().map(|(id, _)| *id) != Some(event.attempt) { + errors.push(format!("{prefix}: failed send outside current attempt")); + } + // Healthy EOF also leaves is_pipe_broken() true. Only a failed + // send plus this flag establishes the terminal destination. + if event.writer_broken { + self.first_failed_send.get_or_insert_with(|| event.clone()); + } + if let Some(attempt) = self.attempts.get_mut(&event.attempt) { + attempt.failed_sends.push(event); + } + } + "read_finished" => { + if self.first_failed_send.is_some() && event.outcome.as_deref() != Some("error") { + errors.push(format!("{prefix}: terminal send failure was not returned")); + } + let last = self.attempts.last_key_value().map_or(0, |(id, _)| *id); + if self.started.is_none() + || event.attempt != last + || !matches!(event.outcome.as_deref(), Some("ok" | "error")) + || event.error_code.is_none() + { + errors.push(format!("{prefix}: invalid read_finished")); + } + if event.outcome.as_deref() == Some("error") + && event.writer_broken + && self.first_failed_send.is_none() + { + errors.push(format!( + "{prefix}: error with broken writer but no observed failed send" + )); + } + self.finished.get_or_insert(event); + } + _ => errors.push(format!("{prefix}: unknown event {}", event.kind)), + } + } +} + +fn load_trace(path: &Path, errors: &mut Vec) -> BTreeMap { + let mut reads = BTreeMap::new(); + let contents = match fs::read(path) { + Ok(contents) => contents, + Err(error) => { + errors.push(format!("reading {}: {error}", path.display())); + return reads; + } + }; + if let Err(error) = std::str::from_utf8(&contents) { + errors.push(format!("trace contains invalid UTF-8: {error}")); + } + let contents = String::from_utf8_lossy(&contents); + let mut sequence = 0_u64; + let mut elapsed = 0.0; + let mut footer = false; + for (line, contents) in contents.lines().enumerate() { + let prefix = format!("trace line {}", line + 1); + if footer { + errors.push(format!("{prefix}: record after trace_finished")); + } + let value: Value = match serde_json::from_str(contents) { + Ok(value) => value, + Err(error) => { + errors.push(format!("{prefix}: {error}")); + continue; + } + }; + if value["sequence"].as_u64().is_none() + || value["sequence"].as_u64() != sequence.checked_add(1) + { + errors.push(format!("{prefix}: missing or nonconsecutive sequence")); + } + sequence = value["sequence"].as_u64().unwrap_or(sequence); + match value["elapsed_s"].as_f64() { + Some(time) if time.is_finite() && time >= elapsed => elapsed = time, + _ => errors.push(format!("{prefix}: missing or decreasing elapsed_s")), + } + if value["event"] == "trace_finished" { + footer = true; + if value.get("capture_error") != Some(&Value::Null) { + errors.push(format!("{prefix}: missing or non-null capture_error")); + } + continue; + } + match serde_json::from_value::(value) { + Ok(event) => reads + .entry(event.read_id) + .or_insert_with(Read::new) + .observe(event, errors), + Err(error) => errors.push(format!("{prefix}: {error}")), + } + } + if !footer { + errors.push("missing trace_finished marker".into()); + } + for (id, read) in &reads { + if read.started.is_none() || read.finished.is_none() { + errors.push(format!("read {id}: incomplete lifecycle")); + } + } + reads +} + +#[derive(Deserialize)] +struct HttpIdentity { + request_id: u64, + read_id: u64, + attempt: u64, + sent: u64, + writer_broken: bool, + bytes: u64, + blob: String, +} + +fn attach_request(reads: &mut BTreeMap, value: &Value) -> Result<(), String> { + let identity: HttpIdentity = + serde_json::from_value(value.clone()).map_err(|e| e.to_string())?; + let read = reads.get_mut(&identity.read_id).ok_or("unknown read_id")?; + if read + .started + .as_ref() + .and_then(|started| started.blob.as_deref()) + != Some(identity.blob.as_str()) + { + return Err("HTTP blob differs from read_started".into()); + } + let attempt = read + .attempts + .get_mut(&identity.attempt) + .ok_or("unknown attempt")?; + if identity.sent != attempt.started.sent + || identity.writer_broken != attempt.started.writer_broken + { + return Err("SDK identity differs from attempt_started snapshot".into()); + } + // No timestamp restriction: the SDK can keep an old request alive after + // read_finished, and server and worker clocks are different boundaries. + attempt.body_bytes = attempt + .body_bytes + .checked_add(identity.bytes) + .ok_or("body byte overflow")?; + attempt.http_requests.push(value.clone()); + Ok(()) +} + +/// Reads are the account's GET ledger for exactly one worker process. Missing +/// identities and unfinished traces stay visible and make `complete` false. +pub(crate) fn analyze(path: &Path, http_reads: &[Value]) -> Value { + let mut errors = Vec::new(); + let mut reads = load_trace(path, &mut errors); + let mut request_ids = BTreeSet::new(); + let mut uncorrelated = Vec::new(); + for request in http_reads { + let parsed = serde_json::from_value::(request.clone()); + let result = match parsed { + Ok(identity) if !request_ids.insert(identity.request_id) => { + Err("duplicate request_id".into()) + } + Ok(_) => attach_request(&mut reads, request), + Err(error) => Err(error.to_string()), + }; + if let Err(error) = result { + errors.push(format!("HTTP request {}: {error}", request["request_id"])); + uncorrelated.push(request.clone()); + } + } + let new_attempts = reads + .values() + .flat_map(|r| r.attempts.values()) + .filter(|a| a.after_failed_send_sequence.is_some()) + .count(); + let failed_reads = reads + .values() + .filter(|r| r.first_failed_send.is_some()) + .count(); + json!({ + "complete": errors.is_empty(), + "errors": errors, + "failed_destination_reads": failed_reads, + "new_attempts_after_failed_send": new_attempts, + "reads": reads, + "uncorrelated_requests": uncorrelated, + "byte_boundary": "successful response payload yielded by account Body to Hyper, grouped by exact SDK outer attempt; not confirmed receipt or billed egress", + "ordering_boundary": "worker capture sequence; server and worker elapsed times are not compared", + }) +} + +/// Compact causal chain, with earlier SDK attempts kept separate from new ones. +pub(crate) fn explain(report: &Value) -> String { + let mut lines = vec![format!( + "complete={} failed_destination_reads={} new_attempts_after_failed_send={}", + report["complete"], + report["failed_destination_reads"], + report["new_attempts_after_failed_send"], + )]; + if let Some(reads) = report["reads"].as_object() { + for (id, read) in reads { + if read["first_failed_send"].is_null() { + continue; + } + let failed = &read["first_failed_send"]; + let mut first_forbidden = None; + let mut requests = 0; + let mut bytes = 0_u128; + let mut earlier_bytes = 0_u128; + if let Some(attempts) = read["attempts"].as_object() { + for attempt in attempts.values() { + if attempt["after_failed_send_sequence"].is_null() { + earlier_bytes += u128::from(attempt["body_bytes"].as_u64().unwrap_or(0)); + } else { + let id = attempt["started"]["attempt"].as_u64().unwrap_or(0); + first_forbidden = Some(first_forbidden.map_or(id, |old: u64| old.min(id))); + requests += attempt["http_requests"].as_array().map_or(0, Vec::len); + bytes += u128::from(attempt["body_bytes"].as_u64().unwrap_or(0)); + } + } + } + let delta = read["finished"]["sent"] + .as_u64() + .zip(failed["sent"].as_u64()) + .and_then(|(finish, failed)| finish.checked_sub(failed)); + lines.push(format!( + "read={id} failed_send(seq={},attempt={},sent={}) -> first_new_attempt={first_forbidden:?}, GETs={requests}, body_bytes={bytes}; earlier_attempt_body_bytes={earlier_bytes}; sent_delta={delta:?}, finish={}/{}", + failed["sequence"], failed["attempt"], failed["sent"], read["finished"]["outcome"], read["finished"]["error_code"], + )); + } + } + if report["complete"] != true { + lines.push(format!("incomplete evidence: {}", report["errors"])); + } + lines.join("; ") +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU64, Ordering}; + + use tracing_subscriber::prelude::*; + + use super::*; + + static FILE_ID: AtomicU64 = AtomicU64::new(0); + + struct TraceFile(std::path::PathBuf); + + impl TraceFile { + fn new() -> Self { + Self(std::env::temp_dir().join(format!( + "azure-witness-{}-{}.jsonl", + std::process::id(), + FILE_ID.fetch_add(1, Ordering::Relaxed), + ))) + } + + fn write(events: &[Value], footer: bool) -> Self { + let file = Self::new(); + let mut text = String::new(); + for event in events { + text.push_str(&event.to_string()); + text.push('\n'); + } + if footer { + text.push_str( + &json!({ + "sequence": events.len() + 1, + "elapsed_s": events.len() + 1, + "event": "trace_finished", + "capture_error": null, + }) + .to_string(), + ); + text.push('\n'); + } + fs::write(&file.0, text).unwrap(); + file + } + } + + impl Drop for TraceFile { + fn drop(&mut self) { + drop(fs::remove_file(&self.0)); + } + } + + fn event(sequence: u64, event: &str, attempt: u64, sent: u64, broken: bool) -> Value { + json!({ + "sequence": sequence, "elapsed_s": sequence, "event": event, + "read_id": 9, "attempt": attempt, "sent": sent, "writer_broken": broken, + "blob": "payload", "stage": "data", "outcome": "error", "error_code": "Unavailable", + }) + } + + fn forbidden_retry() -> Vec { + vec![ + event(1, "read_started", 0, 0, false), + event(2, "attempt_started", 1, 0, false), + event(3, "send_failed", 1, 4, true), + event(4, "attempt_started", 2, 4, true), + event(5, "read_finished", 2, 4, true), + ] + } + + fn request(id: u64, attempt: u64, sent: u64, broken: bool, bytes: u64) -> Value { + json!({ + "request_id": id, "read_id": 9, "attempt": attempt, + "sent": sent, "writer_broken": broken, "bytes": bytes, + "blob": "payload", "offset": sent, "length": bytes, + "status": 206, "start_s": 999, "end_s": 1000, + }) + } + + #[test] + fn final_throttle_does_not_hide_prior_forbidden_retry() { + let file = TraceFile::write(&forbidden_retry(), true); + let mut throttle = request(1, 2, 4, true, 0); + throttle["status"] = 503.into(); + let report = analyze(&file.0, &[request(0, 1, 0, false, 32), throttle]); + assert_eq!(report["complete"], true); + assert_eq!(report["failed_destination_reads"], 1); + assert_eq!(report["new_attempts_after_failed_send"], 1); + assert_eq!( + report["reads"]["9"]["attempts"]["2"]["after_failed_send_sequence"], + 3 + ); + assert!(explain(&report).contains("first_new_attempt=Some(2)")); + } + + #[test] + fn legitimate_retries_before_first_failed_send_are_not_violations() { + let events = vec![ + event(1, "read_started", 0, 0, false), + event(2, "attempt_started", 1, 0, false), + event(3, "attempt_started", 2, 0, false), + event(4, "send_failed", 2, 4, true), + event(5, "read_finished", 2, 4, true), + ]; + let file = TraceFile::write(&events, true); + let mut throttle = request(0, 1, 0, false, 0); + throttle["status"] = 503.into(); + let report = analyze(&file.0, &[throttle, request(1, 2, 0, false, 32)]); + assert_eq!(report["complete"], true); + assert_eq!(report["failed_destination_reads"], 1); + assert_eq!(report["new_attempts_after_failed_send"], 0); + } + + #[test] + fn old_sdk_tasks_keep_their_original_attempt_identity() { + let file = TraceFile::write(&forbidden_retry(), true); + let report = analyze( + &file.0, + &[ + request(0, 2, 4, true, 8), + // Observed after the later attempt and after read_finished. The + // arbitrary server clock cannot turn this into a new retry. + request(1, 1, 0, false, 32), + ], + ); + assert_eq!(report["complete"], true); + assert_eq!(report["new_attempts_after_failed_send"], 1); + assert_eq!(report["reads"]["9"]["attempts"]["1"]["body_bytes"], 32); + assert_eq!(report["reads"]["9"]["attempts"]["2"]["body_bytes"], 8); + } + + #[test] + fn healthy_eof_and_nonterminal_send_failure_do_not_anchor_a_violation() { + let mut events = forbidden_retry(); + events[2]["writer_broken"] = false.into(); + events[3]["writer_broken"] = false.into(); + events[4]["outcome"] = "ok".into(); + events[4]["error_code"] = "".into(); + let file = TraceFile::write(&events, true); + let report = analyze(&file.0, &[]); + assert_eq!(report["complete"], true); + assert_eq!(report["failed_destination_reads"], 0); + assert_eq!(report["new_attempts_after_failed_send"], 0); + } + + #[test] + fn missing_evidence_never_becomes_a_clean_pass() { + let file = TraceFile::write(&forbidden_retry(), true); + let mut missing_identity = request(0, 1, 0, false, 32); + missing_identity["read_id"] = Value::Null; + let mut bad_snapshot = request(0, 1, 4, false, 32); + bad_snapshot["writer_broken"] = true.into(); + for request in [missing_identity, request(0, 3, 4, true, 32), bad_snapshot] { + let report = analyze(&file.0, &[request]); + assert_eq!(report["complete"], false); + assert_eq!(report["new_attempts_after_failed_send"], 1); + assert_eq!(report["uncorrelated_requests"].as_array().unwrap().len(), 1); + } + let no_footer = TraceFile::write(&forbidden_retry(), false); + assert_eq!(analyze(&no_footer.0, &[])["complete"], false); + let no_finish = TraceFile::write(&forbidden_retry()[..4], true); + let report = analyze(&no_finish.0, &[]); + assert_eq!(report["complete"], false); + assert_eq!(report["new_attempts_after_failed_send"], 1); + let mut hole = forbidden_retry(); + hole.remove(1); + let hole = TraceFile::write(&hole, true); + assert_eq!(analyze(&hole.0, &[])["complete"], false); + } + + #[test] + fn unexplained_closure_or_post_failure_delivery_is_incomplete() { + let mut events = forbidden_retry(); + events.remove(2); + for (index, event) in events.iter_mut().enumerate() { + event["sequence"] = (index + 1).into(); + } + let file = TraceFile::write(&events, true); + let report = analyze(&file.0, &[]); + assert_eq!(report["complete"], false); + assert_eq!(report["new_attempts_after_failed_send"], 0); + + let events = vec![ + event(1, "read_started", 0, 0, false), + event(2, "attempt_started", 1, 0, false), + event(3, "read_finished", 1, 4, true), + ]; + let file = TraceFile::write(&events, true); + assert_eq!(analyze(&file.0, &[])["complete"], false); + + let mut events = forbidden_retry(); + events[4]["sent"] = 8.into(); + let file = TraceFile::write(&events, true); + let report = analyze(&file.0, &[]); + assert_eq!(report["complete"], false); + assert_eq!(report["new_attempts_after_failed_send"], 1); + } + + #[test] + fn terminal_send_cannot_finish_writable_or_successful() { + for (field, value) in [("writer_broken", json!(false)), ("outcome", json!("ok"))] { + let mut events = forbidden_retry(); + events[4][field] = value; + let file = TraceFile::write(&events, true); + assert_eq!(analyze(&file.0, &[])["complete"], false); + } + } + + #[test] + #[cfg(target_os = "linux")] + fn capture_write_error_survives_until_finish() { + let mut state = Capture { + writer: BufWriter::with_capacity( + 0, + OpenOptions::new().write(true).open("/dev/full").unwrap(), + ), + started: Instant::now(), + sequence: 0, + error: None, + finished: false, + }; + state.write(Map::new()); + assert!(state.error.is_some()); + let handle = CaptureHandle(Arc::new(Mutex::new(state))); + assert!(handle.finish().is_err()); + } + + #[test] + fn capture_orders_real_events_and_requires_explicit_finish() { + let file = TraceFile::new(); + let (layer, handle) = capture(&file.0).unwrap(); + tracing::subscriber::with_default(tracing_subscriber::registry().with(layer), || { + tracing::debug!(target: "unrelated", event = "ignored"); + tracing::debug!(target: TARGET, event = "read_started", read_id = 9_u64, + attempt = 0_u64, sent = 0_u64, writer_broken = false, blob = %"payload"); + // A permit failure can finish before any SDK download starts. + tracing::debug!(target: TARGET, event = "read_finished", read_id = 9_u64, + attempt = 0_u64, sent = 0_u64, writer_broken = false, + outcome = "error", error_code = %"ResourceExhausted"); + }); + handle.finish().unwrap(); + let report = analyze(&file.0, &[]); + assert_eq!(report["complete"], true, "{report}"); + assert_eq!(report["reads"]["9"]["started"]["sequence"], 1); + assert_eq!(report["reads"]["9"]["finished"]["sequence"], 2); + } +} diff --git a/examples/azure_fleet/workload.rs b/examples/azure_fleet/workload.rs new file mode 100644 index 00000000000..e0e14d5cde4 --- /dev/null +++ b/examples/azure_fleet/workload.rs @@ -0,0 +1,311 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Generate the entire workload before either version runs. + +use std::collections::BTreeSet; + +use bytes::Bytes; +use nativelink_proto::build::bazel::remote::execution::v2::command::EnvironmentVariable; +use nativelink_proto::build::bazel::remote::execution::v2::platform::Property; +use nativelink_proto::build::bazel::remote::execution::v2::{ + Action, Command, Directory, FileNode, Platform, +}; +use nativelink_store::ac_utils::compute_buf_digest; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::DigestHasherFunc; +use prost::Message; +use serde::{Deserialize, Serialize}; + +pub(crate) const MIB: u64 = 1024 * 1024; +pub(crate) const NODES: usize = 4; + +// Keyed draws keep the workload identical when asynchronous execution reorders. +pub(crate) fn uniform(seed: u64, key: u64) -> f64 { + let mut x = seed.wrapping_add(key.wrapping_mul(0x9e37_79b9_7f4a_7c15)); + x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + (f64::from(u32::try_from((x ^ (x >> 31)) >> 32).unwrap()) + 0.5) / 4_294_967_296.0 +} + +fn lognormal(seed: u64, key: u64, median: f64, sigma: f64) -> f64 { + let z = (-2.0 * uniform(seed, key * 2).ln()).sqrt() + * (core::f64::consts::TAU * uniform(seed, key * 2 + 1)).cos(); + + median * (sigma * z).exp() +} + +#[expect( + clippy::cast_possible_truncation, + reason = "rounded draw is clamped to the supplied u32 bounds before conversion" +)] +fn lognormal_count(seed: u64, key: u64, median: f64, sigma: f64, min: u32, max: u32) -> u32 { + lognormal(seed, key, median, sigma) + .round() + .clamp(f64::from(min), f64::from(max)) as u32 +} + +#[derive(Clone)] +pub(crate) struct Blob { + pub(crate) digest: DigestInfo, + pub(crate) data: Bytes, +} + +impl Blob { + fn new(data: Bytes) -> Self { + let digest = compute_buf_digest(&data, &mut DigestHasherFunc::Sha256.hasher()); + Self { digest, data } + } + + pub(crate) fn name(&self) -> String { + format!("{}-{}", self.digest.packed_hash(), self.digest.size_bytes()) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct Metadata { + pub(crate) hash: String, + pub(crate) size: u64, + pub(crate) hex: String, +} + +#[derive(Clone, Debug, Serialize)] +pub(crate) struct Job { + pub(crate) id: u32, + pub(crate) node: usize, + pub(crate) arrival_s: f64, + pub(crate) inputs: Vec, + pub(crate) execution_s: f64, + pub(crate) hash: String, + pub(crate) size: u64, +} + +impl Job { + pub(crate) fn digest(&self) -> DigestInfo { + DigestInfo::try_new(&self.hash, self.size).unwrap() + } +} + +pub(crate) struct Workload { + pub(crate) blobs: Vec, + pub(crate) metadata: Vec, + pub(crate) warmup: Vec, + pub(crate) jobs: Vec, + pub(crate) sentinel: Vec, + pub(crate) catalogue_bytes: Vec, +} + +impl Workload { + pub(crate) fn new(seed: u64) -> Self { + let mut workload = Self { + blobs: Vec::new(), + metadata: Vec::new(), + warmup: Vec::new(), + jobs: Vec::new(), + sentinel: Vec::new(), + catalogue_bytes: Vec::new(), + }; + + workload.populate_catalogue(seed); + workload.add_cache_probes(seed); + workload.add_fill_burst(seed); + workload.add_arrivals(seed); + workload.jobs.sort_by(|left, right| { + left.arrival_s + .total_cmp(&right.arrival_s) + .then(left.id.cmp(&right.id)) + }); + + workload + } + + fn populate_catalogue(&mut self, seed: u64) { + // Declared stress distribution with explicit fixture size parameters. + for i in 0_u32..8 { + let units = lognormal_count(seed, 100 + u64::from(i), 4.0, 0.55, 2, 6); + // A declared large-input probe must exceed the three buffered store + // channels; otherwise Azure can finish before ENOSPC reaches it. + let size = if i == 1 { + 128 * MIB + } else { + u64::from(units) * 4 * MIB + }; + + self.catalogue_bytes.push(size); + self.blobs.push(Blob::new(Bytes::from(vec![ + u8::try_from(i + 1).unwrap(); + usize::try_from(size).unwrap() + ]))); + } + } + + fn add_cache_probes(&mut self, seed: u64) { + // Both probes use the same input. The sentinel must reuse each original + // worker's warm cache after the failed-fill phase has recovered. + for node in 0..NODES { + let job = self.job(seed, u32::try_from(node).unwrap(), node, 0.0, vec![0]); + self.warmup.push(job); + + let job = self.job(seed, 10 + u32::try_from(node).unwrap(), node, 0.0, vec![0]); + self.sentinel.push(job); + } + } + + fn add_fill_burst(&mut self, seed: u64) { + // Two simultaneous offers per node expose the actual fill coalescer. + for node in 0..NODES { + for copy in 0_u32..2 { + let job = self.job( + seed, + 100 + u32::try_from(node).unwrap() * 2 + copy, + node, + 0.0, + vec![1], + ); + self.jobs.push(job); + } + } + } + + fn add_arrivals(&mut self, seed: u64) { + let zipf: Vec = (1..=8).map(|rank| f64::from(rank).powf(-1.1)).collect(); + let total: f64 = zipf.iter().sum(); + let mut time = 0.1; + + for i in 0_u32..32 { + let hazard = -uniform(seed, 1000 + u64::from(i)).ln(); + time = next_arrival(time, hazard); + if time >= 3.0 { + break; + } + + let files = lognormal_count(seed, 2000 + u64::from(i), 2.0, 0.5, 1, 4); + let inputs = choose_inputs(seed, i, files, &zipf, total); + let job = self.job( + seed, + 200 + i, + usize::try_from(i).unwrap() % NODES, + time, + inputs, + ); + self.jobs.push(job); + } + } + + fn message(&mut self, message: &impl Message) -> DigestInfo { + let blob = Blob::new(Bytes::from(message.encode_to_vec())); + self.metadata.push(Metadata { + hash: blob.digest.packed_hash().to_string(), + size: blob.digest.size_bytes(), + hex: hex::encode(&blob.data), + }); + + let digest = blob.digest; + self.blobs.push(blob); + digest + } + + fn job(&mut self, seed: u64, id: u32, node: usize, arrival_s: f64, inputs: Vec) -> Job { + let root = self.message(&Directory { + files: inputs + .iter() + .map(|&i| FileNode { + name: format!("input{i}"), + digest: Some(self.blobs[i].digest.into()), + is_executable: true, + ..Default::default() + }) + .collect(), + ..Default::default() + }); + + let execution_s = lognormal(seed, 10_000 + u64::from(id), 0.02, 0.4).clamp(0.005, 0.1); + let command = self.message(&Command { + arguments: vec![ + "/usr/bin/env".into(), + "sh".into(), + "-c".into(), + format!("sleep {execution_s:.6}; wc -c input* > out"), + format!("job-{id}"), + ], + environment_variables: vec![EnvironmentVariable { + name: "PATH".into(), + value: std::env::var("PATH").unwrap_or_default(), + }], + output_files: vec!["out".into()], + platform: Some(Platform { + properties: vec![Property { + name: "fleet_node".into(), + value: node.to_string(), + }], + }), + ..Default::default() + }); + + let action = self.message(&Action { + command_digest: Some(command.into()), + input_root_digest: Some(root.into()), + do_not_cache: true, + ..Default::default() + }); + + Job { + id, + node, + arrival_s, + inputs, + execution_s, + hash: action.packed_hash().to_string(), + size: action.size_bytes(), + } + } +} + +// Spend exponential hazard across the rate change at one second. Restarting +// the draw at that boundary would change the seeded arrival distribution. +fn next_arrival(time: f64, hazard: f64) -> f64 { + if time >= 1.0 { + return time + hazard / 2.0; + } + + let remaining = (1.0 - time) * 10.0; + if hazard <= remaining { + time + hazard / 10.0 + } else { + 1.0 + (hazard - remaining) / 2.0 + } +} + +fn choose_inputs(seed: u64, job: u32, files: u32, weights: &[f64], total: f64) -> Vec { + // Sampling can choose the same blob twice. An action's input tree contains + // each chosen blob once, in stable order, so use a set before encoding it. + let mut inputs = BTreeSet::new(); + for file in 0..files { + let key = 3000 + u64::from(job) * 8 + u64::from(file); + let mut pick = uniform(seed, key) * total; + let mut chosen = weights.len() - 1; + + for (index, weight) in weights.iter().enumerate() { + pick -= weight; + if pick <= 0.0 { + chosen = index; + break; + } + } + + inputs.insert(chosen); + } + + inputs.into_iter().collect() +} diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index a7715695201..5066c5b3b5b 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -8,6 +8,11 @@ load( "rust_test_suite", ) +exports_files( + ["tests/fixtures/azure_download_response.json"], + visibility = ["//examples/azure_fleet:__pkg__"], +) + # Definitions for the glibc-only symbol names AWS-LC ends up referencing when # it is linked against musl (see aws_lc_musl_shims.c). Empty everywhere else. # Linked into aws-lc-sys through the crate annotation in MODULE.bazel; @@ -143,6 +148,7 @@ rust_test_suite( srcs = [ "tests/ac_utils_test.rs", "tests/azure_blob_store_test.rs", + "tests/azure_download_egress_test.rs", "tests/azurite_store_test.rs", "tests/cache_metrics_remove_metrics_test.rs", "tests/cache_metrics_store_test.rs", @@ -173,6 +179,7 @@ rust_test_suite( "tests/verify_store_test.rs", ], compile_data = [ + "tests/fixtures/azure_download_response.json", "tests/mongo_runner/mod.rs", "tests/mongo_runner/process.rs", "tests/mongo_runner/extractor.rs", diff --git a/nativelink-store/src/azure_blob_store.rs b/nativelink-store/src/azure_blob_store.rs index ca51df422d4..f8c6d310bda 100644 --- a/nativelink-store/src/azure_blob_store.rs +++ b/nativelink-store/src/azure_blob_store.rs @@ -17,6 +17,7 @@ use core::time::Duration; use core::{cmp, mem}; use std::borrow::Cow; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use async_trait::async_trait; use azure_core::credentials::TokenCredential; @@ -76,6 +77,37 @@ const DEFAULT_BLOB_ENDPOINT_SUFFIX: &str = "blob.core.windows.net"; /// Successful reads in a row before the adaptive read limit grows by one. const READ_LIMIT_GROW_AFTER: usize = 16; +static NEXT_DOWNLOAD_READ: AtomicU64 = AtomicU64::new(1); + +/// Observation metadata carried by the SDK context of one outer download. +/// +/// The SDK copies this context to every partition, including tasks that outlive +/// the consuming stream. Read IDs are process-local; attempts start at one. +/// Ordinary clients ignore it. Diagnostic HTTP policies can use it to attribute +/// requests without changing retry decisions or inferring ownership from time. +#[derive(Clone, Copy, Debug)] +pub struct AzureDownloadAttempt { + pub read_id: u64, + pub attempt: u64, + pub sent: u64, + pub writer_broken: bool, +} + +impl AzureDownloadAttempt { + fn record_failed_send(self, writer: &DropCloserWriteHalf, stage: &'static str) { + event!( + target: "nativelink_store::azure_blob_store::download", + Level::DEBUG, + event = "send_failed", + read_id = self.read_id, + attempt = self.attempt, + sent = writer.get_bytes_written(), + writer_broken = writer.is_pipe_broken(), + stage, + ); + } +} + #[derive(MetricsComponent)] pub struct AzureBlobStore { client: Arc, @@ -984,8 +1016,23 @@ where }; let max_bytes = length; - self.retrier - .retry(unfold(writer, move |writer| { + let read_id = NEXT_DOWNLOAD_READ.fetch_add(1, Ordering::Relaxed); + let attempts = AtomicU64::new(0); + let attempts = &attempts; + event!( + target: "nativelink_store::azure_blob_store::download", + Level::DEBUG, + event = "read_started", + read_id, + attempt = 0_u64, + sent = writer.get_bytes_written(), + writer_broken = writer.is_pipe_broken(), + blob = %blob_path, + ); + + let result = self + .retrier + .retry(unfold(&mut *writer, move |writer| { // Resumes where the writer stopped, so a retry continues the // read instead of starting it again and appending a second // copy of what the consumer already has. The S3 store does @@ -1018,12 +1065,31 @@ where } }; + // Allocate identity only when an SDK download will start. + // A failed filesystem-permit acquisition is not a GET. + let observation = AzureDownloadAttempt { + read_id, + attempt: attempts.fetch_add(1, Ordering::Relaxed) + 1, + sent: writer.get_bytes_written(), + writer_broken: writer.is_pipe_broken(), + }; + event!( + target: "nativelink_store::azure_blob_store::download", + Level::DEBUG, + event = "attempt_started", + read_id, + attempt = observation.attempt, + sent = observation.sent, + writer_broken = observation.writer_broken, + ); + // The error carries the wait a throttled response asked for. let result: Result<(), (Error, Option)> = async { - let options = BlobClientDownloadOptions { + let mut options = BlobClientDownloadOptions { range, ..Default::default() }; + options.method_options.context.insert(observation); let response = client .blob_client(&blob_path) .download(Some(options)) @@ -1075,6 +1141,7 @@ where None => chunk, }; writer.send(chunk).await.map_err(|e| { + observation.record_failed_send(writer, "data"); ( make_err!( Code::Aborted, @@ -1106,6 +1173,7 @@ where )); } writer.send_eof().map_err(|e| { + observation.record_failed_send(writer, "eof"); ( make_err!(Code::Aborted, "Failed to send EOF to writer: {e:?}"), None, @@ -1132,7 +1200,20 @@ where } } })) - .await + .await; + + event!( + target: "nativelink_store::azure_blob_store::download", + Level::DEBUG, + event = "read_finished", + read_id, + attempt = attempts.load(Ordering::Relaxed), + sent = writer.get_bytes_written(), + writer_broken = writer.is_pipe_broken(), + outcome = if result.is_ok() { "ok" } else { "error" }, + error_code = %result.as_ref().err().map_or_else(String::new, |e| format!("{:?}", e.code)), + ); + result } fn inner_store(&self, _digest: Option) -> &'_ dyn StoreDriver { diff --git a/nativelink-store/tests/azure_blob_store_test.rs b/nativelink-store/tests/azure_blob_store_test.rs index f7c7a650c1f..cf781ff07dc 100644 --- a/nativelink-store/tests/azure_blob_store_test.rs +++ b/nativelink-store/tests/azure_blob_store_test.rs @@ -672,6 +672,138 @@ async fn test_get_with_retries() -> Result<(), Error> { Ok(()) } +struct DestinationFailureCase { + name: &'static str, + retries: usize, + initial_origin_failures: usize, + final_origin_failure: bool, +} + +async fn observe_destination_failure( + case: &DestinationFailureCase, +) -> Result<(bool, String), Error> { + const VALUE: &[u8] = b"test_content"; + + let mock = MockTransport::new(); + for attempt in 0..=case.retries { + if attempt < case.initial_origin_failures + || (case.final_origin_failure && attempt == case.retries) + { + mock.push_download(error_response(StatusCode::ServiceUnavailable)); + } else { + mock.push_download(download_ok(StatusCode::Ok, VALUE.to_vec())); + } + } + + let store = create_test_store_with_common( + Arc::clone(&mock), + CommonObjectSpec { + // Count each 503 against the stated attempt budget. The separate + // wall-clock throttle budget would otherwise admit extra attempts. + throttle_retry_budget_s: 0.0, + retry: Retry { + max_retries: case.retries, + delay: 0.0, + jitter: 0.0, + ..Default::default() + }, + ..Default::default() + }, + )?; + let digest = DigestInfo::try_new(TEST_HASH, VALUE.len())?; + + // Closing the reader is only observable by the writer on its first send. + // Origin failures before that send still permit their ordinary retries. + let (mut writer, reader) = make_buf_channel_pair(); + drop(reader); + let result = store.get_part(digest, &mut writer, 0, Some(12)).await; + + // One small range means one GET per outer attempt, with no speculative + // partition requests. The first usable origin response must end the read. + let expected_downloads = case.initial_origin_failures + 1; + let expected_range = format!("bytes=0-{}", VALUE.len() - 1); + let ranges = mock.download_ranges(); + let bytes_written = writer.get_bytes_written(); + let broken = writer.is_pipe_broken(); + let preserved_delivery_error = result.as_ref().is_err_and(|error| { + error.code == Code::Aborted && error.to_string().contains("Failed to send data to writer") + }); + let passed = preserved_delivery_error + && broken + && bytes_written == 0 + && ranges.len() == expected_downloads + && ranges + .iter() + .all(|range| range.as_deref() == Some(expected_range.as_str())); + + Ok(( + passed, + format!( + "{}: passed={passed}, retry_budget={}, expected_downloads={expected_downloads}, \ + ranges={ranges:?}, bytes_written={bytes_written}, broken={broken}, result={result:?}", + case.name, case.retries, + ), + )) +} + +/// Once delivery fails, another origin read cannot repair the same writer. +/// Negative controls distinguish an exhausted budget from an actual retry +/// after closure. The final error alone cannot identify that forbidden retry. +#[nativelink_test(start_paused = true)] +async fn a_failed_download_destination_is_not_retried() -> Result<(), Error> { + let cases = [ + DestinationFailureCase { + name: "no retry available after failed send", + retries: 0, + initial_origin_failures: 0, + final_origin_failure: false, + }, + DestinationFailureCase { + name: "one retry available after failed send", + retries: 1, + initial_origin_failures: 0, + final_origin_failure: false, + }, + DestinationFailureCase { + name: "six retries available after failed send", + retries: 6, + initial_origin_failures: 0, + final_origin_failure: false, + }, + DestinationFailureCase { + name: "legitimate 503 before failed send", + retries: 6, + initial_origin_failures: 1, + final_origin_failure: false, + }, + DestinationFailureCase { + name: "final 503 masks a forbidden retry", + retries: 1, + initial_origin_failures: 0, + final_origin_failure: true, + }, + DestinationFailureCase { + name: "first failed send on the last allowed attempt", + retries: 6, + initial_origin_failures: 6, + final_origin_failure: false, + }, + ]; + + // Run the entire matrix before asserting so RED shows both negative + // controls alongside the excess requests and the masking final error. + let mut all_passed = true; + let mut observations = Vec::new(); + for case in &cases { + let (passed, observation) = observe_destination_failure(case).await?; + all_passed &= passed; + observations.push(observation); + } + assert!(all_passed, "{}", observations.join("\n")); + + Ok(()) +} + // Paused tokio time: the only thing that advances the clock is the store's // own sleep, so the elapsed time is exactly what the retrier waited. #[nativelink_test(start_paused = true)] @@ -1209,6 +1341,64 @@ async fn throttled_reads_halve_the_read_limit_and_successes_regrow_it() -> Resul Ok(()) } +/// A failed downstream send must not count as an answered read when deciding +/// whether to grow the adaptive limit. The next healthy read still gets credit. +#[nativelink_test(start_paused = true)] +async fn failed_downstream_send_does_not_grow_the_adaptive_read_limit() -> Result<(), Error> { + const VALUE: &[u8] = b"test_content"; + + let mock = MockTransport::new(); + mock.push_download(error_response(StatusCode::ServiceUnavailable)); + let store = create_test_store_with_common( + Arc::clone(&mock), + CommonObjectSpec { + retry: Retry { + // The first throttle can use the normal throttle budget, but + // a failed send gets no retries. This isolates limit accounting. + max_retries: 0, + delay: 0.0, + jitter: 0.0, + ..Default::default() + }, + max_concurrent_reads: 16, + adaptive_read_concurrency: true, + min_concurrent_reads: 2, + ..Default::default() + }, + )?; + let digest = DigestInfo::try_new(TEST_HASH, VALUE.len())?; + + // The throttle cuts the limit to eight. Stop one successful download short + // of the sixteen needed for regrowth, then exercise a failed destination. + for _ in 0..15 { + mock.push_download(download_ok(StatusCode::Ok, VALUE.to_vec())); + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, VALUE); + } + assert_eq!(store.read_limit(), Some(8)); + + mock.push_download(download_ok(StatusCode::Ok, VALUE.to_vec())); + let (mut writer, reader) = make_buf_channel_pair(); + drop(reader); + let error = store + .get_part(digest, &mut writer, 0, None) + .await + .unwrap_err(); + assert_eq!(error.code, Code::Aborted); + assert!(error.to_string().contains("Failed to send data to writer")); + assert!(writer.is_pipe_broken()); + assert_eq!(writer.get_bytes_written(), 0); + assert_eq!( + store.read_limit(), + Some(8), + "a failed destination must not earn origin-success credit" + ); + + mock.push_download(download_ok(StatusCode::Ok, VALUE.to_vec())); + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, VALUE); + assert_eq!(store.read_limit(), Some(9)); + Ok(()) +} + /// Existence checks are outside the download pool: a throttled `HEAD` /// does not cut the limit and an answered one does not grow it. #[nativelink_test(start_paused = true)] diff --git a/nativelink-store/tests/azure_download_egress_test.rs b/nativelink-store/tests/azure_download_egress_test.rs new file mode 100644 index 00000000000..bd79171f297 --- /dev/null +++ b/nativelink-store/tests/azure_download_egress_test.rs @@ -0,0 +1,958 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Real store + SDK + retrier; only HTTP is simulated. The response distribution +//! comes from measured Azure reads (see the adjacent calibration fixture). +//! A declared token bucket causes 503s, then recovers. After recovery, compare +//! a complete read with a reader dropped after its first SDK partition. +//! +//! Each HTTP request carries the actual `NativeLink` download attempt from the +//! SDK context. No attempt may start with its writer already broken. Old SDK +//! tasks retain their original attempt when they finish or request another +//! partition; post-close bytes alone do not identify a new outer retry. +//! Retry counts and overlapping bytes remain independent secondary checks. +//! No copied store, Python, cloud access, or model of the expected amplification +//! is involved. Paused time preserves production retry delays. The thirty +//! seeded pairs are reproducible stress cases, not production confidence +//! intervals; scheduling and SDK parallelism still affect measured byte counts. + +use core::pin::Pin; +use core::task::{Context, Poll}; +use core::time::Duration; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::{Arc, Mutex}; + +use azure_core::http::headers::Headers; +use azure_core::http::policies::{Policy, PolicyResult, TransportPolicy}; +use azure_core::http::{ + AsyncRawResponse, Context as AzureContext, HttpClient, Method, Request, RetryOptions, + StatusCode, Transport, Url, +}; +use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions}; +use bytes::Bytes; +use futures::{Future, Stream}; +use nativelink_config::stores::{CommonObjectSpec, ExperimentalAzureSpec, Retry}; +use nativelink_error::Code; +use nativelink_macro::nativelink_test; +use nativelink_store::azure_blob_store::{AzureBlobStore, AzureDownloadAttempt}; +use nativelink_util::buf_channel::make_buf_channel_pair; +use nativelink_util::common::DigestInfo; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use nativelink_util::store_trait::StoreLike; +use serde_json::{Value, json}; +use tokio::time::Instant; + +const PARTITION: u32 = 4 * 1024 * 1024; +const CHUNK: u32 = 64 * 1024; +const BLOB: u32 = 32 * PARTITION; +const RECOVERY: Duration = Duration::from_secs(2); +const QUIET_WINDOW: Duration = Duration::from_millis(100); +const SEEDS: u32 = 30; + +type TestStore = Arc MockInstantWrapped>>; + +// The production download context is per call and is cloned by the SDK for +// its partition requests. Only this test policy puts it on HTTP headers. +#[derive(Debug)] +struct AttemptHeaders; + +#[async_trait::async_trait] +impl Policy for AttemptHeaders { + async fn send( + &self, + context: &AzureContext, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + if let Some(attempt) = context.value::() { + request.insert_header("x-nativelink-test-read", attempt.read_id.to_string()); + request.insert_header("x-nativelink-test-attempt", attempt.attempt.to_string()); + request.insert_header("x-nativelink-test-sent", attempt.sent.to_string()); + request.insert_header( + "x-nativelink-test-broken", + attempt.writer_broken.to_string(), + ); + } + next[0].send(context, request, &next[1..]).await + } +} + +fn request_attempt(request: &Request) -> AzureDownloadAttempt { + let header = |name: &str| { + request + .headers() + .iter() + .find(|(key, _)| key.as_str().eq_ignore_ascii_case(name)) + .unwrap_or_else(|| panic!("GET lost its NativeLink attempt context: {name}")) + .1 + .as_str() + }; + AzureDownloadAttempt { + read_id: header("x-nativelink-test-read").parse().unwrap(), + attempt: header("x-nativelink-test-attempt").parse().unwrap(), + sent: header("x-nativelink-test-sent").parse().unwrap(), + writer_broken: header("x-nativelink-test-broken").parse().unwrap(), + } +} + +#[nativelink_test] +async fn attempt_headers_preserve_an_older_context_after_a_newer_request() { + let account = Account::new(17); + let transport: Arc = Arc::new(TransportPolicy::new(Transport::new(Arc::new( + AzureTransport(account.clone()), + )))); + let original = AzureContext::new().with_value(AzureDownloadAttempt { + read_id: 41, + attempt: 1, + sent: 0, + writer_broken: false, + }); + let newer = original.clone().with_value(AzureDownloadAttempt { + read_id: 41, + attempt: 2, + sent: u64::from(PARTITION), + writer_broken: true, + }); + let mut request = Request::new( + Url::parse("https://fixture.blob.core.windows.net/cas/context").unwrap(), + Method::Get, + ); + request.insert_header("range", format!("bytes=0-{}", PARTITION - 1)); + + // Reuse the request to make stale-header bugs visible. The third request + // still belongs to attempt one even though attempt two was seen first. + for context in [&original, &newer, &original] { + let response = AttemptHeaders + .send(context, &mut request, core::slice::from_ref(&transport)) + .await + .unwrap(); + drop(response); + } + + let state = account.state.lock().unwrap(); + let ledger = &state.warmup; + ledger.assert_conserved(); + assert_eq!( + ledger + .requests + .iter() + .map(|request| request.attempt) + .collect::>(), + vec![1, 2, 1], + ); + assert_eq!(ledger.attempts[&(41, 1)].gets, 2); + assert_eq!(ledger.attempts[&(41, 1)].sent, 0); + assert!(!ledger.attempts[&(41, 1)].writer_broken); + assert_eq!(ledger.attempts[&(41, 2)].gets, 1); + assert_eq!(ledger.attempts[&(41, 2)].sent, u64::from(PARTITION)); + assert!(ledger.attempts[&(41, 2)].writer_broken); + assert_eq!(ledger.retries(), 1); + assert_eq!(ledger.jitter_draws, 0); +} + +// SplitMix64: each (seed, range, occurrence) gets the same latency quantile in +// both arms, independent of HTTP task polling order. The duration also depends +// on observed overlap, which need not match across arms. Retry jitter uses a +// distinct stream and the production multiplier formula. +fn uniform(seed: u64, index: u64) -> f64 { + let mut x = seed.wrapping_add(index.wrapping_mul(0x9e37_79b9_7f4a_7c15)); + x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + f64::from(u32::try_from((x ^ (x >> 31)) >> 32).unwrap()) / 4_294_967_296.0 +} + +#[derive(Debug)] +struct Timing { + intercept: f64, + size_power: f64, + concurrency_power: f64, + residuals: Vec, +} + +impl Timing { + fn measured() -> Self { + let fixture: Value = + serde_json::from_str(include_str!("fixtures/azure_download_response.json")).unwrap(); + let fit = &fixture["baseline_ms"]; + let residuals: Vec = fixture["log_residual_quantiles"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap()) + .collect(); + assert_eq!(residuals.len(), 129); + assert!(residuals.windows(2).all(|p| p[0] <= p[1])); + Self { + intercept: fit["intercept"].as_f64().unwrap(), + size_power: fit["size_power"].as_f64().unwrap(), + concurrency_power: fit["concurrency_power"].as_f64().unwrap(), + residuals, + } + } + + fn baseline_ms(&self, inflight: u32) -> f64 { + assert!((1..=32).contains(&inflight), "outside measured overlap"); + // Fixed 4 MiB ranges: 1 + requested_MiB = 5. + self.intercept + * 5.0_f64.powf(self.size_power) + * (1.0 + f64::from(inflight)).powf(self.concurrency_power) + } + + fn duration(&self, inflight: u32, quantile: f64) -> Duration { + let position = quantile * 128.0; + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "quantile is in [0,1), so the index is in 0..128" + )] + let index = position as usize; + let fraction = position - f64::from(u32::try_from(index).unwrap()); + let log_residual = (self.residuals[index + 1] - self.residuals[index]) + .mul_add(fraction, self.residuals[index]); + Duration::from_secs_f64(self.baseline_ms(inflight) * log_residual.exp() / 1000.0) + } +} + +#[derive(Debug, Default)] +struct AttemptLedger { + sent: u64, + writer_broken: bool, + gets: u32, + refused: u32, + bytes: u64, + after_close: u64, + repeated: u64, +} + +#[derive(Debug)] +struct RequestLedger { + read_id: u64, + attempt: u64, + offset: u32, + end: u32, + admitted: bool, + bytes: u64, +} + +#[derive(Debug, Default)] +struct Ledger { + gets: u32, + refused: u32, + admitted: u64, + bytes: u64, + after_close: u64, + repeated: u64, + blocks: BTreeSet, + // 100 ms bins expose retry-driven bursts after the initial transfer. + bins: BTreeMap, + // An admitted range occurrence is not an outer NativeLink attempt. + range_occurrences: BTreeMap, + jitter_draws: u32, + // An attempt may own many HTTP requests, including requests that start + // after a newer attempt. Neither timestamps nor a current-attempt counter + // can recover that ownership; the immutable SDK context supplies it. + attempts: BTreeMap<(u64, u64), AttemptLedger>, + requests: Vec, +} + +impl Ledger { + fn retries(&self) -> usize { + self.attempts + .keys() + .filter(|(_, attempt)| *attempt > 1) + .count() + } + + fn request_started(&mut self, identity: &AzureDownloadAttempt, offset: u32, end: u32) -> usize { + let attempt = self + .attempts + .entry((identity.read_id, identity.attempt)) + .or_insert_with(|| AttemptLedger { + sent: identity.sent, + writer_broken: identity.writer_broken, + ..Default::default() + }); + assert_eq!(attempt.sent, identity.sent); + assert_eq!(attempt.writer_broken, identity.writer_broken); + attempt.gets += 1; + + let request_id = self.requests.len(); + self.requests.push(RequestLedger { + read_id: identity.read_id, + attempt: identity.attempt, + offset, + end, + admitted: false, + bytes: 0, + }); + request_id + } + + fn assert_conserved(&self) { + assert_eq!(self.bytes, self.bins.values().sum::()); + assert_eq!( + self.bytes, + u64::try_from(self.blocks.len()).unwrap() * u64::from(CHUNK) + self.repeated + ); + assert!(self.bytes <= self.admitted); + assert_eq!( + self.admitted, + u64::from(self.gets - self.refused) * u64::from(PARTITION) + ); + assert_eq!( + self.bytes, + self.requests.iter().map(|r| r.bytes).sum::() + ); + assert_eq!( + self.bytes, + self.attempts.values().map(|a| a.bytes).sum::() + ); + assert_eq!( + self.gets, + self.attempts.values().map(|a| a.gets).sum::() + ); + assert_eq!( + self.refused, + self.attempts.values().map(|a| a.refused).sum::() + ); + assert_eq!(usize::try_from(self.gets).unwrap(), self.requests.len()); + for request in &self.requests { + assert!(request.bytes <= u64::from(request.end - request.offset)); + assert!(request.admitted || request.bytes == 0); + } + } +} + +#[derive(Debug)] +struct State { + tokens: f64, + refilled: Instant, + active: u32, + peak_active: u32, + probe: bool, + closed: bool, + probe_start: Instant, + // Assign a logical read to warmup/probe only on its first request. Later + // partition requests keep that assignment across a phase transition. + read_phases: BTreeMap, + warmup: Ledger, + read: Ledger, +} + +impl State { + const fn ledger(&mut self) -> &mut Ledger { + self.ledger_for(self.probe) + } + + const fn ledger_for(&mut self, probe: bool) -> &mut Ledger { + if probe { + &mut self.read + } else { + &mut self.warmup + } + } +} + +#[derive(Debug)] +struct Account { + seed: u64, + start: Instant, + timing: Timing, + range_body: Bytes, + state: Mutex, +} + +impl Account { + fn new(seed: u64) -> Arc { + let now = Instant::now(); + Arc::new(Self { + seed, + start: now, + timing: Timing::measured(), + // Each range has identical contents; the store's digest size is real. + range_body: Bytes::from(vec![42; usize::try_from(PARTITION).unwrap()]), + state: Mutex::new(State { + tokens: f64::from(8 * PARTITION), + refilled: now, + active: 0, + peak_active: 0, + probe: false, + closed: false, + probe_start: now, + read_phases: BTreeMap::new(), + warmup: Ledger::default(), + read: Ledger::default(), + }), + }) + } + + fn jitter(&self, nominal: Duration) -> Duration { + let mut state = self.state.lock().unwrap(); + let phase = u64::from(state.probe); + let index = state.ledger().jitter_draws; + // A throttle can draw jitter even when its wait exceeds the budget. + // Count draws for RNG indexing; SDK identities count actual retries. + state.ledger().jitter_draws += 1; + #[expect( + clippy::cast_possible_truncation, + reason = "production jitter uses f32" + )] + let u = uniform(self.seed, 10_000 + phase * 100 + u64::from(index)) as f32; + nominal.mul_f32(0.5_f32.mul_add(u - 0.5, 1.0)) + } + + fn refill(&self, state: &mut State, now: Instant) { + let elapsed = now.duration_since(self.start).as_secs_f64(); + let previous = state.refilled.duration_since(self.start).as_secs_f64(); + let reference_rate = f64::from(BLOB) / (self.timing.baseline_ms(32) / 1000.0); + let low = elapsed.min(2.0) - previous.min(2.0); + let high = (elapsed - 2.0).max(0.0) - (previous - 2.0).max(0.0); + + // Integrate both rates if a request straddles recovery. The budget is + // a declared fixture policy, not a fitted Azure account limit. + let burst = if elapsed < 2.0 { 8 * PARTITION } else { BLOB }; + state.tokens = reference_rate + .mul_add(0.1_f64.mul_add(low, 1.2 * high), state.tokens) + .min(f64::from(burst)); + state.refilled = now; + } + + async fn quiesce(&self) { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let before = { + let state = self.state.lock().unwrap(); + ( + state.active, + state.warmup.gets + state.read.gets, + state.warmup.bytes + state.read.bytes, + ) + }; + if before.0 != 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + continue; + } + // SDK tasks spawned before the store returned may not have + // entered HTTP yet. Require no requests or body progress over + // a quiet interval, not merely one zero-active snapshot. + tokio::time::sleep(QUIET_WINDOW).await; + let state = self.state.lock().unwrap(); + let after = ( + state.active, + state.warmup.gets + state.read.gets, + state.warmup.bytes + state.read.bytes, + ); + if after == before { + break; + } + } + }) + .await + .expect("HTTP requests and SDK bodies must reach an observed quiet interval"); + } +} + +struct Body { + account: Arc, + probe: bool, + phase_start: Instant, + request_id: usize, + attempt_key: (u64, u64), + offset: u32, + sent: u32, + start: Instant, + duration: Duration, + sleep: Option>>, +} + +impl Stream for Body { + type Item = azure_core::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.sent == PARTITION { + return Poll::Ready(None); + } + let end = this.sent + CHUNK; + let sleep = this.sleep.get_or_insert_with(|| { + Box::pin(tokio::time::sleep_until( + this.start + this.duration.mul_f64(f64::from(end) / f64::from(PARTITION)), + )) + }); + if sleep.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + this.sleep = None; + let chunk = this + .account + .range_body + .slice(usize::try_from(this.sent).unwrap()..usize::try_from(end).unwrap()); + let mut state = this.account.state.lock().unwrap(); + let closed = state.closed; + let bin = this.phase_start.elapsed().as_millis() / 100; + // This body keeps its originating request even when a later outer + // attempt is active. Never charge it to the latest attempt or phase. + let ledger = state.ledger_for(this.probe); + // Charge bytes when this HTTP body yields them to the SDK. Admission + // alone is not a transfer, and yielding does not prove cache delivery. + ledger.bytes += u64::from(CHUNK); + if closed { + ledger.after_close += u64::from(CHUNK); + } + let repeated = !ledger.blocks.insert(this.offset + this.sent); + if repeated { + ledger.repeated += u64::from(CHUNK); + } + *ledger.bins.entry(bin).or_default() += u64::from(CHUNK); + ledger.requests[this.request_id].bytes += u64::from(CHUNK); + let attempt = ledger.attempts.get_mut(&this.attempt_key).unwrap(); + attempt.bytes += u64::from(CHUNK); + if closed { + attempt.after_close += u64::from(CHUNK); + } + if repeated { + attempt.repeated += u64::from(CHUNK); + } + this.sent = end; + Poll::Ready(Some(Ok(chunk))) + } +} + +impl Drop for Body { + fn drop(&mut self) { + self.account.state.lock().unwrap().active -= 1; + } +} + +#[derive(Debug)] +struct AzureTransport(Arc); + +impl AzureTransport { + fn requested_range(request: &Request) -> (u32, u32) { + let range = request + .headers() + .iter() + .find(|(key, _)| { + key.as_str().eq_ignore_ascii_case("range") + || key.as_str().eq_ignore_ascii_case("x-ms-range") + }) + .unwrap() + .1 + .as_str(); + let (start, end) = range + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let offset: u32 = start.parse().unwrap(); + let end = (end.parse::().unwrap() + 1).min(BLOB); + assert_eq!(end - offset, PARTITION, "outside calibrated range size"); + (offset, end) + } + + fn download(&self, request: &Request, mut headers: Headers) -> AsyncRawResponse { + let account = &self.0; + let (offset, end) = Self::requested_range(request); + let identity = request_attempt(request); + let attempt_key = (identity.read_id, identity.attempt); + + // Keep refill, admission and the occurrence draw in one critical + // section so concurrent SDK requests observe the same ordering. + let mut state = account.state.lock().unwrap(); + let now = Instant::now(); + account.refill(&mut state, now); + let current_phase = (state.probe, state.probe_start); + let (probe, phase_start) = *state + .read_phases + .entry(identity.read_id) + .or_insert(current_phase); + let phase = u64::from(probe); + state.ledger_for(probe).gets += 1; + let request_id = state + .ledger_for(probe) + .request_started(&identity, offset, end); + + // Refusals come from the shared budget or active-body limit, never a + // random throttle draw. Refused ranges spend no response quantile. + if state.tokens < f64::from(PARTITION) || state.active >= 32 { + state.ledger_for(probe).refused += 1; + state + .ledger_for(probe) + .attempts + .get_mut(&attempt_key) + .unwrap() + .refused += 1; + headers.insert("x-ms-error-code", "ServerBusy"); + headers.insert("retry-after", "1"); + return AsyncRawResponse::from_bytes( + StatusCode::ServiceUnavailable, + headers, + b"ServerBusyFixture capacity exhausted" + .to_vec(), + ); + } + + state.tokens -= f64::from(PARTITION); + state.active += 1; + state.peak_active = state.peak_active.max(state.active); + let inflight = state.active; + let ledger = state.ledger_for(probe); + ledger.admitted += u64::from(PARTITION); + ledger.requests[request_id].admitted = true; + let occurrence = ledger.range_occurrences.entry(offset).or_default(); + let quantile = uniform( + account.seed, + phase * 1_000 + u64::from(*occurrence) * 32 + u64::from(offset / PARTITION), + ); + *occurrence += 1; + let duration = account.timing.duration(inflight, quantile); + drop(state); + + headers.insert("content-length", PARTITION.to_string()); + headers.insert( + "content-range", + format!("bytes {offset}-{}/{BLOB}", end - 1), + ); + AsyncRawResponse::new( + StatusCode::PartialContent, + headers, + Box::pin(Body { + account: account.clone(), + probe, + phase_start, + request_id, + attempt_key, + offset, + sent: 0, + start: now, + duration, + sleep: None, + }), + ) + } +} + +#[async_trait::async_trait] +impl HttpClient for AzureTransport { + async fn execute_request(&self, request: &Request) -> azure_core::Result { + let mut headers = Headers::new(); + headers.insert("last-modified", "Thu, 01 Jan 1970 00:00:00 GMT"); + headers.insert("etag", "\"unchanged\""); + + if request.method() == Method::Head { + headers.insert("content-length", BLOB.to_string()); + return Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers, + Vec::new(), + )); + } + + assert_eq!(request.method(), Method::Get); + Ok(self.download(request, headers)) + } +} + +fn create_store(account: &Arc) -> TestStore { + let mut options = BlobContainerClientOptions::default(); + options + .client_options + .per_call_policies + .push(Arc::new(AttemptHeaders)); + // Only NativeLink retries in this fixture. The real SDK still owns the + // partition requests and the tasks that can outlive the consuming read. + options.client_options.retry = RetryOptions::none(); + options.client_options.transport = + Some(Transport::new(Arc::new(AzureTransport(account.clone())))); + let client = BlobContainerClient::new( + Url::parse("https://fixture.blob.core.windows.net/cas").unwrap(), + None, + Some(options), + ) + .unwrap(); + let spec = ExperimentalAzureSpec { + account_name: "fixture".into(), + container: "cas".into(), + common: CommonObjectSpec { + // No separate throttle-budget retry path. + throttle_retry_budget_s: 0.0, + retry: Retry { + max_retries: 6, + delay: 0.3, + jitter: 0.5, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() + }; + let jitter = account.clone(); + AzureBlobStore:: MockInstantWrapped>::new_with_client_and_jitter( + &spec, + client, + Arc::new(move |delay| jitter.jitter(delay)), + MockInstantWrapped::default, + ) + .unwrap() +} + +async fn warm_up(store: &TestStore, account: &Account, digest: DigestInfo) { + assert_eq!(store.has(digest).await.unwrap(), Some(u64::from(BLOB))); + + // Exercise the cap first, then complete a real SDK read after recovery. + let warmup = tokio::time::timeout( + Duration::from_mins(1), + store.get_part_unchunked(digest, 0, None), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(warmup.len(), usize::try_from(BLOB).unwrap()); + drop(warmup); + account.quiesce().await; + tokio::time::sleep_until(account.start + RECOVERY + Duration::from_secs(1)).await; + tokio::time::sleep(Duration::from_secs(1)).await; + + // Separate the initial throttling from the assertion phase. Its bucket is + // now large enough for a whole blob, and its counters start independently. + { + let mut state = account.state.lock().unwrap(); + assert!(state.warmup.refused > 0, "the cap must actually throttle"); + state.probe = true; + state.probe_start = Instant::now(); + } +} + +struct ProbeResult { + delivered: u64, + returned: Duration, + bytes_written: u64, +} + +async fn read_probe( + store: &TestStore, + account: &Account, + digest: DigestInfo, + close: bool, +) -> ProbeResult { + let (mut writer, mut reader) = make_buf_channel_pair(); + let read = store.get_part(digest, &mut writer, 0, None); + let consume = async { + let mut delivered = 0_u64; + loop { + let chunk = reader.recv().await.unwrap(); + if chunk.is_empty() { + break; + } + delivered += u64::try_from(chunk.len()).unwrap(); + if close { + // Mark the consumer boundary immediately before dropping it. + // Existing SDK tasks may still yield bytes after this point. + account.state.lock().unwrap().closed = true; + drop(reader); + break; + } + } + delivered + }; + let (result, delivered) = tokio::time::timeout(Duration::from_secs(90), async { + tokio::join!(read, consume) + }) + .await + .expect("bounded production read"); + let returned = account.state.lock().unwrap().probe_start.elapsed(); + if close { + assert_eq!(result.unwrap_err().code, Code::Aborted); + assert!(writer.is_pipe_broken()); + assert_eq!(delivered, u64::from(PARTITION)); + } else { + result.unwrap(); + assert_eq!(delivered, u64::from(BLOB)); + } + ProbeResult { + delivered, + returned, + bytes_written: writer.get_bytes_written(), + } +} + +fn attempt_rows(ledger: &Ledger) -> Vec { + ledger + .attempts + .iter() + .map(|(&(read_id, attempt), entry)| { + json!({ + "read_id": read_id, + "attempt": attempt, + "sent_at_start": entry.sent, + "writer_broken_at_start": entry.writer_broken, + "gets": entry.gets, + "refused": entry.refused, + "body_bytes": entry.bytes, + "after_close_bytes": entry.after_close, + "repurchased_bytes": entry.repeated, + }) + }) + .collect() +} + +fn request_rows(ledger: &Ledger) -> Vec { + ledger + .requests + .iter() + .enumerate() + .map(|(request_id, request)| { + json!({ + "request_id": request_id, + "read_id": request.read_id, + "attempt": request.attempt, + "range_start": request.offset, + "range_end_exclusive": request.end, + "admitted": request.admitted, + "body_bytes": request.bytes, + }) + }) + .collect() +} + +fn report( + account: &Account, + seed: u64, + close: bool, + probe: &ProbeResult, + process_cpus: usize, +) -> Value { + let state = account.state.lock().unwrap(); + for ledger in [&state.warmup, &state.read] { + ledger.assert_conserved(); + } + assert!(state.read.repeated <= state.read.after_close); + assert!(probe.delivered <= state.read.bytes); + if !close { + assert_eq!(state.read.bytes, u64::from(BLOB)); + } + + let forbidden: Vec<_> = state + .read + .attempts + .iter() + .filter(|(_, attempt)| attempt.writer_broken) + .collect(); + // Once a send has broken this writer, no later attempt can advance it. + // This is separate from reader delivery: the channel may have buffered a + // larger prefix before its consumer received one partition and closed. + for (_, attempt) in &forbidden { + assert_eq!(attempt.sent, probe.bytes_written); + } + let forbidden_bytes: u64 = forbidden.iter().map(|(_, attempt)| attempt.bytes).sum(); + let first_forbidden = forbidden.first().map(|((read_id, attempt), entry)| { + json!({ + "read_id": read_id, + "attempt": attempt, + "sent_at_start": entry.sent, + "writer_bytes_written_at_return": probe.bytes_written, + "additional_written_bytes": probe.bytes_written - entry.sent, + "body_bytes": entry.bytes, + }) + }); + + json!({ + "seed": seed, + "closed": close, + "delivered_bytes": probe.delivered, + "writer_bytes_written": probe.bytes_written, + "body_bytes": state.read.bytes, + "after_close_bytes": state.read.after_close, + "repurchased_bytes": state.read.repeated, + "gets": state.read.gets, + "refused": state.read.refused, + "retries": state.read.retries(), + "jitter_draws": state.read.jitter_draws, + "return_seconds": probe.returned.as_secs_f64(), + "warmup_refused": state.warmup.refused, + "peak_inflight": state.peak_active, + "bytes_by_100ms_bin": state.read.bins, + "process_cpus_at_start": process_cpus, + "sdk_parallelism_from_cpu_snapshot": process_cpus.clamp(8, 96), + "fixture_active_body_limit": 32, + "observed_quiet_window_ms": QUIET_WINDOW.as_millis(), + "forbidden_attempts": forbidden.len(), + "forbidden_attempt_body_bytes": forbidden_bytes, + "max_observed_attempt": state.read.attempts.keys().map(|(_, attempt)| *attempt).max(), + "first_forbidden_attempt": first_forbidden, + "attempts": attempt_rows(&state.read), + "requests": request_rows(&state.read), + "warmup_attempts": attempt_rows(&state.warmup), + "attempt_identity_boundary": "AzureBlobStore context at download start, copied into each SDK HTTP request", + }) +} + +async fn run(seed: u64, close: bool, process_cpus: usize) -> Value { + let account = Account::new(seed); + let store = create_store(&account); + let digest = DigestInfo::try_new(&"00".repeat(32), BLOB).unwrap(); + + warm_up(&store, &account, digest).await; + let probe = read_probe(&store, &account, digest, close).await; + account.quiesce().await; + + report(&account, seed, close, &probe, process_cpus) +} + +#[nativelink_test(start_paused = true)] +#[expect( + clippy::print_stdout, + reason = "print every seed and byte ledger for replay and comparison" +)] +async fn a_closed_destination_does_not_repurchase_input_ranges() { + let mut violations = Vec::new(); + // The real SDK chooses its default from process-visible parallelism. The + // fixture's independent 32-body admission limit keeps fitted overlap in + // range; it does not replace or fix the SDK's concurrency. + let process_cpus = std::thread::available_parallelism().map_or(1, core::num::NonZeroUsize::get); + println!( + "process_cpus={process_cpus}, sdk_default_parallelism={}", + process_cpus.clamp(8, 96) + ); + for index in 0..SEEDS { + let seed = 2_026_101_700 + u64::from(index); + // Rotate arm order even with virtual time; draws remain paired by key. + for close in if index % 2 == 0 { + [false, true] + } else { + [true, false] + } { + let row = run(seed, close, process_cpus).await; + println!("{row}"); + let repurchased = row["repurchased_bytes"].as_u64().unwrap(); + let retries = row["retries"].as_u64().unwrap(); + let forbidden = row["forbidden_attempts"].as_u64().unwrap(); + if forbidden != 0 || repurchased != 0 || retries != 0 { + violations.push(json!({ + "seed": seed, + "closed": close, + "forbidden_attempts": forbidden, + "forbidden_attempt_body_bytes": row["forbidden_attempt_body_bytes"], + "first_forbidden_attempt": row["first_forbidden_attempt"], + "repurchased_bytes": repurchased, + "retries": retries, + })); + } + } + } + assert!( + violations.is_empty(), + "a download started with an already-broken writer, or a secondary retry/range invariant failed: {violations:#?}" + ); +} diff --git a/nativelink-store/tests/fixtures/azure_download_response.json b/nativelink-store/tests/fixtures/azure_download_response.json new file mode 100644 index 00000000000..395119e7f6e --- /dev/null +++ b/nativelink-store/tests/fixtures/azure_download_response.json @@ -0,0 +1,184 @@ +{ + "baseline_ms": { + "concurrency_power": 0.094254164, + "intercept": 6.3339399, + "size_power": 1.5779125 + }, + "calibration_observations": 1728, + "calibration_runs": 4, + "capture_boundary": "throttle_proxy: client/proxy complete-body time; includes network, not server-only latency", + "capture_runs": 23, + "correction": { + "date": "2026-10-17", + "policy": "Canonical capture lineage owns snapshots before overlap calculation and split assignment; request ID plus original start time identifies a body within a lineage. Conflicting copies fail ingestion.", + "removed_duplicate_bodies": 3522, + "removed_training_copies_of_evaluation_bodies": 1929, + "scope": "Previous simulator measurements remain observations under their archived timing profile; this correction changes calibration, not those measured byte counts.", + "superseded_fit_sha256": "a203faae06cf4445136ad6436077a4cd0e722b6c4817be7b9f3339382df0e41b" + }, + "description": "Conditional log-duration fit plus empirical log-residual quantiles. No production traffic, identifiers, or credentials.", + "fixture_policy": "Token bucket, limits, fixed object size, consumer closure, and capacity recovery are explicit fixture stress parameters. The fit calibrates response timing only.", + "goodness_of_fit": { + "bootstrap_seed_range_inclusive": [ + 20261016, + 20262015 + ], + "caution": "Retrospective evaluation after snapshot deduplication, using full residual ECDF before compression to 129 quantiles. Five previously inspected evaluation runs, not new prospective data or ten repetitions of one condition.", + "ci95": [ + 0.8301186125570175, + 0.9205140876181208 + ], + "equal_run_coverage": 0.8762129204671598, + "ks_distance_descriptive": 0.1809682102525919, + "ks_iid_p_value_suppressed": "requests are correlated; no independent-request p-value is valid", + "nominal": 0.9, + "test": "held-out whole-run bootstrap of nominal 90% prediction-interval coverage" + }, + "inflight": [ + 1, + 32 + ], + "log_residual_quantiles": [ + -1.0120708231107503, + -0.8276217778621519, + -0.764566114416728, + -0.7022179838422641, + -0.666642216558905, + -0.6307590191666995, + -0.5985297377976262, + -0.5775257777601428, + -0.5544888997729018, + -0.5233451703051331, + -0.49871300041098027, + -0.47682309044496707, + -0.45435192037656247, + -0.432374626149159, + -0.41086489207507926, + -0.39958967322534517, + -0.3898130681026921, + -0.3787115265807577, + -0.3691953014246838, + -0.35205104790244096, + -0.34898928544831653, + -0.3291928289814483, + -0.329188154107261, + -0.31809128745951387, + -0.3097715066706588, + -0.3097715066706588, + -0.29949931826454446, + -0.2907246958423979, + -0.2802647344803544, + -0.2720326884403991, + -0.26236319256186835, + -0.2541183825848587, + -0.25368170656740574, + -0.25075372423309034, + -0.24040868348270578, + -0.23566451701069902, + -0.23566451701069902, + -0.2185977740821465, + -0.21796620968536828, + -0.21796202690628785, + -0.20057569296378652, + -0.2005715822963996, + -0.2005715822963996, + -0.18348244403720848, + -0.1834784030372223, + -0.1834784030372223, + -0.1684851132287708, + -0.1666724970191229, + -0.1644890458001107, + -0.15014436762010686, + -0.15014045911813767, + -0.1446986411971303, + -0.133884981477566, + -0.13388113601146667, + -0.1266346998537145, + -0.1178819544072125, + -0.1178819544072125, + -0.10678041288527806, + -0.10213472094059917, + -0.10213099565091177, + -0.09219582927747626, + -0.0866316233380866, + -0.08662795535644907, + -0.08034491533741185, + -0.07136159473382708, + -0.07136159473382708, + -0.06026005321189265, + -0.056321237214701725, + -0.05137282376270404, + -0.04150738733126319, + -0.038685732721525014, + -0.026909707005306274, + -0.02690951804099484, + -0.0199544118662772, + -0.01252181719186396, + -0.012222486111887452, + 0.001661810396177188, + 0.005822775609119327, + 0.01565038432287924, + 0.022041929841191153, + 0.029442708508219617, + 0.02944597451286679, + 0.043047389587585805, + 0.04609909523790387, + 0.05646946483511428, + 0.05647264375253336, + 0.06971690846282286, + 0.06971690846282286, + 0.08081844998475729, + 0.08279054456715373, + 0.08960593439427482, + 0.09569360143750583, + 0.10843169861326096, + 0.11173113762454881, + 0.12100957543265789, + 0.13211111695459232, + 0.1334341558385317, + 0.1457033510409218, + 0.15275761179947822, + 0.15782383580690773, + 0.1698020102548119, + 0.18158772964225836, + 0.19273490605730048, + 0.20017770351401742, + 0.21016808121076422, + 0.21632138081843344, + 0.22762011414949157, + 0.2387926110890417, + 0.2493442472341606, + 0.2599432990771929, + 0.2611596933913968, + 0.2819078023669141, + 0.2917716560533641, + 0.30333162893233645, + 0.3136961313929687, + 0.3338754767445462, + 0.35393288040314475, + 0.37657264291070563, + 0.38638446363844214, + 0.40404596776523993, + 0.4303314791862549, + 0.44206019922799034, + 0.4664750416170151, + 0.5077905307393755, + 0.5430047671927323, + 0.5886571711612961, + 0.6456829957919367, + 0.7562638828258352, + 1.8987232332591546 + ], + "observations": 61288, + "range_bytes": [ + 4194304, + 4194304 + ], + "residual_sampling": "Linear interpolation at uniformly spaced probabilities 0..1. Residuals pooled from validation runs after normalizing capture boundary, size, and overlap. Independent replay draws do not reproduce within-run correlation.", + "schema": 1, + "source_capture_sha256": "357871ceda2fff932bc8be0f51e8a61d740fde0480d22598f1a1d8dc5266f894", + "source_data_sha256": "482d2b93cf21b7a5425294ce0175cf9f98d0e4d9fc4b53106320b887ef5cef45", + "source_export_sha256": "5da9956c6d44933b3fc1c1df1a8884314005fdcbbb3945a88497da847903cf0a", + "source_fit_sha256": "d3ab0a69324fe56e3d3b1f9abb3ab2be928d73d1d007ec9044fbe89b02f0c012", + "source_requests_sha256": "eba7f4aeca22b426243354251ae0018a1cc3b91f3ab8ad79e56c483fa3b1f91c" +}