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\":[]}}", 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") { 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-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 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/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index 24ba13cd75b..cd2561b6ea1 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -1842,6 +1842,7 @@ impl 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/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/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(()) +} 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(()) +} 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(); 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)); + } +}