diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index c0832d0ef15..90a2b439508 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -493,3 +493,10 @@ Flybody MuJoCo connectome simulatable +<<<<<<< HEAD +[Rr]eachability +recursed +recursing +======= +unary +>>>>>>> origin/b7r6/shard-ring-dead-peer-failover diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 764cc5d9ce7..41d4e0a0820 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -1,17 +1,23 @@ name: CI / Integration # 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 diff --git a/Cargo.lock b/Cargo.lock index ae21f4472d8..a4bfc4d6b52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4604,6 +4604,7 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "digest 0.10.7", + "sha2-asm", ] [[package]] @@ -4617,6 +4618,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sha2-asm" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab" +dependencies = [ + "cc", +] + [[package]] name = "sharded-slab" version = "0.1.7" diff --git a/Cargo.toml b/Cargo.toml index 0f20bfb61f3..97074514c7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,7 +65,10 @@ rand = { version = "0.9.0", default-features = false, features = [ rustls-pki-types = { version = "1.13.1", features = [ "std", ], default-features = false } -sha2 = { version = "0.10.8", default-features = false } +# `asm` unlocks the RustCrypto hardware SHA-2 backends (aarch64 is gated behind +# it) with runtime HWCAP detection — hardware SHA on-core, software fallback +# elsewhere. +sha2 = { version = "0.10.8", default-features = false, features = ["asm"] } tokio = { version = "1.52.2", features = [ "fs", "io-util", diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 938bdaae8be..2fa1188c61b 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1152,6 +1152,7 @@ "sha1_0.10.6": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"},{\"name\":\"sha1-asm\",\"optional\":true,\"req\":\"^0.5\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"}],\"features\":{\"asm\":[\"sha1-asm\"],\"compress\":[],\"default\":[\"std\"],\"force-soft\":[],\"loongarch64_asm\":[],\"oid\":[\"digest/oid\"],\"std\":[\"digest/std\"]}}", "sha1_0.11.0": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.3\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.11\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.11\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"digest/alloc\"],\"default\":[\"alloc\",\"oid\"],\"oid\":[\"digest/oid\"],\"zeroize\":[\"digest/zeroize\"]}}", "sha1_smol_1.0.1": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"openssl\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.4\"},{\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"}],\"features\":{\"alloc\":[],\"std\":[\"alloc\"]}}", + "sha2-asm_0.6.4": "{\"dependencies\":[{\"kind\":\"build\",\"name\":\"cc\",\"req\":\"^1.0\"}],\"features\":{}}", "sha2_0.10.9": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"},{\"name\":\"sha2-asm\",\"optional\":true,\"req\":\"^0.6.1\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"}],\"features\":{\"asm\":[\"sha2-asm\"],\"asm-aarch64\":[\"asm\"],\"compress\":[],\"default\":[\"std\"],\"force-soft\":[],\"force-soft-compact\":[],\"loongarch64_asm\":[],\"oid\":[\"digest/oid\"],\"std\":[\"digest/std\"]}}", "sha2_0.11.0": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1\"},{\"name\":\"cpufeatures\",\"req\":\"^0.3\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.11\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.11\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"digest/alloc\"],\"default\":[\"alloc\",\"oid\"],\"oid\":[\"digest/oid\"],\"zeroize\":[\"digest/zeroize\"]}}", "sharded-slab_0.1.7": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.3\"},{\"kind\":\"dev\",\"name\":\"indexmap\",\"req\":\"^1\"},{\"name\":\"lazy_static\",\"req\":\"^1\"},{\"features\":[\"checkpoint\"],\"name\":\"loom\",\"optional\":true,\"req\":\"^0.5\",\"target\":\"cfg(loom)\"},{\"features\":[\"checkpoint\"],\"kind\":\"dev\",\"name\":\"loom\",\"req\":\"^0.5\",\"target\":\"cfg(loom)\"},{\"kind\":\"dev\",\"name\":\"memory-stats\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"slab\",\"req\":\"^0.4.2\"}],\"features\":{}}", diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 1b4aac223fe..348e40c0084 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; #[cfg(feature = "dev-schema")] @@ -20,7 +20,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use tracing::warn; -use crate::schedulers::SchedulerSpec; +use crate::schedulers::{ExperimentalSimpleSchedulerBackend, SchedulerSpec}; use crate::serde_utils::{ convert_boolean_with_shellexpand, convert_data_size_with_shellexpand, convert_duration_with_shellexpand, convert_numeric_with_shellexpand, @@ -148,6 +148,17 @@ pub struct CasStoreConfig { /// and behavior is identical to when this option did not exist. #[serde(default)] pub experimental_chunking: Option, + + /// Per-blob timeout in seconds applied inside `BatchReadBlobs`. Each + /// blob read in a batch races this deadline so a single stalled read + /// cannot pin the whole batch until the client's overall deadline. + /// Cold reads through a slow backing tier (e.g. an object store that + /// populates a local cache on read) can legitimately take longer than + /// the default for large blobs; raise this value for such deployments. + /// + /// Default: 30 seconds. + #[serde(default)] + pub batch_read_per_blob_timeout_s: Option, } #[derive(Deserialize, Serialize, Debug, Clone)] @@ -2007,6 +2018,347 @@ impl CasConfig { } Ok(()) } + + /// Offline structural validation of every store/scheduler name reference + /// in this config. This is a pure, read-only check intended for + /// `nativelink --check`: it constructs no stores, binds no sockets, opens + /// no connections, and creates no directories. + /// + /// Every store name referenced by a worker or by a server service, and + /// every scheduler name referenced by a service, must resolve to a store + /// declared in `stores` / a scheduler declared in `schedulers`. Store + /// references made *inside* a declared store's spec (the `ref_store` + /// wrappers such as `fast_slow`, `dedup`, `completeness_checking`, ...) + /// are also resolved, but only for the transitive closure of stores that + /// are actually wired to a worker or service. Store definitions that are + /// never referenced by any consumer (for example a pure store-catalog + /// sample config) are intentionally not walked, so their illustrative + /// placeholder references are never flagged. + /// + /// All problems are collected and reported together, each with a path + /// such as `servers[0].services.cas[main].cas_store references + /// undefined store 'CAS_STORE'`, rather than failing on the first one. + /// + /// Duplicate store names and duplicate scheduler names are also reported, + /// because [`crate::stores`] / the store manager resolve names last-wins + /// and would otherwise silently drop the earlier definition. + /// + /// # Errors + /// + /// Returns `Err(Code::InvalidArgument)` if any referenced store/scheduler + /// name is undeclared, or if a store/scheduler name is declared more than + /// once. + pub fn validate_references(&self) -> Result<(), Error> { + let mut problems: Vec = Vec::new(); + + // Declared registries. Collision on insert is a duplicate name. + let mut store_names: HashSet<&str> = HashSet::new(); + for store in &self.stores { + if !store_names.insert(store.name.as_str()) { + problems.push(format!("duplicate store name '{}'", store.name)); + } + } + let mut scheduler_names: HashSet<&str> = HashSet::new(); + for scheduler in self.schedulers.iter().flatten() { + if !scheduler_names.insert(scheduler.name.as_str()) { + problems.push(format!("duplicate scheduler name '{}'", scheduler.name)); + } + } + + // Only validate a reference category when its registry is non-empty. + // A config that declares zero stores (or zero schedulers) is a + // fragment or legacy-format sample whose backing definitions live + // elsewhere; every reference would then trivially be "unresolved" and + // flagging them would be a false positive on an otherwise-valid + // sample. Real deployments always declare their stores, so genuine + // typos are still caught. + let mut checker = ReferenceChecker { + check_stores: !store_names.is_empty(), + check_schedulers: !scheduler_names.is_empty(), + store_names, + scheduler_names, + problems, + used_stores: Vec::new(), + }; + + for (worker_idx, worker) in self.workers.iter().flatten().enumerate() { + let WorkerConfig::Local(local) = worker; + checker.store( + &format!("workers[{worker_idx}].local.cas_fast_slow_store"), + &local.cas_fast_slow_store, + ); + if let Some(ac_store) = &local.upload_action_result.ac_store { + checker.store( + &format!("workers[{worker_idx}].local.upload_action_result.ac_store"), + ac_store, + ); + } + if let Some(historical) = &local.upload_action_result.historical_results_store { + checker.store( + &format!( + "workers[{worker_idx}].local.upload_action_result.historical_results_store" + ), + historical, + ); + } + } + + if let Some(origin_events) = &self.experimental_origin_events { + checker.store( + "experimental_origin_events.publisher.store", + &origin_events.publisher.store, + ); + } + + for (server_idx, server) in self.servers.iter().enumerate() { + let Some(services) = &server.services else { + continue; + }; + let prefix = format!("servers[{server_idx}].services"); + for entry in services.cas.iter().flatten() { + let name = &entry.instance_name; + checker.store(&format!("{prefix}.cas[{name}].cas_store"), &entry.cas_store); + if let Some(index_store) = entry + .experimental_chunking + .as_ref() + .and_then(|chunking| chunking.index_store.as_ref()) + { + checker.store( + &format!("{prefix}.cas[{name}].experimental_chunking.index_store"), + index_store, + ); + } + } + for entry in services.ac.iter().flatten() { + let name = &entry.instance_name; + checker.store(&format!("{prefix}.ac[{name}].ac_store"), &entry.ac_store); + } + for entry in services.execution.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.execution[{name}].cas_store"), + &entry.cas_store, + ); + checker.scheduler( + &format!("{prefix}.execution[{name}].scheduler"), + &entry.scheduler, + ); + } + for entry in services.capabilities.iter().flatten() { + if let Some(remote_execution) = &entry.remote_execution { + let name = &entry.instance_name; + checker.scheduler( + &format!("{prefix}.capabilities[{name}].remote_execution.scheduler"), + &remote_execution.scheduler, + ); + } + } + for entry in services.bytestream.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.bytestream[{name}].cas_store"), + &entry.cas_store, + ); + } + for entry in services.fetch.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.fetch[{name}].fetch_store"), + &entry.fetch_store, + ); + } + for entry in services.push.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.push[{name}].push_store"), + &entry.push_store, + ); + } + if let Some(worker_api) = &services.worker_api { + checker.scheduler( + &format!("{prefix}.worker_api.scheduler"), + &worker_api.scheduler, + ); + } + if let Some(bep) = &services.experimental_bep { + checker.store(&format!("{prefix}.experimental_bep.store"), &bep.store); + } + } + + // Scheduler specs themselves reference stores (and nest further + // schedulers); walk every declared scheduler so those references are + // resolved and their stores seed the `ref_store` reachability walk. + for scheduler in self.schedulers.iter().flatten() { + checker.scheduler_spec(&format!("schedulers[{}]", scheduler.name), &scheduler.spec); + } + + // Resolve `ref_store` references reachable from the wired-up stores. + if checker.check_stores { + let store_specs: HashMap<&str, &StoreSpec> = self + .stores + .iter() + .map(|store| (store.name.as_str(), &store.spec)) + .collect(); + checker.walk_used_stores(&store_specs); + } + + if checker.problems.is_empty() { + Ok(()) + } else { + Err(make_err!( + Code::InvalidArgument, + "configuration reference validation failed:\n - {}", + checker.problems.join("\n - ") + )) + } + } +} + +/// Recursively collects every `ref_store` name that appears anywhere inside a +/// single [`StoreSpec`] tree (the wrapper stores nest other `StoreSpec`s). The +/// match is exhaustive on purpose so a newly added wrapping store forces this +/// to be revisited. +fn collect_ref_store_names<'a>(spec: &'a StoreSpec, out: &mut Vec<&'a StoreRefName>) { + match spec { + StoreSpec::RefStore(ref_spec) => out.push(&ref_spec.name), + StoreSpec::CacheMetrics(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::Verify(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::CompletenessChecking(inner) => { + collect_ref_store_names(&inner.backend, out); + collect_ref_store_names(&inner.cas_store, out); + } + StoreSpec::Compression(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::Dedup(inner) => { + collect_ref_store_names(&inner.index_store, out); + collect_ref_store_names(&inner.content_store, out); + } + StoreSpec::ExistenceCache(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::FastSlow(inner) => { + collect_ref_store_names(&inner.fast, out); + collect_ref_store_names(&inner.slow, out); + } + StoreSpec::Shard(inner) => { + for shard in &inner.stores { + collect_ref_store_names(&shard.store, out); + } + } + StoreSpec::SizePartitioning(inner) => { + collect_ref_store_names(&inner.lower_store, out); + collect_ref_store_names(&inner.upper_store, out); + } + // Leaf stores and stores whose backends are concrete (non-`StoreSpec`) + // specs carry no nested `ref_store` names. + StoreSpec::Memory(_) + | StoreSpec::ExperimentalCloudObjectStore(_) + | StoreSpec::OntapS3ExistenceCache(_) + | StoreSpec::Filesystem(_) + | StoreSpec::Grpc(_) + | StoreSpec::RedisStore(_) + | StoreSpec::Noop(_) + | StoreSpec::ExperimentalMongo(_) => {} + } +} + +/// Accumulates unresolved-reference problems while walking a [`CasConfig`]. +struct ReferenceChecker<'a> { + store_names: HashSet<&'a str>, + scheduler_names: HashSet<&'a str>, + check_stores: bool, + check_schedulers: bool, + problems: Vec, + /// Declared store names reached from a worker/service, used to seed the + /// `ref_store` reachability walk. + used_stores: Vec<&'a str>, +} + +impl<'a> ReferenceChecker<'a> { + /// Records that `name` (found at `path`) must resolve to a declared store. + fn store(&mut self, path: &str, name: &'a str) { + if !self.check_stores { + return; + } + if self.store_names.contains(name) { + self.used_stores.push(name); + } else { + self.problems + .push(format!("{path} references undefined store '{name}'")); + } + } + + /// Records that `name` (found at `path`) must resolve to a declared + /// scheduler. + fn scheduler(&mut self, path: &str, name: &str) { + if !self.check_schedulers { + return; + } + if !self.scheduler_names.contains(name) { + self.problems + .push(format!("{path} references undefined scheduler '{name}'")); + } + } + + /// Walks a [`SchedulerSpec`] tree, resolving every store reference it + /// carries and recursing into nested schedulers. The match is exhaustive + /// on purpose so a newly added scheduler variant forces this to be + /// revisited (same design as [`collect_ref_store_names`]). + fn scheduler_spec(&mut self, path: &str, spec: &'a SchedulerSpec) { + match spec { + SchedulerSpec::Simple(simple) => match &simple.experimental_backend { + Some(ExperimentalSimpleSchedulerBackend::Redis(redis)) => { + self.store( + &format!("{path}.simple.experimental_backend.redis.redis_store"), + &redis.redis_store, + ); + } + Some(ExperimentalSimpleSchedulerBackend::Memory) | None => {} + }, + SchedulerSpec::CacheLookup(inner) => { + self.store(&format!("{path}.cache_lookup.ac_store"), &inner.ac_store); + self.scheduler_spec(&format!("{path}.cache_lookup.scheduler"), &inner.scheduler); + } + SchedulerSpec::PropertyModifier(inner) => { + self.scheduler_spec( + &format!("{path}.property_modifier.scheduler"), + &inner.scheduler, + ); + } + SchedulerSpec::HistoricalResource(inner) => { + self.scheduler_spec( + &format!("{path}.historical_resource.scheduler"), + &inner.scheduler, + ); + } + // Carries no store references and no nested scheduler. + SchedulerSpec::Grpc(_) => {} + } + } + + /// Walks the transitive closure of stores reachable from + /// [`Self::used_stores`], resolving each nested `ref_store` against the + /// declared store set. A visited set bounds the walk against `ref_store` + /// cycles. + fn walk_used_stores(&mut self, store_specs: &HashMap<&'a str, &'a StoreSpec>) { + let mut visited: HashSet<&'a str> = HashSet::new(); + while let Some(name) = self.used_stores.pop() { + if !visited.insert(name) { + continue; + } + let Some(&spec) = store_specs.get(name) else { + continue; + }; + let mut refs: Vec<&'a StoreRefName> = Vec::new(); + collect_ref_store_names(spec, &mut refs); + for ref_name in refs { + if self.store_names.contains(ref_name.as_str()) { + self.used_stores.push(ref_name.as_str()); + } else { + self.problems.push(format!( + "store '{name}' references undefined store '{ref_name}'" + )); + } + } + } + } } #[cfg(test)] @@ -2290,4 +2642,354 @@ mod tests { assert!(logs_contain("instance_name=\"opt-out\"")); assert!(logs_contain(CasConfig::ZSTD_COMPRESSION_DOCS_URL)); } + + // ---------------------------------------------------------------------- + // `validate_references` (offline `nativelink --check`) tests. + // ---------------------------------------------------------------------- + + fn parse_cas_config(json5: &str) -> CasConfig { + serde_json5::from_str(json5).expect("valid CasConfig json5") + } + + /// A fully-resolved config that exercises worker, service and nested + /// `ref_store` reference sites. `WRAPPED` wraps a `ref_store` to `CAS`, + /// so the reachability walk must resolve it. + const VALID_CONFIG: &str = r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "AC", memory: {} }, + { name: "WRAPPED", fast_slow: { + fast: { ref_store: { name: "CAS" } }, + slow: { noop: {} }, + } }, + ], + schedulers: [{ name: "SCHED", simple: {} }], + workers: [{ + local: { + worker_api_endpoint: { uri: "grpc://127.0.0.1:50061" }, + cas_fast_slow_store: "WRAPPED", + upload_action_result: { ac_store: "AC" }, + work_directory: "/tmp/work", + platform_properties: {}, + }, + }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "main", cas_store: "CAS" }], + ac: [{ instance_name: "main", ac_store: "AC" }], + execution: [{ instance_name: "main", cas_store: "CAS", scheduler: "SCHED" }], + capabilities: [{ instance_name: "main", remote_execution: { scheduler: "SCHED" } }], + bytestream: [{ instance_name: "main", cas_store: "CAS" }], + worker_api: { scheduler: "SCHED" }, + }, + }], + }"#; + + #[test] + fn validate_references_accepts_resolved_config() { + parse_cas_config(VALID_CONFIG) + .validate_references() + .expect("every reference resolves"); + } + + #[test] + fn validate_references_flags_undefined_service_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { cas: [{ instance_name: "main", cas_store: "CAS_TYPO" }] }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + assert_eq!(err.code, Code::InvalidArgument); + let message = err.messages.join("\n"); + assert!( + message.contains("CAS_TYPO"), + "error must name the bad store ref: {message}" + ); + assert!( + message.contains("servers[0].services.cas[main].cas_store"), + "error must name the reference site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_scheduler() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ name: "SCHED", simple: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + execution: [{ + instance_name: "main", + cas_store: "CAS", + scheduler: "SCHED_TYPO", + }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + assert_eq!(err.code, Code::InvalidArgument); + let message = err.messages.join("\n"); + assert!( + message.contains("SCHED_TYPO"), + "error must name the bad scheduler ref: {message}" + ); + assert!( + message.contains("scheduler"), + "error must name the reference site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_worker_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + workers: [{ + local: { + worker_api_endpoint: { uri: "grpc://127.0.0.1:50061" }, + cas_fast_slow_store: "WORKER_TYPO", + work_directory: "/tmp/work", + platform_properties: {}, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("WORKER_TYPO") + && message.contains("workers[0].local.cas_fast_slow_store"), + "error must name the bad worker store ref and site: {message}" + ); + } + + #[test] + fn validate_references_flags_nested_ref_store_typo() { + // `WRAPPED` is wired to a service, so its nested `ref_store` typo must + // be resolved and flagged via the reachability walk. + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "WRAPPED", fast_slow: { + fast: { ref_store: { name: "CAS_TYPO" } }, + slow: { noop: {} }, + } }, + ], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { cas: [{ instance_name: "main", cas_store: "WRAPPED" }] }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("CAS_TYPO") && message.contains("store 'WRAPPED'"), + "error must name the nested bad ref and its owning store: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_scheduler_redis_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ + name: "SCHED", + simple: { + experimental_backend: { + redis: { redis_store: "REDIS_TYPO" }, + }, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("REDIS_TYPO") + && message + .contains("schedulers[SCHED].simple.experimental_backend.redis.redis_store"), + "error must name the bad redis store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_cache_lookup_ac_store() { + // The nested scheduler inside `cache_lookup` must be recursed into, + // and its `ac_store` reference resolved. + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ + name: "SCHED", + property_modifier: { + modifications: [], + scheduler: { + cache_lookup: { + ac_store: "AC_TYPO", + scheduler: { simple: {} }, + }, + }, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("AC_TYPO") + && message.contains( + "schedulers[SCHED].property_modifier.scheduler.cache_lookup.ac_store" + ), + "error must name the nested bad ac_store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_chunking_index_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ + instance_name: "main", + cas_store: "CAS", + experimental_chunking: { index_store: "INDEX_TYPO" }, + }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("INDEX_TYPO") + && message + .contains("servers[0].services.cas[main].experimental_chunking.index_store"), + "error must name the bad index store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_duplicate_store_name() { + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "CAS", memory: {} }, + ], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("duplicate store name 'CAS'"), + "error must report the duplicate store: {message}" + ); + } + + #[test] + fn validate_references_flags_duplicate_scheduler_name() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [ + { name: "SCHED", simple: {} }, + { name: "SCHED", simple: {} }, + ], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("duplicate scheduler name 'SCHED'"), + "error must report the duplicate scheduler: {message}" + ); + } + + #[test] + fn validate_references_collects_all_problems() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ name: "SCHED", simple: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "main", cas_store: "BAD_CAS" }], + ac: [{ instance_name: "main", ac_store: "BAD_AC" }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + // Reporting is not fail-fast: both bad references appear together. + assert!( + message.contains("BAD_CAS") && message.contains("BAD_AC"), + "error must collect every unresolved reference: {message}" + ); + } + + #[test] + fn validate_references_skips_empty_store_registry() { + // A fragment/legacy sample with no declared stores must not be + // flagged: every reference would trivially be unresolved, which would + // be a false positive on an otherwise-valid sample. + let cfg = parse_cas_config( + r#"{ + stores: [], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "", cas_store: "CAS_MAIN_STORE" }], + execution: [{ + instance_name: "", + cas_store: "WORKER_STORE", + scheduler: "MAIN_SCHEDULER", + }], + }, + }], + }"#, + ); + cfg.validate_references() + .expect("empty registries defer to boot, not a false positive"); + } + + #[test] + fn validate_references_ignores_unreferenced_store_internal_refs() { + // A pure store-catalog config (no servers, no workers) may carry + // illustrative placeholder `ref_store` names; because nothing wires + // those stores to a consumer, their internal refs are not walked and + // must not be flagged. + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "REAL", memory: {} }, + { name: "SHOWCASE", ref_store: { name: "PLACEHOLDER_NOT_DECLARED" } }, + ], + servers: [], + }"#, + ); + cfg.validate_references() + .expect("unreferenced catalog stores are not walked"); + } } diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 81b6179f929..0a48f575f03 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -844,6 +844,32 @@ pub struct FilesystemSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub max_concurrent_writes: usize, + /// Maximum number of concurrent background file-deletion operations. + /// When an entry leaves the store, its backing file is unlinked by a + /// background task. Without a bound, an eviction storm (cache wipe or a + /// flood of overwrites) can spawn these without limit: memory grows with the + /// queued task frames and disk/inode reclamation lags behind what the async + /// runtime can schedule. Bounding it applies backpressure the same way + /// `max_concurrent_writes` does — once this many deletes are in flight, a + /// further delete runs synchronously inline in the evicting path instead of + /// spawning, which throttles eviction to delete throughput and never leaks + /// a file. + /// + /// Bounded by default, unlike `max_concurrent_writes`: an unbounded delete + /// backlog is a latent memory/inode blowup, and leaving the protection off + /// by default just moves that risk from a high-context code review to a + /// 4 AM config change nobody remembers to make. The default is generous + /// enough that normal steady-state eviction never reaches it (so the inline + /// path never fires in ordinary operation); it only engages under a genuine + /// storm, which is exactly when the backpressure is wanted. + /// A value of 0 is the explicit escape hatch for unlimited. + /// Default: 1024. + #[serde( + default = "default_max_concurrent_deletes", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub max_concurrent_deletes: usize, + /// When true, advise the kernel to drop the page cache for each blob after /// it is written or read (`posix_fadvise` with `POSIX_FADV_DONTNEED`). On /// real filesystems this takes a globally serialized, all-CPU kernel path @@ -1398,6 +1424,42 @@ pub struct CommonObjectSpec { )] pub multipart_max_concurrent_uploads: Option, + /// Target size of each part in a multipart upload, growing only as + /// needed to stay under the provider's 10,000-part ceiling. Parts sized + /// at the 5 MiB provider floor turn a multi-GiB blob into thousands of + /// tiny concurrent requests, which churns HTTP connection pools; a + /// larger part keeps the request count low. Each queued or in-flight + /// part is buffered in memory, so one multipart upload buffers up to + /// roughly `(2 * multipart_max_concurrent_uploads + 1) * + /// multipart_part_size`; `max_concurrent_multipart_uploads` bounds the + /// aggregate across uploads. Values are clamped to the provider's + /// 5 MiB minimum and 5 GiB maximum part sizes. + /// + /// Currently honored by the S3 and ONTAP S3 providers. + /// + /// Default: 64MB. + #[serde( + default, + deserialize_with = "convert_optional_data_size_with_shellexpand" + )] + pub multipart_part_size: Option, + + /// Store-wide cap on the number of blobs concurrently uploaded via the + /// multipart path. Each multipart upload can buffer roughly + /// `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`, + /// so this bounds the store's aggregate upload buffer (about 5 GiB at + /// the defaults). Uploads over the 5 MiB multipart threshold wait for a + /// slot; smaller blobs are unaffected. + /// + /// Currently honored by the S3 and ONTAP S3 providers. + /// + /// Default: 4. + #[serde( + default, + deserialize_with = "convert_optional_numeric_with_shellexpand" + )] + pub max_concurrent_multipart_uploads: Option, + /// Allow unencrypted HTTP connections. Only use this for local testing. /// /// Default: false @@ -1502,6 +1564,8 @@ impl Default for CommonObjectSpec { consider_expired_after_s: 0, max_retry_buffer_per_request: None, multipart_max_concurrent_uploads: None, + multipart_part_size: None, + max_concurrent_multipart_uploads: None, insecure_allow_http: false, disable_http2: false, throttle_retry_delay_s: default_throttle_retry_delay_s(), @@ -1628,12 +1692,18 @@ pub struct GrpcSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub connections_per_endpoint: usize, - /// Maximum time (seconds) allowed for a single RPC request (e.g. a - /// `ByteStream.Write` call) before it is cancelled. + /// Maximum time (seconds) allowed for a single RPC request before it + /// is cancelled with `DeadlineExceeded` (retryable). This applies to + /// every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, + /// `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment + /// of read streams through their first response chunk. The body of an + /// established read stream is not bounded by this deadline. /// /// A value of 0 (the default) disables the per-RPC timeout. Dead /// connections are still detected by the HTTP/2 and TCP keepalive - /// mechanisms configured on each endpoint. + /// mechanisms configured on each endpoint, and endpoints that cannot + /// be connected to at all fail fast with `Unavailable` rather than + /// queuing requests indefinitely. /// /// For large uploads (multi-GB), either leave this at 0 or set it /// large enough to accommodate the full transfer time. @@ -1797,6 +1867,14 @@ const fn default_read_batching_dispatch_slots() -> usize { 4 } +/// Generous enough that normal eviction never reaches it (so the synchronous +/// inline delete path never fires in ordinary operation), low enough that the +/// in-flight backlog stays bounded to a few MB under a storm. Tune per +/// deployment; 0 disables the bound. See `FilesystemSpec::max_concurrent_deletes`. +const fn default_max_concurrent_deletes() -> usize { + 1024 +} + const fn default_read_batching_max_queued_bytes() -> u64 { 32 * 1024 * 1024 // 32 MiB. } diff --git a/nativelink-config/tests/json5_test.rs b/nativelink-config/tests/json5_test.rs index f1dadde10ab..3f2b4ccdd24 100644 --- a/nativelink-config/tests/json5_test.rs +++ b/nativelink-config/tests/json5_test.rs @@ -38,8 +38,11 @@ fn test_example_parsing() { if !config_file.contains(".json5") { continue; } - CasConfig::try_from_json5_file(&config_file) + let config = CasConfig::try_from_json5_file(&config_file) .unwrap_or_else(|e| panic!("Error while reading {config_file}: {e}")); + config + .validate_references() + .unwrap_or_else(|e| panic!("Unresolved reference in {config_file}: {e}")); found_at_least_one_entry = true; } diff --git a/nativelink-redis-tester/src/dynamic_fake_redis.rs b/nativelink-redis-tester/src/dynamic_fake_redis.rs index 3b602c597a9..10e4bc6b189 100644 --- a/nativelink-redis-tester/src/dynamic_fake_redis.rs +++ b/nativelink-redis-tester/src/dynamic_fake_redis.rs @@ -156,12 +156,13 @@ impl FakeRedisBackend { }; // Lazy implementation making assumptions. assert_eq!( - args[load_offset..load_offset + 4], + args[load_offset..load_offset + 5], vec![ OwnedFrame::BulkString(b"LOAD".to_vec()), - OwnedFrame::BulkString(b"2".to_vec()), + OwnedFrame::BulkString(b"3".to_vec()), OwnedFrame::BulkString(b"data".to_vec()), - OwnedFrame::BulkString(b"version".to_vec()) + OwnedFrame::BulkString(b"version".to_vec()), + OwnedFrame::BulkString(b"@__key".to_vec()) ] ); let mut results = vec![Value::Int(0)]; @@ -178,12 +179,13 @@ impl FakeRedisBackend { let desc = matches!(args.get(i + 3), Some(OwnedFrame::BulkString(d)) if d == b"DESC"); Some((field, desc)) }); - let sorted = |mut rows: Vec>| { + let sorted = |mut rows: Vec<(String, HashMap)>| { if let Some((field, desc)) = &sort_spec { - let key = |r: &HashMap| match r.get(field) { - Some(Value::BulkString(b)) => b.clone(), - _ => Vec::new(), - }; + let key = + |r: &(String, HashMap)| match r.1.get(field) { + Some(Value::BulkString(b)) => b.clone(), + _ => Vec::new(), + }; rows.sort_by_key(key); if *desc { rows.reverse(); @@ -199,16 +201,20 @@ impl FakeRedisBackend { .table .lock() .unwrap() - .values() - .filter(|f| f.contains_key("data") && f.contains_key("version")) - .cloned() + .iter() + .filter(|(_, f)| { + f.contains_key("data") && f.contains_key("version") + }) + .map(|(k, f)| (k.clone(), f.clone())) .collect(); - for fields in sorted(rows) { + for (key, fields) in sorted(rows) { results.push(Value::Array(vec![ Value::BulkString(b"data".to_vec()), fields["data"].clone(), Value::BulkString(b"version".to_vec()), fields["version"].clone(), + Value::BulkString(b"__key".to_vec()), + Value::BulkString(key.into_bytes()), ])); } } else { @@ -225,14 +231,14 @@ impl FakeRedisBackend { .table .lock() .unwrap() - .values() - .filter(|f| { + .iter() + .filter(|(_, f)| { f.get(field) == Some(&Value::BulkString(value.as_bytes().to_vec())) }) - .cloned() + .map(|(k, f)| (k.clone(), f.clone())) .collect(); - for fields in sorted(rows) { + for (key, fields) in sorted(rows) { { let mut record = vec![ Value::BulkString(b"data".to_vec()), @@ -245,6 +251,8 @@ impl FakeRedisBackend { record.push(Value::BulkString(b"version".to_vec())); record.push(version.clone()); } + record.push(Value::BulkString(b"__key".to_vec())); + record.push(Value::BulkString(key.into_bytes())); results.push(Value::Array(record)); } } @@ -354,6 +362,38 @@ impl FakeRedisBackend { } } + "EVAL" => { + // Only the delete-lost-record script is run via + // EVAL; emulate its semantics: delete KEYS[1] only + // if the liveness guard KEYS[2] does not exist and + // the record still holds version ARGV[1]. + let script = + str::from_utf8(args[0].as_bytes().expect("Script not bytes")).unwrap(); + assert!(script.contains("DEL"), "Unexpected EVAL script: {script:?}"); + assert_eq!(args[1], OwnedFrame::BulkString(b"2".to_vec())); + let key = str::from_utf8(args[2].as_bytes().expect("Key not bytes")) + .expect("Key cannot be parsed as string") + .to_string(); + let guard_key = + str::from_utf8(args[3].as_bytes().expect("Guard key not bytes")) + .expect("Guard key cannot be parsed as string") + .to_string(); + let expected_version = + args[4].as_bytes().expect("Version not bytes").to_vec(); + let mut table = self.table.lock().unwrap(); + if table.contains_key(&guard_key) { + Value::Int(-1) + } else if table.get(&key).and_then(|f| f.get("version")) + == Some(&Value::BulkString(expected_version)) + { + table.remove(&key); + debug!(%key, "Deleted key via delete-lost-record"); + Value::Int(1) + } else { + Value::Int(0) + } + } + "HMSET" => { let mut values = HashMap::new(); assert_eq!( diff --git a/nativelink-redis-tester/src/fake_redis.rs b/nativelink-redis-tester/src/fake_redis.rs index 95a49685b87..4f38592314f 100644 --- a/nativelink-redis-tester/src/fake_redis.rs +++ b/nativelink-redis-tester/src/fake_redis.rs @@ -15,6 +15,8 @@ use core::fmt::Write; use core::hash::BuildHasher; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use nativelink_util::background_spawn; use redis::Value; @@ -292,3 +294,51 @@ pub async fn make_fake_redis_with_responses u16 { make_fake_redis_with_multiple_responses(vec![responses]).await } + +/// Like [`make_fake_redis_with_responses`], except that the first request +/// matching a key of `first_responses` gets that key's response instead, once. +/// Every later request gets its normal response. Use this to test that a +/// caller retries after a transient error. +pub async fn make_fake_redis_with_first_responses< + B: BuildHasher + Clone + Send + 'static + Sync, +>( + responses: HashMap, + first_responses: HashMap, +) -> u16 { + let first_responses: Vec<(String, String, Arc)> = first_responses + .into_iter() + .map(|(key, value)| (key, value, Arc::new(AtomicBool::new(false)))) + .collect(); + let handler = move |buf: &[u8]| -> String { + let str_buf = String::from_utf8_lossy(buf).into_owned(); + for (key, value, used) in &first_responses { + if str_buf.starts_with(key) && !used.swap(true, Ordering::SeqCst) { + info!("Responding once to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + for (key, value) in &responses { + if str_buf.starts_with(key) { + info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + warn!( + "Unknown command: {}", + str_buf.chars().take(1000).collect::() + ); + String::new() + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>(); + background_spawn!("listener", async move { + fake_redis_internal(listener, listener_ready_tx, vec![handler]).await; + }); + listener_ready_rx + .await + .expect("Expected successful listener boot"); + info!(port, "Fake redis booted"); + port +} diff --git a/nativelink-redis-tester/src/lib.rs b/nativelink-redis-tester/src/lib.rs index 896d384df96..65ec4874674 100644 --- a/nativelink-redis-tester/src/lib.rs +++ b/nativelink-redis-tester/src/lib.rs @@ -20,7 +20,8 @@ mod read_only_redis; pub use dynamic_fake_redis::{FakeRedisBackend, SubscriptionManagerNotify}; pub use fake_redis::{ add_lua_script, add_to_response, add_to_response_raw, fake_redis_sentinel_master_stream, - fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_responses, + fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_first_responses, + make_fake_redis_with_responses, }; pub use pubsub::MockPubSub; pub use read_only_redis::ReadOnlyRedis; diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 25077de27cc..3edfa9274d7 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -74,11 +74,14 @@ rust_test_suite( "tests/admin_test.rs", "tests/backpressure_test.rs", "tests/cache_lookup_scheduler_test.rs", + "tests/client_keepalive_clobber_test.rs", + "tests/client_keepalive_failed_write_test.rs", "tests/dispatch_ack_test.rs", "tests/escalation_test.rs", "tests/execution_active_count_test.rs", "tests/fleet_capabilities_test.rs", "tests/historical_resource_scheduler_test.rs", + "tests/kill_revoked_redispatch_race_test.rs", "tests/live_memory_veto_test.rs", "tests/liveness_clock_step_test.rs", "tests/placement_test.rs", @@ -94,6 +97,7 @@ rust_test_suite( "tests/unsatisfiable_action_test.rs", "tests/worker_capability_index_test.rs", "tests/worker_registry_test.rs", + "tests/worker_retry_reservation_test.rs", "tests/worker_running_operations_report_test.rs", ], compile_data = [ diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 64d7439e449..c7c5f78ea66 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -371,6 +371,14 @@ struct ApiWorkerSchedulerImpl { /// about it without waiting for the next periodic exchange. local_fleet_change_notify: Arc, + /// The next dispatch generation; incremented under this struct's mutex + /// for every dispatch, so each dispatch of any operation to any worker + /// carries a unique, strictly increasing identity. Scheduler-wide rather + /// than per-worker: a per-worker counter would restart when a worker + /// re-registers under the same id, letting a stale kill snapshot from + /// the old incarnation match a fresh dispatch on the new one. + next_dispatch_generation: u64, + /// Whether an action with a given property shape could run on some /// connected worker when that worker is idle. `None` means it could. This /// depends only on what workers registered with, so it holds until the @@ -546,8 +554,11 @@ impl ApiWorkerSchedulerImpl { fn add_worker(&mut self, worker: Worker) -> Result<(), Error> { let worker_id = worker.id.clone(); let platform_properties = worker.platform_properties.clone(); - // A replacement is one out and one in, so the gauge only moves for a - // genuinely new worker. + // The caller evicts a live entry under this id before admitting the + // new incarnation (see `ApiWorkerScheduler::add_worker`), so `put` + // does not normally replace anything. The branch below is kept as a + // defensive path: a replacement is one out and one in, so the gauge + // only moves for a genuinely new worker. let replaced = self.workers.put(worker_id.clone(), worker); // Add to capability index for fast matching @@ -1657,12 +1668,17 @@ impl ApiWorkerSchedulerImpl { action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { + // Allocated under the scheduler mutex: a re-dispatch always gets a + // strictly greater generation than any snapshot taken before it. + self.next_dispatch_generation += 1; + let dispatch_generation = self.next_dispatch_generation; if let Some(worker) = self.workers.get_mut(&worker_id) { let notify_worker_result = worker .notify_update(WorkerUpdate::RunAction(Box::new(( operation_id.clone(), action_info.clone(), dispatched_at, + dispatch_generation, )))) .await; @@ -1939,13 +1955,24 @@ impl ApiWorkerSchedulerImpl { &mut self, worker_id: &WorkerId, operation_id: OperationId, + snapshot_dispatch_generation: u64, ) -> Result<(), Error> { let Some(worker) = self.workers.get_mut(worker_id) else { // Gone between the snapshot and now; its actions were requeued. return Ok(()); }; // Already told, or finished in the meantime; nothing more to send. - if !worker.running_action_infos.contains_key(&operation_id) + let Some(pending_action_info) = worker.running_action_infos.get(&operation_id) else { + return Ok(()); + }; + // A different generation means the snapshot's dispatch was requeued + // and the operation legitimately re-dispatched to this worker while + // the sweep ran lock-free; the revocation evidence belongs to the + // old dispatch, not this one. The generation, not `dispatched_at`, + // is the dispatch's identity: the timestamp has second resolution, + // so a requeue-and-redispatch inside one wall-clock second carries + // the same `dispatched_at` and would be killed by mistake. + if pending_action_info.dispatch_generation != snapshot_dispatch_generation || worker.is_kill_requested(&operation_id) { return Ok(()); @@ -2075,6 +2102,14 @@ pub struct ApiWorkerScheduler { /// Woken when a worker joins or leaves this scheduler. local_fleet_change_notify: Arc, + + /// Test-only rendezvous: `remove_timedout_workers` parks here between its + /// liveness snapshot and re-acquiring the eviction lock, at an existing + /// await boundary, so tests can deterministically interleave a keepalive. + /// `.0` is signalled when the sweep reaches the boundary, `.1` lets it + /// proceed. + #[cfg(test)] + test_evict_gate: std::sync::OnceLock>, } impl ApiWorkerScheduler { @@ -2148,6 +2183,7 @@ impl ApiWorkerScheduler { peer_fleet_known_since: None, peer_fleet_last_refresh: None, local_fleet_change_notify: local_fleet_change_notify.clone(), + next_dispatch_generation: 0, static_verdicts: HashMap::new(), }), platform_property_manager, @@ -2160,6 +2196,8 @@ impl ApiWorkerScheduler { local_fleet_change_notify, capacity_generation, scheduler_id, + #[cfg(test)] + test_evict_gate: std::sync::OnceLock::new(), }) } @@ -2217,6 +2255,21 @@ impl ApiWorkerScheduler { load: Option, keepalive: bool, ) -> Result<(), Error> { + // TL1: stamp the independent registry liveness signal on RECEIPT, before + // taking `inner.lock()`. The registry is the second liveness source that + // `remove_timedout_workers` consults ("alive if EITHER source"), but it + // was previously stamped only AFTER the inner.lock block below — which + // holds the lock across an await (the unacked-dispatch sweep). A punctual + // keepalive whose inner.lock acquisition is contended, or whose task is + // starved on a saturated runtime, then refreshed NEITHER liveness source + // before the eviction sweep read them, falsely evicting a live worker. The + // registry write is a lightweight RwLock insert (no await-under-lock, no + // store, no sweep), so stamping it first decouples the monotone liveness + // FACT from the coordinated state write; it can no longer be starved. + let now = UNIX_EPOCH + Duration::from_secs(timestamp); + self.worker_registry + .update_worker_heartbeat(worker_id, now) + .await; { let mut inner = self.inner.lock().await; inner @@ -2258,10 +2311,6 @@ impl ApiWorkerScheduler { } } } - let now = UNIX_EPOCH + Duration::from_secs(timestamp); - self.worker_registry - .update_worker_heartbeat(worker_id, now) - .await; Ok(()) } @@ -2538,6 +2587,36 @@ impl WorkerScheduler for ApiWorkerScheduler { "Received request to add worker while shutting down" )); } + // A worker re-registering under an id the pool still holds is a new + // incarnation of the process, not the one the pool knows. Silently + // replacing the entry dropped the old `Worker` together with its + // `running_action_infos`: nothing was requeued, the old channel never + // heard a Disconnect, and its operations stayed Executing against + // this id until they timed out and were requeued to another worker + // while the original process might still be running them (double + // execution). Evict the old incarnation exactly like a disconnect + // first, then admit the new one. A failure to requeue must not block + // the reconnecting worker, so it is logged and the add proceeds. + if inner.workers.contains(&worker_id) { + let evict_res = inner + .immediate_evict_worker( + &worker_id, + make_err!( + Code::Internal, + "Worker {worker_id} re-registered; evicting the old incarnation" + ), + true, + WorkerDisconnectReason::Disconnected, + ) + .await; + if let Err(err) = evict_res { + warn!( + ?worker_id, + ?err, + "Failed to requeue everything the replaced worker incarnation was running" + ); + } + } let result = inner .add_worker(worker) .err_tip(|| "Error while adding worker, removing from pool"); @@ -2748,10 +2827,41 @@ impl WorkerScheduler for ApiWorkerScheduler { return Ok(()); } + // Test-only hold at this pre-existing await boundary (the inner lock + // was dropped after the snapshot and is only re-acquired below); lets + // a test run a concurrent keepalive here deterministically. + #[cfg(test)] + if let Some(gate) = self.test_evict_gate.get() { + gate.0.notify_one(); + gate.1.notified().await; + } + let mut inner = self.inner.lock().await; let mut result = Ok(()); for (worker_id, kill_overdue) in &worker_ids_to_remove { + // The staleness decision above was made from a snapshot taken + // before this lock was (re-)acquired; the registry read is an + // await boundary every sweep crosses with the lock released. A + // keepalive processed in that window refreshes + // `last_update_timestamp` under this same mutex, so re-validate + // liveness here and skip eviction if the worker heartbeated past + // the threshold. Kill-overdue evictions are intentionally exempt: + // keepalives from a worker that will not acknowledge a kill are + // exactly the wedged-worker symptom being handled. + if !*kill_overdue + && inner + .workers + .peek(worker_id) + .is_some_and(|worker| worker.last_update_timestamp > timeout_threshold) + { + trace!( + ?worker_id, + timeout_threshold, + "Worker heartbeated between staleness snapshot and eviction lock; keeping" + ); + continue; + } let err = if *kill_overdue { make_err!( Code::Internal, @@ -2824,7 +2934,7 @@ impl WorkerScheduler for ApiWorkerScheduler { async fn kill_revoked_operations(&self) -> Result<(), Error> { let (worker_state_manager, running) = { let inner = self.inner.lock().await; - let running: Vec<(WorkerId, OperationId)> = inner + let running: Vec<(WorkerId, OperationId, u64)> = inner .workers .iter() .flat_map(|(worker_id, worker)| { @@ -2841,7 +2951,16 @@ impl WorkerScheduler for ApiWorkerScheduler { .filter(|(_, pending_action_info)| { pending_action_info.kill_requested_at.is_none() }) - .map(|(operation_id, _)| (worker_id.clone(), operation_id.clone())) + // `dispatch_generation` identifies the dispatch + // this snapshot saw; the kill is only sent if the + // entry still carries it when the lock is re-taken. + .map(|(operation_id, pending_action_info)| { + ( + worker_id.clone(), + operation_id.clone(), + pending_action_info.dispatch_generation, + ) + }) }) .collect(); (inner.worker_state_manager.clone(), running) @@ -2849,8 +2968,8 @@ impl WorkerScheduler for ApiWorkerScheduler { // On store-backed deployments each check is a network round-trip, // so run them lock-free with bounded concurrency. - let revoked: Vec<(WorkerId, OperationId)> = futures::stream::iter(running) - .map(|(worker_id, operation_id)| { + let revoked: Vec<(WorkerId, OperationId, u64)> = futures::stream::iter(running) + .map(|(worker_id, operation_id, dispatch_generation)| { let worker_state_manager = worker_state_manager.clone(); async move { match worker_state_manager @@ -2858,7 +2977,7 @@ impl WorkerScheduler for ApiWorkerScheduler { .await { Ok(true) => None, - Ok(false) => Some((worker_id, operation_id)), + Ok(false) => Some((worker_id, operation_id, dispatch_generation)), // Only kill on positive evidence; try again next pass. Err(err) => { warn!( @@ -2883,10 +3002,13 @@ impl WorkerScheduler for ApiWorkerScheduler { // Re-check lock-free so the scheduler mutex is never held across // store I/O; the checks above may be stale by the time we get here. - // The remaining TOCTOU window is fine: worker_notify_kill_operation - // re-guards with contains_key + is_kill_requested under the lock. + // The remaining TOCTOU window is closed under the lock: + // worker_notify_kill_operation only kills the entry whose + // `dispatch_generation` matches the snapshot, so an operation + // requeued and re-dispatched to the same worker while the sweep ran + // lock-free is left alone. let mut confirmed = Vec::new(); - for (worker_id, operation_id) in revoked { + for (worker_id, operation_id, dispatch_generation) in revoked { // A result of Ok(true) or Err(_) means the operation was // reassigned to this worker after the first check, or the state // manager went quiet: nothing to kill on this pass. @@ -2895,7 +3017,7 @@ impl WorkerScheduler for ApiWorkerScheduler { .await .is_ok_and(|executing| !executing); if revoked { - confirmed.push((worker_id, operation_id)); + confirmed.push((worker_id, operation_id, dispatch_generation)); } } @@ -2905,10 +3027,10 @@ impl WorkerScheduler for ApiWorkerScheduler { let mut inner = self.inner.lock().await; let mut result = Ok(()); - for (worker_id, operation_id) in confirmed { + for (worker_id, operation_id, dispatch_generation) in confirmed { result = result.merge( inner - .worker_notify_kill_operation(&worker_id, operation_id) + .worker_notify_kill_operation(&worker_id, operation_id, dispatch_generation) .await, ); } @@ -3023,3 +3145,285 @@ mod peer_census_trusted_tests { )); } } + +#[cfg(test)] +mod timedout_worker_eviction_race_tests { + use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, + }; + use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::UpdateForWorker; + + use super::*; + use crate::worker_registry::WorkerRegistry; + + #[derive(Debug, Default)] + struct NoopWorkerStateManager; + + impl MetricsComponent for NoopWorkerStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl WorkerStateManager for NoopWorkerStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + Ok(()) + } + + async fn is_executing_on_worker( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + Ok(true) + } + } + + /// repro: `remove_timedout_workers` evicts a worker whose keepalive + /// arrived between the liveness snapshot and the eviction lock. + /// + /// The sweep snapshots `last_update_timestamp` under the inner lock, + /// drops the lock across the registry await, then re-locks and calls + /// `immediate_evict_worker` without re-checking the timestamp. A + /// keepalive that lands in that window (a real await boundary crossed by + /// every sweep) refreshes the worker under the same mutex, yet the + /// worker is still evicted. The test parks the sweep at that boundary + /// via a test-only gate, delivers the keepalive, and asserts the worker + /// survives; on current code the worker is evicted despite being + /// provably alive. + #[nativelink_macro::nativelink_test] + async fn keepalive_between_snapshot_and_evict_lock_still_evicts() { + const WORKER_TIMEOUT_S: u64 = 100; + const OLD_TS: WorkerTimestamp = 1_000; + const NOW_TS: WorkerTimestamp = 2_000; + // Fresh relative to NOW_TS: NOW_TS - FRESH_TS < WORKER_TIMEOUT_S. + const FRESH_TS: WorkerTimestamp = 1_999; + + let registry = Arc::new(WorkerRegistry::new()); + let scheduler = ApiWorkerScheduler::new( + Arc::new(NoopWorkerStateManager), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + None, + Arc::new(Notify::new()), + WORKER_TIMEOUT_S, + 10, + 0, + registry, + None, + false, + Duration::from_secs(15), + "test-scheduler".to_string(), + ); + let gate = Arc::new((Notify::new(), Notify::new())); + scheduler + .test_evict_gate + .set(gate.clone()) + .expect("gate set once"); + + let worker_id = WorkerId::from(String::from("worker1")); + // Keep rx alive so eviction's Disconnect send would succeed. + let (tx, _rx) = mpsc::channel::(16); + let worker = Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + OLD_TS, + 1, + ); + scheduler.add_worker(worker).await.expect("add worker"); + + // Start the sweep; it snapshots local_alive=false (OLD_TS is past + // the timeout), finds the registry stale too, and parks at the gate + // just before re-acquiring the inner lock to evict. + let sweep_scheduler = scheduler.clone(); + let sweep = nativelink_util::background_spawn!("timedout_sweep", async move { + sweep_scheduler.remove_timedout_workers(NOW_TS).await + }); + gate.0.notified().await; + + // The worker's keepalive arrives now: it acquires the inner lock, + // advances last_update_timestamp to FRESH_TS, and updates the + // registry heartbeat. The worker is provably alive. + scheduler + .worker_keep_alive_received(&worker_id, FRESH_TS, None) + .await + .expect("keepalive"); + assert!(scheduler.contains_worker_for_test(&worker_id).await); + + // Let the sweep proceed to its eviction lock. + gate.1.notify_one(); + sweep.await.expect("join").expect("sweep result"); + + // A correct sweep re-checks liveness under the lock and keeps the + // worker; current code evicts it unconditionally. + assert!( + scheduler.contains_worker_for_test(&worker_id).await, + "worker was evicted by remove_timedout_workers even though its \ + keepalive (timestamp {FRESH_TS}, well within the {WORKER_TIMEOUT_S}s \ + timeout at now={NOW_TS}) was fully processed before the eviction \ + lock was taken" + ); + } +} + +#[cfg(test)] +mod keepalive_eviction_race_tests { + //! Deterministic reproduction of the check-then-act race in + //! `remove_timedout_workers`: a keepalive that lands after the sweep's + //! unlocked registry check but before it re-acquires the inner lock is + //! never re-consulted, so a demonstrably live worker is evicted. + + use nativelink_config::schedulers::WorkerAllocationStrategy; + + use super::{ + ApiWorkerScheduler, Arc, Duration, Error, HashMap, MetricFieldData, MetricKind, + MetricPublishKnownKindData, MetricsComponent, Notify, OperationId, PlatformProperties, + UNIX_EPOCH, UpdateOperationType, Worker, WorkerId, WorkerScheduler, WorkerStateManager, + async_trait, mpsc, + }; + use crate::platform_property_manager::PlatformPropertyManager; + use crate::worker_registry::WorkerRegistry; + + #[derive(Debug)] + struct NoopWorkerStateManager; + + impl MetricsComponent for NoopWorkerStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl WorkerStateManager for NoopWorkerStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + Ok(()) + } + + async fn is_executing_on_worker( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + Ok(true) + } + } + + const WORKER_TIMEOUT_S: u64 = 100; + /// Worker registered (locally and in the registry) at t=0. + const REGISTERED_AT: u64 = 0; + /// The sweep and the keepalive both happen at t=200, i.e. one full + /// timeout past expiry, so both liveness sources initially say "dead". + const NOW: u64 = 200; + + #[nativelink_macro::nativelink_test] + async fn keepalive_between_registry_check_and_eviction_lock_is_ignored() { + let registry = Arc::new(WorkerRegistry::new()); + let scheduler = ApiWorkerScheduler::new( + Arc::new(NoopWorkerStateManager), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + None, + Arc::new(Notify::new()), + WORKER_TIMEOUT_S, + 60, // unacknowledged_kill_timeout_s + 0, // dispatch_ack_timeout_s + registry.clone(), + None, + false, // has_peers + Duration::from_secs(15), // record_ttl + "test-scheduler".to_string(), + ); + + let worker_id = WorkerId("stale_then_alive_worker".to_string()); + let (tx, _rx) = mpsc::channel(64); + scheduler + .add_worker(Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + REGISTERED_AT, + 4, + )) + .await + .expect("add_worker failed"); + + // Gate 1: hold the registry lock so the sweep, after snapshotting the + // (stale) local state and RELEASING the inner lock, parks at its + // unlocked `is_worker_alive` await. + let registry_guard = registry.write_lock_for_test().await; + + let sweep = nativelink_util::background_spawn!("timedout_sweep", { + let scheduler = scheduler.clone(); + async move { scheduler.remove_timedout_workers(NOW).await } + }); + // Let the sweep run: it takes the inner lock, snapshots the worker as + // locally stale, releases the lock, and parks on the registry read. + for _ in 0..8 { + tokio::task::yield_now().await; + } + + // Gate 2: take the inner lock (currently free - the sweep released it) + // so that after its registry read the sweep parks again before it can + // evict. This models inner-lock contention on a busy scheduler. + let mut inner_guard = scheduler.inner.lock().await; + + // Release gate 1: the sweep reads the registry (heartbeat still t=0, + // so "dead"), queues the worker for eviction, and parks on the inner + // lock we hold. + drop(registry_guard); + for _ in 0..8 { + tokio::task::yield_now().await; + } + + // The keepalive lands NOW, while the sweep is parked awaiting the + // inner lock. These are exactly the two steps of + // `worker_keep_alive_received` -> `refresh_worker`: refresh the local + // lifetime under the inner lock, then update the registry heartbeat. + inner_guard + .refresh_lifetime(&worker_id, NOW, None, true) + .expect("refresh_lifetime failed"); + registry + .update_worker_heartbeat(&worker_id, UNIX_EPOCH + Duration::from_secs(NOW)) + .await; + + // Release gate 2: the sweep re-acquires the inner lock and applies its + // stale decision without re-checking either liveness source. + drop(inner_guard); + sweep + .await + .expect("sweep panicked") + .expect("remove_timedout_workers failed"); + + // The worker heartbeated at t=NOW - both the local timestamp and the + // registry say it is alive - yet the stale sweep evicted it. + assert!( + scheduler.contains_worker_for_test(&worker_id).await, + "live worker (keepalive at t={NOW}, timeout {WORKER_TIMEOUT_S}s) \ + was evicted by a stale timeout decision" + ); + } +} diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 3d6ced82639..49c76ab4f91 100644 --- a/nativelink-scheduler/src/cache_lookup_scheduler.rs +++ b/nativelink-scheduler/src/cache_lookup_scheduler.rs @@ -231,8 +231,15 @@ impl CacheLookupScheduler { let inflight_cache_checks = self.inflight_cache_checks.clone(); // We need this spawn because we are returning a stream and this spawn will populate the stream's data. background_spawn!("cache_lookup_scheduler_add_action", async move { - // If our spawn ever dies, we will remove the action from the inflight_cache_checks map. - let _scope_guard = scope_guard; + // If our spawn ever dies before we take ownership of the map entry, remove the + // action from the inflight_cache_checks map (drop-safety). We defuse this guard + // the moment the cache check returns (below): from that point we remove the entry + // ourselves on every path with no intervening `.await`, so the guard's stale + // remove must NOT fire afterwards — a concurrent same-key request may have + // re-inserted a fresh entry in the window between our remove and this spawn's + // exit, and letting the guard run would orphan it (its oneshot tx is dropped, + // surfacing a spurious "tx hung up" to that client and forcing a retry). + let scope_guard = scope_guard; let unique_key = match &action_info.unique_qualifier { ActionUniqueQualifier::Cacheable(unique_key) => unique_key, @@ -254,6 +261,11 @@ impl CacheLookupScheduler { action_info.unique_qualifier.digest_function(), ) .await; + // The cache check has returned; from here every path removes our own map entry + // with no intervening `.await` (no drop point), so defusing the drop-guard now is + // leak-free and prevents its stale remove from clobbering a concurrent same-key + // re-insertion after we remove our entry. + scopeguard::ScopeGuard::into_inner(scope_guard); match maybe_action_result { Ok(action_result) => { let maybe_pending_txs = { @@ -303,10 +315,10 @@ impl CacheLookupScheduler { } Err(err) => { // A NotFound just means we need to execute the action. Any - // other lookup failure (an unavailable AC backend, a decode - // fault) degrades to execution too: a cache fault must never - // fail the request, only cost a redundant execution. The - // error itself still reaches operators via the log. + // other lookup failure (an unavailable AC backend, a timeout) + // degrades to execution too: a cache fault must never fail + // the request, only cost a redundant execution. The error + // itself still reaches operators via the log. if err.code != Code::NotFound { warn!(?err, "Cache lookup failed; falling through to execution"); } diff --git a/nativelink-scheduler/src/memory_awaited_action_db.rs b/nativelink-scheduler/src/memory_awaited_action_db.rs index 384b0237277..40ccb410f37 100644 --- a/nativelink-scheduler/src/memory_awaited_action_db.rs +++ b/nativelink-scheduler/src/memory_awaited_action_db.rs @@ -19,14 +19,13 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use async_lock::Mutex; -use futures::{FutureExt, Stream}; +use futures::{FutureExt, Stream, stream}; use nativelink_config::stores::EvictionPolicy; use nativelink_error::{Code, Error, ResultExt, error_if, make_err}; use nativelink_metric::MetricsComponent; use nativelink_util::action_messages::{ ActionInfo, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, }; -use nativelink_util::chunked_stream::ChunkedStream; use nativelink_util::evicting_map::{EvictingMap, LenEntry}; use nativelink_util::instant_wrapper::InstantWrapper; use nativelink_util::metrics::{ @@ -637,6 +636,20 @@ impl I + Clone + Send + Sync> AwaitedActionDbI } new_awaited_action.increment_version(); + // Client keep-alive stamps are written in place (see + // `handle_action_events` and `try_subscribe`) without bumping the + // version, so a read-modify-write cycle that snapshotted before + // the stamp still passes the version check above. Merge the stamp + // forward so such a write can never move the keep-alive backwards + // and cause the client-timeout sweep to kill an operation whose + // client is alive and polling. + if old_awaited_action.last_client_keepalive_timestamp() + > new_awaited_action.last_client_keepalive_timestamp() + { + new_awaited_action + .update_client_keep_alive(old_awaited_action.last_client_keepalive_timestamp()); + } + error_if!( old_awaited_action.action_info().unique_qualifier != new_awaited_action.action_info().unique_qualifier, @@ -979,28 +992,25 @@ impl I + Clone + Send + Sync + 'static> Awaite .await } - fn get_all_awaited_actions( + /// The listing is read whole, under one hold of the lock taken here, + /// before the stream exists, and the stream then only hands out what + /// was read. The lock must not be taken inside the stream: a reader + /// that prefetches rows (the filter pass reads 64 at a time) would be + /// left with a lock wait parked inside the stream while it hands rows + /// on, and if the reader then writes, on the same task, it waits for + /// the same lock; the release goes to the parked wait, which nothing + /// polls, and the scheduler stops for good. + async fn get_all_awaited_actions( &self, - ) -> impl Future>, Error>> - { - std::future::ready(Ok(ChunkedStream::new( - Bound::Unbounded, - Bound::Unbounded, - move |start, end, mut output| async move { - let inner = self.inner.lock().await; - let mut maybe_new_start = None; - - for (operation_id, item) in - inner.get_awaited_actions_range(start.as_ref(), end.as_ref()) - { - output.push_back(item); - maybe_new_start = Some(operation_id); - } - - Ok(maybe_new_start - .map(|new_start| ((Bound::Excluded(new_start.clone()), end), output))) - }, - ))) + ) -> Result> + Send, Error> { + let listed: Vec> = { + let inner = self.inner.lock().await; + inner + .get_awaited_actions_range(Bound::Unbounded, Bound::Unbounded) + .map(|(_, item)| Ok(item)) + .collect() + }; + Ok(stream::iter(listed)) } async fn get_by_operation_id( @@ -1010,52 +1020,32 @@ impl I + Clone + Send + Sync + 'static> Awaite Ok(self.inner.lock().await.get_by_operation_id(operation_id)) } - fn get_range_of_actions( + /// Read whole under one hold of the lock, like `get_all_awaited_actions`, + /// and for the same reason. + async fn get_range_of_actions( &self, state: SortedAwaitedActionState, start: Bound, end: Bound, desc: bool, - ) -> impl Future> + Send, Error>> - { - std::future::ready(Ok(ChunkedStream::new( - start, - end, - move |start, end, mut output| async move { - let inner = self.inner.lock().await; - let mut done = true; - let mut new_start = start.as_ref(); - let mut new_end = end.as_ref(); - - let iterator = inner - .get_range_of_actions(state, (start.as_ref(), end.as_ref())) - .map(|res| res.err_tip(|| "In AwaitedActionDb::get_range_of_actions")); - - // TODO(palfrey) This should probably use the `.left()/right()` pattern, - // but that doesn't exist in the std or any libraries we use. - if desc { - for result in iterator.rev() { - let (sorted_awaited_action, item) = - result.err_tip(|| "In AwaitedActionDb::get_range_of_actions")?; - output.push_back(item); - new_end = Bound::Excluded(sorted_awaited_action); - done = false; - } - } else { - for result in iterator { - let (sorted_awaited_action, item) = - result.err_tip(|| "In AwaitedActionDb::get_range_of_actions")?; - output.push_back(item); - new_start = Bound::Excluded(sorted_awaited_action); - done = false; - } - } - if done { - return Ok(None); - } - Ok(Some(((new_start.cloned(), new_end.cloned()), output))) - }, - ))) + ) -> Result> + Send, Error> { + let listed: Vec> = { + let inner = self.inner.lock().await; + let iterator = inner + .get_range_of_actions(state, (start.as_ref(), end.as_ref())) + .map(|res| { + res.map(|(_, item)| item) + .err_tip(|| "In AwaitedActionDb::get_range_of_actions") + }); + // TODO(palfrey) This should probably use the `.left()/right()` pattern, + // but that doesn't exist in the std or any libraries we use. + if desc { + iterator.rev().collect() + } else { + iterator.collect() + } + }; + Ok(stream::iter(listed)) } async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 69afdd02272..e7fef86f245 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/nativelink-scheduler/src/simple_scheduler.rs @@ -49,7 +49,7 @@ use opentelemetry::baggage::BaggageExt; use opentelemetry::context::{Context, FutureExt as OtelFutureExt}; use opentelemetry_semantic_conventions::attribute::ENDUSER_ID; use parking_lot::Mutex; -use tokio::sync::{Notify, mpsc}; +use tokio::sync::{Notify, mpsc, oneshot}; use tokio::time::Duration; use tracing::{debug, error, info, info_span, warn}; use uuid::Uuid; @@ -87,6 +87,11 @@ const FLEET_RECORD_TTL_INTERVALS: u32 = 3; /// version-conflict backoff sleeps a single pass can incur. const MAX_UNSATISFIABLE_FAILS_PER_PASS: usize = 16; +/// Test-only: milliseconds to stall inside the `Arc::new_cyclic` closure +/// right after the matching-engine task is spawned. See the seam below. +#[cfg(test)] +static TEST_CONSTRUCTOR_PAUSE_MS: AtomicUsize = AtomicUsize::new(0); + /// A backlog larger than `floor * DIVISOR` raises the cap to `backlog / /// DIVISOR`, so a large single-shape burst drains in ~DIVISOR back-to-back /// passes rather than `backlog / floor` of them (whose tail could outlast a @@ -680,15 +685,28 @@ impl SimpleScheduler { action_state.client_operation_id.clone() }; - // Tell the matching engine that the operation is being assigned to a worker. + // Tell the matching engine that the operation is being assigned + // to a worker, validating that the requirements this worker was + // chosen for still hold. A peer that requeued the action with a + // raised reservation (memory escalation) between our read and + // now must not have our stale, smaller placement dispatched at + // it; the accept rejects such a placement with Aborted. let assign_result = matching_engine_state_manager - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation( + &operation_id, + Ok(&worker_id), + Some(&action_info.inner.platform_properties), + ) .await .err_tip(|| "Failed to assign operation in do_try_match"); if let Err(err) = assign_result { if err.code == Code::Aborted { - // If the operation was aborted, it means that the operation was - // cancelled due to another operation being assigned to the worker. + // Aborted here means either another operation was + // assigned to the worker first, or our placement is + // stale because the action's requirements changed after + // we chose the worker. Either way nothing was sent; the + // action stays queued and is rematched next pass against + // its current requirements. return Ok(Dispatch::LostRace); } // Any other error is a real error. @@ -1459,10 +1477,25 @@ impl SimpleScheduler { }) }; + // The matching task is spawned inside `Arc::new_cyclic`, during which + // `weak_inner.upgrade()` returns `None` (strong count is still zero). + // The task treats a failed upgrade as "shutting down", so if it gets + // polled before the closure finishes (e.g. `task_change_notify` + // already holds a permit from a pubsub event that raced construction), + // it would consume that permit, see `None`, and exit permanently. + // Gate the loop on this signal, sent after `new_cyclic` returns, so + // the task cannot observe a partially constructed scheduler. + let (constructed_tx, constructed_rx) = oneshot::channel::<()>(); let action_scheduler = Arc::new_cyclic(move |weak_self| -> Self { let weak_inner = weak_self.clone(); let task_worker_matching_spawn = spawn!("simple_scheduler_task_worker_matching", async move { + // Wait for construction to complete. An `Err` means the + // sender was dropped without sending, which cannot happen: + // it is sent unconditionally right after `new_cyclic`. + if constructed_rx.await.is_err() { + return; + } let mut last_match_successful = true; let mut worker_match_logging_last: Option = None; // Break out of the loop only when the inner is dropped. @@ -1540,6 +1573,19 @@ impl SimpleScheduler { // Unreachable. }); + // Test-only timing seam: widens the window between spawning the + // matching task and the end of `Arc::new_cyclic` (during which + // `weak_inner.upgrade()` returns `None` because the strong count + // is still zero). This changes no logic, only timing, so a test + // can deterministically exercise the constructor race. + #[cfg(test)] + { + let pause_ms = TEST_CONSTRUCTOR_PAUSE_MS.load(Ordering::Acquire); + if pause_ms > 0 { + std::thread::sleep(Duration::from_millis(pause_ms as u64)); + } + } + let worker_match_logging_interval = match spec.worker_match_logging_interval_s { // -1 or 0 means disabled (0 used to cause expensive logging on every call) -1 | 0 => None, @@ -1576,6 +1622,9 @@ impl SimpleScheduler { now_fn: scheduler_now_fn, } }); + // Release the matching-engine task now that `self` is fully + // constructed and weak upgrades succeed. + let _ = constructed_tx.send(()); (action_scheduler, worker_scheduler_clone) } } @@ -1754,3 +1803,86 @@ mod unsatisfiable_fail_cap_tests { assert!(unsatisfiable_fail_cap(over) > MAX_UNSATISFIABLE_FAILS_PER_PASS); } } + +#[cfg(test)] +mod constructor_race_tests { + use core::sync::atomic::Ordering; + + use nativelink_config::schedulers::SimpleSpec; + use nativelink_config::stores::EvictionPolicy; + use nativelink_util::instant_wrapper::MockInstantWrapped; + use tokio::sync::{Notify, mpsc}; + use tokio::time::Duration; + + use super::{Arc, SimpleScheduler, TEST_CONSTRUCTOR_PAUSE_MS}; + use crate::memory_awaited_action_db::MemoryAwaitedActionDb; + + /// Repro for the spawn-then-configure race in `SimpleScheduler::new*`: + /// the matching-engine task is spawned inside `Arc::new_cyclic`, holding + /// only a `Weak` whose `upgrade()` returns `None` until the closure + /// finishes. If `task_change_notify` already holds a permit (as the + /// Redis pubsub path can arrange before the constructor runs), the task's + /// first poll completes instantly, sees `upgrade() == None`, and + /// interprets "not yet constructed" as "shutting down": it returns and + /// the scheduler never matches anything again. + #[expect( + clippy::disallowed_methods, + reason = "the race needs a multi-thread runtime so the spawned task \ + is polled while the constructor thread is blocked" + )] + #[test] + fn matching_engine_survives_prestored_notify_permit_during_construction() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + rt.block_on(async { + // Widen the in-closure window so the freshly spawned task is + // guaranteed to get its first poll while the strong count is 0. + TEST_CONSTRUCTOR_PAUSE_MS.store(500, Ordering::Release); + + let task_change_notify = Arc::new(Notify::new()); + // A permit stored before construction, exactly like + // pull_task_change_subscriber's notify_one() on a Redis pubsub + // event that arrives before SimpleScheduler::new. + task_change_notify.notify_one(); + + let db = MemoryAwaitedActionDb::new( + &EvictionPolicy::default(), + task_change_notify.clone(), + MockInstantWrapped::default, + ); + let (tx, mut rx) = mpsc::unbounded_channel::<()>(); + let (_scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec::default(), + db, + move || { + // Runs after every matching pass; proves the engine is alive. + let tx = tx.clone(); + async move { + let _ = tx.send(()); + } + }, + task_change_notify.clone(), + MockInstantWrapped::default, + None, + ); + TEST_CONSTRUCTOR_PAUSE_MS.store(0, Ordering::Release); + + // Nudge it again post-construction; a live engine must run a + // matching pass for either the pre-stored permit or this one. + task_change_notify.notify_one(); + // `Ok(None)` means the channel closed: the matching task (which + // owns the only sender) returned. `Err(_)` means it is alive but + // never ran a pass. Either way it is dead or wedged. + let matched = tokio::time::timeout(Duration::from_secs(2), rx.recv()).await; + assert!( + matched == Ok(Some(())), + "matching engine never ran a pass: it exited during \ + Arc::new_cyclic because weak_inner.upgrade() was None \ + (strong count still 0), mistaking construction for shutdown" + ); + }); + } +} diff --git a/nativelink-scheduler/src/simple_scheduler_state_manager.rs b/nativelink-scheduler/src/simple_scheduler_state_manager.rs index 656f7b68826..c53bc6da5e3 100644 --- a/nativelink-scheduler/src/simple_scheduler_state_manager.rs +++ b/nativelink-scheduler/src/simple_scheduler_state_manager.rs @@ -673,6 +673,7 @@ where Code::Unavailable, "The worker finished this operation but its result never reached the scheduler" )), + None, ) .await { @@ -1147,6 +1148,7 @@ where self.assign_operation( operation_id, Err(make_err!(Code::DeadlineExceeded, "Operation {cause}")), + None, ) .await } @@ -1156,6 +1158,7 @@ where operation_id: &OperationId, maybe_worker_id: Option<&WorkerId>, update: UpdateOperationType, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error> { let update_type_str = match &update { UpdateOperationType::KeepAlive => "KeepAlive", @@ -1289,6 +1292,31 @@ where warn!(state = ?awaited_action.state(), "Action already assigned"); return Err(make_err!(Code::Aborted, "Action already assigned")); } + // Placement validation: the matcher chose this worker for a + // snapshot of the action's requirements. If the action's + // current requirements no longer match that snapshot — an + // escalation or other requeue raised the reservation after + // the worker was picked — the placement is stale. Accepting + // it would flip the action to Executing while the matcher + // dispatches the old, smaller reservation to the worker, + // silently bypassing the escalation. Reject with Aborted so + // the matcher treats it as a lost race and the action is + // rematched against its current requirements. + if stage == &ActionStage::Executing + && let Some(expected) = expected_platform_properties + && &awaited_action.action_info().platform_properties != expected + { + warn!( + %operation_id, + expected = ?expected, + current = ?awaited_action.action_info().platform_properties, + "Rejecting stale placement: requirements changed since the worker was chosen" + ); + return Err(make_err!( + Code::Aborted, + "Stale placement for {operation_id}: action requirements changed since the worker was chosen; rematch required" + )); + } stage.clone() } UpdateOperationType::UpdateWithError(err) => { @@ -1761,7 +1789,7 @@ where worker_id: &WorkerId, update: UpdateOperationType, ) -> Result<(), Error> { - self.inner_update_operation(operation_id, Some(worker_id), update) + self.inner_update_operation(operation_id, Some(worker_id), update, None) .await } @@ -1821,6 +1849,7 @@ where &self, operation_id: &OperationId, worker_id_or_reason_for_unassign: Result<&WorkerId, Error>, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error> { let (maybe_worker_id, update) = match worker_id_or_reason_for_unassign { Ok(worker_id) => ( @@ -1829,8 +1858,13 @@ where ), Err(err) => (None, UpdateOperationType::UpdateWithError(err)), }; - self.inner_update_operation(operation_id, maybe_worker_id, update) - .await + self.inner_update_operation( + operation_id, + maybe_worker_id, + update, + expected_platform_properties, + ) + .await } async fn fail_queued_operation( diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 01d486798ee..a21c7cb18c0 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -56,6 +56,13 @@ const MAX_RETRIES_FOR_CLIENT_KEEPALIVE: u32 = 8; /// Use separate non-versioned Redis key for client keepalives. const USE_SEPARATE_CLIENT_KEEPALIVE_KEY: bool = true; +/// How long to wait before retrying a client keepalive write that failed +/// (e.g. a transient store error). The local `last_known_keepalive_ts` is +/// only advanced on a successful write, so failures naturally re-arm the +/// keepalive gate; this delay bounds the retry frequency so a hard store +/// outage does not spin the subscriber loop. +const CLIENT_KEEPALIVE_RETRY_DELAY: Duration = Duration::from_secs(1); + /// How often the actions in each stage are recounted for /// `execution.active.count`. const ACTIVE_COUNT_REFRESH_INTERVAL: Duration = Duration::from_secs(15); @@ -309,6 +316,10 @@ where loop { // This is set if the maybe_last_state doesn't match the state in the store. let mut maybe_changed_action = None; + // Set when a client keepalive store write fails so the sleep + // below becomes a bounded retry delay instead of a full + // CLIENT_KEEPALIVE_DURATION park. + let mut keepalive_write_failed = false; let last_known_keepalive_ts = self.last_known_keepalive_ts.load(Ordering::Acquire); // Only a subscriber that stands for a client may say a client is @@ -348,18 +359,27 @@ where ) .await; - if let Err(e) = update_result { - warn!( - ?self.subscription_key, - ?e, - "Failed to update client keepalive (non-versioned)" - ); + match update_result { + Ok(_) => { + // Only advance the local timestamp once the store + // actually holds the new keepalive. Advancing it on + // failure would suppress retries for a full + // CLIENT_KEEPALIVE_DURATION while the store's stamp + // keeps aging, letting a peer replica retire this + // live client's action as abandoned. + self.last_known_keepalive_ts + .store(now_ts, Ordering::Release); + } + Err(e) => { + keepalive_write_failed = true; + warn!( + ?self.subscription_key, + ?e, + "Failed to update client keepalive (non-versioned), will retry" + ); + } } - // Update local timestamp - self.last_known_keepalive_ts - .store(now_ts, Ordering::Release); - // Check if state changed (for unreliable subscription managers) if self.maybe_last_stage.is_some() { let awaited_action = Self::inner_get_awaited_action( @@ -430,12 +450,19 @@ where Some(core::mem::discriminant(&changed_action.state().stage)); return Ok(changed_action); } - // Determine the sleep time based on the last client keep alive. - let sleep_time = CLIENT_KEEPALIVE_DURATION - .checked_sub( - I::from_secs(self.last_known_keepalive_ts.load(Ordering::Acquire)).elapsed(), - ) - .unwrap_or(Duration::from_millis(100)); + // Determine the sleep time based on the last successful client + // keep alive write. If the last write failed, wake after a + // bounded retry delay instead of parking for the full period. + let sleep_time = if keepalive_write_failed { + CLIENT_KEEPALIVE_RETRY_DELAY + } else { + CLIENT_KEEPALIVE_DURATION + .checked_sub( + I::from_secs(self.last_known_keepalive_ts.load(Ordering::Acquire)) + .elapsed(), + ) + .unwrap_or(Duration::from_millis(100)) + }; tokio::select! { result = &mut changed_fut => { result?; @@ -662,6 +689,19 @@ impl SchedulerIndexProvider for SearchStateToAwaitedAction { fn index_value(&self) -> Cow<'_, str> { Cow::Borrowed(self.0) } + + /// A queued record whose data cannot be decoded may be removed, but + /// only while no client keepalive exists for it: the keepalive is + /// written with a TTL of the client action timeout, so its absence + /// means no client has been listening for at least that long. Scoped + /// to the queued listing; records in any other stage are left alone. + fn lost_record_liveness_guard_key(index_value: &str, record_key: &str) -> Option { + if index_value != get_state_prefix(SortedAwaitedActionState::Queued) { + return None; + } + let operation_id = record_key.strip_prefix(OPERATION_ID_TO_AWAITED_ACTION_KEY_PREFIX)?; + Some(format!("{CLIENT_KEEPALIVE_KEY_PREFIX}{operation_id}")) + } } impl SchedulerStoreDecodeTo for SearchStateToAwaitedAction { type DecodeOutput = AwaitedAction; @@ -1003,15 +1043,40 @@ where tokio::pin!(stream); // This index search is by prefix. In particular, an unscoped key // is a prefix of its scoped variants and must not join those jobs. + // + // The index retains completed records alongside any live successor + // for the same qualifier, and the backend does not order results by + // state. Stopping at the first exact match can therefore hand back a + // finished record while a joinable live one is also present, forcing + // an unnecessary recreate. Scan every exact match and prefer a + // joinable (neither finished nor abandoned) candidate; keep the + // first exact match only as a fallback when none is joinable. + let mut fallback: Option = None; while let Some(candidate) = stream .try_next() .await .err_tip(|| "In RedisAwaitedActionDb::try_subscribe")? { - if &candidate.action_info().unique_qualifier == unique_qualifier { + if &candidate.action_info().unique_qualifier != unique_qualifier { + continue; + } + let abandoned = self + .executing_action_is_abandoned( + &candidate, + no_event_action_timeout, + (self.now_fn)().now(), + ) + .await; + if !candidate.state().stage.is_finished() && !abandoned { maybe_awaited_action = Some(candidate); break; } + if fallback.is_none() { + fallback = Some(candidate); + } + } + if maybe_awaited_action.is_none() { + maybe_awaited_action = fallback; } if maybe_awaited_action.is_some() { break; diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index cf6f65ddb38..6481f3da032 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -17,7 +17,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use nativelink_error::{Code, Error, ResultExt}; +use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_metric::MetricsComponent; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ ActionResourceUsage, ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, @@ -54,8 +54,8 @@ pub struct ActionInfoWithProps { #[derive(Debug)] pub enum WorkerUpdate { /// Requests that the worker begin executing this action, dispatched at - /// this scheduler time. - RunAction(Box<(OperationId, ActionInfoWithProps, WorkerTimestamp)>), + /// this scheduler time with this scheduler-wide dispatch generation. + RunAction(Box<(OperationId, ActionInfoWithProps, WorkerTimestamp, u64)>), /// Request that the worker is no longer in the pool and may discard any jobs. Disconnect, @@ -79,6 +79,15 @@ pub struct PendingActionInfoData { /// The worker's last-seen timestamp when the dispatch was sent. #[metric(help = "When this operation was dispatched to the worker.")] pub dispatched_at: WorkerTimestamp, + /// The identity of this dispatch: a scheduler-wide monotonic counter + /// allocated under the scheduler mutex at send time. `dispatched_at` + /// has second resolution and so cannot distinguish a requeue and + /// re-dispatch inside one wall-clock second; the generation is strictly + /// increasing per dispatch, so a snapshot of it names exactly one + /// dispatch. Scheduler-wide (not per-worker) so a worker that + /// disconnects and re-registers under the same id can never reissue a + /// generation an old snapshot still holds. + pub dispatch_generation: u64, /// The worker said it took the action. A worker that does not speak /// the acknowledgement never sets this, which is why the /// unacknowledged sweep is opt-in. @@ -285,9 +294,14 @@ impl Worker { pub async fn notify_update(&mut self, worker_update: WorkerUpdate) -> Result<(), Error> { match worker_update { WorkerUpdate::RunAction(action) => { - let (operation_id, action_info, dispatched_at) = *action; - self.run_action(operation_id, action_info, dispatched_at) - .await + let (operation_id, action_info, dispatched_at, dispatch_generation) = *action; + self.run_action( + operation_id, + action_info, + dispatched_at, + dispatch_generation, + ) + .await } WorkerUpdate::Disconnect => { self.metrics.notify_disconnect.inc(); @@ -415,7 +429,21 @@ impl Worker { operation_id: OperationId, action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, + dispatch_generation: u64, ) -> Result<(), Error> { + // Reject-if-exists: a same-worker retry can dispatch an op whose + // earlier attempt is still resident here (e.g. an execution-deadline + // requeue that left the ledger entry in place). Overwriting the + // op-keyed entry would deduct the platform budget a second time while + // completion only refunds the surviving entry, leaking the first + // reservation permanently. Refuse the redundant dispatch instead. + if self.running_action_infos.contains_key(&operation_id) { + return Err(make_err!( + Code::AlreadyExists, + "Worker {} already holds a reservation for operation {operation_id}; refusing redundant dispatch", + self.id, + )); + } let tx = &mut self.tx; let worker_platform_properties = &mut self.platform_properties; let running_action_infos = &mut self.running_action_infos; @@ -447,6 +475,7 @@ impl Worker { kill_requested_at: None, missed_reports: 0, dispatched_at, + dispatch_generation, accepted: false, last_usage: None, }, diff --git a/nativelink-scheduler/src/worker_registry.rs b/nativelink-scheduler/src/worker_registry.rs index 08cbcce3da1..4e645a4d672 100644 --- a/nativelink-scheduler/src/worker_registry.rs +++ b/nativelink-scheduler/src/worker_registry.rs @@ -67,10 +67,22 @@ impl WorkerRegistry { } } - /// Updates the heartbeat timestamp for a worker. + /// Advances a worker's heartbeat to `now`. Monotone: a late or out-of-order + /// keepalive never regresses `last_seen`. The liveness signal is the MAX of + /// receipt times — a monotone fact — and must not move backward and + /// manufacture a false `Stale` (pairs with the TL1 fix in + /// `ApiWorkerScheduler::refresh_worker`, which now stamps this on receipt, + /// possibly ahead of `refresh_lifetime`'s own out-of-order rejection). pub async fn update_worker_heartbeat(&self, worker_id: &WorkerId, now: SystemTime) { let mut workers = self.workers.write().await; - workers.insert(worker_id.clone(), now); + workers + .entry(worker_id.clone()) + .and_modify(|prev| { + if now > *prev { + *prev = now; + } + }) + .or_insert(now); trace!(?worker_id, now = %humantime::format_rfc3339(now), "FLOW: Worker heartbeat updated in registry"); } @@ -126,6 +138,15 @@ impl WorkerRegistry { self.check_liveness(worker_id, timeout, now).await == WorkerLiveness::Alive } + /// Test-only: hold the internal write lock so that concurrent liveness + /// reads park at a deterministic await point. + #[cfg(test)] + pub(crate) async fn write_lock_for_test( + &self, + ) -> async_lock::RwLockWriteGuard<'_, HashMap> { + self.workers.write().await + } + pub async fn get_worker_last_seen(&self, worker_id: &WorkerId) -> Option { let workers = self.workers.read().await; workers.get(worker_id).copied() @@ -234,4 +255,40 @@ mod tests { .await ); } + + // TL1 regression (monotone half): an out-of-order / late keepalive must not + // regress `last_seen` and manufacture a false `Stale`. Fails WITHOUT the + // monotone fix (plain `insert` overwrites with the older timestamp, so the + // liveness check past the old deadline reads Stale). + #[nativelink_test] + async fn heartbeat_is_monotone_never_regresses() { + let registry = WorkerRegistry::new(); + let w = WorkerId::from(String::from("w")); + let timeout = Duration::from_secs(5); + let t0 = SystemTime::now(); + let t_late = t0.checked_add(Duration::from_secs(10)).unwrap(); + + registry.register_worker(&w, t0).await; + registry.update_worker_heartbeat(&w, t_late).await; // advance + registry.update_worker_heartbeat(&w, t0).await; // out-of-order (older) — must NOT regress + + assert_eq!( + registry.get_worker_last_seen(&w).await, + Some(t_late), + "an older keepalive regressed last_seen — manufactures false Stale (TL1)" + ); + // A check 1s past the OLD deadline (t0 + timeout) must still be Alive, + // because the real last_seen is t_late. + let past_old_deadline = t0 + .checked_add(timeout) + .unwrap() + .checked_add(Duration::from_secs(1)) + .unwrap(); + assert_eq!( + registry + .check_liveness(&w, timeout, past_old_deadline) + .await, + WorkerLiveness::Alive + ); + } } diff --git a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index cbff20b9867..dc0a295ae47 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -32,12 +32,14 @@ use nativelink_scheduler::cache_lookup_scheduler::CacheLookupScheduler; use nativelink_scheduler::mock_scheduler::MockActionScheduler; use nativelink_store::memory_store::MemoryStore; use nativelink_util::action_messages::{ - ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, + ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, }; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; -use nativelink_util::operation_state_manager::{ClientStateManager, OperationFilter}; +use nativelink_util::operation_state_manager::{ + ActionStateResult, ClientStateManager, OperationFilter, +}; use nativelink_util::store_trait::{ RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; @@ -48,11 +50,10 @@ use tokio::{self}; use tokio_stream::StreamExt; use utils::scheduler_utils::{TokioWatchActionStateResult, make_base_action_info}; -/// An AC store whose reads always fail with a non-`NotFound` backend error. +/// An AC store whose operations always fail with a non-`NotFound` backend error. #[derive(Debug, MetricsComponent)] struct AlwaysErrStore { - #[metric(help = "Error code every operation fails with")] - code: u32, + code: Code, } #[async_trait] @@ -66,7 +67,7 @@ impl StoreDriver for AlwaysErrStore { _keys: &[StoreKey<'_>], _results: &mut [Option], ) -> Result<(), Error> { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } async fn update( @@ -75,7 +76,7 @@ impl StoreDriver for AlwaysErrStore { _reader: DropCloserReadHalf, _size_info: UploadSizeInfo, ) -> Result { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } async fn get_part( @@ -85,7 +86,7 @@ impl StoreDriver for AlwaysErrStore { _offset: u64, _length: Option, ) -> Result<(), Error> { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } fn inner_store(&self, _key: Option) -> &dyn StoreDriver { @@ -199,14 +200,14 @@ async fn add_action_handles_skip_cache() -> Result<(), Error> { Ok(()) } -/// A faulting AC backend (unavailable Redis, decode error) must degrade the +/// A faulting AC backend (an unavailable Redis, a timeout) must degrade the /// cache lookup to execution, not fail the Execute request: a cache fault /// costs a redundant execution, never a failed build. #[nativelink_test] async fn ac_fault_falls_through_to_execution() -> Result<(), Error> { let mock_scheduler = Arc::new(MockActionScheduler::new()); let ac_store = Store::new(Arc::new(AlwaysErrStore { - code: Code::Unavailable as u32, + code: Code::Unavailable, })); let cache_scheduler = CacheLookupScheduler::new(ac_store, mock_scheduler.clone())?; let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 123)); @@ -218,14 +219,20 @@ async fn ac_fault_falls_through_to_execution() -> Result<(), Error> { last_transition_timestamp: SystemTime::now(), })); let client_operation_id = OperationId::default(); - let (add_result, _) = join!( - cache_scheduler.add_action(client_operation_id.clone(), action_info.clone()), - mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( - client_operation_id, - action_info, - forward_watch_channel_rx - )))) - ); + // If the fault failed the request instead, the mock would never see the + // action and `expect_add_action` would wait forever; bound it. + let (add_result, _) = tokio::time::timeout(std::time::Duration::from_secs(5), async { + join!( + cache_scheduler.add_action(client_operation_id.clone(), action_info.clone()), + mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + forward_watch_channel_rx + )))) + ) + }) + .await + .expect("the AC fault was not forwarded to execution"); add_result.expect("AC backend fault must fall through to execution, not fail the client"); Ok(()) } @@ -253,3 +260,95 @@ async fn find_by_client_operation_id_call_passed() -> Result<(), Error> { ); Ok(()) } + +fn queued_result( + client_operation_id: OperationId, + action_info: Arc, +) -> Box { + let (_tx, rx) = watch::channel(Arc::new(ActionState { + client_operation_id: client_operation_id.clone(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::now(), + })); + Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + rx, + )) +} + +// A leader's cache-check task removes its in-flight entry before forwarding a +// miss to the inner scheduler, and that forward is an await point. A second +// request for the same key arriving in that window inserts a fresh entry. The +// leader's drop guard must not remove that fresh entry when the leader's task +// exits, or the second request's sender is dropped and it fails with +// "tx hung up". The default test runtime is single-threaded, so the +// interleaving below is deterministic. +#[nativelink_test] +async fn leader_drop_guard_does_not_orphan_a_later_same_key_request() -> Result<(), Error> { + let context = make_cache_scheduler()?; + let cache_scheduler = Arc::new(context.cache_scheduler); + let mock_scheduler = context.mock_scheduler; + // Nothing in the AC, so both requests miss and are forwarded. + let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 77)); + + // Request A becomes the leader, misses, removes its entry and parks in the + // forward to the mock scheduler, which has not answered yet. + let op_a = OperationId::default(); + let request_a = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_a, action_info) = (op_a.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_a, action_info) + .await + .map(|_| ()) + } + }); + for _ in 0..50 { + tokio::task::yield_now().await; + } + + // Request B for the same key arrives while A is parked: A's entry is gone, + // so B inserts its own. + let op_b = OperationId::default(); + let request_b = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_b, action_info) = (op_b.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_b, action_info) + .await + .map(|_| ()) + } + }); + + // Answer A's forward. A's task then finishes while B's entry is in the map. + let (forwarded_a, _) = mock_scheduler + .expect_add_action(Ok(queued_result(op_a.clone(), action_info.clone()))) + .await; + assert_eq!(forwarded_a, op_a); + + // B must still be forwarded and answered, not orphaned. + let serve_b = tokio::spawn({ + let mock_scheduler = mock_scheduler.clone(); + let action_info = action_info.clone(); + async move { + mock_scheduler + .expect_add_action(Ok(queued_result(op_b, action_info))) + .await + } + }); + let result_b = tokio::time::timeout(std::time::Duration::from_secs(5), request_b) + .await + .expect("request B never completed") + .expect("request B panicked"); + assert!( + result_b.is_ok(), + "a same-key request that arrived while the leader was forwarding was orphaned: {result_b:?}" + ); + serve_b.await.expect("serving B panicked"); + request_a.await.expect("request A panicked")?; + Ok(()) +} diff --git a/nativelink-scheduler/tests/client_keepalive_clobber_test.rs b/nativelink-scheduler/tests/client_keepalive_clobber_test.rs new file mode 100644 index 00000000000..bea779f915c --- /dev/null +++ b/nativelink-scheduler/tests/client_keepalive_clobber_test.rs @@ -0,0 +1,156 @@ +// 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. + +//! Repro: client keep-alive writes bypass the version CAS in the memory +//! backend, so a concurrent read-modify-write (a worker update) that +//! snapshotted the action before the keep-alive silently reverts +//! `last_client_keepalive_timestamp` — and passes the version check, so the +//! `Code::Aborted` retry defense in `SimpleSchedulerStateManager` never +//! fires. Repeated indefinitely, the client-timeout sweep kills an action +//! whose client is alive and sending keep-alives. + +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::SystemTime; + +use mock_instant::thread_local::MockClock; +use nativelink_config::stores::EvictionPolicy; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::awaited_action_db::{ + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, CLIENT_KEEPALIVE_DURATION, +}; +use nativelink_scheduler::memory_awaited_action_db::MemoryAwaitedActionDb; +use nativelink_util::action_messages::{ + ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; + +const INSTANCE_NAME: &str = "keepalive_clobber_instance"; + +fn make_action_info() -> Arc { + let insert = SystemTime::UNIX_EPOCH + Duration::from_secs(100); + Arc::new(ActionInfo { + command_digest: DigestInfo::new([7; 32], 1), + input_root_digest: DigestInfo::new([7; 32], 2), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: insert, + insert_timestamp: insert, + unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: DigestInfo::new([7; 32], 3), + }), + }) +} + +/// Extract `last_client_keepalive_timestamp` from the Debug rendering of an +/// `AwaitedAction` (the field has no public accessor). +fn keepalive_of(awaited_action: &AwaitedAction) -> String { + let debug = format!("{awaited_action:?}"); + let field = "last_client_keepalive_timestamp: "; + let start = debug.find(field).expect("field in Debug output") + field.len(); + let rest = &debug[start..]; + let end = rest.find('}').expect("SystemTime Debug closes with }") + 1; + rest[..end].to_string() +} + +#[nativelink_test] +async fn client_keepalive_survives_concurrent_same_version_update() -> Result<(), Error> { + let db = MemoryAwaitedActionDb::new( + &EvictionPolicy::default(), + Arc::new(Notify::new()), + MockInstantWrapped::default, + ); + + let client_id = OperationId::from("client-keepalive"); + // The subscriber returned here carries client_info: its changed() loop is + // what emits ActionEvent::ClientKeepAlive. + let mut client_subscriber = db + .add_action( + client_id.clone(), + make_action_info(), + Duration::from_mins(1), + ) + .await?; + let operation_id = client_subscriber.borrow().await?.operation_id().clone(); + + // The "worker" side: snapshot the action (version N, keepalive K0), the + // way SimpleSchedulerStateManager does before update_awaited_action. + let worker_subscriber = db + .get_by_operation_id(&operation_id) + .await? + .expect("operation must exist"); + let stale_snapshot = worker_subscriber.borrow().await?; + let keepalive_k0 = keepalive_of(&stale_snapshot); + + // Advance the mock clock so (a) the client's changed() loop decides it is + // time to send a keep-alive, and (b) the stamped timestamp K1 differs + // from K0. + MockClock::advance(CLIENT_KEEPALIVE_DURATION + Duration::from_secs(1)); + + // Drive the client subscriber: changed() immediately sends + // ActionEvent::ClientKeepAlive, then parks awaiting a watch change. + let changed_task = tokio::spawn(async move { client_subscriber.changed().await }); + + // Wait (deterministically, bounded) for the event-handler task to stamp + // the keep-alive into the stored AwaitedAction. + let keepalive_k1 = { + let mut stamped = None; + for _ in 0..1000 { + tokio::task::yield_now().await; + let current = worker_subscriber.borrow().await?; + let ts = keepalive_of(¤t); + if ts != keepalive_k0 { + stamped = Some(ts); + break; + } + } + stamped.expect("client keep-alive was never stamped into the db") + }; + changed_task.abort(); + + // The keep-alive did NOT bump the version, so the worker's stale + // pre-keep-alive snapshot still passes the version CAS... + let update_result = db.update_awaited_action(stale_snapshot).await; + + // ...and send_replace overwrites the whole struct, reverting the + // keep-alive to K0. Correct behavior would be either rejecting the stale + // write with Code::Aborted (version bump on keep-alive) or merging the + // keep-alive timestamp. Either way, after an accepted update the stored + // keep-alive must still be K1. + let after = worker_subscriber.borrow().await?; + let keepalive_after = keepalive_of(&after); + assert!( + update_result.is_err() || keepalive_after == keepalive_k1, + "Lost update: stale same-version write was accepted (result: {update_result:?}) and \ + reverted last_client_keepalive_timestamp from {keepalive_k1} back to {keepalive_after} \ + (K0 was {keepalive_k0}). The Aborted-retry defense in \ + SimpleSchedulerStateManager can never fire, so the client-timeout \ + sweep can kill an action with a live client." + ); + // Redundant with the assert above, but gives a crisp diff on failure. + if update_result.is_ok() { + assert_eq!(keepalive_after, keepalive_k1); + } + Ok(()) +} diff --git a/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs new file mode 100644 index 00000000000..d783faf2bc2 --- /dev/null +++ b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs @@ -0,0 +1,260 @@ +// Copyright 2024 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. + +// Repro for: a failed client-keepalive store write in +// `OperationSubscriber::changed` (USE_SEPARATE_CLIENT_KEEPALIVE_KEY branch) +// still advances the local `last_known_keepalive_ts`, so the retry gate +// (`elapsed() > CLIENT_KEEPALIVE_DURATION`) is suppressed for a full +// CLIENT_KEEPALIVE_DURATION even though the shared store's ck_* key was +// never written. Peer scheduler replicas judge client liveness solely from +// the store, so a couple of consecutive failed writes leave the store +// timestamp stale past `client_action_timeout` and the peer retires the +// action as abandoned while the client is actively listening. + +use core::sync::atomic::{AtomicUsize, Ordering}; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::SystemTime; + +use bytes::Bytes; +use futures::{Stream, stream}; +use nativelink_error::{Code, Error, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::awaited_action_db::{ + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, +}; +use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb; +use nativelink_util::action_messages::{ + ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::store_trait::{ + SchedulerCurrentVersionProvider, SchedulerIndexProvider, SchedulerStore, + SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider, + SchedulerSubscription, SchedulerSubscriptionManager, StoreKey, +}; +use tokio::sync::Notify; + +const INSTANCE_NAME: &str = "foo"; + +struct PendingSubscription; +impl SchedulerSubscription for PendingSubscription { + async fn changed(&mut self) -> Result<(), Error> { + futures::future::pending().await + } +} + +struct PendingSubscriptionManager; +impl SchedulerSubscriptionManager for PendingSubscriptionManager { + type Subscription = PendingSubscription; + fn subscribe(&self, _key: K) -> Result + where + K: SchedulerStoreKeyProvider, + { + Ok(PendingSubscription) + } + fn is_reliable() -> bool { + true + } +} + +/// A store holding one queued `AwaitedAction` whose `ck_*` (client +/// keepalive) writes always fail — a Redis blip / dead shard. Every other +/// operation succeeds. +struct KeepaliveWriteFailsStore { + operation_id: OperationId, + encoded_action: Bytes, + keepalive_write_attempts: AtomicUsize, +} + +impl SchedulerStore for KeepaliveWriteFailsStore { + type SubscriptionManager = PendingSubscriptionManager; + + fn subscription_manager( + &self, + ) -> impl Future, Error>> { + std::future::ready(Ok(Arc::new(PendingSubscriptionManager))) + } + + async fn update_data(&self, data: T, _expiry: Option) -> Result, Error> + where + T: SchedulerStoreDataProvider + + SchedulerStoreKeyProvider + + SchedulerCurrentVersionProvider + + Send, + { + let StoreKey::Str(key) = data.get_key() else { + panic!("unexpected non-string scheduler key"); + }; + if key.starts_with("ck_") { + self.keepalive_write_attempts.fetch_add(1, Ordering::SeqCst); + return Err(make_err!( + Code::Unavailable, + "injected transient store failure for client keepalive write" + )); + } + Ok(Some(1)) + } + + fn search_by_index_prefix( + &self, + _index: K, + ) -> impl Future< + Output = Result< + impl Stream::DecodeOutput, Error>> + Send, + Error, + >, + > + where + K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send, + ::DecodeOutput: Send, + { + std::future::ready(Ok(stream::empty())) + } + + async fn count_by_index_prefix(&self, _index: K) -> Result + where + K: SchedulerIndexProvider + Send, + { + Ok(0) + } + + async fn get_and_decode( + &self, + key: K, + ) -> Result::DecodeOutput>, Error> + where + K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send, + { + let StoreKey::Str(key) = key.get_key() else { + panic!("unexpected non-string scheduler key"); + }; + if key.starts_with("cid_") { + let bytes = Bytes::from(serde_json::to_vec(&self.operation_id).unwrap()); + return Ok(Some(K::decode(0, bytes)?)); + } + if key.starts_with("aa_") { + return Ok(Some(K::decode(1, self.encoded_action.clone())?)); + } + if key.starts_with("ck_") { + // No keepalive write ever landed in the store. + return Ok(None); + } + panic!("unexpected key {key}"); + } +} + +fn make_queued_action(operation_id: &OperationId) -> AwaitedAction { + let action_info = Arc::new(ActionInfo { + command_digest: DigestInfo::zero_digest(), + input_root_digest: DigestInfo::zero_digest(), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: SystemTime::UNIX_EPOCH, + insert_timestamp: SystemTime::UNIX_EPOCH, + unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: DigestInfo::zero_digest(), + }), + }); + // The action's embedded keepalive timestamp is at the epoch: from the + // shared store's point of view the client has not checked in for a very + // long time, exactly as after a couple of failed ck_ writes. + AwaitedAction::new(operation_id.clone(), action_info, SystemTime::UNIX_EPOCH) +} + +/// A live client's subscriber whose keepalive store write fails must retry +/// the write well before peer replicas' `client_action_timeout` elapses; +/// otherwise the store timestamp stays stale and +/// `sweep_abandoned_queued_actions` / `apply_filter_predicate` on another +/// replica retires the action as "no more clients listening" out from under +/// a client that is actively listening. +/// +/// Uses `start_paused` tokio time: the subscriber's sleep between keepalive +/// attempts is virtual, so 30 virtual seconds (3x `CLIENT_KEEPALIVE_DURATION`) +/// elapse instantly, while the retry gate — which reads the *local* +/// `last_known_keepalive_ts` through real wall-clock `SystemTime::elapsed` — +/// sees essentially zero elapsed real time after the failed write advanced +/// it. This exposes exactly the bug: the local timestamp is advanced even +/// though the store write failed, so no retry happens for a full +/// `CLIENT_KEEPALIVE_DURATION`. +type TestDb = StoreAwaitedActionDb< + KeepaliveWriteFailsStore, + fn() -> OperationId, + SystemTime, + fn() -> SystemTime, +>; + +fn new_op_id() -> OperationId { + OperationId::from("unused-new-operation") +} + +#[nativelink_test(start_paused = true)] +async fn failed_keepalive_write_must_be_retried() -> Result<(), Error> { + let operation_id = OperationId::from("live-client-operation"); + let action = make_queued_action(&operation_id); + let store = Arc::new(KeepaliveWriteFailsStore { + operation_id: operation_id.clone(), + encoded_action: Bytes::from(serde_json::to_vec(&action).unwrap()), + keepalive_write_attempts: AtomicUsize::new(0), + }); + + let now_fn: fn() -> SystemTime = SystemTime::now; + let op_id_fn: fn() -> OperationId = new_op_id; + let db: TestDb = StoreAwaitedActionDb::new( + store.clone(), + Arc::new(Notify::new()), + now_fn, + op_id_fn, + 60, + 60, + false, + ) + .await + .expect("construct test db"); + + let mut subscriber = db + .get_awaited_action_by_id(&operation_id) + .await? + .expect("subscriber for live client"); + + // Drive the client's changed() loop for 30 virtual seconds — three full + // CLIENT_KEEPALIVE_DURATION periods. The subscription never fires + // (pending) and the action never changes, so the only thing the loop can + // do is maintain the keepalive. + let changed_fut = subscriber.changed(); + tokio::pin!(changed_fut); + tokio::select! { + result = &mut changed_fut => { + panic!("changed() must not resolve: nothing changed; got {result:?}"); + } + () = tokio::time::sleep(Duration::from_secs(30)) => {} + } + + let n = store.keepalive_write_attempts.load(Ordering::SeqCst); + assert!( + n >= 2, + "keepalive write failed but was never retried within 3x \ + CLIENT_KEEPALIVE_DURATION (attempts = {n}); the local \ + last_known_keepalive_ts was advanced despite the store Err, so the \ + shared-store keepalive stays stale and a peer replica will retire \ + this live client's action as abandoned (DeadlineExceeded)" + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs b/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs new file mode 100644 index 00000000000..8e3460aeafe --- /dev/null +++ b/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs @@ -0,0 +1,283 @@ +// Copyright 2024 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. + +//! Deterministic reproduction of the requeue-then-redispatch TOCTOU in +//! `kill_revoked_operations` (`api_worker_scheduler.rs`): the sweep drops the +//! inner mutex between its lock-free `is_executing_on_worker` confirmation +//! and the final locked `worker_notify_kill_operation` loop. If the matcher +//! re-dispatches the same operation to the same worker inside that gap, the +//! re-guard (`contains_key` + `is_kill_requested`) passes on the fresh +//! entry (`kill_requested_at: None`) and the kill lands on live work. + +mod utils { + pub(crate) mod scheduler_utils; +} + +use core::sync::atomic::{AtomicUsize, Ordering}; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use nativelink_config::schedulers::WorkerAllocationStrategy; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, +}; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker; +use nativelink_scheduler::api_worker_scheduler::ApiWorkerScheduler; +use nativelink_scheduler::platform_property_manager::PlatformPropertyManager; +use nativelink_scheduler::worker::{ActionInfoWithProps, Worker}; +use nativelink_scheduler::worker_registry::WorkerRegistry; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_util::action_messages::{OperationId, WorkerId}; +use nativelink_util::common::DigestInfo; +use nativelink_util::operation_state_manager::{Decline, UpdateOperationType, WorkerStateManager}; +use nativelink_util::origin_event::OriginMetadata; +use nativelink_util::platform_properties::PlatformProperties; +use nativelink_util::spawn; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +/// A state manager scripted to answer `is_executing_on_worker` truthfully +/// for the interleaving under test, with rendezvous points so the test can +/// act inside the sweep's lock-free windows. +struct ScriptedStateManager { + op_x: OperationId, + x_calls: AtomicUsize, + y_calls: AtomicUsize, + /// Sweep reached X's first (snapshot-verification) check; the snapshot + /// of running operations has been taken by now. + x1_reached: Notify, + /// Test releases X's first check (after the decline requeued X). + x1_gate: Notify, + /// Test releases Y's first check. + y1_gate: Notify, + /// X's confirmation check has returned false; `confirmed` now contains + /// X and the sweep still holds no lock. + x2_done: Notify, + /// Test releases Y's confirmation check (the await point that models + /// the gap before the sweep re-acquires the inner mutex). + y2_gate: Notify, +} + +impl MetricsComponent for ScriptedStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl WorkerStateManager for ScriptedStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + // The decline requeues the operation in the real state manager; + // nothing to record here. + Ok(()) + } + + async fn is_executing_on_worker( + &self, + operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + if *operation_id == self.op_x { + if self.x_calls.fetch_add(1, Ordering::SeqCst) == 0 { + self.x1_reached.notify_one(); + self.x1_gate.notified().await; + // Truthful: X was declined and requeued. + } else { + // Confirmation check. Truthful: X is still queued (the + // re-dispatch has not happened yet -- it happens right + // after this returns, in the gap before the sweep + // takes the inner lock). + self.x2_done.notify_one(); + } + Ok(false) + } else if self.y_calls.fetch_add(1, Ordering::SeqCst) == 0 { + self.y1_gate.notified().await; + // Truthful: Y was declined and requeued. + Ok(false) + } else { + // Confirmation check for Y: the sweep awaits here after + // X's confirmation returned. This is the window in + // which the matcher re-dispatches X. + self.y2_gate.notified().await; + // Y got reassigned elsewhere: skip it this pass. + Ok(true) + } + } +} + +fn make_action(digest_byte: u8) -> ActionInfoWithProps { + ActionInfoWithProps { + inner: make_base_action_info( + std::time::UNIX_EPOCH, + DigestInfo::new([digest_byte; 32], 512), + ), + platform_properties: PlatformProperties::default(), + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + } +} + +#[nativelink_test] +async fn kill_revoked_operations_must_not_kill_a_redispatched_operation() -> Result<(), Error> { + let worker_id = WorkerId("worker1".to_string()); + let op_x = OperationId::default(); + let op_y = OperationId::default(); + + let scripted = Arc::new(ScriptedStateManager { + op_x: op_x.clone(), + x_calls: AtomicUsize::new(0), + y_calls: AtomicUsize::new(0), + x1_reached: Notify::new(), + x1_gate: Notify::new(), + y1_gate: Notify::new(), + x2_done: Notify::new(), + y2_gate: Notify::new(), + }); + + let scheduler = ApiWorkerScheduler::new( + scripted.clone(), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + None, // reservation_hold + Arc::new(Notify::new()), + 100, // worker_timeout_s + 100, // unacknowledged_kill_timeout_s + 100, // dispatch_ack_timeout_s + Arc::new(WorkerRegistry::new()), + None, + false, + Duration::from_mins(1), + "test-scheduler".to_string(), + ); + + let (tx, mut rx) = mpsc::channel(16); + scheduler + .add_worker(Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + 0, + 4, + )) + .await?; + // Drain the initial connection message. + drop(rx.recv().await); + + // Dispatch X and Y to the worker. + scheduler + .worker_notify_run_action(worker_id.clone(), op_x.clone(), make_action(1), 0) + .await?; + scheduler + .worker_notify_run_action(worker_id.clone(), op_y.clone(), make_action(2), 0) + .await?; + // Drain the two StartAction messages. + drop(rx.recv().await); + drop(rx.recv().await); + + // The sweep runs concurrently, exactly as the worker_api_server + // interval task does. + let sweep = spawn!("kill_revoked_sweep", { + let scheduler = scheduler.clone(); + async move { scheduler.kill_revoked_operations().await } + }); + + // The sweep has snapshotted (W, X) and (W, Y) with + // kill_requested_at: None and reached its lock-free checks. + scripted.x1_reached.notified().await; + + // The worker declines both operations (e.g. the unacknowledged sweep or + // load backpressure): complete_action removes them from + // running_action_infos and the state manager requeues them. + scheduler + .update_action( + &worker_id, + &op_x, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "test".to_string(), + }), + ) + .await?; + scheduler + .update_action( + &worker_id, + &op_y, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "test".to_string(), + }), + ) + .await?; + + // Release the first-round checks, X strictly before Y so the sweep's + // `revoked` list is ordered [X, Y]. + scripted.x1_gate.notify_one(); + for _ in 0..16 { + tokio::task::yield_now().await; + } + scripted.y1_gate.notify_one(); + + // Wait until X's confirmation check has returned false (X is in + // `confirmed`) while the sweep is parked on Y's confirmation check, + // holding no lock. This is the TOCTOU window. + scripted.x2_done.notified().await; + + // The matching engine re-dispatches X to the same worker: a fresh, + // legitimate dispatch with kill_requested_at: None. The wall-clock + // timestamp is the SAME second as the original dispatch: `dispatched_at` + // has second resolution, so a requeue-and-redispatch inside one second + // is indistinguishable by timestamp. Only a per-dispatch generation + // tells the two dispatches apart. + scheduler + .worker_notify_run_action(worker_id.clone(), op_x.clone(), make_action(1), 0) + .await?; + match rx.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(_)) => {} + v => panic!("Expected StartAction for the re-dispatch, got: {v:?}"), + } + + // Let the sweep finish: it re-acquires the inner lock and runs + // worker_notify_kill_operation for X. + scripted.y2_gate.notify_one(); + sweep.await.unwrap()?; + + // The fresh dispatch must not be killed: the state manager considers X + // legitimately executing on this worker again. On buggy code the + // re-guard passes on the NEW entry and a KillOperationRequest lands on + // live work (and the worker's subsequent report is dropped by + // update_action's is_kill_requested check). + match rx.try_recv() { + Err(mpsc::error::TryRecvError::Empty) => {} + Ok(update) => panic!( + "The freshly re-dispatched operation was killed by the stale sweep: {:?}", + update.update + ), + Err(err) => panic!("Worker channel unexpectedly closed: {err:?}"), + } + + Ok(()) +} diff --git a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs index c13978bb319..bc7f733fc2a 100644 --- a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs @@ -35,6 +35,7 @@ use nativelink_scheduler::awaited_action_db::{ }; use nativelink_scheduler::default_scheduler_factory::retain_completed_for_s; use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::simple_scheduler_state_manager::SimpleSchedulerStateManager; use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb; use nativelink_scheduler::worker::Worker; use nativelink_scheduler::worker_scheduler::WorkerScheduler; @@ -723,3 +724,95 @@ async fn unset_retention_is_the_shared_default_not_forever() -> Result<(), Error assert_eq!(retain_completed_for_s(30), 30); Ok(()) } + +/// A queued index entry whose backing record cannot be decoded (for example +/// serialization drift across a rolling upgrade) has no TTL and, before the +/// fix, no deletion path: every sweep listed it, skipped it and counted it +/// as an orphan, forever. The sweep's listing must remove it, and must leave +/// a healthy queued action alone. +#[nativelink_test] +async fn sweep_deletes_lost_queued_record() -> Result<(), Error> { + const CORRUPT_KEY: &str = "aa_corrupt-op"; + + let backend: FakeRedisBackend = FakeRedisBackend::new(); + let port = backend.clone().run().await; + let store = RedisStore::new_standard(RedisSpec { + addresses: vec![format!("redis://127.0.0.1:{port}")], + experimental_pub_sub_channel: Some("lost-record-sweep".to_string()), + ..Default::default() + }) + .await?; + backend.set_subscription_manager(store.subscription_manager().await.unwrap()); + let notify = Arc::new(Notify::new()); + let db = StoreAwaitedActionDb::new( + store, + notify, + MockInstantWrapped::default, + OperationId::default, + 60, + 60, + false, + ) + .await?; + + // A healthy queued action whose client is still alive. + let healthy = make_awaited_action("healthy-op"); + db.add_action( + OperationId::String("healthy-client-op".to_string()), + healthy.action_info().clone(), + Duration::from_mins(1), + ) + .await?; + + // The lost record: listed under the queued index, version intact, but + // its data does not decode. + backend.table.lock().unwrap().insert( + CORRUPT_KEY.to_string(), + HashMap::from([ + ( + "data".to_string(), + Value::BulkString(b"{this is not an AwaitedAction".to_vec()), + ), + ("version".to_string(), Value::BulkString(b"1".to_vec())), + ("state".to_string(), Value::BulkString(b"queued".to_vec())), + ( + "sort_key".to_string(), + Value::BulkString(b"0000000000000000".to_vec()), + ), + ]), + ); + + let state_manager = SimpleSchedulerStateManager::new( + 0_usize, + Duration::from_mins(1), + Duration::from_mins(1), + Duration::from_mins(10), + db, + MockInstantWrapped::default, + None, + ); + + // Several sweeps: without the fix the corrupt record survives every one + // of them (skip-and-count, never delete). + for _ in 0..3 { + state_manager.sweep_abandoned_queued_actions().await?; + } + + assert!( + !backend.table.lock().unwrap().contains_key(CORRUPT_KEY), + "an undecodable queued record must be deleted by the sweep's listing", + ); + assert_eq!( + backend + .table + .lock() + .unwrap() + .keys() + .filter(|k| k.starts_with("aa_")) + .count(), + 1, + "the healthy queued action must survive the sweep", + ); + + Ok(()) +} diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index 1686aa6f203..d08a10bdcd0 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -48,7 +48,7 @@ use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; use nativelink_util::instant_wrapper::MockInstantWrapped; use nativelink_util::operation_state_manager::{ - ActionStateResult, ClientStateManager, MatchingEngineStateManager, OperationFilter, + ActionStateResult, ClientStateManager, Escalation, MatchingEngineStateManager, OperationFilter, OperationStageFlags, OrderDirection, UpdateOperationType, WorkerStateManager, }; use nativelink_util::platform_properties::PlatformProperties; @@ -99,12 +99,19 @@ fn make_system_time(add_time: u64) -> SystemTime { } fn action_info(started: SystemTime) -> ActionInfo { + action_info_with_properties(started, HashMap::default()) +} + +fn action_info_with_properties( + started: SystemTime, + platform_properties: HashMap, +) -> ActionInfo { let action_digest = DigestInfo::zero_digest(); ActionInfo { command_digest: action_digest, input_root_digest: action_digest, timeout: Duration::ZERO, - platform_properties: HashMap::default(), + platform_properties, priority: 0, load_timestamp: SystemTime::UNIX_EPOCH, insert_timestamp: started, @@ -600,7 +607,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { ); state_mgr - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation(&operation_id, Ok(&worker_id), None) .await?; assert!( state_mgr @@ -618,6 +625,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { .assign_operation( &operation_id, Err(make_err!(Code::DeadlineExceeded, "timed out")), + None, ) .await?; assert!( @@ -626,7 +634,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { .await? ); state_mgr - .assign_operation(&operation_id, Ok(&other_worker_id)) + .assign_operation(&operation_id, Ok(&other_worker_id), None) .await?; assert!( !state_mgr @@ -1269,7 +1277,7 @@ async fn fail_queued_operation_leaves_an_executing_action_alone() -> Result<(), let (operation_id, client) = queued_operation_id(state_mgr.as_ref()).await?; let worker_id = WorkerId::from(String::from("worker")); state_mgr - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation(&operation_id, Ok(&worker_id), None) .await?; let failed = state_mgr @@ -1298,6 +1306,160 @@ async fn fail_queued_operation_ignores_an_unknown_action() -> Result<(), Error> Ok(()) } +/// A database wrapper that shares the inner db so the test can act as a +/// concurrent updater (e.g. a worker-event task) racing the state manager. +struct SharedDb(Arc); + +impl MetricsComponent for SharedDb { + fn publish( + &self, + kind: MetricKind, + field_metadata: MetricFieldData, + ) -> Result { + self.0.publish(kind, field_metadata) + } +} + +impl AwaitedActionDb for SharedDb { + type Subscriber = T::Subscriber; + + async fn get_awaited_action_by_id( + &self, + client_operation_id: &OperationId, + ) -> Result, Error> { + self.0.get_awaited_action_by_id(client_operation_id).await + } + + async fn get_all_awaited_actions( + &self, + ) -> Result> + Send, Error> { + self.0.get_all_awaited_actions().await + } + + async fn get_by_operation_id( + &self, + operation_id: &OperationId, + ) -> Result, Error> { + self.0.get_by_operation_id(operation_id).await + } + + async fn get_range_of_actions( + &self, + state: SortedAwaitedActionState, + start: Bound, + end: Bound, + desc: bool, + ) -> Result> + Send, Error> { + self.0.get_range_of_actions(state, start, end, desc).await + } + + async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { + self.0.update_awaited_action(new_awaited_action).await + } + + async fn add_action( + &self, + client_operation_id: OperationId, + action_info: Arc, + no_event_action_timeout: Duration, + ) -> Result { + self.0 + .add_action(client_operation_id, action_info, no_event_action_timeout) + .await + } +} + +fn cacheable_action_info(started: SystemTime) -> ActionInfo { + let mut info = action_info(started); + let ActionUniqueQualifier::Uncacheable(key) = info.unique_qualifier.clone() else { + unreachable!(); + }; + info.unique_qualifier = ActionUniqueQualifier::Cacheable(key); + info +} + +/// Repro for: `try_subscribe`'s initial keep-alive stamp can be reverted by an +/// in-flight versioned update, letting the client-timeout sweep kill an +/// operation right after a new client subscribed. +/// +/// `try_subscribe` stamps the keep-alive with `send_if_modified` without +/// bumping the version, so an updater holding a pre-stamp `borrow()` snapshot +/// at the same version passes the version check in `update_awaited_action` +/// and `send_replace` wholesale-reverts the keep-alive timestamp. The sweep +/// then retires the operation even though a client just attached. +#[nativelink_test] +async fn a_fresh_subscribers_keepalive_survives_an_in_flight_update() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let db = Arc::new(memory_awaited_action_db_factory( + 0, + &task_change_notify, + MockInstantWrapped::default, + )); + let state_mgr = SimpleSchedulerStateManager::new( + 5, + WORKER_TIMEOUT, + Duration::from_mins(5), // client_action_timeout + MAX_EXECUTING, + SharedDb(db.clone()), + MockInstantWrapped::default, + Some(Arc::new(WorkerRegistry::new())), + ); + + // Client 1 queues a cacheable action at t0, then goes idle: its + // stored client keep-alive stays at t0. + let client1_id = OperationId::default(); + let _client1 = state_mgr + .add_action( + client1_id.clone(), + Arc::new(cacheable_action_info(make_system_time(0))), + ) + .await?; + + // A concurrent task (e.g. a worker-event handler) snapshots the + // action -- version N, keep-alive t0. + let subscriber = db + .get_awaited_action_by_id(&client1_id) + .await? + .expect("operation exists"); + let in_flight_snapshot = subscriber.borrow().await?; + + // t0+4min (within the 5min client timeout). + MockClock::advance(Duration::from_mins(4)); + + // A new client subscribes to the same cacheable action: try_subscribe + // stamps the keep-alive to now (t0+4min) so it is not swept before its + // first poll. + let client2_id = OperationId::default(); + let client2 = state_mgr + .add_action( + client2_id.clone(), + Arc::new(cacheable_action_info(make_system_time(0))), + ) + .await?; + + // The in-flight updater resumes and writes its snapshot back. The + // keep-alive stamp never bumped the version, so this passes the + // version check -- and reverts the keep-alive to t0. + db.update_awaited_action(in_flight_snapshot).await?; + + // t0+6min: client 1's keep-alive (t0) is expired, but client 2 + // attached only 2 minutes ago, so the operation must not be retired. + MockClock::advance(Duration::from_mins(2)); + let retired = state_mgr.sweep_abandoned_queued_actions().await?; + assert_eq!( + retired, 0, + "an operation a client subscribed to 2 minutes ago must not be swept" + ); + let (state, _origin_metadata) = client2.as_state().await?; + assert_eq!( + state.stage, + ActionStage::Queued, + "the freshly-subscribed client's operation must still be alive" + ); + Ok(()) +} + #[nativelink_test] async fn fail_queued_operation_ignores_a_completed_action() -> Result<(), Error> { MockClock::set_time(Duration::from_secs(NOW_TIME)); @@ -1364,7 +1526,7 @@ async fn queue_and_assign_with_client( let (state, _) = listed.as_state().await?; let operation_id = state.client_operation_id.clone(); state_mgr - .assign_operation(&operation_id, Ok(worker_id)) + .assign_operation(&operation_id, Ok(worker_id), None) .await?; Ok((operation_id, subscriber)) } @@ -1778,3 +1940,190 @@ async fn the_listing_reads_records_concurrently_in_the_stores_order() -> Result< ); Ok(()) } + +/// N3: matching and assignment must agree on the action's requirements. +/// +/// Two schedulers (two matchers) share one action DB and memory escalation +/// is enabled. Matcher A reads action X while it needs 4 GiB and picks a +/// 4 GiB worker for it. Before A releases that placement, a peer runs X, +/// hits a memory-limit kill, and requeues X escalated to 8 GiB. When A then +/// releases its stale 4 GiB placement, the assignment must reject it: its +/// only guard used to be double-assignment, so it accepted the queued record +/// and flipped X to Executing while A dispatched the parked 4 GiB reservation +/// to the worker — the DB said 8 GiB but `StartExecute` and the worker's +/// reservation carried 4 GiB, silently bypassing the escalation. The accept +/// now validates the placement-relevant requirements the worker was chosen +/// for and rejects the stale proposal, leaving X queued at 8 GiB to be +/// rematched. A fresh placement made against the current 8 GiB requirements +/// is accepted, so stored, matched, and assigned requirements all agree. +#[nativelink_test] +async fn assignment_rejects_a_stale_placement_after_escalation() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = state_manager(Arc::new(WorkerRegistry::new())); + + let four_gib = HashMap::from([("memory_kb".to_string(), "4000".to_string())]); + let eight_gib = HashMap::from([("memory_kb".to_string(), "8000".to_string())]); + + let client = state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info_with_properties( + make_system_time(0), + four_gib.clone(), + )), + ) + .await?; + + // Matcher A reads X and picks a 4 GiB worker. This is the snapshot it + // will release against later; nothing is committed yet. + let operation_id = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the queued operation") + .as_state() + .await? + .0 + .client_operation_id + .clone(); + let matcher_a_snapshot = four_gib.clone(); + let worker_a = WorkerId::from(String::from("matcher-a-worker-4gib")); + + // A peer runs X on its own worker and the action is killed for memory; + // the peer requeues X escalated to 8 GiB. X is Queued again, needing more. + let peer_worker = WorkerId::from(String::from("peer-worker")); + state_mgr + .assign_operation(&operation_id, Ok(&peer_worker), Some(&four_gib)) + .await?; + state_mgr + .update_operation( + &operation_id, + &peer_worker, + UpdateOperationType::UpdateWithEscalation(Escalation { + property: "memory_kb".to_string(), + value: 8000, + reason: make_err!(Code::FailedPrecondition, "killed for memory"), + cpu: None, + }), + ) + .await?; + + // The store now records the escalated requirement. + let stored = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the requeued operation") + .as_action_info() + .await? + .0; + assert_eq!( + stored.platform_properties, eight_gib, + "the escalation raised the stored reservation to 8 GiB" + ); + + // Matcher A now releases its stale 4 GiB placement against its 4 GiB + // worker. Without the fix the assignment accepts this and flips X to + // Executing at 4 GiB, bypassing the escalation. It must be rejected. + let stale = state_mgr + .assign_operation(&operation_id, Ok(&worker_a), Some(&matcher_a_snapshot)) + .await + .expect_err("a stale placement must be rejected, not silently accepted"); + assert_eq!( + stale.code, + Code::Aborted, + "a rejected stale placement is a lost race, not a hard error: {stale}" + ); + assert_eq!( + client.as_state().await?.0.stage, + ActionStage::Queued, + "the rejected placement leaves X queued for a rematch, not Executing" + ); + + // A fresh placement made against the current 8 GiB requirement is + // accepted: stored, matched, and assigned requirements all agree. + state_mgr + .assign_operation(&operation_id, Ok(&worker_a), Some(&eight_gib)) + .await?; + assert_eq!( + client.as_state().await?.0.stage, + ActionStage::Executing, + "a placement matching the current 8 GiB requirement is accepted" + ); + let assigned = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Executing, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the assigned operation") + .as_action_info() + .await? + .0; + assert_eq!( + assigned.platform_properties, eight_gib, + "the dispatched action carries the escalated 8 GiB reservation" + ); + Ok(()) +} + +/// The memory store's listing is read whole, under one hold of the store's +/// lock taken when the listing is requested, and the stream then only hands +/// out what was read. The lock must not be taken inside the stream: the +/// filter pass reads rows 64 at a time, so while it hands rows on it keeps +/// asking the listing for more, and a lock wait that asking starts is +/// parked inside the stream, polled again only when the next row is asked +/// for. Once that wait has been starved past async-lock's fairness +/// threshold the lock is reserved for it, and when the matcher then writes +/// for a row it was handed, on the same task, it waits for the same lock: +/// nothing polls the parked wait, and the scheduler stops for good. Reading +/// lazily was how that wait got into the stream; a row added after the +/// listing was requested is therefore not in it. +#[nativelink_test] +async fn the_listing_is_read_whole_when_it_is_requested() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = state_manager(Arc::new(WorkerRegistry::new())); + // Kept so the clients count as waiting. + let _clients = add_prioritised_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + order_by_priority_direction: Some(OrderDirection::Desc), + ..Default::default() + }, + ) + .await?; + // Added after the listing was requested, before any of it was read. + let _late = state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info(make_system_time(0))), + ) + .await?; + + let listed: Vec<_> = listed.collect().await; + assert_eq!( + listed.len(), + 3, + "the listing holds what was there when it was requested" + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs index cd4237f029f..1f6f481cb40 100644 --- a/nativelink-scheduler/tests/simple_scheduler_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_test.rs @@ -1545,6 +1545,109 @@ async fn update_action_sends_completed_result_to_client_test() -> Result<(), Err Ok(()) } +/// A worker re-registering under an id the pool still holds must be treated +/// as a disconnect of the old incarnation: what it was running is requeued, +/// its channel hears Disconnect, and the new incarnation can pick the work +/// up and complete it. Silent replacement used to drop the old `Worker` with +/// its `running_action_infos`, leaving the operation Executing against an +/// incarnation that never held it (double execution after timeout). +#[nativelink_test] +async fn worker_reregistration_evicts_old_incarnation_and_requeues_test() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + + let task_change_notify = Arc::new(Notify::new()); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec::default(), + memory_awaited_action_db_factory( + 0, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + let action_digest = DigestInfo::new([99u8; 32], 512); + + let mut rx_from_old_worker = + setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?; + let insert_timestamp = make_system_time(1); + let mut action_listener = + setup_action(&scheduler, action_digest, HashMap::new(), insert_timestamp).await?; + + let operation_id = { + // Other tests check full data. We only care if we got StartAction. + match rx_from_old_worker.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(start_execute)) => { + assert_eq!( + action_listener.changed().await.unwrap().0.stage, + ActionStage::Executing + ); + OperationId::from(start_execute.operation_id) + } + v => panic!("Expected StartAction, got : {v:?}"), + } + }; + + // The same worker id re-registers on a fresh channel, as a restarted + // worker process does. + let mut rx_from_new_worker = + setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?; + scheduler.do_try_match_for_test().await?; + + { + // The old incarnation's channel must hear a Disconnect. + let msg_for_worker = rx_from_old_worker.recv().await.unwrap(); + assert_eq!( + msg_for_worker, + UpdateForWorker { + update: Some(update_for_worker::Update::Disconnect(())) + } + ); + } + { + // The operation must be requeued and re-dispatched to the new + // incarnation, not left Executing against the dropped one. + match rx_from_new_worker.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(start_execute)) => { + assert_eq!( + OperationId::from(start_execute.operation_id), + operation_id, + "expected the same operation re-dispatched to the new incarnation" + ); + assert_eq!(start_execute.worker_id, worker_id.to_string()); + } + v => panic!("Expected StartAction on the new channel, got : {v:?}"), + } + } + + // The new incarnation completes the action and the client sees it. + let action_result = ActionResult::default(); + scheduler + .update_action( + &worker_id, + &operation_id, + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed( + action_result.clone(), + )), + ) + .await?; + + // The listener may see intermediate transitions from the requeue; wait + // for the final one. + let mut stage = action_listener.changed().await.unwrap().0.stage.clone(); + for _ in 0..8 { + if stage == ActionStage::Completed(action_result.clone()) { + break; + } + stage = action_listener.changed().await.unwrap().0.stage.clone(); + } + assert_eq!(stage, ActionStage::Completed(action_result)); + + Ok(()) +} + #[nativelink_test] async fn update_action_sends_completed_result_after_disconnect() -> Result<(), Error> { let worker_id = WorkerId("worker_id".to_string()); diff --git a/nativelink-scheduler/tests/store_awaited_action_db_test.rs b/nativelink-scheduler/tests/store_awaited_action_db_test.rs index dbe76f3fc1a..31dc0c4bd90 100644 --- a/nativelink-scheduler/tests/store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/store_awaited_action_db_test.rs @@ -35,7 +35,8 @@ use nativelink_scheduler::store_awaited_action_db::{ }; use nativelink_scheduler::worker_registry::{ORPHANED_ACTION_TIMEOUT, WorkerRegistry}; use nativelink_util::action_messages::{ - ActionInfo, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, WorkerId, + ActionInfo, ActionResult, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, + WorkerId, }; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; @@ -443,6 +444,177 @@ async fn try_subscribe_skips_lookup_for_uncacheable_qualifier() -> Result<(), Er Ok(()) } +/// Fake store whose single `search_by_index_prefix` returns a fixed list of +/// already-encoded candidates, in order. Unlike `EventuallyVisibleStore` this +/// exercises the multi-candidate selection: the unique-qualifier index retains +/// completed records alongside a live successor, and the backend does not order +/// by state, so `try_subscribe` must scan past a finished candidate to join a +/// live one rather than stopping at the first exact match. +struct MultiCandidateStore { + encoded: Vec, +} + +impl MultiCandidateStore { + fn new(actions: &[&AwaitedAction]) -> Self { + Self { + encoded: actions + .iter() + .map(|a| { + Bytes::from(serde_json::to_vec(a).expect("serialize AwaitedAction for fake")) + }) + .collect(), + } + } +} + +impl SchedulerStore for MultiCandidateStore { + type SubscriptionManager = PendingSubscriptionManager; + + fn subscription_manager( + &self, + ) -> impl Future, Error>> { + std::future::ready(Ok(Arc::new(PendingSubscriptionManager))) + } + + fn update_data( + &self, + _data: T, + _expiry: Option, + ) -> impl Future, Error>> + where + T: SchedulerStoreDataProvider + + SchedulerStoreKeyProvider + + SchedulerCurrentVersionProvider + + Send, + { + std::future::ready(Ok(Some(1))) + } + + fn search_by_index_prefix( + &self, + _index: K, + ) -> impl Future< + Output = Result< + impl Stream::DecodeOutput, Error>> + Send, + Error, + >, + > + where + K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send, + ::DecodeOutput: Send, + { + let items: Vec::DecodeOutput, Error>> = self + .encoded + .iter() + .map(|bytes| K::decode(1, bytes.clone())) + .collect(); + std::future::ready(Ok(stream::iter(items))) + } + + async fn count_by_index_prefix(&self, _index: K) -> Result + where + K: SchedulerIndexProvider + Send, + { + Ok(0) + } + + async fn get_and_decode( + &self, + _key: K, + ) -> Result::DecodeOutput>, Error> + where + K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send, + { + todo!("not exercised by try_subscribe") + } +} + +async fn build_multi_db( + store: Arc, +) -> StoreAwaitedActionDb< + MultiCandidateStore, + fn() -> OperationId, + MockInstantWrapped, + fn() -> MockInstantWrapped, +> { + fn new_op_id() -> OperationId { + OperationId::from("new-operation") + } + let now_fn: fn() -> MockInstantWrapped = MockInstantWrapped::default; + let op_id_fn: fn() -> OperationId = new_op_id; + StoreAwaitedActionDb::new( + store, + Arc::new(Notify::new()), + now_fn, + op_id_fn, + 60, + 60, + false, + ) + .await + .expect("construct test db") +} + +fn make_completed_awaited_action(operation_id: &str) -> AwaitedAction { + let now = SystemTime::UNIX_EPOCH; + let mut action = AwaitedAction::new( + OperationId::from(operation_id), + make_cacheable_action_info(), + now, + ); + let mut state = action.state().as_ref().clone(); + state.stage = ActionStage::Completed(ActionResult::default()); + action.worker_set_state(Arc::new(state), now); + action +} + +fn make_queued_awaited_action(operation_id: &str) -> AwaitedAction { + // A fresh `AwaitedAction` is Queued and has no worker, so it is neither + // finished nor abandoned — i.e. joinable. + AwaitedAction::new( + OperationId::from(operation_id), + make_cacheable_action_info(), + SystemTime::UNIX_EPOCH, + ) +} + +/// Regression for the N7 selection defect: with a completed record indexed +/// ahead of a live successor for the same qualifier, `try_subscribe` must join +/// the live one, not return the completed record (which forces a recreate). +#[nativelink_test] +async fn try_subscribe_prefers_live_candidate_over_completed() -> Result<(), Error> { + let completed = make_completed_awaited_action("completed-operation"); + let live = make_queued_awaited_action("live-operation"); + let qualifier = live.action_info().unique_qualifier.clone(); + // Completed record first, live successor second — the order the backend can + // legitimately return. + let store = Arc::new(MultiCandidateStore::new(&[&completed, &live])); + let db = build_multi_db(store).await; + + let result = db + .try_subscribe( + &OperationId::from("client-prefer-live"), + &qualifier, + Duration::from_mins(1), + 0, + ) + .await? + .expect("a joinable live candidate is present"); + + assert_eq!( + *result.operation_id(), + OperationId::from("live-operation"), + "try_subscribe must join the live successor, not the completed record \ + (which would be recreated as a brand-new operation)", + ); + assert_ne!( + *result.operation_id(), + OperationId::from("new-operation"), + "joining a live candidate must not recreate a fresh operation", + ); + Ok(()) +} + // --------------------------------------------------------------------------- // `try_subscribe` abandonment check. // diff --git a/nativelink-scheduler/tests/worker_retry_reservation_test.rs b/nativelink-scheduler/tests/worker_retry_reservation_test.rs new file mode 100644 index 00000000000..c7cd4a8f41c --- /dev/null +++ b/nativelink-scheduler/tests/worker_retry_reservation_test.rs @@ -0,0 +1,134 @@ +// 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. + +use std::collections::HashMap; +use std::time::{Duration, UNIX_EPOCH}; + +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::worker::{ActionInfoWithProps, Worker, WorkerUpdate}; +use nativelink_util::action_messages::{OperationId, WorkerId}; +use nativelink_util::common::DigestInfo; +use nativelink_util::origin_event::OriginMetadata; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; + +mod utils { + pub(crate) mod scheduler_utils; +} +use utils::scheduler_utils::make_base_action_info; + +const NOW_TIME: u64 = 10_000; + +fn minimum(mem: u64) -> PlatformProperties { + PlatformProperties::new(HashMap::from([( + "mem".to_string(), + PlatformPropertyValue::Minimum(mem), + )])) +} + +fn action_needing(mem: u64) -> ActionInfoWithProps { + ActionInfoWithProps { + inner: make_base_action_info( + UNIX_EPOCH + Duration::from_secs(NOW_TIME), + DigestInfo::new([7; 32], 512), + ), + platform_properties: minimum(mem), + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + } +} + +fn mem_of(props: &PlatformProperties) -> u64 { + match props.properties.get("mem") { + Some(PlatformPropertyValue::Minimum(v)) => *v, + _ => panic!("expected a Minimum 'mem' property"), + } +} + +/// N2 regression: a same-worker retry (the earlier attempt of the same op is +/// still resident on the worker) must not double-deduct or leak the worker's +/// resource budget. +/// +/// WITHOUT THE FIX: `Worker::run_action` unconditionally re-runs +/// `reduce_platform_properties` and overwrites the op-keyed entry, so the +/// worker's `mem` budget is deducted a SECOND time (10 -> 6 -> 2) even though +/// the op already holds one reservation. Since completion (`complete_action`) +/// only refunds the single surviving entry, those extra 4 units are leaked +/// forever. This test asserts that re-dispatching an op the worker already +/// holds does not further reduce the budget — it FAILS on the unfixed code +/// (asserts 6, finds 2). +#[nativelink_test] +async fn same_worker_retry_conserves_resource_budget() -> Result<(), Error> { + let (tx, _rx) = mpsc::channel(64); + let mut worker = Worker::new( + WorkerId("w".to_string()), + minimum(10), + tx, + NOW_TIME, + // Two inflight slots so a retry has spare capacity while the first + // attempt is still resident. + 2, + ); + + let op = OperationId::default(); + + // First dispatch of op reserves 4 of the worker's 10 mem. + worker + .notify_update(WorkerUpdate::RunAction(Box::new(( + op.clone(), + action_needing(4), + NOW_TIME, + 0, // dispatch generation + )))) + .await?; + assert_eq!( + mem_of(&worker.platform_properties), + 6, + "first dispatch should reserve 4 of 10" + ); + + // Same op dispatched again to the SAME worker while the earlier attempt is + // still resident (same-worker retry). The worker must not create a second + // live reservation for an op it already holds. + let retry_result = worker + .notify_update(WorkerUpdate::RunAction(Box::new(( + op.clone(), + action_needing(4), + NOW_TIME, + 0, // dispatch generation + )))) + .await; + // The worker may accept or reject the redundant dispatch; either way the + // invariant below (no double reservation) must hold. + drop(retry_result); + + // Exactly one live reservation should exist for the op. + assert_eq!( + worker.running_action_infos.len(), + 1, + "op must have exactly one live reservation on the worker" + ); + + // The op still holds exactly one reservation, so the budget must be + // unchanged from the first dispatch. A second deduction here is a leak: + // completion only refunds the single surviving entry. + assert_eq!( + mem_of(&worker.platform_properties), + 6, + "re-dispatching an already-held op must not deduct the budget again" + ); + Ok(()) +} diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 953c2bc6e5a..fd3d4832896 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -173,6 +173,10 @@ impl ChunkingInstance { #[derive(Debug)] pub struct CasServer { stores: HashMap, + /// Per-instance per-blob deadline applied inside `BatchReadBlobs`. + /// Defaults to `BATCH_PER_BLOB_TIMEOUT`; configurable via + /// `batch_read_per_blob_timeout_s`. + batch_read_per_blob_timeouts: HashMap, remote_cache_compression_instances: RemoteCacheCompressionInstances, chunking_instances: HashMap, chunking_metrics: ChunkingMetrics, @@ -180,7 +184,9 @@ pub struct CasServer { type GetTreeStream = Pin> + Send + 'static>>; -/// Per-blob deadline applied inside `BatchReadBlobs` / `BatchUpdateBlobs`. +/// Default per-blob deadline applied inside `BatchReadBlobs` / +/// `BatchUpdateBlobs`. For `BatchReadBlobs` this can be overridden per +/// instance via `batch_read_per_blob_timeout_s`. const BATCH_PER_BLOB_TIMEOUT: Duration = Duration::from_secs(30); /// Maximum size of a single chunk accepted in a `SpliceBlob` request. @@ -209,11 +215,20 @@ impl CasServer { remote_cache_compression_instances: &RemoteCacheCompressionInstances, ) -> Result { let mut stores = HashMap::with_capacity(configs.len()); + let mut batch_read_per_blob_timeouts = HashMap::with_capacity(configs.len()); let mut chunking_instances = HashMap::new(); for config in configs { let store = store_manager.get_store(&config.cas_store).ok_or_else(|| { make_input_err!("'cas_store': '{}' does not exist", config.cas_store) })?; + batch_read_per_blob_timeouts.insert( + config.instance_name.clone(), + config + .batch_read_per_blob_timeout_s + .map_or(BATCH_PER_BLOB_TIMEOUT, |secs| { + Duration::from_secs(secs.into()) + }), + ); if let Some(chunking_config) = &config.experimental_chunking { let avg_chunk_size_bytes = chunking_config .validated_avg_chunk_size_bytes() @@ -273,6 +288,7 @@ impl CasServer { } Ok(Self { stores, + batch_read_per_blob_timeouts, remote_cache_compression_instances: remote_cache_compression_instances.clone(), chunking_instances, chunking_metrics: ChunkingMetrics::default(), @@ -411,6 +427,11 @@ impl CasServer { } let store_ref = &store; + let per_blob_timeout = self + .batch_read_per_blob_timeouts + .get(instance_name) + .copied() + .unwrap_or(BATCH_PER_BLOB_TIMEOUT); let remote_cache_compression_enabled = self .remote_cache_compression_instances .enabled_for(instance_name); @@ -429,7 +450,7 @@ impl CasServer { // Apply a per-blob deadline so one slow read does not // make the whole batch hit the client's overall deadline. let result = match tokio::time::timeout( - BATCH_PER_BLOB_TIMEOUT, + per_blob_timeout, store_ref.get_part_unchunked(digest_copy, 0, None), ) .await @@ -438,7 +459,7 @@ impl CasServer { Err(_elapsed) => Err(make_err!( Code::DeadlineExceeded, "BatchReadBlobs per-blob timeout ({} s) elapsed for digest {}", - BATCH_PER_BLOB_TIMEOUT.as_secs(), + per_blob_timeout.as_secs(), digest_copy, )), }; diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index 055bfe9a884..3ee5477c25f 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -78,6 +78,7 @@ fn make_cas_server(store_manager: &StoreManager) -> Result { config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), experimental_chunking: None, + batch_read_per_blob_timeout_s: None, }, }], store_manager, @@ -92,6 +93,7 @@ fn make_cas_server_with_zstd(store_manager: &StoreManager) -> Result, _key: StoreKey<'_>, - _writer: &mut DropCloserWriteHalf, + writer: &mut DropCloserWriteHalf, _offset: u64, _length: Option, ) -> Result<(), Error> { tokio::time::sleep(self.delay).await; - Ok(()) + if !self.data.is_empty() { + writer.send(self.data.clone()).await?; + } + writer.send_eof() } fn inner_store(&self, _digest: Option) -> &'_ dyn StoreDriver { @@ -760,14 +767,23 @@ impl StoreDriver for StallStore { default_health_status_indicator!(StallStore); fn make_cas_server_with_stall_store(delay: Duration) -> Result { + make_cas_server_with_stall_store_and_timeout(delay, Bytes::new(), None) +} + +fn make_cas_server_with_stall_store_and_timeout( + delay: Duration, + data: Bytes, + batch_read_per_blob_timeout_s: Option, +) -> Result { let store_manager = Arc::new(StoreManager::new()); - store_manager.add_store("main_cas", Store::new(Arc::new(StallStore { delay })))?; + store_manager.add_store("main_cas", Store::new(Arc::new(StallStore { delay, data })))?; CasServer::new( &[WithInstanceName { instance_name: INSTANCE_NAME.to_string(), config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), experimental_chunking: None, + batch_read_per_blob_timeout_s, }, }], &store_manager, @@ -858,6 +874,47 @@ async fn batch_read_blobs_per_blob_timeout_returns_deadline_exceeded() Ok(()) } +#[nativelink_test(start_paused = true)] +async fn batch_read_blobs_respects_configured_per_blob_timeout() +-> Result<(), Box> { + const VALUE: &str = "1"; + + // The store stalls for 2 minutes — well beyond the 30 s default that + // would fail the read — but the instance is configured with a 3 minute + // per-blob timeout, so the read must succeed. + let cas_server = make_cas_server_with_stall_store_and_timeout( + Duration::from_mins(2), + Bytes::from_static(VALUE.as_bytes()), + Some(180), + )?; + + let digest = Digest { + hash: HASH1.to_string(), + size_bytes: VALUE.len().try_into().unwrap_or(i64::MAX), + }; + let raw_response = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![digest.clone()], + acceptable_compressors: vec![compressor::Value::Identity.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + + let response = raw_response.unwrap().into_inner(); + assert_eq!(response.responses.len(), 1); + let entry = &response.responses[0]; + assert_eq!(entry.digest.as_ref(), Some(&digest)); + let status = entry.status.as_ref().expect("status set"); + assert_eq!( + status.code, + Code::Ok as i32, + "expected Ok with a raised per-blob timeout, got status: {status:?}", + ); + assert_eq!(entry.data, VALUE.as_bytes()); + Ok(()) +} + #[nativelink_test] async fn batch_update_blobs_zstd_compressed() -> Result<(), Box> { let store_manager = make_store_manager().await?; @@ -1061,6 +1118,7 @@ fn make_chunking_cas_server_with_avg( instance_name: INSTANCE_NAME.to_string(), config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), + batch_read_per_blob_timeout_s: None, experimental_chunking: Some(nativelink_config::cas_server::CasChunkingConfig { index_store: Some("chunk_index".to_string()), avg_chunk_size_bytes, @@ -1516,6 +1574,7 @@ async fn chunking_rejects_index_store_same_as_cas_store() -> Result<(), Box Result<(), Box Result<(), Box { - return Err(make_err!( - Code::Internal, - "CompletenessCheckingStore's check_existence_fut ended unexpectedly {r:?}" - )); + // Before `done` is set, this future only finishes when a CAS + // existence check fails. Keep that error's code, so a + // transient CAS fault stays retryable instead of reaching + // the client as Internal. + return Err(match r { + Err(err) => err, + Ok(()) => make_err!( + Code::Internal, + "CompletenessCheckingStore's check_existence_fut ended unexpectedly" + ), + }); } maybe_result = futures.next() => { match maybe_result { Some(Ok(())) => self.complete_entries_counter.inc(), Some(Err((err, i))) => { - self.incomplete_entries_counter.inc(); - // Only a genuine `NotFound` may be laundered into a "missing" - // slot (which callers surface as `Code::NotFound`). That covers - // a truly absent action-result record and a clean completeness - // miss — the intended eviction-coherence behavior (an evicted - // referent must still 404). ANY other code is a real backend - // read/decode failure; collapsing it into `None` would make an - // I/O fault indistinguishable from absence and get served as a - // 404. Propagate it so the caller can tell "genuinely absent" - // from "could not determine". + // Only a `NotFound` becomes a miss: an absent action + // result, or an absent Tree blob. Records that fail to + // decode also arrive as `NotFound` (`ac_utils` maps + // them), so they stay misses too. Any other code is a + // backend failure, such as an unavailable store or a + // timeout, and turning it into a miss would serve an + // outage as a 404 and throw away a valid cached + // result. Propagate it instead. if err.code != Code::NotFound { return Err(err); } + self.incomplete_entries_counter.inc(); state_mux.lock().results[i] = None; } None => { diff --git a/nativelink-store/src/dedup_store.rs b/nativelink-store/src/dedup_store.rs index e3b437ca00a..c2f4a6307c9 100644 --- a/nativelink-store/src/dedup_store.rs +++ b/nativelink-store/src/dedup_store.rs @@ -352,6 +352,55 @@ impl StoreDriver for DedupStore { .await .err_tip(|| "Failed to get_part in content_store in dedup_store")?; + // Content chunks are content-addressed and whole: a chunk must + // contain exactly the number of bytes its index entry declares. + // If the content store returns a chunk of any other size (a + // torn or partial read, or an index/content size disagreement) + // the streaming loop below would either panic on the + // `bytes_to_skip` assert or silently under-deliver and send EOF + // as success. Reject the mismatch here so it surfaces as an + // error instead. + let declared = index_entry.size_bytes(); + let actual = data.len() as u64; + if actual != declared { + // Every blob that references this chunk is unreadable while + // it stays, and nothing else would notice: `has()` sums the + // sizes the content store reports, and `update()` skips any + // chunk the content store already has, so a re-upload would + // keep the bad copy. Drop it, as `VerifyStore` drops a copy + // that fails verification, so `has()` reports those blobs + // missing and the next upload writes the chunk again. This + // is best effort: a content store that does not support + // `remove` keeps the chunk. + match self + .content_store + .as_store_driver_pin() + .remove(index_entry.into()) + .await + { + Ok(removed) => warn!( + %index_entry, + actual, + declared, + removed, + "Dropped a dedup content chunk whose size does not match its digest" + ), + Err(err) => warn!( + %index_entry, + ?err, + "Could not drop a dedup content chunk whose size does not match its digest" + ), + } + // NotFound rather than the DataLoss `VerifyStore` returns: a + // `FastSlowStore` above falls back to its slow tier only on + // NotFound, and the slow tier may hold a good copy. + return Err(make_err!( + Code::NotFound, + "DedupStore content chunk {} is {actual} bytes but its index entry declares {declared}; refusing to serve a torn read as success", + index_entry, + )); + } + Result::<_, Error>::Ok(data) }) .buffered(self.max_concurrent_fetch_per_get); diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index 08d91800559..fd9161f6f70 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -77,7 +77,33 @@ pub struct FastSlowStore { // concurrent writer that has not yet finished pushing to the slow store // is still visible to a concurrent existence check, preventing redundant // duplicate uploads of the same blob. - in_flight_slow_writes: Mutex, Arc>>>, + // + // Each key maps to a single shared slot that *all* concurrent writers of + // that key register on. This is load-bearing for liveness: waiters key on + // the digest, not on any individual upload, so a second upload (e.g. a + // retry under a fresh ByteStream upload id after the first was + // interrupted and parked for resumption) completing the blob wakes + // waiters that subscribed while the first upload was in flight. + in_flight_slow_writes: Mutex, InFlightSlowWrite>>, +} + +/// State of an in-flight slow-store write, broadcast to existence-check +/// waiters through a `watch` channel. +#[derive(Clone, Copy, Debug)] +enum InFlightWriteState { + /// At least one writer is still streaming this key to the slow store. + Pending, + /// The write finished: `Some(size)` if a writer completed successfully, + /// `None` if every writer gave up without completing. + Done(MaybeSize), +} + +/// Per-key slot shared by every concurrent writer of that key. +#[derive(Debug)] +struct InFlightSlowWrite { + tx: tokio::sync::watch::Sender, + /// Number of live [`InFlightSlowWriteGuard`]s registered on this slot. + writers: usize, } // This guard ensures that the populating_digests is cleared even if the future @@ -104,19 +130,22 @@ impl LoaderGuard<'_> { } } -/// Cancel-safe RAII guard that removes an entry from -/// `FastSlowStore::in_flight_slow_writes` when dropped. This ensures the -/// map does not leak entries if the surrounding `update()` future is -/// cancelled before the slow-store write completes. +/// Cancel-safe RAII guard that deregisters a writer from the key's shared +/// slot in `FastSlowStore::in_flight_slow_writes` when dropped. The map entry +/// is removed only when the *last* writer for the key deregisters, so a +/// parked/cancelled upload cannot orphan waiters nor evict a slot a newer +/// upload of the same key is still using. struct InFlightSlowWriteGuard { weak_store: Weak, key: Option>, - write_complete_guard: tokio::sync::OwnedMutexGuard, + tx: tokio::sync::watch::Sender, } impl InFlightSlowWriteGuard { - fn complete(mut self, size: MaybeSize) { - *self.write_complete_guard = size; + /// Marks the write done and wakes all waiters; deregistration happens in + /// `Drop`. + fn complete(self, size: MaybeSize) { + self.tx.send_replace(InFlightWriteState::Done(size)); } } @@ -128,7 +157,27 @@ impl Drop for InFlightSlowWriteGuard { let Some(key) = self.key.take() else { return; }; - store.in_flight_slow_writes.lock().remove(&key); + let mut map = store.in_flight_slow_writes.lock(); + let std::collections::hash_map::Entry::Occupied(mut occupied_entry) = map.entry(key) else { + return; + }; + // Only touch the slot this guard registered on (defense against a + // same-key slot that was replaced out from under us). + if !occupied_entry.get().tx.same_channel(&self.tx) { + return; + } + let slot = occupied_entry.get_mut(); + slot.writers -= 1; + if slot.writers == 0 { + // Last writer for this key. If nobody completed the write, wake + // waiters with "not written" so they do not hang forever. + slot.tx.send_modify(|state| { + if matches!(state, InFlightWriteState::Pending) { + *state = InFlightWriteState::Done(None); + } + }); + occupied_entry.remove(); + } } } @@ -175,18 +224,18 @@ impl FastSlowStore { fn register_in_flight_slow_write(&self, key: StoreKey<'_>) -> InFlightSlowWriteGuard { let owned = key.into_owned(); - let write_complete = Arc::new(tokio::sync::Mutex::new(None)); - let write_complete_guard = write_complete - .clone() - .try_lock_owned() - .expect("Newly created mutex is locked"); - self.in_flight_slow_writes - .lock() - .insert(owned.borrow().into_owned(), write_complete); + let mut map = self.in_flight_slow_writes.lock(); + let slot = map + .entry(owned.borrow().into_owned()) + .or_insert_with(|| InFlightSlowWrite { + tx: tokio::sync::watch::channel(InFlightWriteState::Pending).0, + writers: 0, + }); + slot.writers += 1; InFlightSlowWriteGuard { weak_store: self.weak_self.clone(), key: Some(owned), - write_complete_guard, + tx: slot.tx.clone(), } } @@ -198,6 +247,25 @@ impl FastSlowStore { } } + /// The size to declare to the fast store when filling it with `key`, given + /// the size the slow store's `has()` reported. + /// + /// For a digest there are two answers, and when they agree the size is + /// exact. When they disagree this store cannot tell which is wrong: a slow + /// tier can report a physical size (a filesystem store reports its + /// block-rounded footprint), and an action cache or a dedup index keys an + /// entry by a digest whose size is not the entry's. Declaring either as + /// exact makes a fast store that checks `ExactSize`, as the filesystem + /// store does, refuse a correct blob, so the larger is declared as a bound. + fn fast_store_upload_size(key: &StoreKey<'_>, slow_size: u64) -> UploadSizeInfo { + match Self::digest_size_bytes(key) { + Some(digest_size) if digest_size != slow_size => { + UploadSizeInfo::MaxSize(max(digest_size, slow_size)) + } + _ => UploadSizeInfo::ExactSize(slow_size), + } + } + /// Whether a read should skip dedup and hit the slow store directly. A /// threshold of 0 disables the bypass so every read goes through dedup. fn should_bypass_dedup(&self, key: &StoreKey<'_>) -> bool { @@ -301,23 +369,23 @@ impl FastSlowStore { ); UploadSizeInfo::MaxSize(u64::MAX) } else { - UploadSizeInfo::ExactSize(self - .slow_store - .has(key.borrow()) - .await - .err_tip(|| "Failed to run has() on slow store")? - .ok_or_else(|| { - name_missing_digest( - make_err!( - Code::NotFound, - "Object {} not found in either fast or slow store. \ - If using multiple workers, ensure all workers share the same CAS storage path.", - key.as_str() - ), - &key, - ) - })? - ) + let slow_size = self + .slow_store + .has(key.borrow()) + .await + .err_tip(|| "Failed to run has() on slow store")? + .ok_or_else(|| { + name_missing_digest( + make_err!( + Code::NotFound, + "Object {} not found in either fast or slow store. \ + If using multiple workers, ensure all workers share the same CAS storage path.", + key.as_str() + ), + &key, + ) + })?; + Self::fast_store_upload_size(&key, slow_size) }; let send_range = offset..length.map_or(u64::MAX, |length| length + offset); @@ -560,9 +628,21 @@ impl StoreDriver for FastSlowStore { for (i, (k, result)) in key.iter().zip(results.iter_mut()).enumerate() { if result.is_none() { let owned = k.borrow().into_owned(); - if let Some(maybe_size) = in_flight.get(&owned) { - let maybe_size = maybe_size.clone(); - in_flight_futs.push(async move { (i, *maybe_size.lock().await) }); + if let Some(slot) = in_flight.get(&owned) { + let mut rx = slot.tx.subscribe(); + in_flight_futs.push(async move { + loop { + let state = *rx.borrow_and_update(); + match state { + InFlightWriteState::Done(size) => return (i, size), + InFlightWriteState::Pending => {} + } + if rx.changed().await.is_err() { + // All writers dropped without completing. + return (i, None); + } + } + }); } } } diff --git a/nativelink-store/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index ee26bdfc2da..2d9e3b21a32 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -102,6 +102,17 @@ pub struct SharedContext { // deleted (similar to how it is done in tests). #[metric(help = "Number of active drop spawns")] pub active_drop_spawns: AtomicU64, + /// Count of deletes that ran synchronously inline because the delete + /// semaphore was saturated. A nonzero, growing value means eviction is + /// outrunning background delete throughput — the backpressure signal. + #[metric(help = "Number of deletes performed inline due to delete backpressure")] + pub inline_deletes: AtomicU64, + /// Bounds the number of in-flight background file deletions. `None` means + /// unlimited (the historical behavior). When the permits are exhausted, + /// a drop deletes its file synchronously inline instead of spawning, so + /// an eviction storm applies backpressure rather than spawning without + /// limit. See `FilesystemSpec::max_concurrent_deletes`. + delete_semaphore: Option>, #[metric(help = "Path to the configured temp path")] temp_path: String, #[metric(help = "Path to the configured content path")] @@ -187,6 +198,40 @@ impl Drop for EncodedFilePath { let file_path = self.get_file_path().to_os_string(); let shared_context = self.shared_context.clone(); + + // Bound the number of in-flight background deletes. `drop` is sync, so + // we cannot `.await` a permit the way the write path does; instead we + // `try_acquire_owned` (never blocks) and, when the permits are + // exhausted, delete synchronously inline. That inline path is the + // backpressure: the evicting thread pays the unlink cost directly + // instead of queueing an unbounded number of tasks, and the file is + // always removed (never leaked). `None` keeps the historical + // unbounded-spawn behavior. + let delete_permit = match &shared_context.delete_semaphore { + Some(sem) => { + if let Ok(permit) = Arc::clone(sem).try_acquire_owned() { + Some(permit) + } else { + // Saturated: delete synchronously inline, which is the + // backpressure, then bail out (no spawn). + shared_context + .inline_deletes + .fetch_add(1, Ordering::Relaxed); + match std::fs::remove_file(&file_path) { + Ok(()) => { + debug!(?file_path, "File deleted inline (delete backpressure)"); + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + debug!(?file_path, "File already gone, nothing to delete"); + } + Err(err) => error!(?file_path, ?err, "Failed to delete file inline"), + } + return; + } + } + None => None, + }; + // .fetch_add returns previous value, so we add one to get approximate current value let current_active_drop_spawns = shared_context .active_drop_spawns @@ -198,6 +243,9 @@ impl Drop for EncodedFilePath { "Spawned a filesystem_delete_file" ); background_spawn!("filesystem_delete_file", async move { + // Hold the permit until the delete actually completes so the bound + // reflects deletes in flight, not merely scheduled. + let _delete_permit = delete_permit; match fs::remove_file(&file_path).await { Ok(()) => debug!(?file_path, "File deleted"), // The file already being gone is the desired end state of a @@ -1439,8 +1487,15 @@ impl FilesystemStore { enabled }; + let delete_semaphore = if spec.max_concurrent_deletes > 0 { + Some(Arc::new(Semaphore::new(spec.max_concurrent_deletes))) + } else { + None + }; let shared_context = Arc::new(SharedContext { active_drop_spawns: AtomicU64::new(0), + inline_deletes: AtomicU64::new(0), + delete_semaphore, temp_path: spec.temp_path.clone(), content_path: spec.content_path.clone(), }); diff --git a/nativelink-store/src/gcs_client/mocks.rs b/nativelink-store/src/gcs_client/mocks.rs index f5ba13b5506..33a617a9cc9 100644 --- a/nativelink-store/src/gcs_client/mocks.rs +++ b/nativelink-store/src/gcs_client/mocks.rs @@ -354,7 +354,11 @@ impl GcsOperations for MockGcsOperations { start )); } - core::cmp::min(usize::try_from(e).unwrap_or(usize::MAX), content.len()) + // HTTP byte ranges are inclusive of the end index. + core::cmp::min( + usize::try_from(e).unwrap_or(usize::MAX).saturating_add(1), + content.len(), + ) } else { content.len() }; diff --git a/nativelink-store/src/gcs_store.rs b/nativelink-store/src/gcs_store.rs index 9fe7d161d56..b59ca03e852 100644 --- a/nativelink-store/src/gcs_store.rs +++ b/nativelink-store/src/gcs_store.rs @@ -436,8 +436,15 @@ where return Ok(()); } + if length == Some(0) { + writer.send_eof()?; + return Ok(()); + } let object_path = self.make_object_path(&key); - let end_offset = length.map(|len| offset + len); + // HTTP byte ranges are inclusive on both ends; the last byte index for a + // `len`-byte read from `offset` is `offset + len - 1`. `offset + len` + // over-reads one byte on any range that stops before the object end. + let end_offset = length.map(|len| offset + len - 1); let client = &self.client; let object_path_ref = &object_path; diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 451ad956cd2..61eb26d72ae 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -334,6 +334,32 @@ impl GrpcStore { })) } + /// Bounds `fut` by the configured per-RPC deadline (`rpc_timeout_s`). + /// A timeout maps to `DeadlineExceeded`, which is retryable, so the + /// existing retry machinery can route around a peer that accepted the + /// request but will never answer (e.g. a restarting shard). + async fn with_rpc_timeout( + &self, + context: &'static str, + fut: impl Future> + Send, + ) -> Result { + // Boxed so this wrapper adds a pointer to the caller's future, not a + // second inline copy of `fut` — the RPC futures it wraps are large, + // and embedding them twice measurably deepens poll stacks. + let mut fut = Box::pin(fut); + if self.rpc_timeout.is_zero() { + return fut.as_mut().await; + } + match tokio::time::timeout(self.rpc_timeout, fut.as_mut()).await { + Ok(result) => result, + Err(_elapsed) => Err(make_err!( + Code::DeadlineExceeded, + "{context} timed out after {}s", + self.rpc_timeout.as_secs() + )), + } + } + async fn perform_request(&self, input: I, mut request: F) -> Result where F: FnMut(I) -> Fut + Send + Copy, @@ -345,7 +371,7 @@ impl GrpcStore { .retry(unfold(input, move |input| async move { let input_clone = input.clone(); Some(( - request(input_clone) + self.with_rpc_timeout("GrpcStore RPC", request(input_clone)) .await .map_or_else(RetryResult::Retry, RetryResult::Ok), input, @@ -1454,11 +1480,14 @@ impl GrpcStore { }; let mut stream = match self - .read_internal(ReadRequest { - resource_name, - read_offset: 0, - read_limit: 0, - }) + .with_rpc_timeout( + "GrpcStore::get_part_compressed read", + self.read_internal(ReadRequest { + resource_name, + read_offset: 0, + read_limit: 0, + }), + ) .await { Ok(stream) => stream, @@ -2043,7 +2072,7 @@ impl StoreDriver for GrpcStore { read_limit: local_state.read_limit, }; let mut stream = match self - .read_internal(request) + .with_rpc_timeout("GrpcStore::get_part read", self.read_internal(request)) .await .err_tip(|| "in GrpcStore::get_part()") { @@ -2087,6 +2116,17 @@ impl StoreDriver for GrpcStore { return Some((RetryResult::Err(err), local_state)); } local_state.read_offset += length; + // A bounded read must also shrink its remaining limit by the + // bytes already forwarded. `read_limit` is a count measured + // from `read_offset` (REAPI ByteStream semantics), so a + // mid-stream retry that advances `read_offset` without + // decrementing `read_limit` re-requests the *original* window + // size from the new offset and over-reads by exactly the + // bytes already delivered. `read_limit == 0` means "read to + // end" and must stay unbounded across resumes. + if local_state.read_limit != 0 { + local_state.read_limit = local_state.read_limit.saturating_sub(length); + } } })) .await diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index 8a7bbb0096f..185b9f855b4 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -975,7 +975,7 @@ impl SchedulerStore for ExperimentalMongoStore { let result = match self .scheduler_collection .find_one_and_update(filter, update_doc) - .upsert(true) + .upsert(current_version == 0) .return_document(ReturnDocument::After) .await { diff --git a/nativelink-store/src/ontap_s3_store.rs b/nativelink-store/src/ontap_s3_store.rs index 224fc2df52e..1f85bb3a029 100644 --- a/nativelink-store/src/ontap_s3_store.rs +++ b/nativelink-store/src/ontap_s3_store.rs @@ -50,6 +50,7 @@ use rustls::{ClientConfig, RootCertStore}; use rustls_pki_types::CertificateDer; use rustls_pki_types::pem::PemObject; use sha2::{Digest, Sha256}; +use tokio::sync::Semaphore; use tokio::time::sleep; use tracing::{Level, event}; @@ -71,6 +72,16 @@ const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 20 * 1024 * 1024; // 20MB // Default limit for concurrent part uploads per multipart upload const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; +// Default target size for each multipart part; see the `multipart_part_size` +// config docs. Sizing parts at the 5 MiB floor maximises the part count and +// drowns the endpoint in tiny concurrent PUTs on multi-GiB blobs. +const DEFAULT_MULTIPART_PART_SIZE: u64 = 64 * 1024 * 1024; // 64MB. + +// Default store-wide admission cap on concurrent multipart uploads; see the +// `max_concurrent_multipart_uploads` config docs. A permit is held for the +// whole upload (never per part), so this cannot deadlock the part loop. +const DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS: usize = 4; + #[derive(Debug, MetricsComponent)] pub struct OntapS3Store { s3_client: Arc, @@ -86,6 +97,11 @@ pub struct OntapS3Store { max_retry_buffer_per_request: usize, #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] multipart_max_concurrent_uploads: usize, + #[metric(help = "The target size of each part in a multipart upload")] + multipart_part_size: u64, + // Store-wide admission control for concurrent multipart uploads. + // Bounds aggregate buffered memory. + multipart_upload_slots: Arc, remove_callbacks: Mutex>, } @@ -212,6 +228,19 @@ where .common .multipart_max_concurrent_uploads .unwrap_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS), + multipart_part_size: spec + .common + .multipart_part_size + .map_or(DEFAULT_MULTIPART_PART_SIZE, |size| { + u64::try_from(size).unwrap_or(u64::MAX) + }) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE), + multipart_upload_slots: Arc::new(Semaphore::new( + spec.common + .max_concurrent_multipart_uploads + .unwrap_or(DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS) + .max(1), + )), remove_callbacks: Mutex::new(vec![]), })) } @@ -425,6 +454,16 @@ where } // Handle multipart upload for large files + // Admission control: cap concurrent multipart uploads store-wide so a + // burst of large writes cannot multiply the per-upload buffer into an + // OOM. Held (RAII) for the whole upload; released on drop. + let _multipart_slot = self + .multipart_upload_slots + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "multipart upload semaphore closed: {e}"))?; + let upload_id = &self .retrier .retry(unfold((), move |()| async move { @@ -458,10 +497,23 @@ where })) .await?; - let bytes_per_upload_part = - (max_size / (MIN_MULTIPART_SIZE - 1)).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); + // Size each part at the configured target, growing it only as needed + // to keep the part count within the 10,000-part ceiling for very + // large objects. (The previous formula divided the object size by the + // 5 MiB floor — a part COUNT — and clamped that as a byte size, so + // every blob used 5 MiB parts.) + let bytes_per_upload_part = cmp::max( + self.multipart_part_size, + max_size.div_ceil(MAX_UPLOAD_PARTS as u64), + ) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); let upload_parts = move || async move { + // Together with the gated `recv` in the select loop below, this + // bounds buffered part data for one upload at roughly + // `(2 * multipart_max_concurrent_uploads + 1) * bytes_per_upload_part`: + // once the channel is full and N uploads are in flight, the + // `tx.send` above blocks and the reader stops consuming. let (tx, mut rx) = tokio::sync::mpsc::channel(self.multipart_max_concurrent_uploads); let read_stream_fut = ( @@ -548,7 +600,11 @@ where total_uploaded = result?; }, Some(upload_result) = upload_futures.next() => completed_parts.push(upload_result?), - Some(fut) = rx.recv() => upload_futures.push(fut), + // Gated so the bounded channel exerts real backpressure: + // draining it eagerly into the unbounded `FuturesUnordered` + // would free channel slots and let the reader buffer far + // ahead of the uploads. + Some(fut) = rx.recv(), if upload_futures.len() < self.multipart_max_concurrent_uploads => upload_futures.push(fut), } } @@ -622,8 +678,22 @@ where } let s3_path = &self.make_s3_path(&key); + // A zero-length bounded read has no valid inclusive byte range; satisfy + // it with an immediate EOF (matches `MemoryStore::get_part` semantics). + if length == Some(0) { + writer + .send_eof() + .err_tip(|| "Failed to send zero-length EOF in get_part")?; + return Ok(()); + } + // HTTP byte ranges are inclusive on both ends: reading `length` bytes + // from `offset` ends at byte index `offset + length - 1`. Using + // `offset + length` over-reads one byte on any range that stops before + // the object end. let end_read_byte = length - .map_or(Some(None), |length| Some(offset.checked_add(length))) + .map_or(Some(None), |length| { + Some(offset.checked_add(length).map(|end| end - 1)) + }) .err_tip(|| "Integer overflow protection triggered")?; self.retrier diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 433e07ede6c..52c0e96e508 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -52,8 +52,8 @@ use redis::cluster::ClusterClient; use redis::cluster_async::ClusterConnection; use redis::sentinel::{SentinelClient, SentinelNodeConnectionInfo, SentinelServerType}; use redis::{ - AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, ScanOptions, Script, Value, - pipe, + AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, RedisError, ScanOptions, + Script, Value, pipe, }; use serde::Deserialize; use serde::de::IntoDeserializer; @@ -90,11 +90,19 @@ const DEFAULT_RETRY_DELAY: f32 = 0.1; /// paths so a transient error doesn't fail an otherwise healthy request. const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5; +/// First delay between attempts to subscribe to eviction events after a +/// failure. It doubles after each failure, up to +/// [`EVICTION_SUBSCRIBE_MAX_DELAY`]. +const EVICTION_SUBSCRIBE_INITIAL_DELAY: Duration = Duration::from_millis(50); + +/// Longest delay between attempts to subscribe to eviction events. +const EVICTION_SUBSCRIBE_MAX_DELAY: Duration = Duration::from_secs(5); + /// Whether a Redis error is worth re-resolving the master and retrying: a /// dropped/refused connection, an IO error, or a write that hit a freshly /// demoted replica. Anything else (e.g. a real protocol/logic error) is /// returned as-is. -fn is_retryable_redis_error(err: &redis::RedisError) -> bool { +fn is_retryable_redis_error(err: &RedisError) -> bool { err.is_connection_dropped() || err.is_connection_refusal() || err.is_io_error() @@ -270,6 +278,16 @@ where Ok(()) } + /// The recorded subscriptions that are not in `subscribed`. + fn subscriptions_not_in(&self, subscribed: &HashSet) -> Vec { + self.subscriptions + .lock() + .iter() + .filter(|pattern| !subscribed.contains(*pattern)) + .cloned() + .collect() + } + async fn new(connect_func: Box>) -> Result { let connection_manager = connect_func().await?; let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT); @@ -320,36 +338,27 @@ impl RedisManager for StandardRedisManager let mut connection_manager = (self.connect_func)().await?; let new_uuid = Uuid::new_v4(); self.configure(&mut connection_manager).await?; - let subscriptions = { - let guard = self.subscriptions.lock(); - guard.iter().map(Clone::clone).collect::>() - }; - for subscription in &subscriptions { - connection_manager.psubscribe(subscription).await?; - } - let mut replayed = subscriptions.into_iter().collect::>(); - // Publish the new handle under a brief exclusive lock. Before - // publishing, replay any patterns registered since the snapshot above: - // `psubscribe` records its pattern in `self.subscriptions` *before* it - // fetches a connection, so a pattern missing from this final delta was - // inserted while we hold the write lock, and its `get_connection` will - // block here and observe the new handle — it subscribes on the live - // connection itself. Either way there is no window where a pattern is - // recorded as active but subscribed only on the old connection. - { - let mut guard = self.connection_manager.write().await; - let missing = { - let subs = self.subscriptions.lock(); - subs.iter() - .filter(|s| !replayed.contains(*s)) - .cloned() - .collect::>() - }; - for subscription in missing { + // Subscribe the new connection to every recorded pattern before it is + // published, so a pattern is never recorded as active but missing from + // the live connection. `psubscribe` records its pattern before it calls + // `get_connection`, so a pattern recorded while this runs is either + // picked up by a later pass of this loop, or is recorded after the + // final check, in which case its `get_connection` waits for the write + // lock below and then subscribes on the new connection itself. The + // subscribing happens without the write lock, so `get_connection` + // readers keep using the old handle meanwhile; the lock is held only to + // check that nothing is missing and to publish. + let mut subscribed = HashSet::new(); + loop { + for subscription in self.subscriptions_not_in(&subscribed) { connection_manager.psubscribe(&subscription).await?; - replayed.insert(subscription); + subscribed.insert(subscription); + } + let mut guard = self.connection_manager.write().await; + if self.subscriptions_not_in(&subscribed).is_empty() { + *guard = (connection_manager.clone(), new_uuid); + break; } - *guard = (connection_manager.clone(), new_uuid); } info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master"); Ok((connection_manager, new_uuid)) @@ -532,8 +541,8 @@ pub fn decode_key<'a>( impl RedisStore where - C: ConnectionLike + Clone + Sync, - M: RedisManager + Sync, + C: ConnectionLike + Clone + Send + Sync, + M: RedisManager + Send + Sync, { /// Used for testing when determinism is required. #[expect(clippy::too_many_arguments)] @@ -609,6 +618,86 @@ where } } + /// Subscribes to Redis keyspace events, so evictions and expiries reach the + /// remove callbacks. Succeeds at most once; later calls return at once. + async fn subscribe_to_eviction_events(&self) -> Result<(), Error> { + self.has_remove_callback_subscribe + .get_or_try_init(|| async { + self.warn_if_eviction_events_are_disabled().await; + let (_, uuid) = self.connection_manager.get_connection().await?; + if let Err(err) = self.connection_manager.psubscribe("__key*__:*").await { + // A dropped connection or a failover in progress: re-resolve + // the connection, as the other retry paths in this store do, + // so the next attempt does not reuse a dead one. + if matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) + && let Err(reconnect_err) = self.connection_manager.reconnect(uuid).await + { + debug!(?reconnect_err, "Reconnect after a failed PSUBSCRIBE failed"); + } + return Err(err); + } + Ok(()) + }) + .await?; + Ok(()) + } + + /// Logs if the server's `notify-keyspace-events` setting will not deliver + /// the eviction events the remove callbacks need. Advisory only: the + /// subscription is attempted either way, including when the setting + /// cannot be read (some managed Redis services block `CONFIG`). + async fn warn_if_eviction_events_are_disabled(&self) { + let cfg = async { + let mut client = self.get_client().await?; + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events") + .query_async::>(&mut client.connection_manager) + .await + .map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) + } + .await; + let cfg = match cfg { + Ok(cfg) => cfg, + Err(err) => { + warn!( + ?err, + "Could not read notify-keyspace-events; subscribing to eviction events anyway" + ); + return; + } + }; + if cfg.len() != 1 { + warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); + return; + } + let Some((_, events_cfg)) = cfg.first() else { + return; + }; + if events_cfg.is_empty() { + error!( + "notify-keyspace-events not enabled for Redis, will fail to get remove callbacks" + ); + } else if !events_cfg.contains('K') { + error!( + notify_keyspace_events = events_cfg, + "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events" + ); + } else if !events_cfg.contains('A') { + error!( + notify_keyspace_events = events_cfg, + "notify-keyspace-events does not contain 'A' so we won't get eviction events" + ); + } + // FIXME: Redis events spec appears unreliable, so we subscribe anyways + // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ + // but I'm yet to get reliable eviction events out of that + info!( + notify_keyspace_events = events_cfg, + "Attempting to subscribe to eviction events" + ); + } + async fn get_client(&self) -> Result, Error> { let local_client_permits = self.client_permits.clone(); let remaining = local_client_permits.available_permits(); @@ -1550,67 +1639,55 @@ where fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { debug!(?callback, "New callback"); + // The task holds only a weak reference, so a store that is dropped + // while the subscription is still being retried ends the task instead + // of being kept alive by it. + let weak_self = Arc::downgrade(&self); background_spawn!("remove_callback_subscribe", async move { - self.remove_callbacks.lock().await.push(callback); - // The subscription init is retried until it succeeds: this - // function already told its caller Ok(()) (e.g. - // ExistenceCacheStore, which relies on eviction callbacks for - // invalidation), so a transient Redis error here (e.g. -LOADING - // during a restart/failover) must not permanently disable remove - // callbacks. Previously a one-shot `get_or_try_init` failure was - // only logged, the OnceCell stayed empty, and nothing ever - // retried, so the `__key*__:*` subscription was silently never - // established. - let mut retry_delay = Duration::from_millis(50); + let Some(store) = weak_self.upgrade() else { + return; + }; + store.remove_callbacks.lock().await.push(callback); + drop(store); + // This function has already returned Ok(()) to its caller (an + // `ExistenceCacheStore` relies on these callbacks to invalidate its + // entries), so the subscription is retried until it succeeds. A + // single transient error, such as -LOADING during a restart or + // failover, used to leave remove callbacks off for good. + let mut retry_delay = EVICTION_SUBSCRIBE_INITIAL_DELAY; + let mut failures: u32 = 0; loop { - let local_self = self.clone(); - let init_result = self - .has_remove_callback_subscribe - .get_or_try_init(|| async move { - // The notify-keyspace-events check is advisory: we log - // misconfiguration loudly, but a failure to *read* the - // config must not block the subscription itself. - let cfg_result = async { - let mut client = local_self.get_client().await?; - redis::cmd("CONFIG").arg("GET").arg("notify-keyspace-events").to_owned().query_async::>(&mut client.connection_manager).await.map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) - }.await; - match &cfg_result { - Ok(cfg) => { - if cfg.len() != 1 { - warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); - } else if let Some((_, events_cfg)) = cfg.first() { - if events_cfg.is_empty() { - error!("notify-keyspace-events not enabled for Redis, will fail to get remove callbacks"); - } else if !events_cfg.contains('K') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events"); - } else if !events_cfg.contains('A') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'A' so we won't get eviction events"); - } - // FIXME: Redis events spec appears unreliable, so we subscribe anyways - // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ - // but I'm yet to get reliable eviction events out of that - info!(notify_keyspace_events=events_cfg, "Attempting to subscribe to eviction events"); - } - } - Err(err) => { - error!(?err, "Error while trying to initialise remove_callback_subscribe: could not verify notify-keyspace-events; subscribing anyway"); - } - } - local_self.connection_manager.psubscribe("__key*__:*").await?; - Ok::<(), Error>(()) - }) - .await; - match init_result { - Ok(()) => break, - Err(err) => { - error!( - ?err, - "Error while trying to initialise remove_callback_subscribe; will retry" + let Some(store) = weak_self.upgrade() else { + debug!("RedisStore dropped before it subscribed to eviction events"); + return; + }; + let Err(err) = store.subscribe_to_eviction_events().await else { + if failures > 0 { + info!( + failures, + "Subscribed to Redis eviction events after retrying" ); - sleep(retry_delay).await; - retry_delay = (retry_delay * 2).min(Duration::from_secs(5)); } + return; + }; + drop(store); + failures += 1; + // Log the first failure as an error, then only at powers of + // two, so a permanent failure (for example an ACL that forbids + // PSUBSCRIBE) does not log every few seconds forever. + if failures == 1 { + error!( + ?err, + "Could not subscribe to Redis eviction events; retrying until it succeeds" + ); + } else if failures.is_power_of_two() { + warn!( + ?err, + failures, "Still unable to subscribe to Redis eviction events" + ); } + sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(EVICTION_SUBSCRIBE_MAX_DELAY); } }); Ok(()) @@ -1726,6 +1803,10 @@ const CURSOR_IDLE_MS: u64 = 30_000; const DATA_FIELD_NAME: &str = "data"; /// The name of the field in the Redis hash that stores the version. const VERSION_FIELD_NAME: &str = "version"; +/// The `RediSearch` built-in property holding the document's own key. Loaded in +/// searches so a row whose record cannot be decoded can still be identified +/// and removed rather than sitting in the index forever. +const DOC_KEY_FIELD_NAME: &str = "__key"; #[allow(rustdoc::broken_intra_doc_links)] /// Lua script to set a key if the version matches. @@ -1781,6 +1862,48 @@ return {{ 1, new_version }} " ); +/// Lua script that deletes a record only if its version still matches the +/// version a search listed it at. Used to garbage-collect a record whose +/// `data` cannot be decoded: records are written atomically (delete and +/// re-insert in `LUA_VERSION_SET_SCRIPT`), so an undecodable record is never +/// a write in progress, and no writer can ever update it either, because +/// every update needs the current version from a decoded copy. The version +/// guard makes the removal safe even if that reasoning is ever violated: a +/// record rewritten after the listing is left alone. +pub const LUA_DELETE_LOST_RECORD_SCRIPT: &str = formatcp!( + r" +if redis.call('EXISTS', KEYS[2]) == 1 then + return -1 +end +if redis.call('HGET', KEYS[1], '{VERSION_FIELD_NAME}') == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +" +); + +/// Deletes the lost record at `key`, but only if the liveness guard at +/// `guard_key` does not exist (returns -1 if it does) and the record still +/// holds `version` (returns 0 if it does not). Returns 1 on deletion. +async fn delete_lost_record( + mut connection: C, + key: &str, + guard_key: &str, + version: i64, +) -> Result +where + C: ConnectionLike + Send, +{ + redis::cmd("EVAL") + .arg(LUA_DELETE_LOST_RECORD_SCRIPT) + .arg(2) + .arg(key) + .arg(guard_key) + .arg(version.to_string()) + .query_async(&mut connection) + .await +} + /// How long a freshly created index is given to finish its background /// scan before a read goes ahead regardless, and how often it is asked. const INDEX_READY_TIMEOUT: Duration = Duration::from_secs(5); @@ -2387,7 +2510,11 @@ where format!("@{}:{{ {} }}", K::INDEX_NAME, index_value) }, FtAggregateOptions { - load: vec![DATA_FIELD_NAME.into(), VERSION_FIELD_NAME.into()], + load: vec![ + DATA_FIELD_NAME.into(), + VERSION_FIELD_NAME.into(), + format!("@{DOC_KEY_FIELD_NAME}"), + ], cursor: FtAggregateCursor { count: self.max_count_per_cursor, max_idle: CURSOR_IDLE_MS, @@ -2435,6 +2562,10 @@ where }; let (connection_manager, connect_id) = self.connection_manager.get_connection().await?; + // For garbage-collecting a listed record that fails to decode. + let delete_connection = connection_manager.clone(); + let owned_index_value = index_value.to_string(); + let key_prefix = self.key_prefix.clone(); let stream = match run_ft_aggregate(connection_manager.clone()).await { // A demoted master answers READONLY and a dead/old master drops the // connection or times the command out. Both mean the master moved @@ -2518,7 +2649,11 @@ where Ok(stream) => Ok(stream), }?; - Ok(stream.filter_map(|result| async move { + Ok(stream.filter_map(move |result| { + let delete_connection = delete_connection.clone(); + let owned_index_value = owned_index_value.clone(); + let key_prefix = key_prefix.clone(); + async move { let raw_redis_map = match result { Ok(v) => v, Err(e) => { @@ -2542,6 +2677,7 @@ where let mut redis_map_iter = redis_map.iter(); let mut bytes_data: Option = None; let mut version: Option = None; + let mut doc_key: Option = None; while let Some(key) = redis_map_iter.next() { let value = redis_map_iter.next().unwrap(); let Value::BulkString(k) = key else { @@ -2581,6 +2717,15 @@ where }; version = Some(raw_version); } + DOC_KEY_FIELD_NAME => { + let Ok(str_v) = str::from_utf8(v) else { + return Some(Err(Error::new( + Code::Internal, + format!("Non-utf8 document key from ft_aggregate: {v:?}"), + ))); + }; + doc_key = Some(str_v.to_string()); + } other => { if K::MAYBE_SORT_KEY == Some(other) { // ignore sort keys @@ -2599,11 +2744,50 @@ where format!("Missing '{DATA_FIELD_NAME}' in ft_aggregate, got: {raw_redis_map:?}"), ))); }; - Some( - K::decode(version.unwrap_or(0), found_bytes_data) - .err_tip(|| "In RedisStore::search_by_index_prefix::decode"), - ) - })) + let err = match K::decode(version.unwrap_or(0), found_bytes_data) { + Ok(decoded) => return Some(Ok(decoded)), + Err(err) => err, + }; + // The record is in the index but its data does not decode (for + // example serialization drift across a rolling upgrade). It has + // no TTL and no deletion path, so without removal it is listed, + // skipped and counted by every sweep forever. Records are + // written atomically, so this is never a write in progress. + // Delete it only if the index opted in, its liveness guard says + // nothing is still waiting on it, and its version is unchanged + // since the listing; the row is still reported as lost so the + // caller counts it this one last time. + if err.code == Code::InvalidArgument + && let Some(doc_key) = doc_key + { + let record_key = doc_key.strip_prefix(&key_prefix).unwrap_or(&doc_key); + if let Some(guard_key) = + K::lost_record_liveness_guard_key(&owned_index_value, record_key) + { + let guard_key = format!("{key_prefix}{guard_key}"); + match delete_lost_record( + delete_connection, + &doc_key, + &guard_key, + version.unwrap_or(0), + ) + .await + { + Ok(removed) => warn!( + key = %doc_key, + removed, + "Undecodable record listed; removed if no longer live (1 = removed, 0 = version changed, -1 = liveness guard present)" + ), + Err(del_err) => warn!( + key = %doc_key, + ?del_err, + "Failed to remove a listed record whose data cannot be decoded" + ), + } + } + } + Some(Err(err).err_tip(|| "In RedisStore::search_by_index_prefix::decode")) + }})) } async fn count_by_index_prefix(&self, index: K) -> Result @@ -2719,33 +2903,21 @@ where } #[cfg(test)] -// This module hand-rolls a minimal RESP3 server plus a dedicated -// current-thread runtime so the psubscribe/reconnect interleaving is exact and -// timing-independent; that legitimately needs the otherwise-disallowed tokio -// runtime/spawn primitives. -#[expect( - clippy::disallowed_methods, - reason = "deterministic in-test RESP server needs a hand-built runtime and raw spawns" -)] -#[expect( - clippy::doc_markdown, - reason = "choreography docs reference bare Redis/RESP identifiers (PSUBSCRIBE, Sha1) by name" -)] -#[expect( - trivial_casts, - reason = "explicit Pin> coercion for the connect closure's return type" -)] -mod psubscribe_reconnect_race_tests { +mod psubscribe_reconnect_tests { use core::time::Duration; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; + use nativelink_macro::nativelink_test; + use nativelink_util::{background_spawn, spawn}; use redis::aio::{ConnectionManager, ConnectionManagerConfig}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Notify; + use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + use tokio::time::timeout; - use super::{RedisManager, StandardRedisManager}; + use super::{RedisConnectFn, RedisManager, StandardRedisManager}; /// Reads one CRLF-terminated line from the stream. async fn read_line(stream: &mut TcpStream) -> Option { @@ -2765,7 +2937,7 @@ mod psubscribe_reconnect_race_tests { Some(String::from_utf8_lossy(&line).into_owned()) } - /// Parses one RESP command (array of bulk strings) from the client. + /// Parses one RESP command (an array of bulk strings) from the client. async fn parse_command(stream: &mut TcpStream) -> Option> { let header = read_line(stream).await?; if !header.starts_with('*') { @@ -2783,21 +2955,51 @@ mod psubscribe_reconnect_race_tests { Some(parts) } + /// A `PSUBSCRIBE` reply the fake server holds back on the reconnect + /// connection until the test releases it. + #[derive(Default)] + struct HeldReply { + /// Notified when the `PSUBSCRIBE` arrives. + arrived: Notify, + /// The reply is sent once this is notified. + release: Notify, + } + + #[derive(Default)] struct FakeServerState { - /// Sha1 to return for SCRIPT LOAD (must match the client's own hash). + /// SHA1 to return for `SCRIPT LOAD` (must match the client's own hash). script_hash: String, - /// Patterns each connection (by accept order) was PSUBSCRIBEd to. + /// Patterns each connection, by accept order, was subscribed to. patterns_per_conn: parking_lot::Mutex>>, - /// Signalled when conn #1 (the reconnect connection) receives its - /// first PSUBSCRIBE (the replay of the pre-existing pattern). - replay_psubscribe_arrived: Notify, - /// The server holds conn #1's first PSUBSCRIBE reply until this is - /// notified, keeping reconnect() parked inside its replay loop. - release_replay_reply: Notify, + /// Replies to hold back on connection 1, the one `reconnect` opens, + /// keyed by pattern. + held_on_reconnect: HashMap>, + } + + impl FakeServerState { + fn new(held_patterns: &[&str]) -> Arc { + Arc::new(Self { + script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) + .get_hash() + .to_string(), + held_on_reconnect: held_patterns + .iter() + .map(|pattern| ((*pattern).to_string(), Arc::default())) + .collect(), + ..Self::default() + }) + } + + fn held(&self, pattern: &str) -> &HeldReply { + &self.held_on_reconnect[pattern] + } + + fn patterns_on(&self, conn_idx: usize) -> Vec { + self.patterns_per_conn.lock()[conn_idx].clone() + } } async fn handle_conn(mut stream: TcpStream, conn_idx: usize, state: Arc) { - let mut first_psubscribe_on_this_conn = true; while let Some(cmd) = parse_command(&mut stream).await { let name = cmd .first() @@ -2820,13 +3022,11 @@ mod psubscribe_reconnect_race_tests { "PSUBSCRIBE" => { let pattern = cmd.get(1).cloned().unwrap_or_default(); state.patterns_per_conn.lock()[conn_idx].push(pattern.clone()); - if conn_idx == 1 && first_psubscribe_on_this_conn { - first_psubscribe_on_this_conn = false; - // Tell the test that reconnect() is now parked awaiting - // this reply inside its replay loop, then hold the - // reply until the test releases it. - state.replay_psubscribe_arrived.notify_one(); - state.release_replay_reply.notified().await; + if conn_idx == 1 + && let Some(held) = state.held_on_reconnect.get(&pattern) + { + held.arrived.notify_one(); + held.release.notified().await; } let reply = format!( ">3\r\n$10\r\npsubscribe\r\n${}\r\n{}\r\n:1\r\n", @@ -2842,121 +3042,134 @@ mod psubscribe_reconnect_race_tests { } } - /// Deterministic repro: a `psubscribe` that interleaves with a `reconnect` - /// whose subscription snapshot was taken before the insert lands only on - /// the old (about-to-be-replaced) connection. The new live connection is - /// never subscribed, yet the pattern is recorded as active in - /// `self.subscriptions`, so notifications are silently lost until an - /// unrelated future reconnect replays the set. - /// - /// Choreography (all real await points, no timing races): - /// 1. Manager connects (conn #0) and psubscribes "pre". - /// 2. Task B calls reconnect(): connects conn #1, snapshots - /// subscriptions = {"pre"}, and parks in the replay loop because the - /// fake server withholds conn #1's PSUBSCRIBE "pre" reply. - /// 3. The test (task A) calls psubscribe("pat"): get_connection still - /// returns conn #0 (reconnect has not published yet), the pattern is - /// inserted into `subscriptions`, and PSUBSCRIBE succeeds on conn #0. - /// 4. The server releases the reply; reconnect publishes conn #1. - /// 5. Conn #1 — now the live connection — was never subscribed to "pat". - #[test] - fn psubscribe_lost_across_reconnect() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - runtime.block_on(async { - let state = Arc::new(FakeServerState { - script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) - .get_hash() - .to_string(), - patterns_per_conn: parking_lot::Mutex::new(Vec::new()), - replay_psubscribe_arrived: Notify::new(), - release_replay_reply: Notify::new(), - }); - - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - { - let state = state.clone(); - tokio::spawn(async move { - let mut conn_idx = 0usize; - loop { - let (stream, _) = listener.accept().await.unwrap(); - state.patterns_per_conn.lock().push(Vec::new()); - tokio::spawn(handle_conn(stream, conn_idx, state.clone())); - conn_idx += 1; - } - }); + /// Starts the fake server and returns a manager connected to it, the + /// generation id of its first connection, and the receiver for server + /// pushes, which must stay alive for the length of the test. + async fn connect( + state: &Arc, + ) -> ( + Arc>, + uuid::Uuid, + UnboundedReceiver, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server_state = state.clone(); + background_spawn!("fake_redis_accept", async move { + let mut conn_idx = 0usize; + loop { + let (stream, _) = listener.accept().await.unwrap(); + server_state.patterns_per_conn.lock().push(Vec::new()); + background_spawn!( + "fake_redis_conn", + handle_conn(stream, conn_idx, server_state.clone()) + ); + conn_idx += 1; } + }); - // Keep the push receiver alive for the whole test. - let (push_tx, _push_rx) = tokio::sync::mpsc::unbounded_channel(); - let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); - let connect_func = Box::new(move || { - let url = url.clone(); - let push_tx = push_tx.clone(); - Box::pin(async move { - let client = redis::Client::open(url.as_str()).unwrap(); - let config = ConnectionManagerConfig::new() - .set_number_of_retries(0) - .set_connection_timeout(Some(Duration::from_secs(5))) - .set_response_timeout(Some(Duration::from_secs(5))) - .set_push_sender(push_tx); - ConnectionManager::new_with_config(client, config) - .await - .map_err(Into::into) - }) as core::pin::Pin> - }); - - let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); - - // Step 1: a pre-existing subscription so reconnect's replay loop - // has an await point we can park it on. - manager.psubscribe("pre").await.unwrap(); - let (_conn, uuid0) = manager.get_connection().await.unwrap(); - - // Step 2: start reconnect; it will connect conn #1, snapshot - // subscriptions (= {"pre"}), and park replaying "pre". - let reconnect_task = { - let manager = manager.clone(); - tokio::spawn(async move { manager.reconnect(uuid0).await }) - }; - tokio::time::timeout( - Duration::from_secs(5), - state.replay_psubscribe_arrived.notified(), - ) + let (push_tx, push_rx) = unbounded_channel(); + let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); + let connect_func: Box> = Box::new(move || { + let url = url.clone(); + let push_tx = push_tx.clone(); + Box::pin(async move { + let client = redis::Client::open(url.as_str()).unwrap(); + let config = ConnectionManagerConfig::new() + .set_number_of_retries(0) + .set_connection_timeout(Some(Duration::from_secs(5))) + .set_response_timeout(Some(Duration::from_secs(5))) + .set_push_sender(push_tx); + ConnectionManager::new_with_config(client, config) + .await + .map_err(Into::into) + }) + }); + let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); + let (_, uuid) = manager.get_connection().await.unwrap(); + (manager, uuid, push_rx) + } + + /// A `psubscribe` that runs while a `reconnect` is subscribing its new + /// connection must end up subscribed on that new connection. + /// + /// The interleaving, with every step at a real await point: + /// 1. The manager connects (connection 0) and subscribes to `pre`. + /// 2. `reconnect` opens connection 1 and parks while subscribing it to + /// `pre`, because the fake server holds that reply. + /// 3. The test subscribes to `pat`. `get_connection` still returns + /// connection 0, so `pat` lands there. + /// 4. The server releases the reply and `reconnect` publishes connection 1. + /// + /// Before the fix, `reconnect` only replayed a snapshot taken before step 3, + /// so the live connection never subscribed to `pat` and its events were lost + /// until some later reconnect. + #[nativelink_test] + async fn psubscribe_during_reconnect_lands_on_new_connection() { + let state = FakeServerState::new(&["pre"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } + }); + timeout(Duration::from_secs(5), state.held("pre").arrived.notified()) .await - .expect("reconnect never reached its subscription replay loop"); - - // Step 3: concurrent psubscribe while reconnect is parked - // post-snapshot, pre-publish. It succeeds — on the old connection. - manager.psubscribe("pat").await.unwrap(); - - // Step 4: let reconnect finish and publish conn #1. - state.release_replay_reply.notify_one(); - let (_new_conn, new_uuid) = reconnect_task.await.unwrap().unwrap(); - assert_ne!( - uuid0, new_uuid, - "reconnect must have published a new connection" - ); + .expect("reconnect never started subscribing its new connection"); - // The manager believes "pat" is an active subscription... - let recorded: HashSet = manager.subscriptions.lock().clone(); - assert!( - recorded.contains("pat"), - "sanity: pattern recorded as active" - ); + manager.psubscribe("pat").await.unwrap(); - // ...but the live connection (conn #1) was never subscribed to it. - let conn1_patterns = state.patterns_per_conn.lock()[1].clone(); - assert!( - conn1_patterns.contains(&"pat".to_string()), - "RACE REPRODUCED: pattern \"pat\" is recorded as an active \ - subscription but the live post-reconnect connection was only \ - subscribed to {conn1_patterns:?}; keyspace notifications for \ - \"pat\" are silently lost until the next reconnect", - ); + state.held("pre").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + + let recorded: HashSet = manager.subscriptions.lock().clone(); + assert!( + recorded.contains("pat"), + "pattern must be recorded as active" + ); + let live = state.patterns_on(1); + assert!( + live.contains(&"pat".to_string()), + "pattern \"pat\" is recorded as active but the live connection is only \ + subscribed to {live:?}" + ); + } + + /// While `reconnect` subscribes patterns recorded after its first pass, it + /// must not hold the connection write lock, or every `get_connection` + /// (every Redis operation) waits on those network round trips. + #[nativelink_test] + async fn reconnect_catch_up_does_not_block_get_connection() { + let state = FakeServerState::new(&["pre", "pat"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } }); + timeout(Duration::from_secs(5), state.held("pre").arrived.notified()) + .await + .expect("reconnect never started subscribing its new connection"); + manager.psubscribe("pat").await.unwrap(); + state.held("pre").release.notify_one(); + + // `reconnect` now subscribes connection 1 to `pat`, and the server holds + // that reply. Readers must still get the old connection meanwhile. + timeout(Duration::from_secs(5), state.held("pat").arrived.notified()) + .await + .expect("reconnect never caught up on the pattern recorded mid-reconnect"); + let (_conn, uuid) = timeout(Duration::from_secs(2), manager.get_connection()) + .await + .expect("get_connection blocked while reconnect was subscribing") + .unwrap(); + assert_eq!(uuid, uuid0, "the new connection must not be published yet"); + + state.held("pat").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + assert!(state.patterns_on(1).contains(&"pat".to_string())); } } diff --git a/nativelink-store/src/s3_store.rs b/nativelink-store/src/s3_store.rs index 34e532a542a..617e1a14c5b 100644 --- a/nativelink-store/src/s3_store.rs +++ b/nativelink-store/src/s3_store.rs @@ -51,7 +51,7 @@ use nativelink_util::store_trait::{ RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo, }; use parking_lot::Mutex; -use tokio::sync::mpsc; +use tokio::sync::{Semaphore, mpsc}; use tokio::time::sleep; use tracing::{error, info}; @@ -75,6 +75,17 @@ const MAX_MULTIPART_SIZE: u64 = 5 * 1024 * 1024 * 1024; // 5GB. // Note: Type 'u64' chosen to simplify calculations const MAX_UPLOAD_PARTS: u64 = 10_000; +// Default target size for each multipart part. Deliberately far above the +// 5 MiB floor: sizing parts at the floor turns a multi-GiB blob into +// thousands of tiny PUTs, which churns and poisons the HTTP connection pool +// (aws-smithy "connection never set" storms) and makes large uploads slow and +// flaky — on a lossy link the retries exhaust and the write fails outright. +// 64 MiB keeps a 12 GiB blob at ~192 parts; combined with the concurrency +// limit it bounds in-flight memory to +// `multipart_max_concurrent_uploads * multipart_part_size`. +// Note: If you change this, adjust the docs in the config. +const DEFAULT_MULTIPART_PART_SIZE: u64 = 64 * 1024 * 1024; // 64MB. + // Default max buffer size for retrying upload requests. // Note: If you change this, adjust the docs in the config. const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5MB. @@ -83,6 +94,16 @@ const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5MB. // Note: If you change this, adjust the docs in the config. const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; +// Default store-wide admission cap on concurrent multipart uploads. The +// per-upload concurrency and part size bound the buffer of a single upload +// at roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size` +// (~1.3 GiB at the defaults); without a store-wide limit, N simultaneous +// large writes multiply that footprint unbounded. Four concurrent uploads +// bounds the default worst case near 5 GiB. A permit is held for the whole +// upload (never per part), so this cannot deadlock the inner part loop. +// Note: If you change this, adjust the docs in the config. +const DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS: usize = 4; + #[derive(Debug, MetricsComponent)] pub struct S3Store { s3_client: Arc, @@ -98,6 +119,12 @@ pub struct S3Store { max_retry_buffer_per_request: usize, #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] multipart_max_concurrent_uploads: usize, + #[metric(help = "The target size of each part in a multipart upload")] + multipart_part_size: u64, + // Store-wide admission control for concurrent multipart uploads (see + // DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS). Bounds aggregate buffered + // memory. + multipart_upload_slots: Arc, remove_callbacks: Mutex>, } @@ -212,6 +239,19 @@ where .common .multipart_max_concurrent_uploads .unwrap_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS), + multipart_part_size: spec + .common + .multipart_part_size + .map_or(DEFAULT_MULTIPART_PART_SIZE, |size| { + u64::try_from(size).unwrap_or(u64::MAX) + }) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE), + multipart_upload_slots: Arc::new(Semaphore::new( + spec.common + .max_concurrent_multipart_uploads + .unwrap_or(DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS) + .max(1), + )), remove_callbacks: Mutex::new(Vec::new()), })) } @@ -422,6 +462,20 @@ where .await; } + // Admission control: cap concurrent multipart uploads store-wide so a + // burst of large writes cannot multiply the per-upload buffer + // (`multipart_max_concurrent_uploads * multipart_part_size`) into an + // OOM. The permit is held (RAII) for the whole upload and acquired + // after the single-request early-return above, so small blobs never + // wait on a slot; it is whole-upload admission, never per-part, so it + // cannot deadlock the part loop below. + let _multipart_slot = self + .multipart_upload_slots + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "multipart upload semaphore closed: {e}"))?; + let upload_id = &self .retrier .retry(unfold((), move |()| async move { @@ -458,26 +512,37 @@ where // S3 requires us to upload in parts if the size is greater than 5GB. The part size must be at least // 5MB (except last part) and can have up to 10,000 parts. - // Calculate of number of chunks if we upload in 5MB chucks (min chunk size), clamping to - // 10,000 parts and correcting for lossy integer division. This provides the - let chunk_count = (max_size / MIN_MULTIPART_SIZE).clamp(0, MAX_UPLOAD_PARTS - 1) + 1; + // Size each part at the configured target, growing it only as needed + // to keep the part count within the 10,000-part ceiling for very + // large objects. Sizing at the 5 MiB floor (the old behaviour) + // maximised the part count and drowned providers in tiny concurrent + // PUTs. + let bytes_per_upload_part = cmp::max( + self.multipart_part_size, + max_size.div_ceil(MAX_UPLOAD_PARTS), + ) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); - // Using clamped first approximation of number of chunks, calculate byte count of each - // chunk, excluding last chunk, clamping to min/max upload size 5MB, 5GB. - let bytes_per_upload_part = - (max_size / chunk_count).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); + // Number of parts this yields (last part may be short). + let chunk_count = max_size.div_ceil(bytes_per_upload_part); - // Sanity check before continuing. - if !(MIN_MULTIPART_SIZE..MAX_MULTIPART_SIZE).contains(&bytes_per_upload_part) { + // Sanity check before continuing. Inclusive of MAX_MULTIPART_SIZE: + // the construction-time clamp allows a configured part size of + // exactly the 5 GiB provider maximum. + if !(MIN_MULTIPART_SIZE..=MAX_MULTIPART_SIZE).contains(&bytes_per_upload_part) { return Err(make_err!( Code::FailedPrecondition, - "Failed to calculate file chuck size (min, max, calc): {MIN_MULTIPART_SIZE}, {MAX_MULTIPART_SIZE}, {bytes_per_upload_part}", + "Failed to calculate file chunk size (min, max, calc): {MIN_MULTIPART_SIZE}, {MAX_MULTIPART_SIZE}, {bytes_per_upload_part}", )); } let upload_parts = move || async move { - // This will ensure we only have `multipart_max_concurrent_uploads` * `bytes_per_upload_part` - // bytes in memory at any given time waiting to be uploaded. + // Together with the gated `recv` in the select loop below, this + // bounds buffered part data for one upload at roughly + // `(2 * multipart_max_concurrent_uploads + 1) * bytes_per_upload_part`: + // up to N parts uploading, N queued in the channel, and one being + // read. Once the channel is full and N uploads are in flight, the + // `tx.send` above blocks and the reader stops consuming. let (tx, mut rx) = mpsc::channel(self.multipart_max_concurrent_uploads); let read_stream_fut = async move { @@ -559,7 +624,11 @@ where total_uploaded = result?; }, // Return error or wait for other futures. Some(upload_result) = upload_futures.next() => completed_parts.push(upload_result?), - Some(fut) = rx.recv() => upload_futures.push(fut), + // Gated so the bounded channel exerts real backpressure: + // draining it eagerly into the unbounded `FuturesUnordered` + // would free channel slots and let the reader buffer far + // ahead of the uploads. + Some(fut) = rx.recv(), if upload_futures.len() < self.multipart_max_concurrent_uploads => upload_futures.push(fut), } } @@ -637,8 +706,22 @@ where } let s3_path = &self.make_s3_path(&key); + // A zero-length bounded read has no valid inclusive byte range; satisfy + // it with an immediate EOF (matches `MemoryStore::get_part` semantics). + if length == Some(0) { + writer + .send_eof() + .err_tip(|| "Failed to send zero-length EOF in get_part")?; + return Ok(()); + } + // HTTP byte ranges are inclusive on both ends: reading `length` bytes + // from `offset` ends at byte index `offset + length - 1`. Using + // `offset + length` over-reads one byte on any range that stops before + // the object end. let end_read_byte = length - .map_or(Some(None), |length| Some(offset.checked_add(length))) + .map_or(Some(None), |length| { + Some(offset.checked_add(length).map(|end| end - 1)) + }) .err_tip(|| "Integer overflow protection triggered")?; self.retrier diff --git a/nativelink-store/src/shard_store.rs b/nativelink-store/src/shard_store.rs index 2f249448199..56f644e69e5 100644 --- a/nativelink-store/src/shard_store.rs +++ b/nativelink-store/src/shard_store.rs @@ -15,6 +15,8 @@ use core::hash::Hasher; use core::ops::BitXor; use core::pin::Pin; +use core::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use core::time::Duration; use std::hash::DefaultHasher; use std::sync::Arc; @@ -22,13 +24,39 @@ use async_trait::async_trait; use futures::future::try_join_all; use futures::stream::{FuturesUnordered, TryStreamExt}; use nativelink_config::stores::ShardSpec; -use nativelink_error::{Error, ResultExt, error_if}; +use nativelink_error::{Code, Error, ResultExt, error_if}; use nativelink_metric::MetricsComponent; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; use nativelink_util::store_trait::{ RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; +use tokio::time::Instant; + +/// Number of consecutive connectivity failures after which a shard is +/// considered suspect and skipped as the first read target. +const SUSPECT_FAILURE_THRESHOLD: u32 = 3; + +/// How long a suspect shard is skipped on the read path before it is +/// probed again. +const SUSPECT_COOLDOWN: Duration = Duration::from_secs(5); + +/// Errors that indicate the shard itself is unreachable (as opposed to a +/// definitive answer like `NotFound`), making read failover safe. +const fn is_connectivity_error(err: &Error) -> bool { + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) +} + +/// Read-path health state for one shard. Writes never consult this: a +/// blob's home shard is fixed by the ring, so writes always target it. +#[derive(Debug, Default)] +struct ShardHealth { + /// Consecutive connectivity failures observed on this shard. + consecutive_failures: AtomicU32, + /// Nanoseconds (since `ShardStore` creation) of the most recent + /// connectivity failure. Zero means "never failed". + last_failure_nanos: AtomicU64, +} #[derive(Debug, MetricsComponent)] struct StoreAndWeight { @@ -36,6 +64,7 @@ struct StoreAndWeight { weight: u32, #[metric(help = "The underlying store")] store: Store, + health: ShardHealth, } #[derive(Debug, MetricsComponent)] @@ -47,6 +76,8 @@ pub struct ShardStore { help = "The weights and stores that are used to determine which store to use" )] weights_and_stores: Vec, + /// Anchor for `ShardHealth::last_failure_nanos` timestamps. + created_at: Instant, } impl ShardStore { @@ -85,8 +116,107 @@ impl ShardStore { weights_and_stores: weights .into_iter() .zip(stores) - .map(|(weight, store)| StoreAndWeight { weight, store }) + .map(|(weight, store)| StoreAndWeight { + weight, + store, + health: ShardHealth::default(), + }) .collect(), + created_at: Instant::now(), + })) + } + + fn now_nanos(&self) -> u64 { + u64::try_from(self.created_at.elapsed().as_nanos()).unwrap_or(u64::MAX) + } + + /// A shard is suspect once it has accumulated enough consecutive + /// connectivity failures and its cooldown has not yet elapsed. After + /// the cooldown it is probed again (treated as healthy). + fn is_suspect(&self, store_idx: usize) -> bool { + let health = &self.weights_and_stores[store_idx].health; + if health.consecutive_failures.load(Ordering::Relaxed) < SUSPECT_FAILURE_THRESHOLD { + return false; + } + let last_failure_nanos = health.last_failure_nanos.load(Ordering::Relaxed); + let cooldown_nanos = u64::try_from(SUSPECT_COOLDOWN.as_nanos()).unwrap_or(u64::MAX); + self.now_nanos().saturating_sub(last_failure_nanos) < cooldown_nanos + } + + fn note_success(&self, store_idx: usize) { + self.weights_and_stores[store_idx] + .health + .consecutive_failures + .store(0, Ordering::Relaxed); + } + + fn note_connectivity_failure(&self, store_idx: usize) { + let health = &self.weights_and_stores[store_idx].health; + health.consecutive_failures.fetch_add(1, Ordering::Relaxed); + health + .last_failure_nanos + .store(self.now_nanos(), Ordering::Relaxed); + } + + /// Runs `has_with_results` against the ring starting at `primary_idx`, + /// failing over in ring order on connectivity errors (each shard is + /// tried at most once). Suspect shards are skipped unless every shard + /// is suspect, in which case the primary is probed directly. + async fn has_on_ring( + &self, + primary_idx: usize, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let num_stores = self.weights_and_stores.len(); + let mut last_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + results.fill(None); + last_err = Some(err); + } + Err(err) => { + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(primary_idx); + return Ok(()); + } + Err(err) => { + if is_connectivity_error(&err) { + self.note_connectivity_failure(primary_idx); + } + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + Err(last_err.unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::has_with_results()".to_string(), + ) })) } @@ -167,11 +297,7 @@ impl StoreDriver for ShardStore { if keys.len() == 1 { // Hot path: It is very common to lookup only one key. let store_idx = self.get_store_index(&keys[0]); - let store = &self.weights_and_stores[store_idx].store; - return store - .has_with_results(keys, results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}}"); + return self.has_on_ring(store_idx, keys, results).await; } let mut keys_for_store: Vec<(KeyIdxVec, KeyVec)> = self .weights_and_stores @@ -192,12 +318,9 @@ impl StoreDriver for ShardStore { .into_iter() .enumerate() .map(|(store_idx, (key_idxs, keys))| async move { - let store = &self.weights_and_stores[store_idx].store; let mut inner_results = vec![None; keys.len()]; - store - .has_with_results(&keys, &mut inner_results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}")?; + self.has_on_ring(store_idx, &keys, &mut inner_results) + .await?; Result::<_, Error>::Ok((key_idxs, inner_results)) }) .collect(); @@ -231,11 +354,66 @@ impl StoreDriver for ShardStore { offset: u64, length: Option, ) -> Result<(), Error> { - let store = self.get_store(&key); - store - .get_part(key, writer, offset, length) - .await - .err_tip(|| "In ShardStore::get_part()") + let primary_idx = self.get_store_index(&key); + let num_stores = self.weights_and_stores.len(); + let bytes_written_start = writer.get_bytes_written(); + let mut last_connectivity_err = None; + let mut last_other_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store + .get_part(key.borrow(), &mut *writer, offset, length) + .await + { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + // Failover is only safe if nothing reached the client. + if writer.get_bytes_written() != bytes_written_start { + return Err(err).err_tip(|| "In ShardStore::get_part()"); + } + last_connectivity_err = Some(err); + } + Err(err) if err.code == Code::NotFound && store_idx != primary_idx => { + // A failover shard legitimately may not have the blob; + // keep looking, but never let it mask a connectivity + // error from the shard that owns the key. + self.note_success(store_idx); + last_other_err = Some(err); + } + Err(err) => return Err(err).err_tip(|| "In ShardStore::get_part()"), + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + let result = store.get_part(key, writer, offset, length).await; + match &result { + Ok(()) => self.note_success(primary_idx), + Err(err) if is_connectivity_error(err) => { + self.note_connectivity_failure(primary_idx); + } + Err(_) => {} + } + return result.err_tip(|| "In ShardStore::get_part()"); + } + Err(last_connectivity_err.or(last_other_err).unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::get_part()".to_string(), + ) + })) + .err_tip(|| "In ShardStore::get_part()") } fn inner_store(&self, key: Option) -> &'_ dyn StoreDriver { diff --git a/nativelink-store/tests/completeness_checking_store_test.rs b/nativelink-store/tests/completeness_checking_store_test.rs index ed3dcd2c38e..01ba27d89b1 100644 --- a/nativelink-store/tests/completeness_checking_store_test.rs +++ b/nativelink-store/tests/completeness_checking_store_test.rs @@ -43,6 +43,21 @@ const STDOUT: DigestInfo = DigestInfo::new([5u8; 32], 0); const STDERR: DigestInfo = DigestInfo::new([6u8; 32], 0); async fn setup() -> Result<(Arc, Arc, DigestInfo), Error> { + let (ac_store, cas_store, action_result_digest, _tree_digest) = setup_with_tree().await?; + Ok((ac_store, cas_store, action_result_digest)) +} + +/// Like `setup`, but also returns the digest of the `Tree` the action result's +/// output directory points at. +async fn setup_with_tree() -> Result< + ( + Arc, + Arc, + DigestInfo, + DigestInfo, + ), + Error, +> { let backend_store = Store::new(MemoryStore::new(&MemorySpec::default())); let cas_store = MemoryStore::new(&MemorySpec::default()); let ac_store = @@ -114,7 +129,7 @@ async fn setup() -> Result<(Arc, Arc, Di ) .await?; - Ok((ac_store, cas_store, action_result_digest)) + Ok((ac_store, cas_store, action_result_digest, tree_digest)) } #[nativelink_test] @@ -241,12 +256,14 @@ async fn verify_completeness_get() -> Result<(), Error> { cas_store.remove_entry(OUTPUT_FILE.into()).await; - assert!( - ac_store - .get_part_unchunked(action_result_digest, 0, None) - .await - .is_err(), - ".get() should fail with item missing in CAS", + let err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err(".get() should fail with item missing in CAS"); + assert_eq!( + err.code, + Code::NotFound, + "a missing referent must stay a NotFound miss: {err:?}" ); } @@ -320,14 +337,14 @@ impl StoreDriver for AlwaysErrStore { default_health_status_indicator!(AlwaysErrStore); -/// Regression: a real backend read/decode failure on the action-result record -/// must propagate with its underlying, non-`NotFound` code — it must NOT be -/// collapsed into `Code::NotFound`. A `NotFound` miss tells the client the -/// result does not exist and to re-execute; a backend fault (unavailable -/// Redis, decode failure) is transient and retryable, and serving it as a -/// clean miss silently discards a valid cached result. The genuine -/// eviction-miss path (referent absent) is covered by -/// `verify_has_function_call_checks_cas` and stays `NotFound`. +/// Regression: a backend failure reading the action-result record must +/// propagate with its own code, not be turned into `Code::NotFound`. A +/// `NotFound` miss tells the client the result does not exist and to +/// re-execute; a backend fault (an unavailable Redis, a timeout) is transient +/// and retryable, and serving it as a miss throws away a valid cached result. +/// The misses that must stay `NotFound` are covered by +/// `verify_has_function_call_checks_cas`, `absent_record_is_a_miss` and +/// `absent_tree_is_a_miss`. #[nativelink_test] async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Error> { let make_store = || { @@ -339,7 +356,7 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err }; let some_digest = DigestInfo::new([9u8; 32], 123); - // get_part path (nix narinfo GET / Bazel AC get). + // get_part path (an AC GetActionResult). let get_err = make_store() .get_part_unchunked(some_digest, 0, None) .await @@ -355,7 +372,7 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err "backend fault must propagate with its real code" ); - // has path (nix narinfo HEAD / Bazel AC has). + // has path. let has_err = make_store() .has(some_digest) .await @@ -373,3 +390,85 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err Ok(()) } + +/// An action result that is not in the AC at all is a plain miss. +#[nativelink_test] +async fn absent_record_is_a_miss() -> Result<(), Error> { + let (ac_store, _cas_store, _action_result_digest) = setup().await?; + let absent = DigestInfo::new([8u8; 32], 64); + + assert_eq!(ac_store.has(absent).await?, None); + let err = ac_store + .get_part_unchunked(absent, 0, None) + .await + .expect_err("an absent action result must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +/// An action result whose output directory's `Tree` blob has been evicted from +/// the CAS is incomplete, so it is a miss, not an error. +#[nativelink_test] +async fn absent_tree_is_a_miss() -> Result<(), Error> { + let (ac_store, cas_store, action_result_digest, tree_digest) = setup_with_tree().await?; + cas_store.remove_entry(tree_digest.into()).await; + + assert_eq!(ac_store.has(action_result_digest).await?, None); + let err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err("an action result with an evicted Tree must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +/// A failure checking the referenced outputs in the CAS keeps its own code. It +/// used to reach the client as `Internal`, which retry policies treat as +/// permanent. +#[nativelink_test] +async fn cas_backend_error_propagates_with_its_code() -> Result<(), Error> { + let backend_store = Store::new(MemoryStore::new(&MemorySpec::default())); + let cas_store = Store::new(Arc::new(AlwaysErrStore { + code: Code::Unavailable, + })); + let ac_store = CompletenessCheckingStore::new(backend_store, cas_store); + + // No output directories, so the only CAS call is the existence check of + // the output file, stdout and stderr. + let action_result = ProtoActionResult { + output_files: vec![OutputFile { + digest: Some(OUTPUT_FILE.into()), + ..Default::default() + }], + stdout_digest: Some(STDOUT.into()), + stderr_digest: Some(STDERR.into()), + ..Default::default() + }; + let action_result_digest = serialize_and_upload_message( + &action_result, + ac_store.as_pin(), + &mut DigestHasherFunc::Blake3.hasher(), + ) + .await?; + + let has_err = ac_store + .has(action_result_digest) + .await + .expect_err("a CAS failure must not be swallowed"); + assert_eq!( + has_err.code, + Code::Unavailable, + "unexpected error: {has_err:?}" + ); + + let get_err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err("a CAS failure must not be swallowed"); + assert_eq!( + get_err.code, + Code::Unavailable, + "unexpected error: {get_err:?}" + ); + Ok(()) +} diff --git a/nativelink-store/tests/dedup_store_test.rs b/nativelink-store/tests/dedup_store_test.rs index e72f2905a62..22850e323f4 100644 --- a/nativelink-store/tests/dedup_store_test.rs +++ b/nativelink-store/tests/dedup_store_test.rs @@ -12,11 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use nativelink_config::stores::{DedupSpec, MemorySpec, StoreSpec}; use nativelink_error::{Code, Error, ResultExt}; use nativelink_macro::nativelink_test; use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS; -use nativelink_store::dedup_store::DedupStore; +use nativelink_store::compression_store::WincodeConfig; +use nativelink_store::dedup_store::{DedupIndex, DedupStore}; use nativelink_store::memory_store::MemoryStore; use nativelink_util::common::DigestInfo; use nativelink_util::store_trait::{Store, StoreLike}; @@ -445,3 +448,177 @@ async fn empty_cas_upload_does_not_publish_dedup_index() -> Result<(), Error> { assert_eq!(content.len_for_test(), 0); Ok(()) } + +/// Builds a `DedupStore` whose index for `VALID_HASH1` names one content chunk +/// of `declared_size` bytes, and stores `stored` under that chunk's digest. A +/// `stored` of any other length is a chunk whose bytes disagree with its index +/// entry, as a torn write or a misbehaving content store would leave it. +async fn store_with_one_chunk( + declared_size: usize, + stored: Vec, +) -> Result<(Arc, Arc, DigestInfo, DigestInfo), Error> { + let index_store = MemoryStore::new(&MemorySpec::default()); + let content_store = MemoryStore::new(&MemorySpec::default()); + let store = DedupStore::new( + &make_default_config(), + Store::new(index_store.clone()), + Store::new(content_store.clone()), + )?; + + let chunk_digest = DigestInfo::try_new(VALID_HASH2, declared_size).unwrap(); + content_store + .update_oneshot(chunk_digest, stored.into()) + .await + .err_tip(|| "Failed to write content chunk")?; + + let serialized = wincode::config::serialize( + &DedupIndex { + entries: vec![chunk_digest], + }, + WincodeConfig::new(), + ) + .unwrap(); + let top_digest = DigestInfo::try_new(VALID_HASH1, declared_size).unwrap(); + index_store + .update_oneshot(top_digest, serialized.into()) + .await + .err_tip(|| "Failed to write index entry")?; + Ok((store, content_store, top_digest, chunk_digest)) +} + +// A content chunk must contain exactly the number of bytes its index entry +// declares. A short chunk used to be served as a truncated read that ended in a +// normal EOF. The read must fail with NotFound, and the bad chunk must be gone +// afterwards so `has()` stops reporting the blob as present. +#[nativelink_test] +async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { + let (store, content_store, top_digest, chunk_digest) = + store_with_one_chunk(100, vec![7u8; 10]).await?; + + let result = store.get_part_unchunked(top_digest, 0, Some(100)).await; + let err = result.expect_err("a short chunk must not be served as a successful read"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + + assert_eq!( + content_store.has(chunk_digest).await?, + None, + "the short chunk must be dropped from the content store" + ); + assert_eq!( + store.has(top_digest).await?, + None, + "a blob whose chunk was dropped must no longer be reported as present" + ); + Ok(()) +} + +// A read that starts inside a short chunk used to trip the `bytes_to_skip` +// assert and panic. It must fail like any other read of that chunk. +#[nativelink_test] +async fn get_part_offset_inside_short_chunk_is_error_test() -> Result<(), Error> { + let (store, _content_store, top_digest, _chunk_digest) = + store_with_one_chunk(100, vec![7u8; 10]).await?; + + let result = store.get_part_unchunked(top_digest, 50, Some(10)).await; + let err = result.expect_err("a read inside a short chunk must fail, not panic"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +// A chunk longer than its index entry declares is just as wrong as a short one. +#[nativelink_test] +async fn get_part_long_chunk_is_error_test() -> Result<(), Error> { + let (store, content_store, top_digest, chunk_digest) = + store_with_one_chunk(100, vec![7u8; 150]).await?; + + let result = store.get_part_unchunked(top_digest, 0, None).await; + let err = result.expect_err("a long chunk must not be served as a successful read"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + assert_eq!(content_store.has(chunk_digest).await?, None); + Ok(()) +} + +// End to end: a blob with one corrupted chunk fails to read, is then reported +// missing, and the client's re-upload makes it readable again. Before the bad +// chunk was dropped, `update()` skipped it because the content store still had +// a value under its digest, so the blob stayed unreadable forever. +#[nativelink_test] +async fn short_chunk_is_rewritten_by_the_next_upload_test() -> Result<(), Error> { + let index_store = MemoryStore::new(&MemorySpec::default()); + let content_store = MemoryStore::new(&MemorySpec::default()); + let store = DedupStore::new( + &make_default_config(), + Store::new(index_store.clone()), + Store::new(content_store.clone()), + )?; + + let original_data = make_random_data(MEGABYTE_SZ); + let digest = DigestInfo::try_new(VALID_HASH1, MEGABYTE_SZ).unwrap(); + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to write data to dedup store")?; + + // Truncate one chunk from the middle of the blob in place. + let index_bytes = index_store.get_part_unchunked(digest, 0, None).await?; + let index: DedupIndex = + wincode::config::deserialize(&index_bytes, WincodeConfig::new()).unwrap(); + assert!(index.entries.len() > 2, "expected several chunks"); + let victim = index.entries[index.entries.len() / 2]; + let victim_bytes = content_store.get_part_unchunked(victim, 0, None).await?; + content_store + .update_oneshot(victim, victim_bytes.slice(..victim_bytes.len() / 2)) + .await?; + + let err = store + .get_part_unchunked(digest, 0, None) + .await + .expect_err("a blob with a truncated chunk must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + assert_eq!( + store.has(digest).await?, + None, + "the blob must now be missing" + ); + + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to re-upload data to dedup store")?; + let rt_data = store + .get_part_unchunked(digest, 0, None) + .await + .err_tip(|| "The re-uploaded blob must read back")?; + assert_eq!(rt_data, original_data, "Expected round trip data to match"); + Ok(()) +} + +// Guard against over-correction: a bounded request for MORE bytes than the blob +// actually contains must still succeed and return the available bytes, matching +// the clamp-to-available semantics of the underlying stores. The +// size-consistency check must not turn a legitimate over-length request into an +// error. +#[nativelink_test] +async fn get_part_over_length_request_still_ok_test() -> Result<(), Error> { + const DATA_SIZE: usize = 64 * 1024; + + let store = DedupStore::new( + &make_default_config(), + Store::new(MemoryStore::new(&MemorySpec::default())), + Store::new(MemoryStore::new(&MemorySpec::default())), + )?; + + let original_data = make_random_data(DATA_SIZE); + let digest = DigestInfo::try_new(VALID_HASH1, DATA_SIZE).unwrap(); + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to write data to dedup store")?; + + let rt = store + .get_part_unchunked(digest, 0, Some((DATA_SIZE as u64) * 2)) + .await + .err_tip(|| "Over-length bounded read should still succeed")?; + assert_eq!(rt, original_data, "Expected available bytes to be returned"); + Ok(()) +} diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index dbbac9ee502..6ffdc9d8f75 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -2411,3 +2411,263 @@ async fn incomplete_slow_read_never_populates_memory_or_filesystem() -> Result<( } Ok(()) } + +/// A filesystem store, which refuses a write whose length is not its declared +/// `ExactSize`, with one-byte blocks so `has()` reports the stored length. +async fn make_filesystem_fast_store() -> Result<(Arc, FilesystemSpec), Error> { + let spec = FilesystemSpec { + content_path: make_temp_path("content_path"), + temp_path: make_temp_path("temp_path"), + block_size: 1, + ..Default::default() + }; + Ok((FilesystemStore::::new(&spec).await?, spec)) +} + +/// Reads `key` through a `FastSlowStore` whose fast tier is `fast` and whose +/// slow tier alone holds `data`, then checks the fast tier was filled with it. +async fn check_populates_filesystem_fast_store( + fast: Arc, + fast_spec: FilesystemSpec, + slow: Store, + key: StoreKey<'_>, + data: &Bytes, +) -> Result<(), Error> { + let fast = Store::new(fast); + let store = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_spec), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast.clone(), + slow.clone(), + )); + slow.update_oneshot(key.borrow(), data.clone()).await?; + assert_eq!(fast.has(key.borrow()).await?, None); + + assert_eq!( + store.get_part_unchunked(key.borrow(), 0, None).await?, + data, + "{key}" + ); + assert_eq!( + fast.has(key.borrow()).await?, + Some(data.len() as u64), + "{key}" + ); + assert_eq!(fast.get_part_unchunked(key.borrow(), 0, None).await?, data); + Ok(()) +} + +/// A slow store whose `has()` reports `reported_size` for whatever it holds, +/// the way a tier that reports a physical size does. +#[derive(MetricsComponent)] +struct MisreportingSlowStore { + inner: Arc, + reported_size: u64, +} + +#[async_trait] +impl StoreDriver for MisreportingSlowStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + Pin::new(self.inner.as_ref()) + .has_with_results(keys, results) + .await?; + for result in results.iter_mut().filter(|result| result.is_some()) { + *result = Some(self.reported_size); + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + Pin::new(self.inner.as_ref()) + .update(key, reader, size_info) + .await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + Pin::new(self.inner.as_ref()) + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _digest: Option) -> &'_ dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(MisreportingSlowStore); + +// The populate cannot take a digest's size from the slow tier's `has()` on +// trust: a tier that reports less (a compressed stream) or more (a +// block-rounded footprint) than the blob must still fill the fast tier with +// the blob as it is. +#[nativelink_test] +async fn populate_does_not_trust_slow_store_has_size_for_digest() -> Result<(), Error> { + let data = Bytes::from(make_random_data(3000)); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + for reported_size in [1, 4096] { + let slow = Store::new(Arc::new(MisreportingSlowStore { + inner: MemoryStore::new(&MemorySpec::default()), + reported_size, + })); + let (fast, fast_spec) = make_filesystem_fast_store().await?; + check_populates_filesystem_fast_store(fast, fast_spec, slow, digest.into(), &data).await?; + } + Ok(()) +} + +// An action cache entry is keyed by the action's digest and a dedup index entry +// by its blob's, so a digest's size is not always the size of what it names. +// The slow tier's `has()` is right here and the digest is not, so the populate +// must not declare the digest's size as exact either. +#[nativelink_test] +async fn populate_keeps_entry_whose_size_is_not_its_digests() -> Result<(), Error> { + let data = Bytes::from(make_random_data(3000)); + let digest = DigestInfo::try_new(VALID_HASH, 100)?; + let slow = Store::new(MemoryStore::new(&MemorySpec::default())); + let (fast, fast_spec) = make_filesystem_fast_store().await?; + check_populates_filesystem_fast_store(fast, fast_spec, slow, digest.into(), &data).await +} + +// The shipped shape of the bug: a filesystem slow tier at the default 4 KiB +// block size reports its block-rounded footprint from `has()` (4096 for a +// 17-byte blob), and the populate declares that to the fast tier as an +// `ExactSize`. Since the populate checks the bytes it received against that +// size, every read that has to fill the fast tier from such a slow tier fails +// with DataLoss unless the blob happens to be a block multiple. This is the +// dedup index store in `nativelink-config/examples/filesystem_cas.json5` and +// the action cache in `kubernetes/nativelink/nativelink-config.json5`. +#[nativelink_test] +async fn populate_from_filesystem_slow_tier_at_default_block_size() -> Result<(), Error> { + let slow_spec = FilesystemSpec { + content_path: make_temp_path("slow_content_path"), + temp_path: make_temp_path("slow_temp_path"), + ..Default::default() + }; + let slow = Store::new(FilesystemStore::::new(&slow_spec).await?); + let fast = Store::new(MemoryStore::new(&MemorySpec::default())); + let data = Bytes::from_static(b"seventeen bytes!!"); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + slow.update_oneshot(digest, data.clone()).await?; + // Before #3006 a filesystem store's `has()` reported its block-rounded + // size on disk (4096 here); since then it reports the content length. + // Either way the populate must not refuse the 17 bytes it gets. + assert!(slow.has(digest).await?.is_some()); + let store = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Memory(MemorySpec::default()), + slow: StoreSpec::Filesystem(slow_spec), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast.clone(), + slow.clone(), + )); + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(got, data); + assert_eq!(fast.has(digest).await?, Some(data.len() as u64)); + Ok(()) +} + +// Regression test for N14: an interrupted ByteStream upload is parked for +// resumption, keeping its FastSlowStore::update() future (and in-flight +// slow-write registration) alive. Existence-check waiters that subscribed to +// that in-flight write must be woken when the SAME digest is completed by a +// DIFFERENT upload (fresh upload id), not starve until their client deadline. +#[nativelink_test] +async fn stalled_upload_waiter_released_by_fresh_upload_of_same_digest() -> Result<(), Error> { + let (fast_slow_store, _fast_store, _slow_store) = make_stores(); + let data = make_random_data(100); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + + // Upload A: starts streaming, then stalls forever mid-stream (simulates a + // cancelled/half-closed ByteStream upload parked for resumption). We keep + // `stalled_tx` alive so the update() future stays pending. + let (mut stalled_tx, stalled_rx) = make_buf_channel_pair(); + let store_a = fast_slow_store.clone(); + let upload_a = tokio::spawn(async move { + store_a + .update(digest, stalled_rx, UploadSizeInfo::ExactSize(100)) + .await + }); + stalled_tx.send(Bytes::from(make_random_data(10))).await?; + // Let upload A register its in-flight slow write. + tokio::task::yield_now().await; + + // Waiter: FindMissingBlobs-style existence check that parks on the + // in-flight write of upload A. + let store_w = fast_slow_store.clone(); + let waiter = tokio::spawn(async move { store_w.has(digest).await }); + // Let the waiter subscribe to the in-flight write. + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!waiter.is_finished(), "Waiter should be parked on upload A"); + + // Upload B: completes the same digest under a "fresh upload id". + fast_slow_store + .update_oneshot(digest, data.clone().into()) + .await?; + + // The waiter must be released promptly by upload B's completion. + let result = tokio::time::timeout(Duration::from_secs(2), waiter) + .await + .map_err(|_| { + make_err!( + Code::DeadlineExceeded, + "Waiter hung after same digest was completed by a fresh upload" + ) + })? + .expect("waiter task panicked")?; + assert_eq!( + result, + Some(100), + "Waiter should see the completed blob size" + ); + + // A genuinely-absent digest must still resolve promptly once its only + // (stalled) writer goes away. + drop(stalled_tx); + drop(upload_a.await); + let absent = DigestInfo::try_new(VALID_HASH, 12345)?; + let absent_result = tokio::time::timeout(Duration::from_secs(2), fast_slow_store.has(absent)) + .await + .map_err(|_| make_err!(Code::DeadlineExceeded, "Absent digest check hung"))??; + assert_eq!(absent_result, None); + + Ok(()) +} diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 59f6ac30a33..161b7715262 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -245,11 +245,17 @@ async fn read_file_contents(file_name: &OsStr) -> Result, Error> { } async fn wait_for_no_open_files() -> Result<(), Error> { + // `get_open_files_for_test` reads a process-global open-handle counter, so + // this drains handles opened by *any* concurrently running test in the same + // binary. Under asan (and other slow-execution configs) a 1s budget is too + // tight and this spuriously times out even though the handles do close, so + // give the drain a generous ceiling — the loop still exits the instant the + // count reaches zero, so a healthy run pays nothing for the larger cap. let mut counter = 0; while fs::get_open_files_for_test() != 0 { sleep(Duration::from_millis(1)).await; counter += 1; - if counter > 1000 { + if counter > 10000 { return Err(make_err!( Code::Internal, "Timed out waiting all files to close" @@ -363,6 +369,35 @@ const HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789a const HASH2: &str = "0123456789abcdef000000000000000000020000000000000123456789abcdef"; const VALUE1: &str = "0123456789"; const VALUE2: &str = "9876543210"; + +// Longer, in `read_buffer_size: 1` chunks, than `make_buf_channel_pair`'s +// in-flight capacity, so a reader that has taken only the first byte keeps the +// store's `get` stalled with the file still open. Tests that need an open +// stream on a file use these instead of `VALUE1`/`VALUE2`. +const LONG_VALUE1: &str = concat!( + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", +); +const LONG_VALUE2: &str = concat!( + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", +); const STRING_NAME: &str = "String_Filename"; #[nativelink_test] @@ -523,7 +558,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> } } - let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; + let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?; let content_path = make_temp_path("content_path"); let temp_path = make_temp_path("temp_path"); @@ -543,7 +578,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> ); // Insert data into store. - store.update_oneshot(digest1, VALUE1.into()).await?; + store.update_oneshot(digest1, LONG_VALUE1.into()).await?; let (writer, mut reader) = make_buf_channel_pair(); let store_clone = store.clone(); @@ -561,13 +596,13 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> .err_tip(|| "Error reading first byte")?; assert_eq!( first_byte[0], - VALUE1.as_bytes()[0], + LONG_VALUE1.as_bytes()[0], "Expected first byte to match" ); } // Replace content. - store.update_oneshot(digest1, VALUE2.into()).await?; + store.update_oneshot(digest1, LONG_VALUE2.into()).await?; // Ensure we let any background tasks finish. tokio::task::yield_now().await; @@ -586,7 +621,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> let data = read_file_contents(path.as_os_str()).await?; assert_eq!( &data[..], - VALUE1.as_bytes(), + LONG_VALUE1.as_bytes(), "Expected file content to match" ); } @@ -603,7 +638,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> assert_eq!( &remaining_file_data, - &VALUE1.as_bytes()[1..], + &LONG_VALUE1.as_bytes()[1..], "Expected file content to match" ); @@ -631,8 +666,8 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { } } - let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; - let digest2 = DigestInfo::try_new(HASH2, VALUE2.len())?; + let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?; + let digest2 = DigestInfo::try_new(HASH2, LONG_VALUE2.len())?; let content_path = make_temp_path("content_path"); let temp_path = make_temp_path("temp_path"); @@ -652,7 +687,10 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { ); // Insert data into store. - store.update_oneshot(digest1, VALUE1.into()).await.unwrap(); + store + .update_oneshot(digest1, LONG_VALUE1.into()) + .await + .unwrap(); let mut reader = { let (writer, reader) = make_buf_channel_pair(); @@ -668,7 +706,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { assert!(reader.peek().await.is_ok(), "Could not peek into reader"); // Insert new content. This will evict the old item. - store.update_oneshot(digest2, VALUE2.into()).await?; + store.update_oneshot(digest2, LONG_VALUE2.into()).await?; // Ensure we let any background tasks finish. tokio::task::yield_now().await; @@ -687,7 +725,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { let data = read_file_contents(path.as_os_str()).await?; assert_eq!( &data[..], - VALUE1.as_bytes(), + LONG_VALUE1.as_bytes(), "Expected file content to match" ); } @@ -702,7 +740,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { .await .err_tip(|| "Error reading remaining bytes")?; - assert_eq!(&reader_data, VALUE1, "Expected file content to match"); + assert_eq!(&reader_data, LONG_VALUE1, "Expected file content to match"); loop { if DELETES_FINISHED.load(Ordering::Relaxed) == 1 { @@ -715,6 +753,62 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { check_storage_dir_empty(&temp_path).await } +// Bounding background deletes with `max_concurrent_deletes` must not change the +// correctness of eviction: an evicted entry's backing file is still removed. +// This exercises the bounded (semaphore) delete path — permit acquired, held +// for the lifetime of the delete, then released. The default unbounded path is +// covered by the rest of this suite; the inline-overflow fallback (which fires +// only under saturation) is exercised under real load by the furnace via the +// `inline_deletes` metric. +#[nativelink_test] +async fn bounded_delete_semaphore_still_removes_evicted_files() -> Result<(), Error> { + static DELETES_FINISHED: AtomicU32 = AtomicU32::new(0); + struct LocalHooks {} + impl FileEntryHooks for LocalHooks { + fn on_drop(_file_entry: &Fe) { + DELETES_FINISHED.fetch_add(1, Ordering::Relaxed); + } + } + + let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; + let digest2 = DigestInfo::try_new(HASH2, VALUE2.len())?; + let content_path = make_temp_path("content_path"); + let temp_path = make_temp_path("temp_path"); + + let store = Arc::new( + FilesystemStore::>::new(&FilesystemSpec { + content_path: content_path.clone(), + temp_path: temp_path.clone(), + eviction_policy: Some(EvictionPolicy { + max_count: 1, + ..Default::default() + }), + // The bound under test. With it set, the delete routes through the + // semaphore path rather than the unconditional spawn. + max_concurrent_deletes: 1, + block_size: 1, + read_buffer_size: 1, + ..Default::default() + }) + .await?, + ); + + // Insert, then evict digest1 by inserting digest2 (max_count == 1). + store.update_oneshot(digest1, VALUE1.into()).await?; + store.update_oneshot(digest2, VALUE2.into()).await?; + + // Wait for the eviction's background delete to finish. + loop { + if DELETES_FINISHED.load(Ordering::Relaxed) == 1 { + break; + } + tokio::task::yield_now().await; + } + + // The evicted file must be gone from the temp staging area. + check_storage_dir_empty(&temp_path).await +} + // Test to ensure that if we are holding a reference to `FileEntry` and the contents are // replaced, the `FileEntry` continues to use the old data. // `FileEntry` file contents should be immutable for the lifetime of the object. diff --git a/nativelink-store/tests/gcs_store_test.rs b/nativelink-store/tests/gcs_store_test.rs index fce8d76b6a6..8a7ccfaa45b 100644 --- a/nativelink-store/tests/gcs_store_test.rs +++ b/nativelink-store/tests/gcs_store_test.rs @@ -549,8 +549,8 @@ async fn get_part_with_range() -> Result<(), Error> { assert_eq!(*read_requests[0].0, 6, "Expected start offset to be 6"); assert_eq!( *read_requests[0].1, - Some(11), - "Expected end offset to be Some(11)" + Some(10), + "Expected inclusive end offset to be Some(10)" ); Ok(()) @@ -1043,3 +1043,52 @@ async fn incomplete_upload_never_publishes_gcs_object() -> Result<(), Error> { } Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_is_end_inclusive() -> Result<(), Error> { + // A strict subrange (offset + length < size) must return exactly `length` + // bytes with an *inclusive* end offset. Regression for the off-by-one that + // passed `offset + length` (one byte too many) to the HTTP byte range. + let mock_ops = Arc::new(MockGcsOperations::new()); + let store = create_test_store(mock_ops.clone()).await?; + let digest = DigestInfo::try_new(VALID_HASH1, 11)?; // "hello world" + let store_key: StoreKey = to_store_key(digest); + let object_path = create_object_path(&store_key); + mock_ops + .add_object(&object_path, b"hello world".to_vec()) + .await; + let (mut tx, mut rx) = make_buf_channel_pair(); + + // bytes [2, 5) = "llo" (3 bytes), ending strictly before the object end. + let store_clone = store.clone(); + let get_fut = nativelink_util::spawn!("get_part_subrange_task", async move { + store_clone.get_part(store_key, &mut tx, 2, Some(3)).await + }); + let received_data = rx.consume(Some(100)).await?; + assert_eq!( + received_data.as_ref(), + b"llo", + "strict subrange must return exactly `length` bytes" + ); + get_fut.await??; + + let requests = mock_ops.get_requests().await; + let read_requests: Vec<_> = requests + .iter() + .filter_map(|req| { + if let MockRequest::ReadContent { start, end, .. } = req { + Some((*start, *end)) + } else { + None + } + }) + .collect(); + assert_eq!(read_requests.len(), 1, "Expected one read content request"); + assert_eq!(read_requests[0].0, 2, "Expected start offset 2"); + assert_eq!( + read_requests[0].1, + Some(4), + "Expected inclusive end offset Some(4) for a 3-byte read from offset 2" + ); + Ok(()) +} diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index a11c6fb518e..156dde6e6fa 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -123,6 +123,11 @@ struct FakeStreamServer { fail_mid_upload: Arc>, failure_code: tonic::Code, require_zero_start: bool, + /// Full byte payload the fake `read` serves (sliced from `read_offset`). + read_payload: Arc>, + /// `Some(n)`: the FIRST `read` serves `n` bytes then a retryable error, + /// forcing the client to resume; later reads serve the remainder. + first_read_partial: Option, } impl FakeStreamServer { @@ -136,6 +141,8 @@ impl FakeStreamServer { fail_mid_upload: Arc::new(Mutex::new(0)), failure_code: tonic::Code::Unavailable, require_zero_start: false, + read_payload: Arc::new(RAW_INPUT.as_bytes().to_vec()), + first_read_partial: None, } } @@ -149,10 +156,6 @@ impl FakeStreamServer { type ReadStream = Pin> + Send + 'static>>; -struct ReaderState { - responded: bool, -} - #[tonic::async_trait] impl ByteStream for FakeStreamServer { type ReadStream = ReadStream; @@ -177,19 +180,39 @@ impl ByteStream for FakeStreamServer { } } let read_request = grpc_request.into_inner(); - self.read_requests.lock().await.push(ReadRequestHolder { - request: read_request, - metadata: request_metadata, - }); - - let folded = unfold(ReaderState { responded: false }, async move |state| { - if state.responded { - return None; + let offset = usize::try_from(read_request.read_offset).unwrap_or(0); + let this_read = { + let mut reqs = self.read_requests.lock().await; + reqs.push(ReadRequestHolder { + request: read_request, + metadata: request_metadata, + }); + reqs.len() + }; + let payload = self.read_payload.clone(); + let partial = self.first_read_partial; + + // Step machine: 0 => emit a data chunk, 1 => emit a retryable error + // (only on a partial first read), 2 => EOF. Default behaviour (no + // `first_read_partial`) is one chunk from `read_offset` then EOF. + let folded = unfold((0u8, payload), move |(step, payload)| async move { + match step { + 0 => { + if let (1, Some(f)) = (this_read, partial) { + let start = offset.min(payload.len()); + let end = (offset + f).min(payload.len()); + let chunk = payload[start..end].to_vec(); + return Some((Ok(ReadResponse { data: chunk.into() }), (1u8, payload))); + } + let chunk = payload.get(offset..).unwrap_or(&[]).to_vec(); + Some((Ok(ReadResponse { data: chunk.into() }), (2u8, payload))) + } + 1 => Some(( + Err(Status::unavailable("simulated mid-stream drop")), + (2u8, payload), + )), + _ => None, } - let response = ReadResponse { - data: RAW_INPUT.as_bytes().into(), - }; - Some((Ok(response), ReaderState { responded: true })) }); Ok(Response::new(Box::pin(folded))) } @@ -1108,3 +1131,67 @@ async fn incomplete_oneshot_cas_upload_never_sends_remote_write() -> Result<(), assert!(server.write_requests.lock().await.is_empty()); Ok(()) } + +/// A 10-byte payload, longer than the partial-forward below, so a bounded read +/// has a meaningful remaining window to carry across a resume. +const RESUME_PAYLOAD: &[u8] = b"0123456789"; + +/// Regression for the `GrpcStore` bounded-read resume over-read: a mid-stream +/// retry advanced `read_offset` but left `read_limit` at its original value, so +/// the resume re-requested the full window from the new offset and over-read by +/// exactly the bytes already forwarded. The deterministic witness is the second +/// (resume) `ReadRequest` on the wire: its `read_limit` must shrink by the +/// forwarded byte count. Fails without the fix by construction (limit stays L). +#[nativelink_test] +async fn bounded_read_resume_shrinks_read_limit() -> Result<(), Error> { + const FORWARDED: usize = 3; + let requested_len = RESUME_PAYLOAD.len(); + + let server = FakeStreamServer { + read_payload: Arc::new(RESUME_PAYLOAD.to_vec()), + first_read_partial: Some(FORWARDED), + ..FakeStreamServer::new() + }; + let (server, port) = spawn_bytestream_server(server).await; + let mut spec = test_spec(format!("http://localhost:{port}"), false); + spec.retry.max_retries = 1; + let store = GrpcStore::new(&spec)?; + let digest = DigestInfo::try_new(VALID_HASH, requested_len).unwrap(); + + let (tx, mut rx) = make_buf_channel_pair(); + let get_fut = store.get_part(digest, tx, 0, Some(u64::try_from(requested_len).unwrap())); + let drain_fut = async move { + let mut total = 0usize; + loop { + match rx.recv().await { + Ok(chunk) if !chunk.is_empty() => total += chunk.len(), + _ => break, + } + } + total + }; + let (res, total) = futures::join!(get_fut, drain_fut); + res?; + assert_eq!( + total, requested_len, + "writer must receive exactly the requested bytes across the resume" + ); + + let reqs = server.read_requests.lock().await; + assert!( + reqs.len() >= 2, + "expected a resume read, saw {}", + reqs.len() + ); + let limit = i64::try_from(requested_len).unwrap(); + let fwd = i64::try_from(FORWARDED).unwrap(); + assert_eq!(reqs[0].request.read_offset, 0); + assert_eq!(reqs[0].request.read_limit, limit); + assert_eq!(reqs[1].request.read_offset, fwd); + assert_eq!( + reqs[1].request.read_limit, + limit - fwd, + "resume read_limit must shrink by bytes already forwarded (over-read bug)" + ); + Ok(()) +} diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs index d6133631254..3933ec5b554 100644 --- a/nativelink-store/tests/mongo_store_test.rs +++ b/nativelink-store/tests/mongo_store_test.rs @@ -839,7 +839,10 @@ async fn test_scheduler_store_operations() -> Result<(), Error> { data.version = version1; // Using old version let result = helper.store.update_data(data.clone(), None).await; - assert!(result.is_err(), "Update with old version should fail"); + assert!( + matches!(result, Ok(None)), + "Stale-version update must return retryable Ok(None), got {result:?}" + ); eprintln!("Correctly rejected update with stale version"); } @@ -1056,3 +1059,50 @@ async fn incomplete_upload_preserves_mongo_value() -> Result<(), Error> { } Ok(()) } + +#[nativelink_test] +async fn update_data_stale_version_returns_retryable_none_not_err() -> Result<(), Error> { + // N12 regression: ExperimentalMongoStore::update_data used upsert(true), so a + // version-CAS miss (a concurrent writer already bumped the version) hit the + // unique _id and surfaced as a DuplicateKeyError (Err) instead of the + // retryable Ok(None) the scheduler contract (and the Redis backend) require. + let helper = TestMongoHelper::new(None).await?; + + let mut data = TestSchedulerData { + key: "test:n12_version_cas".to_string(), + content: "v1".to_string(), + version: 0, + }; + + // Seed at version 0 -> creates version 1 (genuine first create still upserts). + let v1 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "create should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on create"))?; + + // One writer wins with the correct version -> version 2. + data.content = "v2".to_string(); + data.version = v1; + let v2 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "correct-version update should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on update"))?; + assert!(v2 > v1, "version should increment ({v1} -> {v2})"); + + // The losing concurrent writer carries the now-stale version. It MUST get a + // retryable Ok(None), not Err. Before the fix this returned Err(Internal) + // (duplicate _id), destroying the retry signal. + data.content = "stale".to_string(); + data.version = v1; + let stale = helper.store.update_data(data.clone(), None).await; + assert!( + matches!(stale, Ok(None)), + "stale-version update must return retryable Ok(None), got {stale:?}" + ); + + Ok(()) +} diff --git a/nativelink-store/tests/oci_store_test.rs b/nativelink-store/tests/oci_store_test.rs index d2eb5cadbd5..1d092465fd1 100644 --- a/nativelink-store/tests/oci_store_test.rs +++ b/nativelink-store/tests/oci_store_test.rs @@ -250,7 +250,7 @@ async fn ranged_get_sends_range_header_path_style() -> Result<(), Error> { .uri(format!( "https://{HOST}/{BUCKET}/{VALID_HASH1}-1000?x-id=GetObject" )) - .header("range", format!("bytes={OFFSET}-{}", OFFSET + LENGTH)) + .header("range", format!("bytes={OFFSET}-{}", OFFSET + LENGTH - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder() diff --git a/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs b/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs index e195c7df687..299453eb209 100644 --- a/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs +++ b/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs @@ -69,6 +69,8 @@ async fn create_test_store(mock_client: StaticReplayClient) -> Result Result<(), Error> { "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject" ) ) - .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH)) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder().status(StatusCode::OK).body(SdkBody::empty()).unwrap() @@ -491,7 +491,7 @@ async fn ensure_empty_string_in_stream_works_test() -> Result<(), Error> { "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{CAS_ENTRY_SIZE}?x-id=GetObject" ) ) - .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE)) + .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE - 1)) .body(SdkBody::empty()) .unwrap(), http::Response @@ -719,6 +719,12 @@ async fn multipart_update_large_cas() -> Result<(), Error> { bucket: BUCKET_NAME.to_string(), vserver_name: VSERVER_NAME.to_string(), endpoint: "https://example.com".to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }), s3_client, @@ -773,3 +779,62 @@ async fn incomplete_small_upload_never_sends_ontap_put() -> Result<(), Error> { } Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_range_is_end_inclusive() -> Result<(), Error> { + // A bounded read ending strictly before the object end must request an + // *inclusive* HTTP byte range of exactly `length` bytes. Regression for the + // off-by-one that requested `offset + length` (one byte too many). The + // StaticReplayClient only serves the exact-match range header, so without + // the fix the store emits `bytes=100-300`, fails to match, and errors. + const AC_ENTRY_SIZE: u64 = 1000; + const OFFSET: u64 = 100; + const LENGTH: u64 = 200; // OFFSET + LENGTH = 300 < AC_ENTRY_SIZE: a strict subrange. + let expected_payload = vec![0x88_u8; usize::try_from(LENGTH).unwrap()]; + let mock_client = StaticReplayClient::new(vec![ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", + )) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(expected_payload.clone())) + .unwrap(), + )]); + let test_config = Builder::new() + .behavior_version(BehaviorVersion::latest()) + .region(Region::from_static(VSERVER_NAME)) + .http_client(mock_client.clone()) + .build(); + let s3_client = aws_sdk_s3::Client::from_conf(test_config); + let store = OntapS3Store::new_with_client_and_jitter( + &(ExperimentalOntapS3Spec { + bucket: BUCKET_NAME.to_string(), + vserver_name: VSERVER_NAME.to_string(), + endpoint: "https://example.com".to_string(), + ..Default::default() + }), + s3_client, + Arc::new(move |_delay| Duration::from_secs(0)), + MockInstantWrapped::default, + )?; + + let got = store + .get_part_unchunked( + DigestInfo::try_new(VALID_HASH1, AC_ENTRY_SIZE)?, + OFFSET, + Some(LENGTH), + ) + .await?; + assert_eq!( + got.len(), + usize::try_from(LENGTH).unwrap(), + "bounded read returned wrong byte count" + ); + assert_eq!(got, Bytes::from(expected_payload)); + mock_client.assert_requests_match(&[]); + Ok(()) +} diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 6b3d43a8135..f8e85711343 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -26,13 +26,13 @@ use nativelink_macro::nativelink_test; use nativelink_redis_tester::{ ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response, add_to_response_raw, fake_redis_sentinel_master_stream, fake_redis_sentinel_stream, fake_redis_stream, - make_fake_redis_with_responses, + make_fake_redis_with_first_responses, make_fake_redis_with_responses, }; use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS; use nativelink_store::existence_cache_store::ExistenceCacheStore; use nativelink_store::redis_store::{ ClusterRedisManager, DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE, DEFAULT_MAX_COUNT_PER_CURSOR, - LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, decode_key, + LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, StandardRedisManager, decode_key, }; use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; @@ -44,6 +44,7 @@ use nativelink_util::store_trait::{ TrueValue, UploadSizeInfo, }; use pretty_assertions::assert_eq; +use redis::aio::ConnectionManager; use redis::{PushInfo, RedisError, Value, make_extension_error}; use redis_test::{MockCmd, MockRedisConnection}; use tokio::sync::mpsc::UnboundedReceiver; @@ -566,6 +567,94 @@ async fn test_large_downloads_are_chunked() -> Result<(), Error> { Ok(()) } +// C10 (read integrity): equal-length in-place replacement tears a chunked read. +// +// `get_part` reads a value larger than `read_chunk_size` with one `GETRANGE` +// per chunk, each a separate round-trip. The value under a key is mutable: +// `update` lands data on a temp key and `RENAME`s it onto the final key, +// discarding whatever was there (CAS content-addressing makes this a non-issue +// for CAS, but the AC and arbitrary-key uses of this same driver mutate values +// legitimately). If a writer replaces the value between two chunk reads with a +// *different value of the same total length*, the reader stitches bytes from +// two different values and returns them as success. +// +// A length guard cannot catch this: STRLEN is identical before, during, and +// after the replacement. This test pins the current (buggy) behavior so the +// hole is captured in-tree and so any future fix has a fail->pass target. +// When a fix lands that establishes stable value identity across chunks, this +// test's final assertion flips from "returns torn bytes as Ok" to "returns an +// error", and the comment/name should change with it. +#[nativelink_test] +async fn c10_equal_length_replacement_tears_chunked_read() -> Result<(), Error> { + const READ_CHUNK_SIZE: usize = 1024; + // Two values of identical total length. Value A is all 'A's, value B all + // 'B's. The reader asks for A (via the first chunk) and then, mid-read, + // the key is swapped to B, so the second chunk comes from B. + let total_len = READ_CHUNK_SIZE + 128; + let value_a = Bytes::from(vec![b'A'; total_len]); + let value_b = Bytes::from(vec![b'B'; total_len]); + assert_eq!( + value_a.len(), + value_b.len(), + "same length is the whole point" + ); + + let digest = DigestInfo::try_new(VALID_HASH1, 1)?; + let real_key = format!("{digest}"); + + let commands = vec![ + // get_part goes straight to GETRANGE per chunk (no up-front STRLEN): + // the existence/length pipeline lives in has_with_results, a separate + // call. STRLEN would be identical for A and B anyway -- the swap is + // indistinguishable by length. + // First chunk: bytes [0, READ_CHUNK_SIZE) of value A. + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(0) + .arg(READ_CHUNK_SIZE.try_into().unwrap_or(i64::MAX) - 1), + Ok(Value::BulkString(value_a.slice(..READ_CHUNK_SIZE).into())), + ), + // *** The value is replaced by B here (equal length). *** + // Second chunk: bytes [READ_CHUNK_SIZE, total_len) now come from B. + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key) + .arg(READ_CHUNK_SIZE.try_into().unwrap_or(i64::MAX)) + .arg(total_len.try_into().unwrap_or(i64::MAX) - 1), + Ok(Value::BulkString(value_b.slice(READ_CHUNK_SIZE..).into())), + ), + ]; + + let store = make_mock_store(commands).await; + + let get_result = store + .get_part_unchunked(digest, 0, Some(total_len as u64)) + .await; + + // The stitched result is neither A nor B: it is the head of A spliced onto + // the tail of B -- a torn blob served as success. This is the C10 hole. + let torn = get_result.expect("current code returns the torn blob as Ok"); + let mut expected_torn = Vec::with_capacity(total_len); + expected_torn.extend_from_slice(&value_a[..READ_CHUNK_SIZE]); + expected_torn.extend_from_slice(&value_b[READ_CHUNK_SIZE..]); + assert_eq!( + torn, + Bytes::from(expected_torn), + "C10: read is torn across an equal-length replacement -- head from A, tail from B", + ); + assert_ne!( + torn, value_a, + "torn blob is not the value the reader asked for" + ); + assert_ne!( + torn, value_b, + "torn blob is not the replacement value either" + ); + + Ok(()) +} + #[nativelink_test] async fn yield_between_sending_packets_in_update() -> Result<(), Error> { let data_p1 = Bytes::from(vec![b'A'; DEFAULT_READ_CHUNK_SIZE + 512]); @@ -1411,9 +1500,10 @@ fn test_search_by_index() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1524,9 +1614,10 @@ fn test_search_by_index_descending() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1581,9 +1672,10 @@ fn test_search_by_index_retries_on_failover() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1684,9 +1776,10 @@ fn test_search_by_index_skips_docs_that_expired_mid_query() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1781,7 +1874,7 @@ fn test_search_by_index_failure() -> Result<(), Error> { "Client: TEST - Client: unexpected command", "Error with ft_create in RedisStore::search_by_index_prefix(test:_content_prefix_sort_key_3e762c15)", "---", "Client: TEST - Client: unexpected command", "Error with second ft_aggregate in RedisStore::search_by_index_prefix(test:_content_prefix_sort_key_3e762c15)"].iter().map(ToString::to_string).collect())); assert!(logs_contain( - "Error calling ft.aggregate e=TEST - Client: unexpected command index=\"test:_content_prefix_sort_key_3e762c15\" query=\"*\" options=FtAggregateOptions { load: [\"data\", \"version\"], cursor: FtAggregateCursor { count: 1500, max_idle: 30000 }, sort_by: [\"@sort_key\"], sort_desc: false } all_args=[\"FT.AGGREGATE\", \"test:_content_prefix_sort_key_3e762c15\", \"*\", \"TIMEOUT\", \"10000\", \"LOAD\", \"2\", \"data\", \"version\", \"WITHCURSOR\", \"COUNT\", \"1500\", \"MAXIDLE\", \"30000\", \"SORTBY\", \"2\", \"@sort_key\", \"ASC\", \"MAX\", \"1000000\"]" + "Error calling ft.aggregate e=TEST - Client: unexpected command index=\"test:_content_prefix_sort_key_3e762c15\" query=\"*\" options=FtAggregateOptions { load: [\"data\", \"version\", \"@__key\"], cursor: FtAggregateCursor { count: 1500, max_idle: 30000 }, sort_by: [\"@sort_key\"], sort_desc: false } all_args=[\"FT.AGGREGATE\", \"test:_content_prefix_sort_key_3e762c15\", \"*\", \"TIMEOUT\", \"10000\", \"LOAD\", \"3\", \"data\", \"version\", \"@__key\", \"WITHCURSOR\", \"COUNT\", \"1500\", \"MAXIDLE\", \"30000\", \"SORTBY\", \"2\", \"@sort_key\", \"ASC\", \"MAX\", \"1000000\"]" )); Ok(()) @@ -1802,9 +1895,10 @@ fn test_search_by_index_swallows_already_exists_from_ft_create() -> Result<(), E .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1886,9 +1980,10 @@ fn test_search_by_index_preserves_other_ft_create_errors() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1966,9 +2061,10 @@ fn test_search_by_index_with_sort_key() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2052,9 +2148,10 @@ fn test_search_by_index_resp3() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2161,9 +2258,10 @@ fn test_search_by_index_skips_int_from_cursor_read() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2252,9 +2350,10 @@ async fn search_by_index_prefix_reads_a_resp3_cursor_page() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2694,7 +2793,10 @@ where Arc::into_inner(RedisStore::new_standard(spec).await.expect("Working spec")).unwrap(); raw_store.replace_temp_name_generator(mock_uuid_generator); let store = Arc::new(raw_store); - store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; timeout(Duration::from_secs(3), async move { loop { @@ -2743,21 +2845,16 @@ async fn callback_for_eviction_good() -> Result<(), Error> { Ok(()) } -// Repro for a fail-silent one-shot init bug: `register_remove_callback` -// returns Ok(()) immediately and performs the CONFIG GET + PSUBSCRIBE lazily -// inside a `background_spawn` via `OnceCell::get_or_try_init`. If that init -// hits a single transient Redis error (here: `-LOADING`), the error is only -// logged, the OnceCell stays uninitialized, and nothing ever retries — so the -// `__key*__:*` subscription is never established and evicted/expired keys -// never fire remove callbacks, silently and permanently. +// `register_remove_callback` returns Ok(()) at once and sets up the +// `__key*__:*` subscription in the background. Reading `notify-keyspace-events` +// first is only advisory, so a server that refuses `CONFIG GET` (here with a +// transient `-LOADING`) must still end up subscribed. This used to leave remove +// callbacks off for the life of the process. #[nativelink_test] -async fn transient_init_error_permanently_disables_remove_callbacks() -> Result<(), Error> { +async fn remove_callback_subscribes_even_if_config_cannot_be_read() -> Result<(), Error> { let redis_span = info_span!("redis"); let mut responses = add_lua_version_script(fake_redis_stream()); - // The one-time init query fails with a transient, retryable server error - // (Redis restarting/failing over). The server is otherwise healthy and - // would happily accept a PSUBSCRIBE. add_to_response_raw( &mut responses, redis::cmd("CONFIG") @@ -2773,6 +2870,125 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< let redis_port = make_fake_redis_with_responses(responses) .instrument(redis_span) .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not read notify-keyspace-events; subscribing to eviction events anyway", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + Ok(()) +} + +// A `PSUBSCRIBE` that fails once with a transient error must be retried until +// the subscription is established. +#[nativelink_test] +async fn remove_callback_subscription_is_retried_after_a_failed_psubscribe() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + vec![Value::Map(vec![( + Value::BulkString(b"notify-keyspace-events".into()), + Value::BulkString(b"KA".into()), + )])], + ); + add_to_response( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + vec![Value::Nil], + ); + let mut first_responses = HashMap::new(); + add_to_response_raw( + &mut first_responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-LOADING Redis is loading the dataset in memory\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_first_responses(responses, first_responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + wait_for_log( + logs_contain, + "Subscribed to Redis eviction events after retrying failures=1", + ) + .await; + Ok(()) +} + +// The subscription retry task holds only a weak reference to the store, so a +// store dropped while its subscription keeps failing (here every PSUBSCRIBE is +// refused) ends the task instead of being kept alive and retried forever. +#[nativelink_test] +async fn remove_callback_retry_stops_when_the_store_is_dropped() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response_raw( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + "-NOPERM this user has no permissions to run the 'config|get' command\r\n".to_string(), + ); + add_to_response_raw( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-NOPERM this user has no permissions to access one of the channels\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_responses(responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + drop(store); + wait_for_log( + logs_contain, + "RedisStore dropped before it subscribed to eviction events", + ) + .await; + Ok(()) +} + +async fn remove_callback_test_store( + redis_port: u16, +) -> Arc>> { let spec = RedisSpec { addresses: vec![format!("redis://127.0.0.1:{redis_port}/")], mode: RedisMode::Standard, @@ -2781,46 +2997,19 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< let mut raw_store = Arc::into_inner(RedisStore::new_standard(spec).await.expect("Working spec")).unwrap(); raw_store.replace_temp_name_generator(mock_uuid_generator); - let store = Arc::new(raw_store); - // The caller (e.g. ExistenceCacheStore at construction, its only call - // site) is told Ok(()) and believes eviction tracking is live. - store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + Arc::new(raw_store) +} - // Deterministic sync point: wait until the background init task has hit - // the transient error and logged it. +/// Waits up to 5 s for `needle` to appear in the calling test's logs. +/// `logs_contain` is the function `#[nativelink_test]` defines in each test. +async fn wait_for_log(logs_contain: impl Fn(&str) -> bool + Send + Sync, needle: &str) { timeout(Duration::from_secs(5), async { - loop { - if logs_contain("Error while trying to initialise remove_callback_subscribe") { - break; - } + while !logs_contain(needle) { sleep(Duration::from_millis(20)).await; } }) .await - .expect("background init never ran/logged its error"); - - // Since register_remove_callback reported success, the keyspace-event - // subscription must (eventually) be established — otherwise remove - // callbacks are silently dead forever. On current code there is no retry - // path at all (the OnceCell is left empty and only another - // register_remove_callback call could re-drive it), so this never happens. - let subscribed = timeout(Duration::from_secs(2), async { - loop { - if logs_contain("new psubscribe complete pattern=\"__key*__:*\"") { - break; - } - sleep(Duration::from_millis(50)).await; - } - }) - .await - .is_ok(); - assert!( - subscribed, - "register_remove_callback returned Ok(()) but a single transient CONFIG GET \ - error permanently disabled the __key*__:* subscription; evicted/expired keys \ - will never invoke remove callbacks (stale ExistenceCacheStore entries)" - ); - Ok(()) + .unwrap_or_else(|_| panic!("timed out waiting for log line: {needle}")); } #[nativelink_test] diff --git a/nativelink-store/tests/s3_store_test.rs b/nativelink-store/tests/s3_store_test.rs index a9f4b7ee185..598e9bafd34 100644 --- a/nativelink-store/tests/s3_store_test.rs +++ b/nativelink-store/tests/s3_store_test.rs @@ -326,7 +326,7 @@ async fn smoke_test_get_part() -> Result<(), Error> { .uri(format!( "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", )) - .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH)) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder() @@ -426,6 +426,98 @@ async fn get_part_simple_retries() -> Result<(), Error> { Ok(()) } +// A blob over the 5 MiB multipart threshold but under the default 64 MiB +// part target must upload as ONE part. Sizing parts at the 5 MiB floor (the +// old behaviour) maximised the part count, drowning providers in tiny +// concurrent PUTs on multi-GiB blobs. +#[nativelink_test] +async fn multipart_default_part_size_uploads_one_part() -> Result<(), Error> { + const MIN_MULTIPART_SIZE: usize = 5 * 1024 * 1024; // 5mb. + const AC_ENTRY_SIZE: usize = MIN_MULTIPART_SIZE * 2 + 50; + + let mut send_data: Vec = Vec::with_capacity(AC_ENTRY_SIZE); + for i in 0..send_data.capacity() { + send_data.push(u8::try_from((i * 3) % 256).expect("modulo 256 always fits in u8")); + } + let digest = DigestInfo::try_new(VALID_HASH1, send_data.len())?; + + let mock_client = StaticReplayClient::new(vec![ + ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?uploads", + )) + .method("POST") + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(concat!( + "", + "Dummy-uploadid", + "" + ))) + .unwrap(), + ), + ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=UploadPart&partNumber=1&uploadId=Dummy-uploadid", + )) + .method("PUT") + .header("content-type", "application/octet-stream") + .header("content-length", AC_ENTRY_SIZE.to_string()) + .body(SdkBody::from(&send_data[..])) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::empty()) + .unwrap(), + ), + ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?uploadId=Dummy-uploadid", + )) + .method("POST") + .header("content-length", "138") + .body(SdkBody::from(concat!( + "", + "1", + "", + ))) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(concat!( + "", + "", + ))) + .unwrap(), + ), + ]); + let test_config = Builder::new() + .behavior_version(BehaviorVersion::latest()) + .region(Region::from_static(REGION)) + .http_client(mock_client.clone()) + .build(); + let s3_client = aws_sdk_s3::Client::from_conf(test_config); + let store = S3Store::new_with_client_and_jitter( + &ExperimentalAwsSpec { + bucket: BUCKET_NAME.to_string(), + ..Default::default() + }, + s3_client, + Arc::new(move |_delay| Duration::from_secs(0)), + MockInstantWrapped::default, + )?; + store + .update_oneshot(digest, send_data.clone().into()) + .await?; + mock_client.assert_requests_match(&[]); + Ok(()) +} + #[nativelink_test] async fn multipart_update_large_cas() -> Result<(), Error> { // Same as in s3_store. @@ -538,6 +630,12 @@ async fn multipart_update_large_cas() -> Result<(), Error> { let store = S3Store::new_with_client_and_jitter( &ExperimentalAwsSpec { bucket: BUCKET_NAME.to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }, s3_client, @@ -564,7 +662,7 @@ async fn ensure_empty_string_in_stream_works_test() -> Result<(), Error> { .uri(format!( "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{CAS_ENTRY_SIZE}?x-id=GetObject", )) - .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE)) + .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder() @@ -890,6 +988,12 @@ async fn multipart_chunk_size_clamp_min() -> Result<(), Error> { let store = S3Store::new_with_client_and_jitter( &ExperimentalAwsSpec { bucket: BUCKET_NAME.to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }, s3_client, @@ -1097,3 +1201,60 @@ async fn incomplete_multipart_upload_aborts_before_s3_commit() -> Result<(), Err assert_eq!(methods, vec!["POST", "DELETE"]); Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_range_is_end_inclusive() -> Result<(), Error> { + // A bounded read ending strictly before the object end must request an + // *inclusive* HTTP byte range of exactly `length` bytes. Regression for the + // off-by-one that requested `offset + length` (one byte too many). The + // StaticReplayClient only serves the exact-match range header, so without + // the fix the store emits `bytes=100-300`, fails to match, and errors. + const AC_ENTRY_SIZE: u64 = 1000; + const OFFSET: u64 = 100; + const LENGTH: u64 = 200; // OFFSET + LENGTH = 300 < AC_ENTRY_SIZE: a strict subrange. + let expected_payload = vec![0x88_u8; usize::try_from(LENGTH).unwrap()]; + let mock_client = StaticReplayClient::new(vec![ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", + )) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(expected_payload.clone())) + .unwrap(), + )]); + let test_config = Builder::new() + .behavior_version(BehaviorVersion::latest()) + .region(Region::from_static(REGION)) + .http_client(mock_client.clone()) + .build(); + let s3_client = aws_sdk_s3::Client::from_conf(test_config); + let store = S3Store::new_with_client_and_jitter( + &ExperimentalAwsSpec { + bucket: BUCKET_NAME.to_string(), + ..Default::default() + }, + s3_client, + Arc::new(move |_delay| Duration::from_secs(0)), + MockInstantWrapped::default, + )?; + + let got = store + .get_part_unchunked( + DigestInfo::try_new(VALID_HASH1, AC_ENTRY_SIZE)?, + OFFSET, + Some(LENGTH), + ) + .await?; + assert_eq!( + got.len(), + usize::try_from(LENGTH).unwrap(), + "bounded read returned wrong byte count" + ); + assert_eq!(got, Bytes::from(expected_payload)); + mock_client.assert_requests_match(&[]); + Ok(()) +} diff --git a/nativelink-store/tests/shard_ring_dead_peer_test.rs b/nativelink-store/tests/shard_ring_dead_peer_test.rs new file mode 100644 index 00000000000..db790cd873a --- /dev/null +++ b/nativelink-store/tests/shard_ring_dead_peer_test.rs @@ -0,0 +1,557 @@ +// 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. + +//! Lab reproduction of the production "shard-ring wedge": a `ShardStore` +//! ring over `GrpcStore` backends (peer CAS shards) where one peer is dead +//! (killed or never started). Before the fix, every CAS RPC that touched +//! the dead shard hung forever: the peer's `ConnectionManager` retried the +//! connect in the background while `connection()` requests queued +//! unbounded, so no error ever reached the retry machinery. The +//! operational workaround was `systemctl restart nativelink` on the +//! frontend. +//! +//! These tests must fail (by generous outer timeout) against the +//! pre-fix code and pass after dead peers yield fast `Unavailable`. + +use core::pin::Pin; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; + +use async_lock::Mutex; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use nativelink_config::stores::{ + GrpcEndpoint, GrpcSpec, MemorySpec, Retry, ShardConfig, ShardSpec, StoreSpec, StoreType, +}; +use nativelink_error::{Code, Error}; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{ + ContentAddressableStorage, ContentAddressableStorageServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::{ + BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, + BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse, GetTreeRequest, + GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse, SplitBlobRequest, SplitBlobResponse, +}; +use nativelink_proto::google::bytestream::byte_stream_server::{ByteStream, ByteStreamServer}; +use nativelink_proto::google::bytestream::{ + QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest, + WriteResponse, +}; +use nativelink_store::grpc_store::GrpcStore; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::shard_store::ShardStore; +use nativelink_util::background_spawn; +use nativelink_util::common::DigestInfo; +use nativelink_util::store_trait::{Store, StoreLike}; +use tokio::time::timeout; +use tonic::transport::Server; +use tonic::transport::server::TcpIncoming; +use tonic::{Request, Response, Status, Streaming}; + +/// Generous bound: far longer than any configured timeout/retry budget in +/// this test, far shorter than "forever". If an operation does not settle +/// within this, it is the production hang. +const GENEROUS_TIMEOUT: Duration = Duration::from_secs(20); + +/// A minimal in-process CAS shard: blobs live in a shared map, +/// `FindMissingBlobs`, `ByteStream.Read` and `ByteStream.Write` behave like +/// a real peer. +#[derive(Debug, Clone)] +struct FakeShardServer { + blobs: Arc>>, +} + +type GetTreeStream = Pin> + Send + 'static>>; + +#[tonic::async_trait] +impl ContentAddressableStorage for FakeShardServer { + type GetTreeStream = GetTreeStream; + + async fn find_missing_blobs( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let blobs = self.blobs.lock().await; + let missing_blob_digests = request + .blob_digests + .into_iter() + .filter(|digest| !blobs.contains_key(&digest.hash)) + .collect(); + Ok(Response::new(FindMissingBlobsResponse { + missing_blob_digests, + })) + } + + #[allow(clippy::unimplemented)] + async fn batch_update_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn batch_read_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn get_tree( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn split_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn splice_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +type ReadStream = Pin> + Send + 'static>>; + +/// Extracts the blob hash from a `ByteStream` resource name of the form +/// `[{instance}/]uploads/{uuid}/blobs/{fn}/{hash}/{size}` or +/// `[{instance}/]blobs/{fn}/{hash}/{size}`. +fn hash_from_resource_name(resource_name: &str) -> String { + let mut components = resource_name.rsplit('/'); + let _size = components.next(); + components.next().unwrap_or_default().to_string() +} + +#[tonic::async_trait] +impl ByteStream for FakeShardServer { + type ReadStream = ReadStream; + + async fn read( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let hash = hash_from_resource_name(&request.resource_name); + let data = self + .blobs + .lock() + .await + .get(&hash) + .cloned() + .ok_or_else(|| Status::not_found(format!("Blob {hash} not seeded")))?; + let offset = usize::try_from(request.read_offset) + .map_err(|_| Status::invalid_argument("Bad read_offset"))?; + if offset > data.len() { + return Err(Status::out_of_range("read_offset past end of blob")); + } + let stream: Self::ReadStream = Box::pin(futures::stream::iter(vec![Ok(ReadResponse { + data: data.slice(offset..), + })])); + Ok(Response::new(stream)) + } + + async fn write( + &self, + grpc_request: Request>, + ) -> Result, Status> { + let mut stream = grpc_request.into_inner(); + let mut hash = String::new(); + let mut data = Vec::new(); + while let Some(request) = stream.next().await { + let request = request?; + if !request.resource_name.is_empty() { + hash = hash_from_resource_name(&request.resource_name); + } + data.extend_from_slice(&request.data); + if request.finish_write { + break; + } + } + let committed_size = + i64::try_from(data.len()).map_err(|_| Status::invalid_argument("Blob too large"))?; + self.blobs.lock().await.insert(hash, Bytes::from(data)); + Ok(Response::new(WriteResponse { committed_size })) + } + + #[allow(clippy::unimplemented)] + async fn query_write_status( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +/// Spawns a live in-process shard peer on `port` (0 for ephemeral) and +/// returns its bound port and blob map. +fn spawn_live_peer_on(port: u16) -> (u16, Arc>>) { + let server = FakeShardServer { + blobs: Arc::new(Mutex::new(HashMap::new())), + }; + let blobs = server.blobs.clone(); + let listener = TcpIncoming::bind(format!("127.0.0.1:{port}").parse().unwrap()).unwrap(); + let port = listener.local_addr().unwrap().port(); + let cas_service = ContentAddressableStorageServer::new(server.clone()); + let bytestream_service = ByteStreamServer::new(server); + background_spawn!("shard_ring_dead_peer_test_server", async move { + Server::builder() + .add_service(cas_service) + .add_service(bytestream_service) + .serve_with_incoming(listener) + .await + .unwrap(); + }); + (port, blobs) +} + +/// Reserves a port with nothing listening on it: the dead peer (killed or +/// never started). Connect attempts fail with `ECONNREFUSED`. +fn dead_peer_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() +} + +fn grpc_spec_for_port(port: u16) -> GrpcSpec { + GrpcSpec { + instance_name: String::new(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 1, + tcp_keepalive_s: 0, + http2_keepalive_interval_s: 0, + http2_keepalive_timeout_s: 0, + }], + store_type: StoreType::Cas, + retry: Retry { + max_retries: 1, + delay: 0.05, + jitter: 0.0, + retry_on_errors: None, + }, + max_concurrent_requests: 0, + connections_per_endpoint: 0, + rpc_timeout_s: 5, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_remote_cache_compression: None, + } +} + +/// Mirrors `ShardStore`'s routing for a two-shard equal-weight ring: the +/// digest's xor-folded key lands in shard 0 iff it is `<= u32::MAX / 2`. +fn shard_index_for(digest: &DigestInfo) -> usize { + let packed_hash = digest.packed_hash(); + let size_bytes = digest.size_bytes().to_le_bytes(); + let mut key: u32 = 0; + for chunk in 0..8 { + key ^= u32::from_le_bytes(packed_hash[chunk * 4..chunk * 4 + 4].try_into().unwrap()); + } + key ^= u32::from_le_bytes(size_bytes[0..4].try_into().unwrap()); + key ^= u32::from_le_bytes(size_bytes[4..8].try_into().unwrap()); + usize::from(key > u32::MAX / 2) +} + +/// Finds a digest (with 64-byte content) that routes to `shard`. +fn digest_for_shard(shard: usize) -> (DigestInfo, Bytes) { + for i in 0..1024usize { + let hash = format!("{i:064x}"); + let content = Bytes::from(format!("{i:0>64}")); + let digest = DigestInfo::try_new(&hash, content.len()).unwrap(); + if shard_index_for(&digest) == shard { + return (digest, content); + } + } + unreachable!("no digest found for shard {shard} in 1024 candidates"); +} + +struct Ring { + shard_store: Arc, + live_shard: usize, + live_blobs: Arc>>, + dead_port: u16, +} + +/// Builds a two-shard equal-weight ring: shard 0 is a live in-process +/// peer, shard 1 is dead (nothing listening). +async fn make_ring() -> Result { + let (live_port, live_blobs) = spawn_live_peer_on(0); + let dead_port = dead_peer_port(); + + let live_store = GrpcStore::new(&grpc_spec_for_port(live_port))?; + let dead_store = GrpcStore::new(&grpc_spec_for_port(dead_port))?; + + // The spec's `store` fields are only used for length/weights by + // `ShardStore::new`; memory placeholders keep the config minimal. + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + vec![Store::new(live_store), Store::new(dead_store)], + )?; + Ok(Ring { + shard_store, + live_shard: 0, + live_blobs, + dead_port, + }) +} + +// Guard for `shard_index_for`: it re-implements `ShardStore`'s xor-fold +// routing, which would silently diverge if upstream ever changes the ring +// hash — the dead-peer tests above would then route probes to the wrong +// shard and test nothing. This builds a REAL `ShardStore` over two +// instrumented `MemoryStore` backends, writes each probe digest through +// it, and asserts the blob landed on the shard `shard_index_for` +// predicted. If the ring hash changes upstream, this fails loudly. +#[nativelink_test] +async fn shard_index_for_matches_real_shard_store_routing() -> Result<(), Error> { + let backends: Vec> = (0..2) + .map(|_| MemoryStore::new(&MemorySpec::default())) + .collect(); + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + backends + .iter() + .map(|backend| Store::new(backend.clone())) + .collect(), + )?; + + for shard in 0..2usize { + let (digest, content) = digest_for_shard(shard); + shard_store.update_oneshot(digest, content).await?; + for (backend_index, backend) in backends.iter().enumerate() { + let present = backend.has(digest).await?.is_some(); + assert_eq!( + present, + backend_index == shard, + "probe digest {digest} predicted for shard {shard} but presence on \ + backend {backend_index} was {present}: `shard_index_for` no longer \ + matches ShardStore's ring hash", + ); + } + } + Ok(()) +} + +// A read routed to the dead shard must settle with an error (which the +// retry machinery upstream can handle) instead of hanging forever. +#[nativelink_test] +async fn read_routed_to_dead_shard_errors_instead_of_hanging() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, _) = digest_for_shard(1 - ring.live_shard); + + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung past the generous timeout: the shard-ring wedge"); + let err = result.expect_err("read from a dead shard cannot succeed"); + assert!( + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded), + "expected Unavailable/DeadlineExceeded from the dead shard, got: {err:?}", + ); + Ok(()) +} + +// The production symptom: a multi-key existence check (FindMissingBlobs) +// fans out to every shard, so ONE dead peer wedged EVERY CAS RPC through +// the frontend. With read failover the dead shard's bucket is retried on +// the live shard, so the whole call succeeds instead of erroring. +#[nativelink_test] +async fn has_with_dead_shard_fails_over_instead_of_erroring() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, live_content) = digest_for_shard(ring.live_shard); + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + // Both blobs are present on the live shard (as a replica would be). + { + let mut blobs = ring.live_blobs.lock().await; + blobs.insert(live_digest.packed_hash().to_string(), live_content.clone()); + blobs.insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + } + + let results = timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .has_many(&[live_digest.into(), dead_digest.into()]), + ) + .await + .expect("multi-key has() with a dead shard hung past the generous timeout")?; + assert_eq!( + results, + vec![ + Some(live_content.len() as u64), + Some(dead_content.len() as u64) + ], + "both keys must resolve: the dead shard's bucket via failover", + ); + Ok(()) +} + +// A single-key has() routed to the dead shard must fail over to the live +// shard that holds the blob instead of surfacing Unavailable. +#[nativelink_test] +async fn has_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("single-key has() routed to the dead shard hung")?; + assert_eq!( + result, + Some(dead_content.len() as u64), + "has() must find the blob on the live shard via failover", + ); + Ok(()) +} + +// A read routed to the dead shard must fail over to the live shard that +// holds a replica of the blob and return its content. +#[nativelink_test] +async fn get_part_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung")?; + assert_eq!( + data, dead_content, + "failover read must return the live shard's replica", + ); + Ok(()) +} + +// Health recovery: after enough consecutive failures the dead shard is +// marked suspect and skipped; once the peer comes back and the cooldown +// elapses it is probed again and serves its own keys directly. +#[nativelink_test] +async fn dead_shard_recovers_after_cooldown() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + + // Trip the suspect threshold: each read fails over to the live shard + // (which lacks the blob), so the call errors, but each attempt counts + // a connectivity failure against the dead shard. + for _ in 0..3 { + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read with dead shard hung"); + assert!(result.is_err(), "blob exists nowhere yet; read must fail"); + } + + // Revive the dead peer on its original port, now holding the blob. + let (_, revived_blobs) = spawn_live_peer_on(ring.dead_port); + revived_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + // Wait out the suspect cooldown so the shard is probed again. + tokio::time::sleep(Duration::from_secs(6)).await; + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read after recovery hung")?; + assert_eq!(data, dead_content, "revived shard must serve its own key"); + + // The success must have reset the failure counter: an immediate + // follow-up (well inside any cooldown) is also served directly. + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("has() after recovery hung")?; + assert_eq!(result, Some(dead_content.len() as u64)); + Ok(()) +} + +// The live shard must keep serving while its ring-mate is dead: routing +// around the dead peer stays possible for keys that don't touch it. +#[nativelink_test] +async fn live_shard_still_serves_while_peer_dead() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, content) = digest_for_shard(ring.live_shard); + + timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .update_oneshot(live_digest, content.clone()), + ) + .await + .expect("write to the live shard hung")?; + assert!( + ring.live_blobs + .lock() + .await + .contains_key(&live_digest.packed_hash().to_string()), + "blob did not land on the live peer", + ); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(live_digest, 0, None), + ) + .await + .expect("read from the live shard hung")?; + assert_eq!(data, content, "live shard returned wrong content"); + Ok(()) +} diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index bc008ef5ea0..fa453f15375 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -106,6 +106,7 @@ rust_test_suite( "tests/cache_size_metrics_test.rs", "tests/channel_body_for_tests_test.rs", "tests/common_test.rs", + "tests/connection_manager_test.rs", "tests/evicting_map_test.rs", "tests/fastcdc_test.rs", "tests/fs_test.rs", diff --git a/nativelink-util/src/buf_channel.rs b/nativelink-util/src/buf_channel.rs index d60f8fa362c..741d1245611 100644 --- a/nativelink-util/src/buf_channel.rs +++ b/nativelink-util/src/buf_channel.rs @@ -27,6 +27,16 @@ use tracing::warn; const ZERO_DATA: Bytes = Bytes::new(); +/// Chunks buffered in flight per channel. Small enough that a fast producer +/// still applies backpressure promptly, but not so small (it was 2) that a +/// momentarily-slow consumer — e.g. a shard CAS mid-forward of a multi-GiB blob +/// — stalls the producer immediately. That instant stall stops the server +/// draining its inbound stream, the HTTP/2 flow-control window empties, and the +/// connection resets mid-stream (dropping large cacheable uploads). tokio's +/// bounded channel does not preallocate, so this capacity costs memory only +/// while data is actually queued (i.e. exactly during a downstream stall). +const CHANNEL_BUFFER_CHUNKS: usize = 32; + /// Create a channel pair that can be used to transport buffer objects around to /// different components. This wrapper is used because the streams give some /// utility like managing EOF in a more friendly way, ensure if no EOF is received @@ -34,11 +44,7 @@ const ZERO_DATA: Bytes = Bytes::new(); /// the number of bytes sent. #[must_use] pub fn make_buf_channel_pair() -> (DropCloserWriteHalf, DropCloserReadHalf) { - // We allow up to 2 items in the buffer at any given time. There is no major - // reason behind this magic number other than thinking it will be nice to give - // a little time for another thread to wake up and consume data if another - // thread is pumping large amounts of data into the channel. - let (tx, rx) = mpsc::channel(2); + let (tx, rx) = mpsc::channel(CHANNEL_BUFFER_CHUNKS); let eof_sent = Arc::new(AtomicBool::new(false)); ( DropCloserWriteHalf { diff --git a/nativelink-util/src/connection_manager.rs b/nativelink-util/src/connection_manager.rs index 753e7cb956a..145f84a790a 100644 --- a/nativelink-util/src/connection_manager.rs +++ b/nativelink-util/src/connection_manager.rs @@ -35,7 +35,7 @@ use crate::retry::{self, Retrier, RetryResult}; #[derive(Debug)] pub struct ConnectionManager { // The channel to request connections from the worker. - worker_tx: mpsc::Sender<(String, oneshot::Sender)>, + worker_tx: mpsc::Sender<(String, oneshot::Sender>)>, } /// The index into `ConnectionManagerWorker::endpoints`. @@ -94,6 +94,21 @@ struct ConnectionManagerWorker { /// The endpoints to establish Channels and the identifier of the last /// connection attempt to that endpoint. endpoints: Vec<(ConnectionIndex, Endpoint)>, + /// Per-endpoint connect health: `true` once a connect retry cycle to + /// that endpoint has failed, cleared when a connect succeeds. When + /// every endpoint is marked failed and no channel is available, the + /// endpoint set is unreachable — requests fail fast with + /// `Unavailable` instead of queuing forever while the background + /// reconnect loop runs (the "dead peer wedges every RPC" hang). + endpoint_connect_failed: Vec, + /// Number of currently established channels, whether sitting in + /// `available_channels` or checked out to in-flight `Connection`s. + /// `available_channels` alone cannot distinguish "everything is down" + /// from "everything is busy": borrowed channels are not in that queue, + /// so a busy-but-healthy pool plus one transient connect failure must + /// not be declared dead. Incremented on connect success, decremented + /// when a channel is torn down on transport error. + live_channels: usize, /// The channel used to communicate between a Connection and the worker. connection_tx: mpsc::UnboundedSender, /// Gates the maximum number of in-flight `Connection` objects. @@ -112,7 +127,7 @@ struct ConnectionManagerWorker { /// Connected channels that are available for use. available_channels: VecDeque, /// Requests for a Channel when available - (reason, request) - waiting_connections: VecDeque<(String, oneshot::Sender)>, + waiting_connections: VecDeque<(String, oneshot::Sender>)>, /// The retry configuration for connecting to an Endpoint, on failure will /// restart the retrier after a 1 second delay. retrier: Retrier, @@ -140,10 +155,11 @@ impl ConnectionManager { // which defeats the object since there would be no backpressure // applied. Therefore it makes sense for this to be unbounded. let (connection_tx, connection_rx) = mpsc::unbounded_channel(); - let endpoints = endpoints + let endpoints: Vec<(ConnectionIndex, Endpoint)> = endpoints .into_iter() .map(|endpoint| (0, endpoint)) .collect(); + let endpoint_connect_failed = vec![false; endpoints.len()]; // Zero means unlimited, which becomes a semaphore holding // Semaphore::MAX_PERMITS. That is a sentinel, not a free-slot count, @@ -159,6 +175,8 @@ impl ConnectionManager { } let worker = ConnectionManagerWorker { endpoints, + endpoint_connect_failed, + live_channels: 0, available_connections: Arc::new(Semaphore::new(max_concurrent_requests)), bounded_connections, connection_tx, @@ -190,7 +208,7 @@ impl ConnectionManager { .await .map_err(|err| make_err!(Code::Unavailable, "Requesting a new connection: {err:?}"))?; rx.await - .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}")) + .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}"))? } } @@ -198,7 +216,7 @@ impl ConnectionManagerWorker { async fn service_requests( mut self, connections_per_endpoint: usize, - mut worker_rx: mpsc::Receiver<(String, oneshot::Sender)>, + mut worker_rx: mpsc::Receiver<(String, oneshot::Sender>)>, mut connection_rx: mpsc::UnboundedReceiver, ) { // Make the initial set of connections, connection failures will be @@ -252,6 +270,13 @@ impl ConnectionManagerWorker { fn handle_connected(&mut self, connection_result: IndexedChannel) { match connection_result { Ok(established_channel) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(established_channel.identifier.endpoint_index) + { + *failed = false; + } + self.live_channels += 1; self.available_channels.push_back(established_channel); self.maybe_available_connection(); } @@ -259,11 +284,52 @@ impl ConnectionManagerWorker { // beginning of the retry period. Never want to be in a // situation where we give up on an Endpoint forever. Err((identifier, _)) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(identifier.endpoint_index) + { + *failed = true; + } + // If this failure means the whole endpoint set is now + // unreachable, waiting requests would hang until a peer + // comes back — fail them fast so the caller's retry/error + // machinery can handle it. + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, Some(identifier.connection_index)); } } } + /// True when no channel is live (available or checked out) and the + /// last connect attempt to every endpoint failed: nothing can service + /// a request until a background reconnect succeeds. + fn all_endpoints_down(&self) -> bool { + self.live_channels == 0 + && !self.endpoint_connect_failed.is_empty() + && self.endpoint_connect_failed.iter().all(|failed| *failed) + } + + fn unavailable_error(&self) -> Error { + make_err!( + Code::Unavailable, + "All {} gRPC endpoint(s) are unreachable (connect attempts failed); failing fast instead of queuing", + self.endpoints.len() + ) + } + + fn fail_requests_if_all_endpoints_down(&mut self) { + if !self.all_endpoints_down() { + return; + } + while let Some((reason, tx)) = self.waiting_connections.pop_front() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing queued connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); + } + } + fn connect_endpoint(&mut self, endpoint_index: usize, connection_index: Option) { let Some((current_connection_index, endpoint)) = self.endpoints.get_mut(endpoint_index) else { @@ -333,7 +399,7 @@ impl ConnectionManagerWorker { } // This must never be made async otherwise the select may cancel it. - fn handle_worker(&mut self, reason: String, tx: oneshot::Sender) { + fn handle_worker(&mut self, reason: String, tx: oneshot::Sender>) { let maybe_permit = self.available_connections.clone().try_acquire_owned().ok(); if let Some(permit) = maybe_permit && let Some(channel) = self.available_channels.pop_front() @@ -341,6 +407,12 @@ impl ConnectionManagerWorker { debug!(reason, "ConnectionManager: request running"); record_connection_acquired("grpc", self.free_connection_slots(), false); self.provide_channel(channel, tx, permit); + } else if self.all_endpoints_down() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); } else { debug!( available_permits = self.available_connections.available_permits(), @@ -357,15 +429,15 @@ impl ConnectionManagerWorker { fn provide_channel( &self, channel: EstablishedChannel, - tx: oneshot::Sender, + tx: oneshot::Sender>, permit: OwnedSemaphorePermit, ) { - drop(tx.send(Connection { + drop(tx.send(Ok(Connection { tx: self.connection_tx.clone(), pending_channel: Some(channel.channel.clone()), channel, _permit: permit, - })); + }))); } fn maybe_available_connection(&mut self) { @@ -413,6 +485,17 @@ impl ConnectionManagerWorker { original_length != self.available_channels.len() }; if should_reconnect { + // The channel this identifier referred to is gone; + // saturate in case of duplicate error reports for the + // same pending channel. + self.live_channels = self.live_channels.saturating_sub(1); + // If this was the last live channel and every endpoint's + // last connect cycle had already failed, the pool just + // became unreachable — flush queued waiters now instead + // of leaving them parked until the reconnect scheduled + // below fails a full connect cycle later (new requests + // already fail fast in `handle_worker` in this state). + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, None); } } diff --git a/nativelink-util/src/operation_state_manager.rs b/nativelink-util/src/operation_state_manager.rs index 4e6e2a30dee..f88f45cfb47 100644 --- a/nativelink-util/src/operation_state_manager.rs +++ b/nativelink-util/src/operation_state_manager.rs @@ -13,6 +13,7 @@ // limitations under the License. use core::pin::Pin; +use std::collections::HashMap; use std::sync::Arc; use std::time::SystemTime; @@ -202,10 +203,19 @@ pub trait MatchingEngineStateManager: Sync + Send + MetricsComponent { ) -> Result, Error>; /// Assign an operation to a worker or unassign it. + /// + /// `expected_platform_properties`, when assigning, is the untyped + /// requirement snapshot the matcher chose this worker for. The accept + /// rejects (with `Aborted`) if the action's current requirements no + /// longer match it — an escalation or other requeue changed the + /// reservation after the worker was picked, so the placement is stale + /// and must be rematched against the current requirements rather than + /// dispatched at the old reservation. `None` skips the check. async fn assign_operation( &self, operation_id: &OperationId, worker_id_or_reason_for_unassign: Result<&WorkerId, Error>, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error>; /// Completes a queued operation with `err`, without retrying it. Does diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index e68d030c670..92b9ad8b890 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -1025,6 +1025,23 @@ pub trait SchedulerIndexProvider { /// The value of the index. fn index_value(&self) -> Cow<'_, str>; + + /// Opts an index into garbage-collecting a listed record whose data + /// exists but cannot be decoded (for example serialization drift across + /// a rolling upgrade). Such a record has no TTL, and no reader or + /// writer can ever use it again, so without removal it is listed, + /// skipped and counted forever. + /// + /// Returning `Some(guard_key)` permits deletion of `record_key` only + /// while `guard_key` does not exist; the guard should be a liveness + /// marker (such as a client keepalive) so a record something still + /// cares about is never removed. `None`, the default, means a record + /// listed by this index is never deleted. Both keys are store keys + /// before any store-level key prefix is applied. + #[must_use] + fn lost_record_liveness_guard_key(_index_value: &str, _record_key: &str) -> Option { + None + } } /// Provides a key to lookup data in the store. diff --git a/nativelink-util/tests/connection_manager_test.rs b/nativelink-util/tests/connection_manager_test.rs index 286ea578311..f293b0201dd 100644 --- a/nativelink-util/tests/connection_manager_test.rs +++ b/nativelink-util/tests/connection_manager_test.rs @@ -161,6 +161,263 @@ async fn aborted_caller_future_does_not_leak_permits() -> Result<(), Error> { Ok(()) } +/// The production "shard-ring wedge": an endpoint that is dead (nothing +/// listening — killed or never started) must yield a fast `Unavailable` +/// error from `connection()` instead of queuing the request forever while +/// the worker retries the connect in the background. Before the fix this +/// test failed: `cm.connection()` never resolved and the outer timeout +/// elapsed. +#[nativelink_test] +async fn dead_endpoint_fails_fast_instead_of_queuing_forever() -> Result<(), Error> { + // Bind a port and immediately drop the listener so nothing is + // listening there — connects get ECONNREFUSED, the "peer killed or + // never started" case observed in production. + let dead_port = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() + }; + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{dead_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + + let cm = ConnectionManager::new(vec![endpoint], 1, 2, Retry::default(), no_jitter()); + + let result = timeout(Duration::from_secs(10), cm.connection("dead-peer".into())) + .await + .expect("connection() to a dead endpoint hung past 10s instead of erroring"); + assert_eq!( + result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected fast Unavailable from a dead endpoint", + ); + Ok(()) +} + +/// The false-negative guard on the fast-fail gating: a busy-but-healthy +/// pool must NOT be declared dead. One endpoint, two connection slots: a +/// single-accept proxy in front of a real gRPC server lets exactly one +/// connect succeed, then drops its listener so the second slot's connect +/// is refused (`ECONNREFUSED`), marking the endpoint's last connect cycle +/// as failed. We hold the one live connection (so `available_channels` is +/// empty) and issue a new `connection()` request: with gating keyed only +/// on `available_channels` + connect health this errored `Unavailable` +/// immediately; with live-channel tracking it queues, and resolves as +/// soon as the held connection is released. +#[nativelink_test] +async fn busy_pool_with_transient_connect_failure_does_not_fail_fast() -> Result<(), Error> { + // Real gRPC server the proxied connection terminates at. + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + // Proxy that accepts exactly ONE connection, pipes it to the backend, + // then drops the listener — every later connect is refused. + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + background_spawn!("single_accept_proxy", async move { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + // Drop the listener so the second connect slot gets ECONNREFUSED. + drop(proxy_listener); + let mut outbound = tokio::net::TcpStream::connect(backend_addr).await.unwrap(); + drop(tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await); + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + Retry::default(), + no_jitter(), + )); + + // Take (and hold) the one connection that could be established. This + // may need a moment while the connect race settles; the winner is the + // proxied slot, the loser marks the endpoint's connect cycle failed. + let held = timeout(Duration::from_secs(10), cm.connection("held-live".into())) + .await + .expect("acquiring the single live connection hung")?; + + // Give the losing connect slot ample time to fail its cycle (default + // retry is a single ~1s-connect-timeout attempt) and set the + // per-endpoint failed flag; pre-fix this is what armed the trap. + tokio::time::sleep(Duration::from_secs(3)).await; + + // A new request must queue, not fail fast: the pool is busy, not dead. + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-while-busy".into()).await }); + tokio::time::sleep(Duration::from_millis(500)).await; + assert_eq!( + queued.is_finished(), + false, + "connection() failed fast while a live connection was checked out — \ + busy-but-healthy pool wrongly declared dead", + ); + + // Releasing the held connection must let the queued request through. + drop(held); + let conn = timeout(Duration::from_secs(5), queued) + .await + .expect("queued request did not resolve after the live connection was released") + .unwrap()?; + drop(conn); + Ok(()) +} + +/// The teardown-to-zero window: when the LAST live channel is torn down +/// by an in-use transport error while every endpoint's last connect cycle +/// had already failed, queued waiters must be flushed with `Unavailable` +/// at that instant — not parked until the freshly scheduled reconnect +/// fails a whole connect cycle later (`handle_worker` already fails new +/// requests fast in that state; queued ones must not be second-class). +/// +/// Setup: one endpoint, two connection slots, both piped through a proxy +/// to a real gRPC server; `max_concurrent_requests` of 2 with both +/// permits held so returned channel clones can never be handed to the +/// queued request. The retrier is configured so a connect cycle to the +/// (killed) endpoint takes ~3s: after severing both connections, an RPC +/// on the second held connection tears its channel down and its +/// reconnect cycle fails, arming the endpoint's connect-failed flag +/// while the first channel is still live. A request is then queued, and +/// an RPC on the first held connection tears the last live channel down +/// to zero. Post-fix the queued request errors `Unavailable` right +/// there; pre-fix nothing flushed it until the next connect cycle failed +/// ~3s later, so the 1.5s deadline below distinguishes the two. +#[nativelink_test] +async fn last_live_channel_teardown_flushes_queued_requests() -> Result<(), Error> { + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + let (kill_tx, kill_rx) = tokio::sync::watch::channel(false); + background_spawn!("teardown_test_proxy", async move { + // Pipe exactly the two initial connections (one per slot). On the + // kill signal, drop the listener first (so reconnects are refused, + // failing their connect cycles) and then sever both pipes. + for _ in 0..2 { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + let mut rx = kill_rx.clone(); + let addr = backend_addr.clone(); + background_spawn!("teardown_test_pipe", async move { + let mut outbound = tokio::net::TcpStream::connect(addr).await.unwrap(); + tokio::select! { + _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound) => {} + _ = rx.changed() => {} + } + }); + } + let mut rx = kill_rx.clone(); + drop(rx.changed().await); + drop(proxy_listener); + // Pipes observe the same signal and sever their connections. + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + // delay 1.5 with one retry makes a failing connect cycle take one + // ~3s backoff sleep between two instant ECONNREFUSED attempts: long + // enough that "flushed at teardown" and "flushed by the next connect + // failure" are cleanly distinguishable below. + let retry = Retry { + max_retries: 1, + delay: 1.5, + ..Default::default() + }; + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + retry, + no_jitter(), + )); + + // Hold both permits (and both live channels). + let held1 = timeout(Duration::from_secs(10), cm.connection("held-1".into())) + .await + .expect("acquiring the first piped connection hung")?; + let held2 = timeout(Duration::from_secs(10), cm.connection("held-2".into())) + .await + .expect("acquiring the second piped connection hung")?; + + // Kill the endpoint: refuse new connects and sever both connections. + tokio::time::sleep(Duration::from_millis(500)).await; + kill_tx.send(true).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + // Tear down the second channel with an in-use error; keep the client + // (and its permit) alive. Its reconnect cycle fails ~3s later, + // arming the endpoint's connect-failed flag while the first channel + // keeps live_channels at 1. + let mut client2 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held2); + let rpc2 = timeout( + Duration::from_secs(5), + client2.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-2".into(), + }), + ) + .await + .expect("RPC on the severed second connection hung"); + assert_eq!(rpc2.is_err(), true, "RPC on a severed connection succeeded"); + tokio::time::sleep(Duration::from_secs(4)).await; + + // Queue a request; it must wait (one channel still live, pool busy). + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-behind-held".into()).await }); + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + queued.is_finished(), + false, + "queued request resolved while a channel was live and permits were held", + ); + + // Tear down the LAST live channel with an in-use error. + let mut client1 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held1); + let rpc1 = timeout( + Duration::from_secs(5), + client1.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-1".into(), + }), + ) + .await + .expect("RPC on the severed first connection hung"); + assert_eq!(rpc1.is_err(), true, "RPC on a severed connection succeeded"); + + // The queued request must be flushed with Unavailable at the + // teardown itself — the next connect-cycle failure is ~3s away. + let queued_result = timeout(Duration::from_millis(1500), queued) + .await + .expect( + "queued request still parked after the last live channel was torn down \ + with all endpoints' connect cycles failed", + ) + .unwrap(); + assert_eq!( + queued_result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected Unavailable for the queued request after teardown-to-zero", + ); + drop(client1); + drop(client2); + Ok(()) +} + #[nativelink_test] async fn extra_request_above_max_blocks_until_a_release() -> Result<(), Error> { const MAX_CONCURRENT: usize = 2; diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 155f28fe3e3..d631248e61f 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -4977,6 +4977,15 @@ pub struct RunningActionsManagerImpl { #[cfg(target_os = "linux")] use_namespaces: UseNamespaces, running_actions: Mutex>>, + /// Kills received for operations that are not yet resident in + /// `running_actions`. A `KillOperationRequest` can arrive while + /// `create_and_add_action` is still awaiting (fetching the Action, waiting + /// for prior cleanup, making the directory) — before the op is inserted. + /// Without a record, the kill looked up nothing, logged, and was forgotten, + /// and the action then started executing unkilled (finding N1). We record + /// the op here under the same lock that guards the `running_actions` + /// insert, so registration can observe an in-transit kill and apply it. + pending_kills: Mutex>, // Note: We don't use Notify because we need to support a .wait_for()-like function, which // Notify does not support. action_done_tx: watch::Sender<()>, @@ -5096,6 +5105,7 @@ impl RunningActionsManagerImpl { max_download_timeout: args.max_download_timeout, timeout_handled_externally: args.timeout_handled_externally, running_actions: Mutex::new(HashMap::new()), + pending_kills: Mutex::new(HashSet::new()), action_done_tx, callbacks, metrics: Arc::new(Metrics { @@ -5539,7 +5549,13 @@ impl RunningActionsManager for RunningActionsManagerImpl { timeout, self.clone(), )); + // True if a kill arrived for this op while its registration was + // still in flight (recorded in `pending_kills` before the op + // became resident). Determined under the `running_actions` lock + // so it cannot race with a concurrent kill. + let killed_in_transit; { + // Lock ordering: `running_actions` before `pending_kills`. let mut running_actions = self.running_actions.lock(); // Check if action already exists and is still alive if let Some(existing_weak) = running_actions.get(&operation_id) @@ -5554,6 +5570,18 @@ impl RunningActionsManager for RunningActionsManagerImpl { running_action .has_manager_entry .store(true, Ordering::Release); + killed_in_transit = self.pending_kills.lock().remove(&operation_id); + } + if killed_in_transit { + // A kill was requested before this op finished registering. + // Apply it now that the action object exists so the kill is + // not lost: `execute()` will observe the consumed kill + // channel and abort instead of running the process. + warn!( + %operation_id, + "Applying kill requested while action registration was in flight", + ); + Self::kill_operation(running_action.clone()).await; } if let Some(config) = &self.execution_configuration.buck2_file_capture && Buck2FileCapture::eligible(capture_metadata.as_ref()) { @@ -5598,12 +5626,26 @@ impl RunningActionsManager for RunningActionsManagerImpl { } async fn kill_operation(&self, operation_id: &OperationId) -> Result<(), Error> { + // Lock ordering: `running_actions` before `pending_kills`. Decide under + // the `running_actions` lock whether the op is resident: if it is not, + // it may still be mid-registration (create_and_add_action fetches the + // Action, waits for prior cleanup, and makes the directory across await + // points before inserting), so record the kill as pending rather than + // dropping it. Registration checks `pending_kills` under this same lock + // and applies the kill, so the request stays effective through the + // registration race (finding N1). let running_action = { let running_actions = self.running_actions.lock(); - running_actions - .get(operation_id) - .and_then(Weak::upgrade) - .ok_or_else(|| make_input_err!("Failed to get running action {operation_id}"))? + let Some(running_action) = running_actions.get(operation_id).and_then(Weak::upgrade) + else { + self.pending_kills.lock().insert(operation_id.clone()); + warn!( + %operation_id, + "Kill received for an operation not yet resident; recorded as pending so registration applies it", + ); + return Ok(()); + }; + running_action }; Self::kill_operation(running_action).await; Ok(()) diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index 07a0abe9941..2020e3bc5cd 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -1740,6 +1740,152 @@ mod tests { Ok(()) } + /// N1 regression: a kill received before an action finishes registering + /// must stay effective through registration. + /// + /// A `KillOperationRequest` can reach the worker while + /// `create_and_add_action` is still awaiting (fetching the Action, waiting + /// for prior cleanup, making the directory) — before the op is inserted + /// into `running_actions`. This test reproduces that ordering by killing + /// the op *before* it is registered. + /// + /// WITHOUT THE FIX: `kill_operation` looked the op up in `running_actions`, + /// found nothing, and returned an Err (the kill was logged and dropped); + /// the action then registered and ran to completion unkilled (exit 0). + /// WITH THE FIX the kill is parked in `pending_kills` and applied at + /// registration, so when the action runs it is killed with `SIGKILL` + /// and reports the killed exit code (9). + /// + /// Two assertions fail without the fix: (1) the pre-registration kill must + /// return Ok (recorded, not rejected), and (2) the resulting action must + /// carry the killed exit code rather than having run normally. + #[nativelink_test] + async fn kill_before_registration_stays_effective() -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + + fn test_monotonic_clock() -> SystemTime { + static CLOCK: AtomicU64 = AtomicU64::new(0); + monotonic_clock(&CLOCK) + } + + let (_, _, cas_store, ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = Arc::new(RunningActionsManagerImpl::new_with_callbacks( + RunningActionsManagerArgs { + root_action_directory, + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + #[cfg(target_os = "linux")] + readonly_input_mounts: Vec::new(), + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + }, + Callbacks { + now_fn: test_monotonic_clock, + sleep_fn: |_duration| Box::pin(future::pending()), + }, + )?); + + let command = Command { + arguments: vec!["sleep".to_string(), "0.2".to_string()], + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let operation_id = OperationId::default(); + + // The kill arrives before the op is resident (models a kill landing + // while create_and_add_action is still mid-flight). This must be + // remembered, not dropped. + running_actions_manager + .kill_operation(&operation_id) + .await + .err_tip(|| "A kill for a not-yet-resident op must be recorded, not rejected")?; + + let execute_request = ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }; + let running_action = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(execute_request), + operation_id: operation_id.to_string(), + queued_timestamp: None, + platform: action.platform.clone(), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + // The in-transit kill must have been applied at registration. Since + // #2846 the kill token is also watched during input fetch, so the + // remembered kill now short-circuits `prepare_action` with `Aborted` + // ("killed while fetching its inputs") before the action ever runs — + // instead of the sleep completing normally (which a dropped in-transit + // kill would allow). Either way the observable proof is that the kill + // took effect: the action never completes successfully. + let err = run_action(running_action) + .await + .expect_err("a kill recorded before registration must abort the action"); + assert_eq!( + err.code, + Code::Aborted, + "a kill applied at registration aborts the action during input fetch; \ + a dropped in-transit kill would instead let it complete (Ok)" + ); + assert!( + err.to_string().contains("killed while fetching its inputs"), + "abort must be the kill path, not an unrelated failure: {err}" + ); + Ok(()) + } + #[nativelink_test] async fn ensure_output_files_full_directories_are_created_test() -> Result<(), Box> { diff --git a/src/bin/nativelink.rs b/src/bin/nativelink.rs index 4a796002009..7ee55120aeb 100644 --- a/src/bin/nativelink.rs +++ b/src/bin/nativelink.rs @@ -128,6 +128,10 @@ struct Args { /// Config file to use. #[clap(value_parser)] config_file: String, + + /// Validate the configuration and exit without starting any services. + #[clap(long)] + check: bool, } trait RoutesExt { @@ -881,14 +885,57 @@ async fn inner_main( Ok(()) } -fn get_config() -> Result { +fn get_config() -> Result<(Args, CasConfig), Error> { let args = Args::parse(); - CasConfig::try_from_json5_file(&args.config_file) + let cfg = CasConfig::try_from_json5_file(&args.config_file) + .err_tip(|| format!("While loading config file {}", args.config_file))?; + Ok((args, cfg)) } fn main() -> Result<(), Box> { install_default_rustls_crypto_provider(); + // Parse args and load the config *before* building any runtime so that + // `--check` can run a fully offline validation and exit without starting + // a serving runtime, binding a socket, connecting to a backend, or + // creating any store directories. + let (args, mut cfg) = get_config()?; + + if args.check { + return match cfg.validate_references() { + Ok(()) => { + let num_schedulers = cfg.schedulers.as_ref().map_or(0, Vec::len); + #[allow( + clippy::print_stdout, + reason = "`--check` reports its success summary on stdout for CI" + )] + { + println!( + "OK: {} — {} stores, {} schedulers, {} servers, all references resolve", + args.config_file, + cfg.stores.len(), + num_schedulers, + cfg.servers.len(), + ); + } + Ok(()) + } + Err(err) => { + #[allow( + clippy::print_stderr, + reason = "`--check` reports validation failures on stderr for CI" + )] + { + eprintln!("FAIL: {} — configuration is invalid", args.config_file); + for message in &err.messages { + eprintln!("{message}"); + } + } + std::process::exit(1); + } + }; + } + // Set QoS to USER_INITIATED on the main thread *before* the tokio // runtime is built so the spawned worker threads inherit P-core // scheduling preference via pthread QoS inheritance on Apple @@ -910,8 +957,6 @@ fn main() -> Result<(), Box> { #[expect(clippy::disallowed_methods, reason = "tracing init on main runtime")] runtime.block_on(async { tokio::spawn(async { init_tracing().await }).await? })?; - let mut cfg = get_config()?; - let global_cfg = if let Some(global_cfg) = &mut cfg.global { if global_cfg.max_open_files == 0 { global_cfg.max_open_files = fs::DEFAULT_OPEN_FILE_LIMIT; diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..97c19799f9e 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (acc4210f3d) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -830,6 +830,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `eviction_policy` | [EvictionPolicy](#evictionpolicy) | — | — | Policy used to evict items out of the store. Failure to set this value will cause items to never be removed from the store causing infinite memory usage. | | `block_size` | integer (uint64) | — | 4kb | The block size of the filesystem for the running machine value is used to determine an entry's actual size on disk consumed For a 4KB block size filesystem, a 1B file actually consumes 4KB | | `max_concurrent_writes` | integer (uint) | — | unlimited | Maximum number of concurrent write operations allowed. Each write involves streaming data to a temp file and calling `sync_all()`, which can saturate disk I/O when many writes happen simultaneously. Limiting concurrency prevents disk saturation from blocking the async runtime. A value of 0 means unlimited (no concurrency limit). | +| `max_concurrent_deletes` | integer (uint) | — | 1024 | Maximum number of concurrent background file-deletion operations. When an entry leaves the store, its backing file is unlinked by a background task. Without a bound, an eviction storm (cache wipe or a flood of overwrites) can spawn these without limit: memory grows with the queued task frames and disk/inode reclamation lags behind what the async runtime can schedule. Bounding it applies backpressure the same way `max_concurrent_writes` does — once this many deletes are in flight, a further delete runs synchronously inline in the evicting path instead of spawning, which throttles eviction to delete throughput and never leaks a file. | | `evict_page_cache` | boolean | — | false | When true, advise the kernel to drop the page cache for each blob after it is written or read (`posix_fadvise` with `POSIX_FADV_DONTNEED`). On real filesystems this takes a globally serialized, all-CPU kernel path that stalls on many-core hosts, and it evicts the page-cache tier that fast-disk deployments rely on. Leave off unless you specifically want to keep this store's I/O out of the page cache. | ## RefSpec @@ -856,7 +857,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `retry` | [Retry](#retry) | — | — | Retry configuration to use when a network request fails. | | `max_concurrent_requests` | integer (uint) | — | `0` | Limit the number of simultaneous upstream requests to this many. A value of zero is treated as unlimited. If the limit is reached the request is queued. | | `connections_per_endpoint` | integer (uint) | — | 1 | The number of connections to make to each specified endpoint to balance the load over multiple TCP connections. | -| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request (e.g. a `ByteStream.Write` call) before it is cancelled. | +| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request before it is cancelled with `DeadlineExceeded` (retryable). This applies to every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment of read streams through their first response chunk. The body of an established read stream is not bounded by this deadline. | | `use_legacy_resource_names` | boolean | — | false | Use legacy `ByteStream` resource name format, omitting the digest function component from the path. | | `headers` | map of string to string | — | — | Static headers to attach to every outgoing gRPC request sent to this store's upstream endpoints. Useful for fixed authentication tokens (e.g. `{"authorization": "Bearer "}`) and other static metadata. | | `forward_headers` | array of string | — | — | Header names to forward from the incoming client request to every outgoing upstream request. The header value is taken from the client request that triggered this store operation. Use this to pass through dynamic credentials such as JWT tokens sent by build clients. | @@ -1126,6 +1127,8 @@ The total delay is additive, so this example produces 9.525 to 15.875 s of total | `consider_expired_after_s` | integer (uint32) | — | 0 | If the number of seconds since the `last_modified` time of the object is greater than this value, the object will not be considered "existing". This allows for external tools to delete objects that have not been uploaded in a long time. If a client receives a `NotFound` the client should re-upload the object. | | `max_retry_buffer_per_request` | integer (uint) | — | 5MB | The maximum buffer size to retain in case of a retryable error during upload. Setting this to zero will disable upload buffering; this means that in the event of a failure during upload, the entire upload will be aborted and the client will likely receive an error. | | `multipart_max_concurrent_uploads` | integer (uint) | — | 10 | Maximum number of concurrent `UploadPart` requests per `MultipartUpload`. | +| `multipart_part_size` | integer (uint) | — | 64MB | Target size of each part in a multipart upload, growing only as needed to stay under the provider's 10,000-part ceiling. Parts sized at the 5 MiB provider floor turn a multi-GiB blob into thousands of tiny concurrent requests, which churns HTTP connection pools; a larger part keeps the request count low. Each queued or in-flight part is buffered in memory, so one multipart upload buffers up to roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`; `max_concurrent_multipart_uploads` bounds the aggregate across uploads. Values are clamped to the provider's 5 MiB minimum and 5 GiB maximum part sizes. | +| `max_concurrent_multipart_uploads` | integer (uint) | — | 4 | Store-wide cap on the number of blobs concurrently uploaded via the multipart path. Each multipart upload can buffer roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`, so this bounds the store's aggregate upload buffer (about 5 GiB at the defaults). Uploads over the 5 MiB multipart threshold wait for a slot; smaller blobs are unaffected. | | `insecure_allow_http` | boolean | — | false | Allow unencrypted HTTP connections. Only use this for local testing. | | `disable_http2` | boolean | — | false | Disable HTTP/2 connections and only use HTTP/1.1. Default client configuration will have HTTP/1.1 and HTTP/2 enabled for connection schemes. HTTP/2 should be disabled if environments have poor support or performance related to HTTP/2. Safe to keep default unless underlying network environment, S3, or GCS API servers specify otherwise. | | `throttle_retry_delay_s` | number (float) | — | 2.0 | Minimum number of seconds to wait after a response that asked the client to slow down (HTTP 503 `ServerBusy` or 429) when it carries no `Retry-After` header.0 | @@ -1491,6 +1494,7 @@ eternal 3600 s), so this maps them to classes without a CAS fetch. | `instance_name` | string | — | `""` | Used when the config references `instance_name` in the protocol. | | `cas_store` | string | Yes | — | The store name referenced in the `stores` map in the main config. This store name referenced here may be reused multiple times. | | `experimental_chunking` | [CasChunkingConfig](#caschunkingconfig) | — | not set — chunking RPCs are rejected, nothing is advertised, | Optional and experimental: enables the REAPI `SplitBlob`/`SpliceBlob` RPCs used by content-defined chunking clients (e.g. Bazel's `--experimental_remote_cache_chunking`). When set, the capabilities service advertises blob split/splice support and `FastCDC` 2020 parameters for this instance. When `cas_store` is a grpc store the RPCs are forwarded to the backend (which must support chunking with matching parameters); otherwise they are served locally. | +| `batch_read_per_blob_timeout_s` | integer (uint32) | — | 30 seconds | Per-blob timeout in seconds applied inside `BatchReadBlobs`. Each blob read in a batch races this deadline so a single stalled read cannot pin the whole batch until the client's overall deadline. Cold reads through a slow backing tier (e.g. an object store that populates a local cache on read) can legitimately take longer than the default for large blobs; raise this value for such deployments. | ## ActionCacheServiceConfig