From 20c288cf36ab130f50335359bee3318d85d2db7c Mon Sep 17 00:00:00 2001 From: b7r6 Date: Sat, 4 Jul 2026 08:09:23 -0400 Subject: [PATCH 1/6] nativelink-store: propagate completeness backend errors instead of swallowing them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CompletenessCheckingStore collapsed every backend failure into a NotFound miss. A miss tells the client the result does not exist and to re-execute; a backend fault (unavailable Redis, undecodable record) is transient and retryable. Laundering the fault as a clean miss silently discards a valid cached result — and behind an AC-serving frontend it turns a 5xx-class condition into a 404. NotFound from the backend still becomes a miss (the genuine eviction path is untouched); any other code now propagates. CacheLookupScheduler is updated in the same change: a cache-lookup fault there must never fail the Execute request, so non-NotFound lookup errors now log a warning and fall through to execution — the same availability the system had when the store laundered them, but visible to operators and correct at the AC service surface. Previously the scheduler's propagate-the-error branch was unreachable for backend faults precisely because the store swallowed them. Tests: an injected always-failing AC backend must surface its real code through both get_part and has (fails against the old store code, which demoted the error to NotFound), and the same injected fault at the scheduler must delegate to execution rather than fail the client. Co-Authored-By: Claude Fable 5 --- .../src/cache_lookup_scheduler.rs | 24 +--- .../tests/cache_lookup_scheduler_test.rs | 105 +++++++++++++- .../src/completeness_checking_store.rs | 19 +-- .../tests/completeness_checking_store_test.rs | 132 +++++++++++++++++- 4 files changed, 251 insertions(+), 29 deletions(-) diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 710eb615ea8..3d6ced82639 100644 --- a/nativelink-scheduler/src/cache_lookup_scheduler.rs +++ b/nativelink-scheduler/src/cache_lookup_scheduler.rs @@ -44,7 +44,7 @@ use parking_lot::{Mutex, MutexGuard}; use scopeguard::guard; use tokio::sync::oneshot; use tonic::{Request, Response}; -use tracing::error; +use tracing::{error, warn}; use crate::known_platform_property_provider::KnownPlatformPropertyProvider; @@ -302,23 +302,13 @@ impl CacheLookupScheduler { return; } Err(err) => { - // NotFound errors just mean we need to execute our action. + // 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. if err.code != Code::NotFound { - let err = err.append("In CacheLookupScheduler::add_action"); - let maybe_pending_txs = { - let mut inflight_cache_checks = inflight_cache_checks.lock(); - // We are ready to resolve the in-flight actions. We remove the - // in-flight actions from the map. - inflight_cache_checks.remove(unique_key) - }; - let Some(pending_txs) = maybe_pending_txs else { - return; // Nobody is waiting for this action anymore. - }; - for (_client_operation_id, pending_tx) in pending_txs { - // Ignore errors here, as the other end may have hung up. - drop(pending_tx.send(Err(err.clone()))); - } - return; + warn!(?err, "Cache lookup failed; falling through to execution"); } } } diff --git a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index c95b3cbface..cbff20b9867 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -19,10 +19,14 @@ mod utils { pub(crate) mod scheduler_utils; } +use core::pin::Pin; + +use async_trait::async_trait; use futures::join; use nativelink_config::stores::MemorySpec; -use nativelink_error::Error; +use nativelink_error::{Code, Error, make_err}; use nativelink_macro::nativelink_test; +use nativelink_metric::MetricsComponent; use nativelink_proto::build::bazel::remote::execution::v2::ActionResult as ProtoActionResult; use nativelink_scheduler::cache_lookup_scheduler::CacheLookupScheduler; use nativelink_scheduler::mock_scheduler::MockActionScheduler; @@ -30,9 +34,13 @@ use nativelink_store::memory_store::MemoryStore; use nativelink_util::action_messages::{ 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::store_trait::{Store, StoreLike}; +use nativelink_util::store_trait::{ + RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, +}; use pretty_assertions::assert_eq; use prost::Message; use tokio::sync::watch; @@ -40,6 +48,68 @@ 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. +#[derive(Debug, MetricsComponent)] +struct AlwaysErrStore { + #[metric(help = "Error code every operation fails with")] + code: u32, +} + +#[async_trait] +impl StoreDriver for AlwaysErrStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + Err(make_err!(Code::Unavailable, "injected AC backend failure")) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + Err(make_err!(Code::Unavailable, "injected AC backend failure")) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(Code::Unavailable, "injected AC backend failure")) + } + + fn inner_store(&self, _key: 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: Arc, + ) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(AlwaysErrStore); + struct TestContext { mock_scheduler: Arc, ac_store: Store, @@ -129,6 +199,37 @@ async fn add_action_handles_skip_cache() -> Result<(), Error> { Ok(()) } +/// A faulting AC backend (unavailable Redis, decode error) 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, + })); + let cache_scheduler = CacheLookupScheduler::new(ac_store, mock_scheduler.clone())?; + let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 123)); + let (_forward_watch_channel_tx, forward_watch_channel_rx) = + watch::channel(Arc::new(ActionState { + client_operation_id: OperationId::default(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + 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 + )))) + ); + add_result.expect("AC backend fault must fall through to execution, not fail the client"); + Ok(()) +} + #[nativelink_test] async fn find_by_client_operation_id_call_passed() -> Result<(), Error> { let context = make_cache_scheduler()?; diff --git a/nativelink-store/src/completeness_checking_store.rs b/nativelink-store/src/completeness_checking_store.rs index 38e054cefa5..9cc282d5d09 100644 --- a/nativelink-store/src/completeness_checking_store.rs +++ b/nativelink-store/src/completeness_checking_store.rs @@ -33,7 +33,6 @@ use nativelink_util::store_trait::{ }; use parking_lot::Mutex; use tokio::sync::Notify; -use tracing::warn; use crate::ac_utils::{get_and_decode_digest, get_size_and_decode_digest}; @@ -308,15 +307,19 @@ impl CompletenessCheckingStore { Some(Ok(())) => self.complete_entries_counter.inc(), Some(Err((err, i))) => { self.incomplete_entries_counter.inc(); - state_mux.lock().results[i] = None; - // Note: Don't return the errors. We just flag the result as - // missing but show a warning if it's not a NotFound. + // 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". if err.code != Code::NotFound { - warn!( - ?err, - "Error checking existence of digest" - ); + return Err(err); } + state_mux.lock().results[i] = None; } None => { // We are done, so flag it done and ensure we notify the diff --git a/nativelink-store/tests/completeness_checking_store_test.rs b/nativelink-store/tests/completeness_checking_store_test.rs index 8766f8b429a..ed3dcd2c38e 100644 --- a/nativelink-store/tests/completeness_checking_store_test.rs +++ b/nativelink-store/tests/completeness_checking_store_test.rs @@ -12,11 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +use core::pin::Pin; use std::sync::Arc; +use async_trait::async_trait; use nativelink_config::stores::MemorySpec; -use nativelink_error::Error; +use nativelink_error::{Code, Error, make_err}; use nativelink_macro::nativelink_test; +use nativelink_metric::MetricsComponent; use nativelink_proto::build::bazel::remote::execution::v2::{ ActionResult as ProtoActionResult, Directory, DirectoryNode, FileNode, OutputDirectory, OutputFile, Tree, @@ -24,9 +27,13 @@ use nativelink_proto::build::bazel::remote::execution::v2::{ use nativelink_store::ac_utils::serialize_and_upload_message; use nativelink_store::completeness_checking_store::CompletenessCheckingStore; use nativelink_store::memory_store::MemoryStore; +use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; -use nativelink_util::store_trait::{Store, StoreLike}; +use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; +use nativelink_util::store_trait::{ + RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, +}; const ROOT_FILE: DigestInfo = DigestInfo::new([0u8; 32], 0); const ROOT_DIRECTORY: DigestInfo = DigestInfo::new([1u8; 32], 0); @@ -245,3 +252,124 @@ async fn verify_completeness_get() -> Result<(), Error> { Ok(()) } + +/// A store driver whose reads always fail with a configurable, non-`NotFound` +/// backend error. Used to prove that a genuine backend read failure on the +/// action-result record is surfaced with its real code rather than laundered +/// into a `NotFound` "miss". +#[derive(Debug, MetricsComponent)] +struct AlwaysErrStore { + code: Code, +} + +#[async_trait] +impl StoreDriver for AlwaysErrStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + Err(make_err!( + self.code, + "injected backend failure in has_with_results" + )) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + Err(make_err!(self.code, "injected backend failure in update")) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(self.code, "injected backend failure in get_part")) + } + + fn inner_store(&self, _key: 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: Arc, + ) -> Result<(), Error> { + Ok(()) + } +} + +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`. +#[nativelink_test] +async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Error> { + let make_store = || { + let ac_store = Store::new(Arc::new(AlwaysErrStore { + code: Code::Unavailable, + })); + let cas_store = Store::new(MemoryStore::new(&MemorySpec::default())); + CompletenessCheckingStore::new(ac_store, cas_store) + }; + let some_digest = DigestInfo::new([9u8; 32], 123); + + // get_part path (nix narinfo GET / Bazel AC get). + let get_err = make_store() + .get_part_unchunked(some_digest, 0, None) + .await + .expect_err("a backend read failure must not be swallowed into a hit"); + assert_ne!( + get_err.code, + Code::NotFound, + "backend fault must not masquerade as a NotFound miss (would 404)" + ); + assert_eq!( + get_err.code, + Code::Unavailable, + "backend fault must propagate with its real code" + ); + + // has path (nix narinfo HEAD / Bazel AC has). + let has_err = make_store() + .has(some_digest) + .await + .expect_err("a backend read failure on has must not be swallowed into a miss"); + assert_ne!( + has_err.code, + Code::NotFound, + "backend fault on has must not masquerade as a NotFound miss (would 404)" + ); + assert_eq!( + has_err.code, + Code::Unavailable, + "backend fault on has must propagate with its real code" + ); + + Ok(()) +} From 56085a463f6ec417f025239a6733b6ebbfdca973 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 14:02:02 +0200 Subject: [PATCH 2/6] fix(evicting_map): stop has() batch from cascading-evicting every queried blob (#2840) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `sizes_for_keys` snapshotted `lru.len()` once for the whole batch and reused that stale length in `should_evict` for every key in the loop. When the map sits one entry over `max_count` — e.g. an action-input lease (or a TTL-pending entry) pins it at `max_count + 1` — the first unleased queried key is popped, bringing the real length back to `max_count`. But iterations 2..N still saw the stale `lru_len == max_count + 1`, so `is_over_count` stayed true and every remaining queried, unleased entry was popped too: remove-callbacks queued, `FileEntryImpl` unref deleting the backing files. The batch then reports `None` for all N digests, the client treats them as absent and re-uploads, and a plain existence check (`has()` → FilesystemStore → `sizes_for_keys`) has destroyed every resident blob it was only asked about. Fix: keep a mutable working count and decrement it as entries are reaped, so the over-count predicate is re-evaluated against the current length each iteration. Eviction now stops the moment the map is back at/under the limit — a size/existence query evicts at most the genuine overage and never cascades. Adds a white-box regression test that pins the map one over `max_count` with a leased entry (bypassing the insert-time eviction pass the way a held lease does) and asserts a batch `sizes_for_keys` over the resident keys leaves all but the single-entry overage present. It fails on the stale-len code (all N reported missing) and passes with the fix. Provenance: adversarial review fleet, highest-severity finding. --- nativelink-util/src/evicting_map.rs | 105 +++++++++++++++++++++++++++- 1 file changed, 104 insertions(+), 1 deletion(-) diff --git a/nativelink-util/src/evicting_map.rs b/nativelink-util/src/evicting_map.rs index a94bfd534a8..8076f9241a0 100644 --- a/nativelink-util/src/evicting_map.rs +++ b/nativelink-util/src/evicting_map.rs @@ -634,7 +634,13 @@ where let (removal_futures, data_to_unref) = { let mut state = self.state.lock(); - let lru_len = state.lru.len(); + // Working copy of the resident count, decremented as entries are + // reaped below. Evaluating `should_evict` against a length + // snapshotted once for the whole batch would keep reporting + // over-count for every remaining key after the first reap brings + // the map back to its limit, cascading into evicting every + // queried entry — reachable from a plain `has()` batch. + let mut lru_len = state.lru.len(); let mut data_to_unref = Vec::new(); let mut removal_futures = Vec::new(); for (key, result) in keys.into_iter().zip(results.iter_mut()) { @@ -656,6 +662,7 @@ where if !is_leased && self.should_evict(lru_len, entry, 0, u64::MAX) { *result = None; if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) { + lru_len -= 1; info!(?key, "Item expired, evicting"); let (data, futures) = state.remove(key.borrow(), &eviction_item, false); @@ -1042,3 +1049,99 @@ where self.state.lock().add_remove_callback(callback); } } + +#[cfg(test)] +mod tests { + use nativelink_config::stores::EvictionPolicy; + use nativelink_macro::nativelink_test; + use pretty_assertions::assert_eq; + + use super::{EvictingMap, EvictionItem, LenEntry}; + use crate::instant_wrapper::MockInstantWrapped; + + #[derive(Clone, Debug, PartialEq, Eq)] + struct TestEntry(u64); + + impl LenEntry for TestEntry { + fn len(&self) -> u64 { + self.0 + } + fn is_empty(&self) -> bool { + self.0 == 0 + } + } + + /// A batch existence check (`has()` → `sizes_for_keys`) over resident + /// keys, run while the map is pinned one entry over `max_count` by a + /// leased entry, must evict at most the genuine overage — never every + /// queried key. + /// + /// The bug: `sizes_for_keys` snapshotted `lru.len()` once for the whole + /// batch and re-used that stale value in `should_evict` for every key. + /// After the first queried entry was reaped (bringing the map back to + /// `max_count`), every remaining queried key still saw the stale + /// over-count and was popped too. A plain `has()` batch would then + /// destroy all N resident blobs and report them all missing. + /// + /// White-box: the pinning entry is planted with a raw leased `State::put` + /// (no eviction pass), modeling the window in which an action-input lease + /// holds the map over its count limit while unleased blobs stay resident. + #[nativelink_test("crate")] + async fn size_check_batch_evicts_at_most_the_overage() { + let map = EvictingMap::::new( + &EvictionPolicy { + max_bytes: 0, + evict_bytes: 0, + max_seconds: 0, + max_count: 3, + }, + MockInstantWrapped::default(), + ); + + map.insert("k1".to_string(), TestEntry(10)).await; + map.insert("k2".to_string(), TestEntry(20)).await; + map.insert("k3".to_string(), TestEntry(30)).await; + + // Pin the map one over `max_count` with a leased entry, bypassing the + // insert-time eviction pass exactly as an already-held lease does. + { + let mut state = map.state.lock(); + state.lease("pinned".to_string()); + let replaced = state.put( + &"pinned".to_string(), + EvictionItem { + seconds_since_anchor: 0, + data: TestEntry(1), + }, + ); + assert!(replaced.is_none()); + state.sum_store_size += 1; + } + assert_eq!(map.len_for_test(), 4); + + let keys = ["k1".to_string(), "k2".to_string(), "k3".to_string()]; + let mut results = [None, None, None]; + map.sizes_for_keys(keys.iter(), &mut results[..], false) + .await; + + // The overage is a single entry, so at most the first queried key may + // be reaped. The rest of the batch must stay resident and report size. + assert_eq!( + results[1], + Some(20), + "k2 must survive a batch size check: {results:?}" + ); + assert_eq!( + results[2], + Some(30), + "k3 must survive a batch size check: {results:?}" + ); + assert_eq!( + map.len_for_test(), + 3, + "only the genuine overage may be evicted by a size query" + ); + assert_eq!(map.size_for_key(&"k2".to_string()).await, Some(20)); + assert_eq!(map.size_for_key(&"k3".to_string()).await, Some(30)); + } +} From f798fa160f78225f059996f945d73334cfa0dcd1 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:28:03 +0100 Subject: [PATCH 3/6] Update Rust crate rustls to v0.23.46 [SECURITY] (#2938) * Update Rust crate rustls to v0.23.46 [SECURITY] * Fix MODULE.bazel.lock --- Cargo.lock | 4 ++-- MODULE.bazel.lock | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8c81817d67b..9c125801680 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4216,9 +4216,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.45" +version = "0.23.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +checksum = "48e13bd8c0e9365c43cfa5c9e8f9ad49d3c8444926c9aac819e0e4dc503c8fdf" dependencies = [ "log", "once_cell", diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 848a85eff59..4075f99ccb2 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1115,7 +1115,7 @@ "rustls-webpki_0.103.13": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aws-lc-rs\",\"optional\":true,\"req\":\"^1.14\"},{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.22\"},{\"kind\":\"dev\",\"name\":\"bencher\",\"req\":\"^0.1.5\"},{\"kind\":\"dev\",\"name\":\"bzip2\",\"req\":\"^0.6\"},{\"kind\":\"dev\",\"name\":\"once_cell\",\"req\":\"^1.17.2\"},{\"default_features\":false,\"name\":\"pki-types\",\"package\":\"rustls-pki-types\",\"req\":\"^1.12\"},{\"default_features\":false,\"features\":[\"aws_lc_rs\"],\"kind\":\"dev\",\"name\":\"rcgen\",\"req\":\"^0.14.2\"},{\"default_features\":false,\"name\":\"ring\",\"optional\":true,\"req\":\"^0.17\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"untrusted\",\"req\":\"^0.9\"},{\"kind\":\"dev\",\"name\":\"x509-parser\",\"req\":\"^0.18.1\"}],\"features\":{\"alloc\":[\"ring?/alloc\",\"pki-types/alloc\"],\"aws-lc-rs\":[\"dep:aws-lc-rs\",\"aws-lc-rs/aws-lc-sys\",\"aws-lc-rs/prebuilt-nasm\"],\"aws-lc-rs-fips\":[\"dep:aws-lc-rs\",\"aws-lc-rs/fips\"],\"aws-lc-rs-unstable\":[\"aws-lc-rs\",\"aws-lc-rs/unstable\"],\"default\":[\"std\"],\"ring\":[\"dep:ring\"],\"std\":[\"alloc\",\"pki-types/std\"]}}", "rustls-webpki_0.103.15": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aws-lc-rs\",\"optional\":true,\"req\":\"^1.18\"},{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.23\"},{\"kind\":\"dev\",\"name\":\"bencher\",\"req\":\"^0.1.5\"},{\"kind\":\"dev\",\"name\":\"bzip2\",\"req\":\"^0.6\"},{\"kind\":\"dev\",\"name\":\"once_cell\",\"req\":\"^1.17.2\"},{\"default_features\":false,\"name\":\"pki-types\",\"package\":\"rustls-pki-types\",\"req\":\"^1.12\"},{\"default_features\":false,\"features\":[\"aws_lc_rs\"],\"kind\":\"dev\",\"name\":\"rcgen\",\"req\":\"^0.14.2\"},{\"default_features\":false,\"name\":\"ring\",\"optional\":true,\"req\":\"^0.17\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"untrusted\",\"req\":\"^0.9\"},{\"kind\":\"dev\",\"name\":\"x509-parser\",\"req\":\"^0.18.1\"}],\"features\":{\"alloc\":[\"ring?/alloc\",\"pki-types/alloc\"],\"aws-lc-rs\":[\"dep:aws-lc-rs\",\"aws-lc-rs/aws-lc-sys\",\"aws-lc-rs/prebuilt-nasm\"],\"aws-lc-rs-fips\":[\"dep:aws-lc-rs\",\"aws-lc-rs/fips\"],\"aws-lc-rs-unstable\":[\"aws-lc-rs\"],\"default\":[\"std\"],\"ring\":[\"dep:ring\"],\"std\":[\"alloc\",\"pki-types/std\"]}}", "rustls_0.23.40": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aws-lc-rs\",\"optional\":true,\"req\":\"^1.14\"},{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.22\"},{\"kind\":\"dev\",\"name\":\"bencher\",\"req\":\"^0.1.5\"},{\"default_features\":false,\"features\":[\"std\"],\"name\":\"brotli\",\"optional\":true,\"req\":\"^8\"},{\"name\":\"brotli-decompressor\",\"optional\":true,\"req\":\"^5.0.0\"},{\"kind\":\"dev\",\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"default_features\":false,\"features\":[\"default-hasher\",\"inline-more\"],\"name\":\"hashbrown\",\"optional\":true,\"req\":\"^0.15\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"name\":\"log\",\"optional\":true,\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"log\",\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"macro_rules_attribute\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"num-bigint\",\"req\":\"^0.4.4\"},{\"default_features\":false,\"features\":[\"alloc\",\"race\"],\"name\":\"once_cell\",\"req\":\"^1.16\"},{\"features\":[\"alloc\"],\"name\":\"pki-types\",\"package\":\"rustls-pki-types\",\"req\":\"^1.12\"},{\"default_features\":false,\"features\":[\"pem\",\"aws_lc_rs\"],\"kind\":\"dev\",\"name\":\"rcgen\",\"req\":\"^0.14\"},{\"name\":\"ring\",\"optional\":true,\"req\":\"^0.17\"},{\"kind\":\"build\",\"name\":\"rustversion\",\"optional\":true,\"req\":\"^1.0.6\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.5.0\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3.6\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"webpki\",\"package\":\"rustls-webpki\",\"req\":\"^0.103.5\"},{\"kind\":\"dev\",\"name\":\"webpki-roots\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"x509-parser\",\"req\":\"^0.17\"},{\"name\":\"zeroize\",\"req\":\"^1.8\"},{\"name\":\"zlib-rs\",\"optional\":true,\"req\":\"^0.6\"}],\"features\":{\"aws-lc-rs\":[\"aws_lc_rs\"],\"aws_lc_rs\":[\"dep:aws-lc-rs\",\"webpki/aws-lc-rs\",\"aws-lc-rs/aws-lc-sys\",\"aws-lc-rs/prebuilt-nasm\"],\"brotli\":[\"dep:brotli\",\"dep:brotli-decompressor\",\"std\"],\"custom-provider\":[],\"default\":[\"aws_lc_rs\",\"logging\",\"prefer-post-quantum\",\"std\",\"tls12\"],\"fips\":[\"aws_lc_rs\",\"aws-lc-rs?/fips\",\"webpki/aws-lc-rs-fips\"],\"logging\":[\"log\"],\"prefer-post-quantum\":[\"aws_lc_rs\"],\"read_buf\":[\"rustversion\",\"std\"],\"ring\":[\"dep:ring\",\"webpki/ring\"],\"std\":[\"webpki/std\",\"pki-types/std\",\"once_cell/std\"],\"tls12\":[],\"zlib\":[\"dep:zlib-rs\"]}}", - "rustls_0.23.45": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aws-lc-rs\",\"optional\":true,\"req\":\"^1.18\"},{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.22\"},{\"kind\":\"dev\",\"name\":\"bencher\",\"req\":\"^0.1.5\"},{\"default_features\":false,\"features\":[\"std\"],\"name\":\"brotli\",\"optional\":true,\"req\":\"^8\"},{\"name\":\"brotli-decompressor\",\"optional\":true,\"req\":\"^5.0.0\"},{\"kind\":\"dev\",\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"default_features\":false,\"features\":[\"default-hasher\",\"inline-more\"],\"name\":\"hashbrown\",\"optional\":true,\"req\":\"^0.15\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"name\":\"log\",\"optional\":true,\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"log\",\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"macro_rules_attribute\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"num-bigint\",\"req\":\"^0.4.4\"},{\"default_features\":false,\"features\":[\"alloc\",\"race\"],\"name\":\"once_cell\",\"req\":\"^1.16\"},{\"features\":[\"alloc\"],\"name\":\"pki-types\",\"package\":\"rustls-pki-types\",\"req\":\"^1.12\"},{\"default_features\":false,\"features\":[\"pem\",\"aws_lc_rs\"],\"kind\":\"dev\",\"name\":\"rcgen\",\"req\":\"^0.14\"},{\"name\":\"ring\",\"optional\":true,\"req\":\"^0.17\"},{\"kind\":\"build\",\"name\":\"rustversion\",\"optional\":true,\"req\":\"^1.0.6\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.5.0\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3.6\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"webpki\",\"package\":\"rustls-webpki\",\"req\":\"^0.103.14\"},{\"kind\":\"dev\",\"name\":\"webpki-roots\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"x509-parser\",\"req\":\"^0.17\"},{\"name\":\"zeroize\",\"req\":\"^1.8\"},{\"name\":\"zlib-rs\",\"optional\":true,\"req\":\"^0.6\"}],\"features\":{\"aws-lc-rs\":[\"aws_lc_rs\"],\"aws_lc_rs\":[\"dep:aws-lc-rs\",\"webpki/aws-lc-rs\",\"aws-lc-rs/aws-lc-sys\",\"aws-lc-rs/prebuilt-nasm\"],\"brotli\":[\"dep:brotli\",\"dep:brotli-decompressor\",\"std\"],\"custom-provider\":[],\"default\":[\"aws_lc_rs\",\"logging\",\"prefer-post-quantum\",\"std\",\"tls12\"],\"fips\":[\"aws_lc_rs\",\"aws-lc-rs?/fips\",\"webpki/aws-lc-rs-fips\"],\"logging\":[\"log\"],\"prefer-post-quantum\":[\"aws_lc_rs\"],\"read_buf\":[\"rustversion\",\"std\"],\"ring\":[\"dep:ring\",\"webpki/ring\"],\"std\":[\"webpki/std\",\"pki-types/std\",\"once_cell/std\"],\"tls12\":[],\"zlib\":[\"dep:zlib-rs\"]}}", + "rustls_0.23.46": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aws-lc-rs\",\"optional\":true,\"req\":\"^1.18\"},{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.22\"},{\"kind\":\"dev\",\"name\":\"bencher\",\"req\":\"^0.1.5\"},{\"default_features\":false,\"features\":[\"std\"],\"name\":\"brotli\",\"optional\":true,\"req\":\"^8\"},{\"name\":\"brotli-decompressor\",\"optional\":true,\"req\":\"^5.0.0\"},{\"kind\":\"dev\",\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"default_features\":false,\"features\":[\"default-hasher\",\"inline-more\"],\"name\":\"hashbrown\",\"optional\":true,\"req\":\"^0.15\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"name\":\"log\",\"optional\":true,\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"log\",\"req\":\"^0.4.8\"},{\"kind\":\"dev\",\"name\":\"macro_rules_attribute\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"num-bigint\",\"req\":\"^0.4.4\"},{\"default_features\":false,\"features\":[\"alloc\",\"race\"],\"name\":\"once_cell\",\"req\":\"^1.16\"},{\"features\":[\"alloc\"],\"name\":\"pki-types\",\"package\":\"rustls-pki-types\",\"req\":\"^1.12\"},{\"default_features\":false,\"features\":[\"pem\",\"aws_lc_rs\"],\"kind\":\"dev\",\"name\":\"rcgen\",\"req\":\"^0.14\"},{\"name\":\"ring\",\"optional\":true,\"req\":\"^0.17\"},{\"kind\":\"build\",\"name\":\"rustversion\",\"optional\":true,\"req\":\"^1.0.6\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.5.0\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3.6\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"webpki\",\"package\":\"rustls-webpki\",\"req\":\"^0.103.14\"},{\"kind\":\"dev\",\"name\":\"webpki-roots\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"x509-parser\",\"req\":\"^0.17\"},{\"name\":\"zeroize\",\"req\":\"^1.8\"},{\"name\":\"zlib-rs\",\"optional\":true,\"req\":\"^0.6\"}],\"features\":{\"aws-lc-rs\":[\"aws_lc_rs\"],\"aws_lc_rs\":[\"dep:aws-lc-rs\",\"webpki/aws-lc-rs\",\"aws-lc-rs/aws-lc-sys\",\"aws-lc-rs/prebuilt-nasm\"],\"brotli\":[\"dep:brotli\",\"dep:brotli-decompressor\",\"std\"],\"custom-provider\":[],\"default\":[\"aws_lc_rs\",\"logging\",\"prefer-post-quantum\",\"std\",\"tls12\"],\"fips\":[\"aws_lc_rs\",\"aws-lc-rs?/fips\",\"webpki/aws-lc-rs-fips\"],\"logging\":[\"log\"],\"prefer-post-quantum\":[\"aws_lc_rs\"],\"read_buf\":[\"rustversion\",\"std\"],\"ring\":[\"dep:ring\",\"webpki/ring\"],\"std\":[\"webpki/std\",\"pki-types/std\",\"once_cell/std\"],\"tls12\":[],\"zlib\":[\"dep:zlib-rs\"]}}", "rustversion_1.0.22": "{\"dependencies\":[{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.49\"}],\"features\":{}}", "ryu_1.0.23": "{\"dependencies\":[{\"default_features\":false,\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.8\",\"target\":\"cfg(not(miri))\"},{\"name\":\"no-panic\",\"optional\":true,\"req\":\"^0.1\"},{\"kind\":\"dev\",\"name\":\"num_cpus\",\"req\":\"^1.8\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"rand_xorshift\",\"req\":\"^0.5\"}],\"features\":{\"small\":[]}}", "safemem_0.3.3": "{\"dependencies\":[],\"features\":{\"default\":[\"std\"],\"std\":[]}}", From 264b7eb28e5fb59311a8348e367d659a58290860 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 09:59:38 -0400 Subject: [PATCH 4/6] Reject short and torn uploads in ByteStream and FilesystemStore, and report QueryWriteStatus honestly (#2881) * fix(service): reject ByteStream finish_write before the declared blob size A ByteStream Write that sets finish_write=true before sending the full size declared in its resource name was accepted as a complete upload: the handler returned OK with committed_size equal to the full declared size, the blob then showed as present in FindMissingBlobs, and Read served the truncated bytes under a digest they do not hash to -- silent CAS content corruption. Both the streaming path (process_client_stream) and the oneshot path (inner_write_oneshot) only guarded the too-many-bytes case and committed on finish_write without checking that the received byte count equalled the declared size. A VerifyStore-fronted write path catches this on EOF, but an unverified store tier (filesystem/memory fast tier) commits the partial. Reject finish_write with InvalidArgument when bytes_received != expected_size in both paths. The compressed path already verifies decoded size and digest. Also fixes a latent same-class defect in max_decoding_message_size_test, whose success case declared a size that did not match the bytes it sent and only passed because of this bug. Co-Authored-By: Claude Opus 4.8 * fix(cas): reject short/torn uploads and report honest write status Unifies the handler-level short finish_write rejection (N13, e5d1b0db) with two further CAS-integrity fixes for the same disclosure (TRA-1373): - filesystem_store: refuse to emplace a blob whose byte count disagrees with an ExactSize upload hint. A torn/aborted upload stream closes the writer half, which arrives at the store as a clean EOF; without this commit-point guard the partial bytes are stored under the full digest key regardless of any handler check (finish_write never arrives on the abort path). - bytestream_server: QueryWriteStatus clamps committed_size to the declared digest size (the filesystem store reports block-rounded on-disk footprint, e.g. 17 -> 4096) and only reports complete once the store holds at least that many bytes, instead of unconditional complete=true. Regression tests: - short_finish_write_is_rejected_and_not_committed (oneshot path, N13) - short_finish_write_is_rejected (streaming multi-message path) - aborted_mid_stream_upload_is_not_committed (EOF without finish_write; only the store-level guard can catch this) - max_decoding_message_size_test updated: it asserted the old short-upload-accepted behavior. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: b7r6 Co-authored-by: Claude Opus 4.8 Co-authored-by: Tom Parker-Shemilt --- nativelink-service/src/bytestream_server.rs | 31 ++++- .../tests/bytestream_server_test.rs | 122 +++++++++++++++++- nativelink-store/src/filesystem_store.rs | 20 ++- .../tests/filesystem_store_test.rs | 67 ++++++++++ 4 files changed, 233 insertions(+), 7 deletions(-) diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index d42c2f6acad..c18cf49dff6 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -944,6 +944,17 @@ impl ByteStreamServer { return Err(make_input_err!("Received more bytes than expected")); } if write_request.finish_write { + // The client declared the blob size in the resource name. If it + // signals completion before sending that many bytes, the upload + // is incomplete and MUST be rejected: committing it would admit a + // truncated artifact into the CAS under a digest it does not hash + // to, which later reads would serve as if it were complete. + let bytes_written = tx.get_bytes_written(); + if bytes_written != expected_size { + return Err(make_input_err!( + "Client finished writing before sending all data. Expected {expected_size} bytes but received {bytes_written}" + )); + } // Gracefully close our stream. tx.send_eof() .err_tip(|| "Failed to send EOF in ByteStream::write")?; @@ -1058,6 +1069,14 @@ impl ByteStreamServer { } if write_request.finish_write { + // Reject an incomplete upload: a finish_write before the declared + // size would otherwise commit a truncated blob under a digest it + // does not hash to. See inner_write for the full rationale. + if bytes_received != expected_size { + return Err(make_input_err!( + "Client finished writing before sending all data. Expected {expected_size} bytes but received {bytes_received}" + )); + } break; } } @@ -1373,9 +1392,17 @@ impl ByteStreamServer { complete: false, })); }; + // `item_size` is the store's reported size; for the filesystem store + // that is the block-rounded on-disk footprint, not the logical byte + // count. Clamp the reported committed_size to the declared digest size + // and only report `complete` once the store holds at least that many + // bytes, so QueryWriteStatus never over-reports nor falsely completes. + let expected_size = digest.size_bytes(); + let complete = item_size >= expected_size; + let committed_size = if complete { expected_size } else { item_size }; Ok(Response::new(QueryWriteStatusResponse { - committed_size: item_size.try_into().unwrap_or(i64::MAX), - complete: true, + committed_size: committed_size.try_into().unwrap_or(i64::MAX), + complete, })) } } diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index 20da09e6bde..972239cdee7 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -357,6 +357,61 @@ pub async fn chunked_stream_receives_all_data() -> Result<(), Box Result<(), Box> { + const SHORT_PREFIX: usize = 5; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + let raw_data = b"12456789abcdefghijk"; + let declared_len = raw_data.len(); + + let (tx, join_handle) = + make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); + + // Declare the full length in the resource name, but only send a short + // prefix and immediately set finish_write=true. + let resource_name = format!( + "{}/uploads/{}/blobs/{}/{}", + INSTANCE_NAME, "4dcec57e-1389-4ab5-b188-4a59f22ceb4b", HASH1, declared_len + ); + let write_request = WriteRequest { + resource_name, + write_offset: 0, + finish_write: true, + data: raw_data[..SHORT_PREFIX].into(), + }; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + + // The write MUST fail (not return OK with a full committed_size). + let server_result = join_handle.await.expect("Failed to join"); + assert!( + server_result.is_err(), + "Short finish_write must be rejected, got Ok: {server_result:?}" + ); + + // The blob MUST NOT be present under its declared digest. + assert!( + store + .has(DigestInfo::try_new(HASH1, declared_len)?) + .await? + .is_none(), + "Truncated blob must not be committed to the store", + ); + + Ok(()) +} + #[nativelink_test] pub async fn sha256_write_without_digest_segment_populates_cache() -> Result<(), Box> { @@ -1848,6 +1903,62 @@ pub async fn test_query_write_status_smoke_test() -> Result<(), Box Result<(), Box> { + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + // Declare a 19-byte blob, but only send 12 bytes across two messages and + // then finish_write (multi-message forces the streaming write path). + let raw_data = b"12456789abcdefghijk"; + let resource_name = format!( + "{}/uploads/{}/blobs/{}/{}", + INSTANCE_NAME, + "4dcec57e-1389-4ab5-b188-4a59f22ceb4b", + HASH1, + raw_data.len() + ); + let (tx, join_handle) = + make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); + + let mut write_request = WriteRequest { + resource_name, + write_offset: 0, + finish_write: false, + data: raw_data[..8].into(), + }; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + // Short final chunk: total 12 < declared 19. + write_request.write_offset = 8; + write_request.data = raw_data[8..12].into(); + write_request.finish_write = true; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + + let server_result = join_handle.await.expect("Failed to join"); + assert!( + server_result.is_err(), + "short finish_write must be rejected, got {server_result:?}" + ); + assert!( + store + .has(DigestInfo::try_new(HASH1, raw_data.len())?) + .await? + .is_none(), + "a rejected short upload must not leave a readable partial blob", + ); + Ok(()) +} + #[nativelink_test] pub async fn max_decoding_message_size_test() -> Result<(), Box> { const MAX_MESSAGE_SIZE: usize = 1024 * 1024; // 1MB. @@ -1876,9 +1987,12 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box Result<(), Box FilesystemStore { mut temp_file: FileSlot, final_key: StoreKey<'static>, mut reader: DropCloserReadHalf, + upload_size: UploadSizeInfo, ) -> Result { let mut data_size = 0; loop { @@ -1860,6 +1861,21 @@ impl FilesystemStore { data_size += data_len as u64; } + // Refuse to commit a blob whose byte count disagrees with its declared + // digest size. A torn or aborted upload stream closes the writer half, + // which arrives here as a normal EOF; without this check the partial + // bytes would be emplaced under the full digest's key and later read + // back OK with a mismatched hash (found by the rechaos chaos monkey). + if let UploadSizeInfo::ExactSize(expected_size) = upload_size + && data_size != expected_size + { + return Err(make_err!( + Code::InvalidArgument, + "Filesystem store received {data_size} bytes but the digest declares \ + {expected_size}; refusing to store a partial blob" + )); + } + let permit = if let Some(sem) = &self.write_semaphore { Some(sem.acquire().await.map_err(|err| { Error::from_std_err(Code::Internal, &err).append("Write semaphore closed") @@ -2122,7 +2138,7 @@ impl StoreDriver for FilesystemStore { self: Pin<&Self>, key: StoreKey<'_>, mut reader: DropCloserReadHalf, - _upload_size: UploadSizeInfo, + upload_size: UploadSizeInfo, ) -> Result { if is_zero_digest(key.borrow()) { // don't need to add, because zero length files are just assumed to exist. @@ -2141,7 +2157,7 @@ impl StoreDriver for FilesystemStore { let (entry, temp_file, temp_full_path) = self.make_temp_file(temp_key).await?; - self.update_file(entry, temp_file, key.into_owned(), reader) + self.update_file(entry, temp_file, key.into_owned(), reader, upload_size) .await .err_tip(|| { format!( diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 27f8bf1833c..3381af555d2 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -3118,3 +3118,70 @@ async fn partial_migration_stays_readable_and_retries_next_startup() -> Result<( Ok(()) } + +// Regression (TRA-1373, aborted-mid-stream variant): a ByteStream Write that is +// torn mid-stream (client abort, injected UNAVAILABLE, etc.) closes the writer +// half, which arrives at the store as a normal EOF with fewer bytes than the +// digest declares -- finish_write is never sent, so a handler-level check can +// never fire. The store itself must refuse to emplace the partial blob at the +// commit point when the upload declared an exact size. +#[nativelink_test] +async fn aborted_mid_stream_upload_is_not_committed() -> Result<(), Error> { + const DECLARED_SIZE: usize = 19; + const PARTIAL: &[u8] = b"12456789abcd"; // 12 bytes < declared 19. + + let content_path = make_temp_path("content_path"); + let temp_path = make_temp_path("temp_path"); + + let store = Box::pin( + FilesystemStore::::new(&FilesystemSpec { + content_path, + temp_path, + read_buffer_size: 1, + ..Default::default() + }) + .await?, + ); + + let digest = DigestInfo::try_new(HASH1, DECLARED_SIZE)?; + + let (mut writer, reader) = make_buf_channel_pair(); + let store_ref = &store; + let update_fut = async move { + store_ref + .update( + digest, + reader, + UploadSizeInfo::ExactSize(DECLARED_SIZE as u64), + ) + .await + }; + // Simulate the torn stream: partial bytes, then the writer half closes + // cleanly (EOF) without ever having sent the declared byte count. + let writer_fut = async move { + writer.send(Bytes::from_static(PARTIAL)).await?; + writer.send_eof()?; + Ok::<_, Error>(()) + }; + + let (update_result, writer_result) = tokio::join!(update_fut, writer_fut); + writer_result?; + + assert!( + update_result.is_err(), + "A torn upload (EOF before the declared size) must not be committed, got {update_result:?}" + ); + assert_eq!( + update_result.unwrap_err().code, + Code::InvalidArgument, + "Expected the store-level size guard to reject the partial blob" + ); + + // The partial blob MUST NOT be present under its declared digest. + assert_eq!( + store.has(digest).await?, + None, + "Partial blob must not be readable under the full digest" + ); + Ok(()) +} From 0cd82fffc2dc9de9e213bf5d754fb05fd2cbf9d5 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 10:30:34 -0400 Subject: [PATCH 5/6] error: map transport-originated Unknown to retryable Unavailable (#2816) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tonic codes a transport failure it cannot classify further (a connection reset or broken pipe mid-stream) as Code::Unknown with message "transport error"; tonic's own Status::from_error maps the shapes it does recognize (e.g. a connect refusal) to Unavailable already. The gRPC status model reserves UNKNOWN for genuinely unmapped errors and codes transport failures UNAVAILABLE. Left as Unknown, such a blip is treated as permanent by REAPI clients — Bazel retries UNAVAILABLE but not UNKNOWN — and by any retry configuration keyed to Unavailable, so a retryable failure becomes terminal and cacheable uploads are silently dropped mid-write. Detection is downcast-first: walk the Status source chain for a concrete tonic::transport::Error. The fallback, for a Status reconstructed from the wire where only the message survives, is anchored to the two shapes that actually occur: tonic's exact "transport error" rendering, and "transport error : " as produced by this crate's own From for Status when a peer re-encodes and forwards one. An app-level Unknown that merely mentions the phrase mid-message is not reclassified. Upstream's #2657 made the same argument for transient Redis failures; this applies it to the adjacent tonic transport surface. --- nativelink-error/src/lib.rs | 45 +++++++++++++- nativelink-error/tests/error_tests.rs | 88 ++++++++++++++++++++++++++- 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/nativelink-error/src/lib.rs b/nativelink-error/src/lib.rs index 71b88435b5d..c8f2bb3f5b6 100644 --- a/nativelink-error/src/lib.rs +++ b/nativelink-error/src/lib.rs @@ -394,10 +394,53 @@ impl From for Error { impl From for Error { fn from(status: tonic::Status) -> Self { - Self::new(status.code(), status.to_string()) + // tonic codes a mid-stream transport/`hyper` failure (a connection reset + // or broken pipe during a large streaming write) as `Code::Unknown` + // with message "transport error". But a broken connection is transient + // and retryable — the gRPC status model reserves `UNKNOWN` for genuinely + // unmapped errors and codes transport failures `UNAVAILABLE`. Left as + // `Unknown`, such a blip falls outside both our own `Retrier`'s and most + // clients' retry sets, so a retryable failure becomes terminal and + // cacheable uploads are silently dropped. Remap transport-originated + // `Unknown` to `Unavailable`. + let code = if status.code() == Code::Unknown && is_transport_status(&status) { + Code::Unavailable + } else { + status.code() + }; + Self::new(code, status.to_string()) } } +/// Whether a `tonic::Status` originates from a transport-layer failure (a +/// reset/broken connection), which tonic surfaces as `Code::Unknown`. +/// +/// tonic's own `Status::from_error` already maps the failures it recognizes +/// (for example a connect refusal) to `Unavailable`; the `Unknown` shape +/// arises for transport errors it does not, and for statuses that crossed a +/// process boundary. Detection is downcast-first: precise when the concrete +/// transport error is still reachable in the source chain; falls back to the +/// message otherwise. +fn is_transport_status(status: &tonic::Status) -> bool { + let mut source = std::error::Error::source(status); + while let Some(err) = source { + if err.is::() { + return true; + } + source = err.source(); + } + // Fallback for a `Status` reconstructed from the wire, where the concrete + // transport error is no longer in the source chain and only the message + // survives. `tonic::transport::Error` renders as exactly "transport + // error" (no detail suffix), and this crate's own `From for + // Status` re-encodes a proxied one as "transport error : ". Anchor to those two shapes — an unanchored `contains` would + // reclassify any app-level `Unknown` whose message merely mentions + // "transport error" as retryable `Unavailable`. + let msg = status.message(); + msg == "transport error" || msg.starts_with("transport error : ") +} + impl From for tonic::Status { fn from(val: Error) -> Self { Self::new(val.code, val.messages.join(" : ")) diff --git a/nativelink-error/tests/error_tests.rs b/nativelink-error/tests/error_tests.rs index 75ded58fc97..21c1370f026 100644 --- a/nativelink-error/tests/error_tests.rs +++ b/nativelink-error/tests/error_tests.rs @@ -12,9 +12,95 @@ // See the License for the specific language governing permissions and // limitations under the License. -use nativelink_error::Error; +use nativelink_error::{Code, Error}; use walkdir::WalkDir; +// A transport-layer failure (mid-stream connection reset) is reported by tonic +// as `Code::Unknown` "transport error", but it is transient and must reach the +// client as the retryable `Unavailable` — otherwise big cacheable uploads are +// silently dropped when a shard resets the connection mid-write. +#[test] +fn transport_unknown_status_maps_to_unavailable() { + let status = tonic::Status::unknown("transport error"); + let err: Error = status.into(); + assert_eq!(err.code, Code::Unavailable); +} + +// A genuinely-unmapped `Unknown` (no transport origin) is preserved as-is. +#[test] +fn genuine_unknown_status_is_preserved() { + let status = tonic::Status::unknown("some application-level failure"); + let err: Error = status.into(); + assert_eq!(err.code, Code::Unknown); +} + +// A transport-error status that passed through this crate's own +// `From for Status` re-encoding (message segments joined with " : ") +// must still be recognized after the source chain is gone — the multi-hop +// proxy case the fallback exists for. +#[test] +fn proxied_transport_error_with_appended_context_maps_to_unavailable() { + let status = tonic::Status::unknown("transport error : while writing to upstream shard"); + let err: Error = status.into(); + assert_eq!(err.code, Code::Unavailable); +} + +// The message match is anchored: an app-level `Unknown` that merely mentions +// "transport error" mid-message must NOT be reclassified as retryable. +#[test] +fn unknown_mentioning_transport_error_mid_message_is_preserved() { + let status = tonic::Status::unknown("upstream proxy saw a transport error and gave up"); + let err: Error = status.into(); + assert_eq!(err.code, Code::Unknown); +} + +// The downcast walk, tested in isolation from the message fallback: a +// `Status` whose SOURCE CHAIN carries a real `tonic::transport::Error` maps +// to `Unavailable` even when its message matches neither anchor. The +// transport error is produced the honest way (a connect to a port with no +// listener, bound-then-dropped so nothing can race onto it) and attached +// via `Status::set_source` — the shape a middleware produces when it wraps +// a transport failure with its own message. Going through tonic's own +// `Status::from_error` instead would not reach our walk: tonic already +// maps the shapes it recognizes to `Unavailable` before we ever see them. +#[expect( + clippy::disallowed_methods, + reason = "obtaining a real transport error requires driving a connect on a runtime" +)] +#[test] +fn transport_error_in_source_chain_maps_to_unavailable() { + let runtime = tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime"); + let transport_err = runtime.block_on(async { + let dead_port = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() + }; + tonic::transport::Endpoint::from_shared(format!("http://127.0.0.1:{dead_port}")) + .unwrap() + .connect_timeout(core::time::Duration::from_secs(5)) + .connect() + .await + .expect_err("connect to a dropped listener must fail") + }); + let mut status = tonic::Status::unknown("middleware wrapped the failure with its own text"); + status.set_source(std::sync::Arc::new(transport_err)); + let err: Error = status.into(); + assert_eq!(err.code, Code::Unavailable); +} + +// Pin the round-trip that produces the proxied shape above: converting an +// `Unavailable` transport Error back to a Status and re-ingesting it must +// stay `Unavailable` with the anchored prefix intact. +#[test] +fn transport_error_round_trip_stays_unavailable() { + let err: Error = tonic::Status::unknown("transport error").into(); + assert_eq!(err.code, Code::Unavailable); + let with_context = err.append("while writing to upstream shard"); + let status: tonic::Status = with_context.into(); + let round_tripped: Error = status.into(); + assert_eq!(round_tripped.code, Code::Unavailable); +} + #[test] fn walkdir_source_error() { for entry in WalkDir::new("/bad/path") { From 89ea706c34b440d02ded72fef602c3649187d5d5 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 11:35:13 -0400 Subject: [PATCH 6/6] redis: close two subscription-lifecycle windows that silently drop events (#2839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two confirmed bugs left the RedisStore keyspace/pattern subscriptions recorded as active while never actually established on the live connection. Both come with a deterministic repro test that fails on the old code and passes on the fix. 1. psubscribe lost across a concurrent reconnect (StandardRedisManager). `psubscribe` used to fetch a connection, then insert the pattern into the tracked `subscriptions` set. If a concurrent `reconnect` snapshotted+replayed that set on the new connection in between, the pattern was PSUBSCRIBEd only on the old, about-to-be-demoted connection and recorded as active — so every PMessage for it was dropped until an unrelated later reconnect replayed the now-complete set. Fix: order registration against the reconnect snapshot. `psubscribe` now inserts the pattern BEFORE fetching a connection, and `reconnect` recomputes the missing delta and subscribes it while holding the `connection_manager` write lock just before publishing the new handle. A racing `psubscribe` then either lands in the replayed snapshot or blocks on the write lock and subscribes on the already-published live connection — no window where a pattern is "active" but subscribed only on a stale connection. 2. transient init error permanently disabled remove callbacks. `register_remove_callback` returned Ok(()) and deferred the CONFIG GET + PSUBSCRIBE "__key*__:*" to a background one-shot `OnceCell::get_or_try_init`. A single transient Redis error (e.g. -LOADING during a restart/failover) was only logged; the OnceCell stayed empty and nothing retried, so eviction/expiry callbacks were silently dead forever (ExistenceCacheStore never invalidated → stale `exists` hits). Fix: retry the init with capped backoff until it succeeds, and demote the advisory notify-keyspace-events CONFIG GET check to non-fatal (log misconfig, subscribe anyway) so a transient read failure can no longer permanently disable the subscription. Both repro tests fail on the unfixed code and pass on the fix (verified both directions on the build fleet). Full `cargo test -p nativelink-store` green, clippy clean, nightly rustfmt applied. Provenance: adversarial review fleet --- nativelink-store/src/redis_store.rs | 366 +++++++++++++++++++-- nativelink-store/tests/redis_store_test.rs | 82 ++++- 2 files changed, 417 insertions(+), 31 deletions(-) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index ef150ba6ccc..f52b199cd6c 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -324,12 +324,31 @@ impl RedisManager for StandardRedisManager let guard = self.subscriptions.lock(); guard.iter().map(Clone::clone).collect::>() }; - for subscription in subscriptions { - connection_manager.psubscribe(&subscription).await?; + for subscription in &subscriptions { + connection_manager.psubscribe(subscription).await?; } - // Publish the new handle under a brief exclusive lock. + 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 { + connection_manager.psubscribe(&subscription).await?; + replayed.insert(subscription); + } *guard = (connection_manager.clone(), new_uuid); } info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master"); @@ -342,10 +361,21 @@ impl RedisManager for StandardRedisManager async fn psubscribe(&self, pattern: &str) -> Result<(), Error> { debug!(pattern, "new psubscribe"); - let mut connection = self.get_connection().await?.0; + // Record the pattern BEFORE fetching a connection. A concurrent + // `reconnect` either sees the pattern in the snapshot it replays under + // its pre-publish write lock (subscribing it on the new connection), + // or it published first — in which case `get_connection` below returns + // the new live connection and we subscribe on it ourselves. Inserting + // after fetching the connection (the previous order) left a window + // where the pattern was recorded as active but PSUBSCRIBEd only on the + // old, about-to-be-replaced connection. let new_subscription = self.subscriptions.lock().insert(String::from(pattern)); if new_subscription { - let result = connection.psubscribe(pattern).await; + let result = async { + let mut connection = self.get_connection().await?.0; + connection.psubscribe(pattern).await.map_err(Error::from) + } + .await; if result.is_err() { self.subscriptions.lock().remove(pattern); } @@ -1471,35 +1501,68 @@ where fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { debug!(?callback, "New callback"); - let local_self = self.clone(); background_spawn!("remove_callback_subscribe", async move { self.remove_callbacks.lock().await.push(callback); - if let Err(err) = local_self.clone().has_remove_callback_subscribe - .get_or_try_init(|| async move { - let mut client = local_self.get_client().await?; - let cfg = 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"))?; - if cfg.len() != 1 { - warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); - return Err(make_input_err!("Got multiple items for CONFIG GET, expected one")); - } - let events_cfg = &cfg.first().ok_or_else(|| make_err!(Code::InvalidArgument, "Only one item"))?.1; - 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"); + // 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); + 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" + ); + sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(Duration::from_secs(5)); } - // 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"); - self.connection_manager.psubscribe("__key*__:*").await?; - Ok::<(), Error>(()) - }) - .await { - error!(?err, "Error while trying to initialise remove_callback_subscribe"); } + } }); Ok(()) } @@ -2597,3 +2660,246 @@ 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 { + use core::time::Duration; + use std::collections::HashSet; + use std::sync::Arc; + + use redis::aio::{ConnectionManager, ConnectionManagerConfig}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::sync::Notify; + + use super::{RedisManager, StandardRedisManager}; + + /// Reads one CRLF-terminated line from the stream. + async fn read_line(stream: &mut TcpStream) -> Option { + let mut line = Vec::new(); + let mut byte = [0u8; 1]; + loop { + if stream.read_exact(&mut byte).await.is_err() { + return None; + } + if byte[0] == b'\n' { + break; + } + if byte[0] != b'\r' { + line.push(byte[0]); + } + } + Some(String::from_utf8_lossy(&line).into_owned()) + } + + /// Parses one RESP command (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('*') { + return Some(vec![header]); + } + let count: usize = header[1..].parse().ok()?; + let mut parts = Vec::with_capacity(count); + for _ in 0..count { + let len_line = read_line(stream).await?; + let len: usize = len_line.strip_prefix('$')?.parse().ok()?; + let mut data = vec![0u8; len + 2]; // include trailing CRLF + stream.read_exact(&mut data).await.ok()?; + parts.push(String::from_utf8_lossy(&data[..len]).into_owned()); + } + Some(parts) + } + + struct FakeServerState { + /// 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_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, + } + + 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() + .map(|c| c.to_ascii_uppercase()) + .unwrap_or_default(); + match name.as_str() { + "HELLO" => { + stream + .write_all( + b"%3\r\n$6\r\nserver\r\n$5\r\nredis\r\n$7\r\nversion\r\n$5\r\n7.4.0\r\n$5\r\nproto\r\n:3\r\n", + ) + .await + .unwrap(); + } + "SCRIPT" => { + let reply = + format!("${}\r\n{}\r\n", state.script_hash.len(), state.script_hash); + stream.write_all(reply.as_bytes()).await.unwrap(); + } + "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; + } + let reply = format!( + ">3\r\n$10\r\npsubscribe\r\n${}\r\n{}\r\n:1\r\n", + pattern.len(), + pattern + ); + stream.write_all(reply.as_bytes()).await.unwrap(); + } + _ => { + stream.write_all(b"+OK\r\n").await.unwrap(); + } + } + } + } + + /// 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; + } + }); + } + + // 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(), + ) + .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" + ); + + // The manager believes "pat" is an active subscription... + let recorded: HashSet = manager.subscriptions.lock().clone(); + assert!( + recorded.contains("pat"), + "sanity: pattern recorded as active" + ); + + // ...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", + ); + }); + } +} diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index b668633a778..cd898f28cde 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -24,7 +24,7 @@ use nativelink_config::stores::{ExistenceCacheSpec, RedisMode, RedisSpec, StoreS use nativelink_error::{Code, Error, ErrorContext, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_redis_tester::{ - ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response, + 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, }; @@ -2717,6 +2717,86 @@ 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. +#[nativelink_test] +async fn transient_init_error_permanently_disables_remove_callbacks() -> 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") + .arg("GET") + .arg("notify-keyspace-events"), + "-LOADING Redis is loading the dataset in memory\r\n".to_string(), + ); + add_to_response( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + vec![Value::Nil], + ); + let redis_port = make_fake_redis_with_responses(responses) + .instrument(redis_span) + .await; + let spec = RedisSpec { + addresses: vec![format!("redis://127.0.0.1:{redis_port}/")], + mode: RedisMode::Standard, + ..Default::default() + }; + 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 {}))?; + + // Deterministic sync point: wait until the background init task has hit + // the transient error and logged it. + timeout(Duration::from_secs(5), async { + loop { + if logs_contain("Error while trying to initialise remove_callback_subscribe") { + break; + } + 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(()) +} + #[nativelink_test] async fn send_eviction_to_subscription_channel() -> Result<(), Error> { let (tx, rx) = tokio::sync::mpsc::unbounded_channel();