diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index c0832d0ef15..c68d4ee5fd3 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -493,3 +493,4 @@ Flybody MuJoCo connectome simulatable +unary diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 81b6179f929..8abaa14c390 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1628,12 +1628,18 @@ pub struct GrpcSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub connections_per_endpoint: usize, - /// Maximum time (seconds) allowed for a single RPC request (e.g. a - /// `ByteStream.Write` call) before it is cancelled. + /// Maximum time (seconds) allowed for a single RPC request before it + /// is cancelled with `DeadlineExceeded` (retryable). This applies to + /// every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, + /// `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment + /// of read streams through their first response chunk. The body of an + /// established read stream is not bounded by this deadline. /// /// A value of 0 (the default) disables the per-RPC timeout. Dead /// connections are still detected by the HTTP/2 and TCP keepalive - /// mechanisms configured on each endpoint. + /// mechanisms configured on each endpoint, and endpoints that cannot + /// be connected to at all fail fast with `Unavailable` rather than + /// queuing requests indefinitely. /// /// For large uploads (multi-GB), either leave this at 0 or set it /// large enough to accommodate the full transfer time. diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index 5742556338f..22116cb06f8 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -166,6 +166,7 @@ rust_test_suite( "tests/redis_store_test.rs", "tests/ref_store_test.rs", "tests/s3_store_test.rs", + "tests/shard_ring_dead_peer_test.rs", "tests/shard_store_test.rs", "tests/size_partitioning_store_test.rs", "tests/store_manager_test.rs", diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 451ad956cd2..3fa7efd31cd 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -334,6 +334,32 @@ impl GrpcStore { })) } + /// Bounds `fut` by the configured per-RPC deadline (`rpc_timeout_s`). + /// A timeout maps to `DeadlineExceeded`, which is retryable, so the + /// existing retry machinery can route around a peer that accepted the + /// request but will never answer (e.g. a restarting shard). + async fn with_rpc_timeout( + &self, + context: &'static str, + fut: impl Future> + Send, + ) -> Result { + // Boxed so this wrapper adds a pointer to the caller's future, not a + // second inline copy of `fut` — the RPC futures it wraps are large, + // and embedding them twice measurably deepens poll stacks. + let mut fut = Box::pin(fut); + if self.rpc_timeout.is_zero() { + return fut.as_mut().await; + } + match tokio::time::timeout(self.rpc_timeout, fut.as_mut()).await { + Ok(result) => result, + Err(_elapsed) => Err(make_err!( + Code::DeadlineExceeded, + "{context} timed out after {}s", + self.rpc_timeout.as_secs() + )), + } + } + async fn perform_request(&self, input: I, mut request: F) -> Result where F: FnMut(I) -> Fut + Send + Copy, @@ -345,7 +371,7 @@ impl GrpcStore { .retry(unfold(input, move |input| async move { let input_clone = input.clone(); Some(( - request(input_clone) + self.with_rpc_timeout("GrpcStore RPC", request(input_clone)) .await .map_or_else(RetryResult::Retry, RetryResult::Ok), input, @@ -1454,11 +1480,14 @@ impl GrpcStore { }; let mut stream = match self - .read_internal(ReadRequest { - resource_name, - read_offset: 0, - read_limit: 0, - }) + .with_rpc_timeout( + "GrpcStore::get_part_compressed read", + self.read_internal(ReadRequest { + resource_name, + read_offset: 0, + read_limit: 0, + }), + ) .await { Ok(stream) => stream, @@ -2043,7 +2072,7 @@ impl StoreDriver for GrpcStore { read_limit: local_state.read_limit, }; let mut stream = match self - .read_internal(request) + .with_rpc_timeout("GrpcStore::get_part read", self.read_internal(request)) .await .err_tip(|| "in GrpcStore::get_part()") { diff --git a/nativelink-store/src/shard_store.rs b/nativelink-store/src/shard_store.rs index 2f249448199..56f644e69e5 100644 --- a/nativelink-store/src/shard_store.rs +++ b/nativelink-store/src/shard_store.rs @@ -15,6 +15,8 @@ use core::hash::Hasher; use core::ops::BitXor; use core::pin::Pin; +use core::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use core::time::Duration; use std::hash::DefaultHasher; use std::sync::Arc; @@ -22,13 +24,39 @@ use async_trait::async_trait; use futures::future::try_join_all; use futures::stream::{FuturesUnordered, TryStreamExt}; use nativelink_config::stores::ShardSpec; -use nativelink_error::{Error, ResultExt, error_if}; +use nativelink_error::{Code, Error, ResultExt, error_if}; use nativelink_metric::MetricsComponent; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; use nativelink_util::store_trait::{ RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; +use tokio::time::Instant; + +/// Number of consecutive connectivity failures after which a shard is +/// considered suspect and skipped as the first read target. +const SUSPECT_FAILURE_THRESHOLD: u32 = 3; + +/// How long a suspect shard is skipped on the read path before it is +/// probed again. +const SUSPECT_COOLDOWN: Duration = Duration::from_secs(5); + +/// Errors that indicate the shard itself is unreachable (as opposed to a +/// definitive answer like `NotFound`), making read failover safe. +const fn is_connectivity_error(err: &Error) -> bool { + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) +} + +/// Read-path health state for one shard. Writes never consult this: a +/// blob's home shard is fixed by the ring, so writes always target it. +#[derive(Debug, Default)] +struct ShardHealth { + /// Consecutive connectivity failures observed on this shard. + consecutive_failures: AtomicU32, + /// Nanoseconds (since `ShardStore` creation) of the most recent + /// connectivity failure. Zero means "never failed". + last_failure_nanos: AtomicU64, +} #[derive(Debug, MetricsComponent)] struct StoreAndWeight { @@ -36,6 +64,7 @@ struct StoreAndWeight { weight: u32, #[metric(help = "The underlying store")] store: Store, + health: ShardHealth, } #[derive(Debug, MetricsComponent)] @@ -47,6 +76,8 @@ pub struct ShardStore { help = "The weights and stores that are used to determine which store to use" )] weights_and_stores: Vec, + /// Anchor for `ShardHealth::last_failure_nanos` timestamps. + created_at: Instant, } impl ShardStore { @@ -85,8 +116,107 @@ impl ShardStore { weights_and_stores: weights .into_iter() .zip(stores) - .map(|(weight, store)| StoreAndWeight { weight, store }) + .map(|(weight, store)| StoreAndWeight { + weight, + store, + health: ShardHealth::default(), + }) .collect(), + created_at: Instant::now(), + })) + } + + fn now_nanos(&self) -> u64 { + u64::try_from(self.created_at.elapsed().as_nanos()).unwrap_or(u64::MAX) + } + + /// A shard is suspect once it has accumulated enough consecutive + /// connectivity failures and its cooldown has not yet elapsed. After + /// the cooldown it is probed again (treated as healthy). + fn is_suspect(&self, store_idx: usize) -> bool { + let health = &self.weights_and_stores[store_idx].health; + if health.consecutive_failures.load(Ordering::Relaxed) < SUSPECT_FAILURE_THRESHOLD { + return false; + } + let last_failure_nanos = health.last_failure_nanos.load(Ordering::Relaxed); + let cooldown_nanos = u64::try_from(SUSPECT_COOLDOWN.as_nanos()).unwrap_or(u64::MAX); + self.now_nanos().saturating_sub(last_failure_nanos) < cooldown_nanos + } + + fn note_success(&self, store_idx: usize) { + self.weights_and_stores[store_idx] + .health + .consecutive_failures + .store(0, Ordering::Relaxed); + } + + fn note_connectivity_failure(&self, store_idx: usize) { + let health = &self.weights_and_stores[store_idx].health; + health.consecutive_failures.fetch_add(1, Ordering::Relaxed); + health + .last_failure_nanos + .store(self.now_nanos(), Ordering::Relaxed); + } + + /// Runs `has_with_results` against the ring starting at `primary_idx`, + /// failing over in ring order on connectivity errors (each shard is + /// tried at most once). Suspect shards are skipped unless every shard + /// is suspect, in which case the primary is probed directly. + async fn has_on_ring( + &self, + primary_idx: usize, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let num_stores = self.weights_and_stores.len(); + let mut last_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + results.fill(None); + last_err = Some(err); + } + Err(err) => { + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(primary_idx); + return Ok(()); + } + Err(err) => { + if is_connectivity_error(&err) { + self.note_connectivity_failure(primary_idx); + } + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + Err(last_err.unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::has_with_results()".to_string(), + ) })) } @@ -167,11 +297,7 @@ impl StoreDriver for ShardStore { if keys.len() == 1 { // Hot path: It is very common to lookup only one key. let store_idx = self.get_store_index(&keys[0]); - let store = &self.weights_and_stores[store_idx].store; - return store - .has_with_results(keys, results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}}"); + return self.has_on_ring(store_idx, keys, results).await; } let mut keys_for_store: Vec<(KeyIdxVec, KeyVec)> = self .weights_and_stores @@ -192,12 +318,9 @@ impl StoreDriver for ShardStore { .into_iter() .enumerate() .map(|(store_idx, (key_idxs, keys))| async move { - let store = &self.weights_and_stores[store_idx].store; let mut inner_results = vec![None; keys.len()]; - store - .has_with_results(&keys, &mut inner_results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}")?; + self.has_on_ring(store_idx, &keys, &mut inner_results) + .await?; Result::<_, Error>::Ok((key_idxs, inner_results)) }) .collect(); @@ -231,11 +354,66 @@ impl StoreDriver for ShardStore { offset: u64, length: Option, ) -> Result<(), Error> { - let store = self.get_store(&key); - store - .get_part(key, writer, offset, length) - .await - .err_tip(|| "In ShardStore::get_part()") + let primary_idx = self.get_store_index(&key); + let num_stores = self.weights_and_stores.len(); + let bytes_written_start = writer.get_bytes_written(); + let mut last_connectivity_err = None; + let mut last_other_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store + .get_part(key.borrow(), &mut *writer, offset, length) + .await + { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + // Failover is only safe if nothing reached the client. + if writer.get_bytes_written() != bytes_written_start { + return Err(err).err_tip(|| "In ShardStore::get_part()"); + } + last_connectivity_err = Some(err); + } + Err(err) if err.code == Code::NotFound && store_idx != primary_idx => { + // A failover shard legitimately may not have the blob; + // keep looking, but never let it mask a connectivity + // error from the shard that owns the key. + self.note_success(store_idx); + last_other_err = Some(err); + } + Err(err) => return Err(err).err_tip(|| "In ShardStore::get_part()"), + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + let result = store.get_part(key, writer, offset, length).await; + match &result { + Ok(()) => self.note_success(primary_idx), + Err(err) if is_connectivity_error(err) => { + self.note_connectivity_failure(primary_idx); + } + Err(_) => {} + } + return result.err_tip(|| "In ShardStore::get_part()"); + } + Err(last_connectivity_err.or(last_other_err).unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::get_part()".to_string(), + ) + })) + .err_tip(|| "In ShardStore::get_part()") } fn inner_store(&self, key: Option) -> &'_ dyn StoreDriver { diff --git a/nativelink-store/tests/shard_ring_dead_peer_test.rs b/nativelink-store/tests/shard_ring_dead_peer_test.rs new file mode 100644 index 00000000000..db790cd873a --- /dev/null +++ b/nativelink-store/tests/shard_ring_dead_peer_test.rs @@ -0,0 +1,557 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Lab reproduction of the production "shard-ring wedge": a `ShardStore` +//! ring over `GrpcStore` backends (peer CAS shards) where one peer is dead +//! (killed or never started). Before the fix, every CAS RPC that touched +//! the dead shard hung forever: the peer's `ConnectionManager` retried the +//! connect in the background while `connection()` requests queued +//! unbounded, so no error ever reached the retry machinery. The +//! operational workaround was `systemctl restart nativelink` on the +//! frontend. +//! +//! These tests must fail (by generous outer timeout) against the +//! pre-fix code and pass after dead peers yield fast `Unavailable`. + +use core::pin::Pin; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; + +use async_lock::Mutex; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use nativelink_config::stores::{ + GrpcEndpoint, GrpcSpec, MemorySpec, Retry, ShardConfig, ShardSpec, StoreSpec, StoreType, +}; +use nativelink_error::{Code, Error}; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{ + ContentAddressableStorage, ContentAddressableStorageServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::{ + BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, + BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse, GetTreeRequest, + GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse, SplitBlobRequest, SplitBlobResponse, +}; +use nativelink_proto::google::bytestream::byte_stream_server::{ByteStream, ByteStreamServer}; +use nativelink_proto::google::bytestream::{ + QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest, + WriteResponse, +}; +use nativelink_store::grpc_store::GrpcStore; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::shard_store::ShardStore; +use nativelink_util::background_spawn; +use nativelink_util::common::DigestInfo; +use nativelink_util::store_trait::{Store, StoreLike}; +use tokio::time::timeout; +use tonic::transport::Server; +use tonic::transport::server::TcpIncoming; +use tonic::{Request, Response, Status, Streaming}; + +/// Generous bound: far longer than any configured timeout/retry budget in +/// this test, far shorter than "forever". If an operation does not settle +/// within this, it is the production hang. +const GENEROUS_TIMEOUT: Duration = Duration::from_secs(20); + +/// A minimal in-process CAS shard: blobs live in a shared map, +/// `FindMissingBlobs`, `ByteStream.Read` and `ByteStream.Write` behave like +/// a real peer. +#[derive(Debug, Clone)] +struct FakeShardServer { + blobs: Arc>>, +} + +type GetTreeStream = Pin> + Send + 'static>>; + +#[tonic::async_trait] +impl ContentAddressableStorage for FakeShardServer { + type GetTreeStream = GetTreeStream; + + async fn find_missing_blobs( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let blobs = self.blobs.lock().await; + let missing_blob_digests = request + .blob_digests + .into_iter() + .filter(|digest| !blobs.contains_key(&digest.hash)) + .collect(); + Ok(Response::new(FindMissingBlobsResponse { + missing_blob_digests, + })) + } + + #[allow(clippy::unimplemented)] + async fn batch_update_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn batch_read_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn get_tree( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn split_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn splice_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +type ReadStream = Pin> + Send + 'static>>; + +/// Extracts the blob hash from a `ByteStream` resource name of the form +/// `[{instance}/]uploads/{uuid}/blobs/{fn}/{hash}/{size}` or +/// `[{instance}/]blobs/{fn}/{hash}/{size}`. +fn hash_from_resource_name(resource_name: &str) -> String { + let mut components = resource_name.rsplit('/'); + let _size = components.next(); + components.next().unwrap_or_default().to_string() +} + +#[tonic::async_trait] +impl ByteStream for FakeShardServer { + type ReadStream = ReadStream; + + async fn read( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let hash = hash_from_resource_name(&request.resource_name); + let data = self + .blobs + .lock() + .await + .get(&hash) + .cloned() + .ok_or_else(|| Status::not_found(format!("Blob {hash} not seeded")))?; + let offset = usize::try_from(request.read_offset) + .map_err(|_| Status::invalid_argument("Bad read_offset"))?; + if offset > data.len() { + return Err(Status::out_of_range("read_offset past end of blob")); + } + let stream: Self::ReadStream = Box::pin(futures::stream::iter(vec![Ok(ReadResponse { + data: data.slice(offset..), + })])); + Ok(Response::new(stream)) + } + + async fn write( + &self, + grpc_request: Request>, + ) -> Result, Status> { + let mut stream = grpc_request.into_inner(); + let mut hash = String::new(); + let mut data = Vec::new(); + while let Some(request) = stream.next().await { + let request = request?; + if !request.resource_name.is_empty() { + hash = hash_from_resource_name(&request.resource_name); + } + data.extend_from_slice(&request.data); + if request.finish_write { + break; + } + } + let committed_size = + i64::try_from(data.len()).map_err(|_| Status::invalid_argument("Blob too large"))?; + self.blobs.lock().await.insert(hash, Bytes::from(data)); + Ok(Response::new(WriteResponse { committed_size })) + } + + #[allow(clippy::unimplemented)] + async fn query_write_status( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +/// Spawns a live in-process shard peer on `port` (0 for ephemeral) and +/// returns its bound port and blob map. +fn spawn_live_peer_on(port: u16) -> (u16, Arc>>) { + let server = FakeShardServer { + blobs: Arc::new(Mutex::new(HashMap::new())), + }; + let blobs = server.blobs.clone(); + let listener = TcpIncoming::bind(format!("127.0.0.1:{port}").parse().unwrap()).unwrap(); + let port = listener.local_addr().unwrap().port(); + let cas_service = ContentAddressableStorageServer::new(server.clone()); + let bytestream_service = ByteStreamServer::new(server); + background_spawn!("shard_ring_dead_peer_test_server", async move { + Server::builder() + .add_service(cas_service) + .add_service(bytestream_service) + .serve_with_incoming(listener) + .await + .unwrap(); + }); + (port, blobs) +} + +/// Reserves a port with nothing listening on it: the dead peer (killed or +/// never started). Connect attempts fail with `ECONNREFUSED`. +fn dead_peer_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() +} + +fn grpc_spec_for_port(port: u16) -> GrpcSpec { + GrpcSpec { + instance_name: String::new(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 1, + tcp_keepalive_s: 0, + http2_keepalive_interval_s: 0, + http2_keepalive_timeout_s: 0, + }], + store_type: StoreType::Cas, + retry: Retry { + max_retries: 1, + delay: 0.05, + jitter: 0.0, + retry_on_errors: None, + }, + max_concurrent_requests: 0, + connections_per_endpoint: 0, + rpc_timeout_s: 5, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_remote_cache_compression: None, + } +} + +/// Mirrors `ShardStore`'s routing for a two-shard equal-weight ring: the +/// digest's xor-folded key lands in shard 0 iff it is `<= u32::MAX / 2`. +fn shard_index_for(digest: &DigestInfo) -> usize { + let packed_hash = digest.packed_hash(); + let size_bytes = digest.size_bytes().to_le_bytes(); + let mut key: u32 = 0; + for chunk in 0..8 { + key ^= u32::from_le_bytes(packed_hash[chunk * 4..chunk * 4 + 4].try_into().unwrap()); + } + key ^= u32::from_le_bytes(size_bytes[0..4].try_into().unwrap()); + key ^= u32::from_le_bytes(size_bytes[4..8].try_into().unwrap()); + usize::from(key > u32::MAX / 2) +} + +/// Finds a digest (with 64-byte content) that routes to `shard`. +fn digest_for_shard(shard: usize) -> (DigestInfo, Bytes) { + for i in 0..1024usize { + let hash = format!("{i:064x}"); + let content = Bytes::from(format!("{i:0>64}")); + let digest = DigestInfo::try_new(&hash, content.len()).unwrap(); + if shard_index_for(&digest) == shard { + return (digest, content); + } + } + unreachable!("no digest found for shard {shard} in 1024 candidates"); +} + +struct Ring { + shard_store: Arc, + live_shard: usize, + live_blobs: Arc>>, + dead_port: u16, +} + +/// Builds a two-shard equal-weight ring: shard 0 is a live in-process +/// peer, shard 1 is dead (nothing listening). +async fn make_ring() -> Result { + let (live_port, live_blobs) = spawn_live_peer_on(0); + let dead_port = dead_peer_port(); + + let live_store = GrpcStore::new(&grpc_spec_for_port(live_port))?; + let dead_store = GrpcStore::new(&grpc_spec_for_port(dead_port))?; + + // The spec's `store` fields are only used for length/weights by + // `ShardStore::new`; memory placeholders keep the config minimal. + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + vec![Store::new(live_store), Store::new(dead_store)], + )?; + Ok(Ring { + shard_store, + live_shard: 0, + live_blobs, + dead_port, + }) +} + +// Guard for `shard_index_for`: it re-implements `ShardStore`'s xor-fold +// routing, which would silently diverge if upstream ever changes the ring +// hash — the dead-peer tests above would then route probes to the wrong +// shard and test nothing. This builds a REAL `ShardStore` over two +// instrumented `MemoryStore` backends, writes each probe digest through +// it, and asserts the blob landed on the shard `shard_index_for` +// predicted. If the ring hash changes upstream, this fails loudly. +#[nativelink_test] +async fn shard_index_for_matches_real_shard_store_routing() -> Result<(), Error> { + let backends: Vec> = (0..2) + .map(|_| MemoryStore::new(&MemorySpec::default())) + .collect(); + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + backends + .iter() + .map(|backend| Store::new(backend.clone())) + .collect(), + )?; + + for shard in 0..2usize { + let (digest, content) = digest_for_shard(shard); + shard_store.update_oneshot(digest, content).await?; + for (backend_index, backend) in backends.iter().enumerate() { + let present = backend.has(digest).await?.is_some(); + assert_eq!( + present, + backend_index == shard, + "probe digest {digest} predicted for shard {shard} but presence on \ + backend {backend_index} was {present}: `shard_index_for` no longer \ + matches ShardStore's ring hash", + ); + } + } + Ok(()) +} + +// A read routed to the dead shard must settle with an error (which the +// retry machinery upstream can handle) instead of hanging forever. +#[nativelink_test] +async fn read_routed_to_dead_shard_errors_instead_of_hanging() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, _) = digest_for_shard(1 - ring.live_shard); + + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung past the generous timeout: the shard-ring wedge"); + let err = result.expect_err("read from a dead shard cannot succeed"); + assert!( + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded), + "expected Unavailable/DeadlineExceeded from the dead shard, got: {err:?}", + ); + Ok(()) +} + +// The production symptom: a multi-key existence check (FindMissingBlobs) +// fans out to every shard, so ONE dead peer wedged EVERY CAS RPC through +// the frontend. With read failover the dead shard's bucket is retried on +// the live shard, so the whole call succeeds instead of erroring. +#[nativelink_test] +async fn has_with_dead_shard_fails_over_instead_of_erroring() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, live_content) = digest_for_shard(ring.live_shard); + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + // Both blobs are present on the live shard (as a replica would be). + { + let mut blobs = ring.live_blobs.lock().await; + blobs.insert(live_digest.packed_hash().to_string(), live_content.clone()); + blobs.insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + } + + let results = timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .has_many(&[live_digest.into(), dead_digest.into()]), + ) + .await + .expect("multi-key has() with a dead shard hung past the generous timeout")?; + assert_eq!( + results, + vec![ + Some(live_content.len() as u64), + Some(dead_content.len() as u64) + ], + "both keys must resolve: the dead shard's bucket via failover", + ); + Ok(()) +} + +// A single-key has() routed to the dead shard must fail over to the live +// shard that holds the blob instead of surfacing Unavailable. +#[nativelink_test] +async fn has_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("single-key has() routed to the dead shard hung")?; + assert_eq!( + result, + Some(dead_content.len() as u64), + "has() must find the blob on the live shard via failover", + ); + Ok(()) +} + +// A read routed to the dead shard must fail over to the live shard that +// holds a replica of the blob and return its content. +#[nativelink_test] +async fn get_part_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung")?; + assert_eq!( + data, dead_content, + "failover read must return the live shard's replica", + ); + Ok(()) +} + +// Health recovery: after enough consecutive failures the dead shard is +// marked suspect and skipped; once the peer comes back and the cooldown +// elapses it is probed again and serves its own keys directly. +#[nativelink_test] +async fn dead_shard_recovers_after_cooldown() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + + // Trip the suspect threshold: each read fails over to the live shard + // (which lacks the blob), so the call errors, but each attempt counts + // a connectivity failure against the dead shard. + for _ in 0..3 { + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read with dead shard hung"); + assert!(result.is_err(), "blob exists nowhere yet; read must fail"); + } + + // Revive the dead peer on its original port, now holding the blob. + let (_, revived_blobs) = spawn_live_peer_on(ring.dead_port); + revived_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + // Wait out the suspect cooldown so the shard is probed again. + tokio::time::sleep(Duration::from_secs(6)).await; + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read after recovery hung")?; + assert_eq!(data, dead_content, "revived shard must serve its own key"); + + // The success must have reset the failure counter: an immediate + // follow-up (well inside any cooldown) is also served directly. + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("has() after recovery hung")?; + assert_eq!(result, Some(dead_content.len() as u64)); + Ok(()) +} + +// The live shard must keep serving while its ring-mate is dead: routing +// around the dead peer stays possible for keys that don't touch it. +#[nativelink_test] +async fn live_shard_still_serves_while_peer_dead() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, content) = digest_for_shard(ring.live_shard); + + timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .update_oneshot(live_digest, content.clone()), + ) + .await + .expect("write to the live shard hung")?; + assert!( + ring.live_blobs + .lock() + .await + .contains_key(&live_digest.packed_hash().to_string()), + "blob did not land on the live peer", + ); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(live_digest, 0, None), + ) + .await + .expect("read from the live shard hung")?; + assert_eq!(data, content, "live shard returned wrong content"); + Ok(()) +} diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index bc008ef5ea0..fa453f15375 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -106,6 +106,7 @@ rust_test_suite( "tests/cache_size_metrics_test.rs", "tests/channel_body_for_tests_test.rs", "tests/common_test.rs", + "tests/connection_manager_test.rs", "tests/evicting_map_test.rs", "tests/fastcdc_test.rs", "tests/fs_test.rs", diff --git a/nativelink-util/src/connection_manager.rs b/nativelink-util/src/connection_manager.rs index 753e7cb956a..145f84a790a 100644 --- a/nativelink-util/src/connection_manager.rs +++ b/nativelink-util/src/connection_manager.rs @@ -35,7 +35,7 @@ use crate::retry::{self, Retrier, RetryResult}; #[derive(Debug)] pub struct ConnectionManager { // The channel to request connections from the worker. - worker_tx: mpsc::Sender<(String, oneshot::Sender)>, + worker_tx: mpsc::Sender<(String, oneshot::Sender>)>, } /// The index into `ConnectionManagerWorker::endpoints`. @@ -94,6 +94,21 @@ struct ConnectionManagerWorker { /// The endpoints to establish Channels and the identifier of the last /// connection attempt to that endpoint. endpoints: Vec<(ConnectionIndex, Endpoint)>, + /// Per-endpoint connect health: `true` once a connect retry cycle to + /// that endpoint has failed, cleared when a connect succeeds. When + /// every endpoint is marked failed and no channel is available, the + /// endpoint set is unreachable — requests fail fast with + /// `Unavailable` instead of queuing forever while the background + /// reconnect loop runs (the "dead peer wedges every RPC" hang). + endpoint_connect_failed: Vec, + /// Number of currently established channels, whether sitting in + /// `available_channels` or checked out to in-flight `Connection`s. + /// `available_channels` alone cannot distinguish "everything is down" + /// from "everything is busy": borrowed channels are not in that queue, + /// so a busy-but-healthy pool plus one transient connect failure must + /// not be declared dead. Incremented on connect success, decremented + /// when a channel is torn down on transport error. + live_channels: usize, /// The channel used to communicate between a Connection and the worker. connection_tx: mpsc::UnboundedSender, /// Gates the maximum number of in-flight `Connection` objects. @@ -112,7 +127,7 @@ struct ConnectionManagerWorker { /// Connected channels that are available for use. available_channels: VecDeque, /// Requests for a Channel when available - (reason, request) - waiting_connections: VecDeque<(String, oneshot::Sender)>, + waiting_connections: VecDeque<(String, oneshot::Sender>)>, /// The retry configuration for connecting to an Endpoint, on failure will /// restart the retrier after a 1 second delay. retrier: Retrier, @@ -140,10 +155,11 @@ impl ConnectionManager { // which defeats the object since there would be no backpressure // applied. Therefore it makes sense for this to be unbounded. let (connection_tx, connection_rx) = mpsc::unbounded_channel(); - let endpoints = endpoints + let endpoints: Vec<(ConnectionIndex, Endpoint)> = endpoints .into_iter() .map(|endpoint| (0, endpoint)) .collect(); + let endpoint_connect_failed = vec![false; endpoints.len()]; // Zero means unlimited, which becomes a semaphore holding // Semaphore::MAX_PERMITS. That is a sentinel, not a free-slot count, @@ -159,6 +175,8 @@ impl ConnectionManager { } let worker = ConnectionManagerWorker { endpoints, + endpoint_connect_failed, + live_channels: 0, available_connections: Arc::new(Semaphore::new(max_concurrent_requests)), bounded_connections, connection_tx, @@ -190,7 +208,7 @@ impl ConnectionManager { .await .map_err(|err| make_err!(Code::Unavailable, "Requesting a new connection: {err:?}"))?; rx.await - .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}")) + .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}"))? } } @@ -198,7 +216,7 @@ impl ConnectionManagerWorker { async fn service_requests( mut self, connections_per_endpoint: usize, - mut worker_rx: mpsc::Receiver<(String, oneshot::Sender)>, + mut worker_rx: mpsc::Receiver<(String, oneshot::Sender>)>, mut connection_rx: mpsc::UnboundedReceiver, ) { // Make the initial set of connections, connection failures will be @@ -252,6 +270,13 @@ impl ConnectionManagerWorker { fn handle_connected(&mut self, connection_result: IndexedChannel) { match connection_result { Ok(established_channel) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(established_channel.identifier.endpoint_index) + { + *failed = false; + } + self.live_channels += 1; self.available_channels.push_back(established_channel); self.maybe_available_connection(); } @@ -259,11 +284,52 @@ impl ConnectionManagerWorker { // beginning of the retry period. Never want to be in a // situation where we give up on an Endpoint forever. Err((identifier, _)) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(identifier.endpoint_index) + { + *failed = true; + } + // If this failure means the whole endpoint set is now + // unreachable, waiting requests would hang until a peer + // comes back — fail them fast so the caller's retry/error + // machinery can handle it. + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, Some(identifier.connection_index)); } } } + /// True when no channel is live (available or checked out) and the + /// last connect attempt to every endpoint failed: nothing can service + /// a request until a background reconnect succeeds. + fn all_endpoints_down(&self) -> bool { + self.live_channels == 0 + && !self.endpoint_connect_failed.is_empty() + && self.endpoint_connect_failed.iter().all(|failed| *failed) + } + + fn unavailable_error(&self) -> Error { + make_err!( + Code::Unavailable, + "All {} gRPC endpoint(s) are unreachable (connect attempts failed); failing fast instead of queuing", + self.endpoints.len() + ) + } + + fn fail_requests_if_all_endpoints_down(&mut self) { + if !self.all_endpoints_down() { + return; + } + while let Some((reason, tx)) = self.waiting_connections.pop_front() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing queued connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); + } + } + fn connect_endpoint(&mut self, endpoint_index: usize, connection_index: Option) { let Some((current_connection_index, endpoint)) = self.endpoints.get_mut(endpoint_index) else { @@ -333,7 +399,7 @@ impl ConnectionManagerWorker { } // This must never be made async otherwise the select may cancel it. - fn handle_worker(&mut self, reason: String, tx: oneshot::Sender) { + fn handle_worker(&mut self, reason: String, tx: oneshot::Sender>) { let maybe_permit = self.available_connections.clone().try_acquire_owned().ok(); if let Some(permit) = maybe_permit && let Some(channel) = self.available_channels.pop_front() @@ -341,6 +407,12 @@ impl ConnectionManagerWorker { debug!(reason, "ConnectionManager: request running"); record_connection_acquired("grpc", self.free_connection_slots(), false); self.provide_channel(channel, tx, permit); + } else if self.all_endpoints_down() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); } else { debug!( available_permits = self.available_connections.available_permits(), @@ -357,15 +429,15 @@ impl ConnectionManagerWorker { fn provide_channel( &self, channel: EstablishedChannel, - tx: oneshot::Sender, + tx: oneshot::Sender>, permit: OwnedSemaphorePermit, ) { - drop(tx.send(Connection { + drop(tx.send(Ok(Connection { tx: self.connection_tx.clone(), pending_channel: Some(channel.channel.clone()), channel, _permit: permit, - })); + }))); } fn maybe_available_connection(&mut self) { @@ -413,6 +485,17 @@ impl ConnectionManagerWorker { original_length != self.available_channels.len() }; if should_reconnect { + // The channel this identifier referred to is gone; + // saturate in case of duplicate error reports for the + // same pending channel. + self.live_channels = self.live_channels.saturating_sub(1); + // If this was the last live channel and every endpoint's + // last connect cycle had already failed, the pool just + // became unreachable — flush queued waiters now instead + // of leaving them parked until the reconnect scheduled + // below fails a full connect cycle later (new requests + // already fail fast in `handle_worker` in this state). + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, None); } } diff --git a/nativelink-util/tests/connection_manager_test.rs b/nativelink-util/tests/connection_manager_test.rs index 286ea578311..f293b0201dd 100644 --- a/nativelink-util/tests/connection_manager_test.rs +++ b/nativelink-util/tests/connection_manager_test.rs @@ -161,6 +161,263 @@ async fn aborted_caller_future_does_not_leak_permits() -> Result<(), Error> { Ok(()) } +/// The production "shard-ring wedge": an endpoint that is dead (nothing +/// listening — killed or never started) must yield a fast `Unavailable` +/// error from `connection()` instead of queuing the request forever while +/// the worker retries the connect in the background. Before the fix this +/// test failed: `cm.connection()` never resolved and the outer timeout +/// elapsed. +#[nativelink_test] +async fn dead_endpoint_fails_fast_instead_of_queuing_forever() -> Result<(), Error> { + // Bind a port and immediately drop the listener so nothing is + // listening there — connects get ECONNREFUSED, the "peer killed or + // never started" case observed in production. + let dead_port = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() + }; + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{dead_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + + let cm = ConnectionManager::new(vec![endpoint], 1, 2, Retry::default(), no_jitter()); + + let result = timeout(Duration::from_secs(10), cm.connection("dead-peer".into())) + .await + .expect("connection() to a dead endpoint hung past 10s instead of erroring"); + assert_eq!( + result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected fast Unavailable from a dead endpoint", + ); + Ok(()) +} + +/// The false-negative guard on the fast-fail gating: a busy-but-healthy +/// pool must NOT be declared dead. One endpoint, two connection slots: a +/// single-accept proxy in front of a real gRPC server lets exactly one +/// connect succeed, then drops its listener so the second slot's connect +/// is refused (`ECONNREFUSED`), marking the endpoint's last connect cycle +/// as failed. We hold the one live connection (so `available_channels` is +/// empty) and issue a new `connection()` request: with gating keyed only +/// on `available_channels` + connect health this errored `Unavailable` +/// immediately; with live-channel tracking it queues, and resolves as +/// soon as the held connection is released. +#[nativelink_test] +async fn busy_pool_with_transient_connect_failure_does_not_fail_fast() -> Result<(), Error> { + // Real gRPC server the proxied connection terminates at. + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + // Proxy that accepts exactly ONE connection, pipes it to the backend, + // then drops the listener — every later connect is refused. + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + background_spawn!("single_accept_proxy", async move { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + // Drop the listener so the second connect slot gets ECONNREFUSED. + drop(proxy_listener); + let mut outbound = tokio::net::TcpStream::connect(backend_addr).await.unwrap(); + drop(tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await); + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + Retry::default(), + no_jitter(), + )); + + // Take (and hold) the one connection that could be established. This + // may need a moment while the connect race settles; the winner is the + // proxied slot, the loser marks the endpoint's connect cycle failed. + let held = timeout(Duration::from_secs(10), cm.connection("held-live".into())) + .await + .expect("acquiring the single live connection hung")?; + + // Give the losing connect slot ample time to fail its cycle (default + // retry is a single ~1s-connect-timeout attempt) and set the + // per-endpoint failed flag; pre-fix this is what armed the trap. + tokio::time::sleep(Duration::from_secs(3)).await; + + // A new request must queue, not fail fast: the pool is busy, not dead. + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-while-busy".into()).await }); + tokio::time::sleep(Duration::from_millis(500)).await; + assert_eq!( + queued.is_finished(), + false, + "connection() failed fast while a live connection was checked out — \ + busy-but-healthy pool wrongly declared dead", + ); + + // Releasing the held connection must let the queued request through. + drop(held); + let conn = timeout(Duration::from_secs(5), queued) + .await + .expect("queued request did not resolve after the live connection was released") + .unwrap()?; + drop(conn); + Ok(()) +} + +/// The teardown-to-zero window: when the LAST live channel is torn down +/// by an in-use transport error while every endpoint's last connect cycle +/// had already failed, queued waiters must be flushed with `Unavailable` +/// at that instant — not parked until the freshly scheduled reconnect +/// fails a whole connect cycle later (`handle_worker` already fails new +/// requests fast in that state; queued ones must not be second-class). +/// +/// Setup: one endpoint, two connection slots, both piped through a proxy +/// to a real gRPC server; `max_concurrent_requests` of 2 with both +/// permits held so returned channel clones can never be handed to the +/// queued request. The retrier is configured so a connect cycle to the +/// (killed) endpoint takes ~3s: after severing both connections, an RPC +/// on the second held connection tears its channel down and its +/// reconnect cycle fails, arming the endpoint's connect-failed flag +/// while the first channel is still live. A request is then queued, and +/// an RPC on the first held connection tears the last live channel down +/// to zero. Post-fix the queued request errors `Unavailable` right +/// there; pre-fix nothing flushed it until the next connect cycle failed +/// ~3s later, so the 1.5s deadline below distinguishes the two. +#[nativelink_test] +async fn last_live_channel_teardown_flushes_queued_requests() -> Result<(), Error> { + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + let (kill_tx, kill_rx) = tokio::sync::watch::channel(false); + background_spawn!("teardown_test_proxy", async move { + // Pipe exactly the two initial connections (one per slot). On the + // kill signal, drop the listener first (so reconnects are refused, + // failing their connect cycles) and then sever both pipes. + for _ in 0..2 { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + let mut rx = kill_rx.clone(); + let addr = backend_addr.clone(); + background_spawn!("teardown_test_pipe", async move { + let mut outbound = tokio::net::TcpStream::connect(addr).await.unwrap(); + tokio::select! { + _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound) => {} + _ = rx.changed() => {} + } + }); + } + let mut rx = kill_rx.clone(); + drop(rx.changed().await); + drop(proxy_listener); + // Pipes observe the same signal and sever their connections. + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + // delay 1.5 with one retry makes a failing connect cycle take one + // ~3s backoff sleep between two instant ECONNREFUSED attempts: long + // enough that "flushed at teardown" and "flushed by the next connect + // failure" are cleanly distinguishable below. + let retry = Retry { + max_retries: 1, + delay: 1.5, + ..Default::default() + }; + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + retry, + no_jitter(), + )); + + // Hold both permits (and both live channels). + let held1 = timeout(Duration::from_secs(10), cm.connection("held-1".into())) + .await + .expect("acquiring the first piped connection hung")?; + let held2 = timeout(Duration::from_secs(10), cm.connection("held-2".into())) + .await + .expect("acquiring the second piped connection hung")?; + + // Kill the endpoint: refuse new connects and sever both connections. + tokio::time::sleep(Duration::from_millis(500)).await; + kill_tx.send(true).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + // Tear down the second channel with an in-use error; keep the client + // (and its permit) alive. Its reconnect cycle fails ~3s later, + // arming the endpoint's connect-failed flag while the first channel + // keeps live_channels at 1. + let mut client2 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held2); + let rpc2 = timeout( + Duration::from_secs(5), + client2.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-2".into(), + }), + ) + .await + .expect("RPC on the severed second connection hung"); + assert_eq!(rpc2.is_err(), true, "RPC on a severed connection succeeded"); + tokio::time::sleep(Duration::from_secs(4)).await; + + // Queue a request; it must wait (one channel still live, pool busy). + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-behind-held".into()).await }); + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + queued.is_finished(), + false, + "queued request resolved while a channel was live and permits were held", + ); + + // Tear down the LAST live channel with an in-use error. + let mut client1 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held1); + let rpc1 = timeout( + Duration::from_secs(5), + client1.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-1".into(), + }), + ) + .await + .expect("RPC on the severed first connection hung"); + assert_eq!(rpc1.is_err(), true, "RPC on a severed connection succeeded"); + + // The queued request must be flushed with Unavailable at the + // teardown itself — the next connect-cycle failure is ~3s away. + let queued_result = timeout(Duration::from_millis(1500), queued) + .await + .expect( + "queued request still parked after the last live channel was torn down \ + with all endpoints' connect cycles failed", + ) + .unwrap(); + assert_eq!( + queued_result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected Unavailable for the queued request after teardown-to-zero", + ); + drop(client1); + drop(client2); + Ok(()) +} + #[nativelink_test] async fn extra_request_above_max_blocks_until_a_release() -> Result<(), Error> { const MAX_CONCURRENT: usize = 2; diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..cdd2255366b 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (b0f3b0067c) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -856,7 +856,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `retry` | [Retry](#retry) | — | — | Retry configuration to use when a network request fails. | | `max_concurrent_requests` | integer (uint) | — | `0` | Limit the number of simultaneous upstream requests to this many. A value of zero is treated as unlimited. If the limit is reached the request is queued. | | `connections_per_endpoint` | integer (uint) | — | 1 | The number of connections to make to each specified endpoint to balance the load over multiple TCP connections. | -| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request (e.g. a `ByteStream.Write` call) before it is cancelled. | +| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request before it is cancelled with `DeadlineExceeded` (retryable). This applies to every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment of read streams through their first response chunk. The body of an established read stream is not bounded by this deadline. | | `use_legacy_resource_names` | boolean | — | false | Use legacy `ByteStream` resource name format, omitting the digest function component from the path. | | `headers` | map of string to string | — | — | Static headers to attach to every outgoing gRPC request sent to this store's upstream endpoints. Useful for fixed authentication tokens (e.g. `{"authorization": "Bearer "}`) and other static metadata. | | `forward_headers` | array of string | — | — | Header names to forward from the incoming client request to every outgoing upstream request. The header value is taken from the client request that triggered this store operation. Use this to pass through dynamic credentials such as JWT tokens sent by build clients. |