From 18b8be700c56de79967cc92284459f63cd66a613 Mon Sep 17 00:00:00 2001 From: Walter Gray Date: Wed, 29 Jul 2026 14:08:23 -0700 Subject: [PATCH 1/2] fix(store): make compressed download failures classify correctly Review follow-ups to the GrpcStore wire-compression work. `get_part_compressed` classified pump-stage failures with `is_retryable_code`, but pump errors come from our own writer and `buf_channel` reports a dropped receiver as `Internal`, which counts as retryable. A consumer that hung up mid-download therefore returned `Ok(Some(forwarded))` and made `get_part` issue a second, full-size identity Read streaming the remainder into a channel nobody was reading. Pump failures are now always terminal: they mean either the consumer went away or the decoder aborted, and the decoder's verdict is already reported by the arm above it since `try_join!` polls decode first. The fallback also treated a late failure from another stage as a reason to resume even when the download had already completed and been verified. The pump now records that it forwarded the decoder's EOF -- which the decoder only sends after the blob passes its size and digest checks -- and a late error after that point resolves as success instead of resuming into an already-closed writer. On a genuine resume, the identity request now asks for exactly the undelivered tail. It previously kept the caller's original `length`, which the entry guard only permits when it covers the whole blob, so the resumed request ran past the end of the blob and relied on the server clamping it. Also: only take the compressed upload path for real digest keys (a string key's `into_digest()` hashes the key itself, which the remote's mandatory digest verification would reject), and skip the background drain spawn when the encoder finished cleanly and the reader is already at EOF. Tests: two new cases in the service suite. The resume path had no coverage at all -- the fake ByteStream now honors `read_offset`/`read_limit` so the resumed range can be asserted byte-exactly, and grew a mid-frame abort mode plus a mixed-entropy payload helper, since `make_content` compresses far enough that multi-megabyte blobs fit in one wire chunk. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/grpc_wire_compression_test.rs | 186 ++++++++++++++++-- nativelink-store/src/grpc_store.rs | 86 ++++++-- 2 files changed, 243 insertions(+), 29 deletions(-) diff --git a/nativelink-service/tests/grpc_wire_compression_test.rs b/nativelink-service/tests/grpc_wire_compression_test.rs index 8ded2f18f..112b174d5 100644 --- a/nativelink-service/tests/grpc_wire_compression_test.rs +++ b/nativelink-service/tests/grpc_wire_compression_test.rs @@ -16,7 +16,7 @@ //! (`GrpcSpec.experimental_remote_cache_compression`) against a real //! compression-enabled `ByteStream`/CAS server. -use core::sync::atomic::{AtomicU64, Ordering}; +use core::sync::atomic::{AtomicI64, AtomicU64, Ordering}; use std::collections::HashMap; use std::sync::Arc; @@ -171,6 +171,30 @@ fn make_content(tag: u8, len: usize) -> (Bytes, DigestInfo) { (Bytes::from(data), digest) } +/// Deterministic content that compresses at roughly 4:1 — one pseudo-random +/// byte in every four, constants elsewhere. `make_content` is compressible +/// enough that even multi-megabyte blobs fit in a single 64KiB wire chunk once +/// encoded, which is useless for tests that must interrupt a compressed stream +/// part-way through. +fn make_mixed_content(seed: u64, len: usize) -> (Bytes, DigestInfo) { + let mut state = seed | 1; + let mut data = vec![0u8; len]; + for (i, b) in data.iter_mut().enumerate() { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *b = if i % 4 == 0 { + u8::try_from((state >> 33) & 0xff).expect("masked to a single byte") + } else { + 0x5a + }; + } + let mut hasher = Sha256::new(); + hasher.update(&data); + let digest = DigestInfo::try_new(&hex::encode(hasher.finalize()), len).unwrap(); + (Bytes::from(data), digest) +} + // A blob above the compression threshold uploads via compressed-blobs and // must land byte-identical (the server decodes and digest-verifies). #[nativelink_test] @@ -283,15 +307,21 @@ struct FakeByteStream { early_complete_write: bool, /// Read behavior: payload served for compressed-blobs reads. compressed_read_payload: Option, - /// Pause the compressed read stream after this many chunks (forces the - /// client feed to lag the decoder, exercising the cascade race). - pause_compressed_read_after_chunks: Option, /// Hold the compressed read stream open forever after this many chunks. /// This ensures a decoder error must cancel the feed rather than waiting /// for the upstream stream to finish. hold_compressed_read_after_chunks: Option, - /// Payload served for identity reads (fallback detector). + /// Fail the compressed read stream with UNAVAILABLE after this many + /// chunks, after letting the decoded prefix settle into the client's + /// writer. Exercises the retryable-interruption resume path. + abort_compressed_read_after_chunks: Option, + /// Payload served for identity reads (fallback detector). Served honoring + /// `read_offset`/`read_limit` so a resume can be checked byte-exactly. identity_read_payload: Option, + /// `read_offset`/`read_limit` of the last identity read. Tests that care + /// initialize these to -1 so "never read" is distinguishable from 0. + identity_read_offset: Arc, + identity_read_limit: Arc, } type ReadStreamT = @@ -305,8 +335,8 @@ impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for Fa &self, request: tonic::Request, ) -> Result, tonic::Status> { - let resource_name = request.into_inner().resource_name; - if resource_name.contains("compressed-blobs") { + let request = request.into_inner(); + if request.resource_name.contains("compressed-blobs") { let payload = self .compressed_read_payload .clone() @@ -315,8 +345,8 @@ impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for Fa .chunks(64 * 1024) .map(Bytes::copy_from_slice) .collect(); - let pause_after = self.pause_compressed_read_after_chunks; let hold_after = self.hold_compressed_read_after_chunks; + let abort_after = self.abort_compressed_read_after_chunks; let stream = futures::stream::unfold((chunks, 0usize), move |(chunks, index)| async move { if index >= chunks.len() { @@ -325,10 +355,17 @@ impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for Fa if Some(index) == hold_after { futures::future::pending::<()>().await; } - if Some(index) == pause_after { - // Hold the wire open long enough that the client - // decoder settles first. - tokio::time::sleep(core::time::Duration::from_millis(500)).await; + if Some(index) == abort_after { + // Let the already-sent prefix finish decoding and + // reach the client's writer, so the resume offset + // reflects real forwarded progress. + tokio::time::sleep(core::time::Duration::from_millis(250)).await; + return Some(( + Err(tonic::Status::unavailable( + "injected mid-stream read failure", + )), + (chunks, index), + )); } let data = chunks[index].clone(); Some((Ok(ReadResponse { data }), (chunks, index + 1))) @@ -337,10 +374,27 @@ impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for Fa return Ok(tonic::Response::new(boxed)); } self.identity_reads.fetch_add(1, Ordering::Relaxed); + self.identity_read_offset + .store(request.read_offset, Ordering::Relaxed); + self.identity_read_limit + .store(request.read_limit, Ordering::Relaxed); let payload = self .identity_read_payload .clone() .ok_or_else(|| tonic::Status::not_found("no identity payload configured"))?; + // Honor the requested range like a real server so a resumed read can + // be checked byte-exactly. `read_limit` of zero means "to the end". + let start = usize::try_from(request.read_offset) + .map_err(|_| tonic::Status::invalid_argument("negative read_offset"))? + .min(payload.len()); + let end = if request.read_limit == 0 { + payload.len() + } else { + let limit = usize::try_from(request.read_limit) + .map_err(|_| tonic::Status::invalid_argument("negative read_limit"))?; + start.saturating_add(limit).min(payload.len()) + }; + let payload = payload.slice(start..end); let chunks: Vec> = payload .chunks(64 * 1024) .map(|c| { @@ -550,3 +604,111 @@ async fn compressed_upload_returns_after_early_completion() producer_result?; Ok(()) } + +// W4: the resume path. A retryable mid-stream failure of a compressed +// download must keep the already-decoded prefix and ask the identity path for +// exactly the undelivered tail — not re-fetch from zero, and not request a +// range that runs past the end of the blob. +#[nativelink_test] +async fn interrupted_compressed_download_resumes_via_identity() +-> Result<(), Box> { + use nativelink_proto::build::bazel::remote::execution::v2::compressor; + + let (content, digest) = make_mixed_content(0x44, 8 * 1024 * 1024); + let frame = + nativelink_util::wire_compression::compress(content.clone(), compressor::Value::Zstd)?; + // The failure has to land mid-frame for this to test anything. + let frame_chunks = frame.len().div_ceil(64 * 1024); + assert!( + frame_chunks >= 4, + "test payload must span several wire chunks, got {frame_chunks}" + ); + let identity_reads = Arc::new(AtomicU64::new(0)); + let identity_read_offset = Arc::new(AtomicI64::new(-1)); + let identity_read_limit = Arc::new(AtomicI64::new(-1)); + let port = start_fake_bytestream_server(FakeByteStream { + identity_reads: identity_reads.clone(), + identity_read_offset: identity_read_offset.clone(), + identity_read_limit: identity_read_limit.clone(), + compressed_read_payload: Some(frame), + // Two 64KiB compressed chunks decode to several full zstd blocks, so + // the prefix forwarded before the failure is always non-empty. + abort_compressed_read_after_chunks: Some(2), + identity_read_payload: Some(content.clone()), + ..Default::default() + }) + .await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let fetched = tokio::time::timeout( + core::time::Duration::from_secs(10), + grpc_store.get_part_unchunked(digest, 0, None), + ) + .await + .expect("interrupted compressed read must not hang")?; + + assert_eq!( + fetched, content, + "resumed read must reproduce the blob byte-for-byte" + ); + assert_eq!( + identity_reads.load(Ordering::Relaxed), + 1, + "the interrupted compressed read must fall back exactly once" + ); + let resume_offset = identity_read_offset.load(Ordering::Relaxed); + let resume_limit = identity_read_limit.load(Ordering::Relaxed); + assert!( + resume_offset > 0, + "the decoded prefix must be kept, not re-fetched: read_offset was {resume_offset}" + ); + assert_eq!( + resume_offset + resume_limit, + i64::try_from(content.len())?, + "the resume range must cover exactly the undelivered tail", + ); + Ok(()) +} + +// W5: a download whose consumer goes away must fail, not trigger a pointless +// identity re-read of the remainder into a channel nobody is reading. +#[nativelink_test] +async fn compressed_download_consumer_hangup_does_not_refetch() +-> Result<(), Box> { + use nativelink_proto::build::bazel::remote::execution::v2::compressor; + use nativelink_util::buf_channel::make_buf_channel_pair; + + let (content, digest) = make_content(0x55, 8 * 1024 * 1024); + let frame = + nativelink_util::wire_compression::compress(content.clone(), compressor::Value::Zstd)?; + let identity_reads = Arc::new(AtomicU64::new(0)); + let port = start_fake_bytestream_server(FakeByteStream { + identity_reads: identity_reads.clone(), + compressed_read_payload: Some(frame), + identity_read_payload: Some(content), + ..Default::default() + }) + .await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (mut writer, reader) = make_buf_channel_pair(); + // Hang up on the read before it can complete. + drop(reader); + let result = tokio::time::timeout( + core::time::Duration::from_secs(10), + grpc_store.get_part(digest, &mut writer, 0, None), + ) + .await + .expect("a hung-up consumer must not hang the read"); + + assert!( + result.is_err(), + "a broken writer must surface as an error, not a silent success" + ); + assert_eq!( + identity_reads.load(Ordering::Relaxed), + 0, + "a writer-side failure must not be retried as an identity read" + ); + Ok(()) +} diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index a40afdb85..377fc8683 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -13,7 +13,7 @@ // limitations under the License. use core::pin::Pin; -use core::sync::atomic::{AtomicU64, Ordering}; +use core::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use core::time::Duration; use std::borrow::Cow; use std::collections::{HashMap, VecDeque}; @@ -1271,19 +1271,22 @@ impl GrpcStore { ?err, "Compressed upload encoder ended early after successful write" ); + // The encoder stopped early, most likely because the + // server completed the write while its response stream + // was still being finalized. It has already released its + // borrow of `reader`, so drain any raw input the producer + // still has in flight. On a clean encode the reader is + // already at EOF and draining would be a no-op, so the + // happy path spawns nothing. + background_spawn!("grpc_store_compressed_upload_drain", async move { + if let Err(err) = reader.drain().await { + debug!( + ?err, + "Compressed upload reader drain failed after early completion" + ); + } + }); } - // The encoder may have stopped because the server completed - // the write while its response stream was still being - // finalized. It has already released its borrow of `reader`, - // so drain any raw input the producer still has in flight. - background_spawn!("grpc_store_compressed_upload_drain", async move { - if let Err(err) = reader.drain().await { - debug!( - ?err, - "Compressed upload reader drain failed after early completion" - ); - } - }); } } Ok(digest.size_bytes()) @@ -1391,6 +1394,12 @@ impl GrpcStore { } }; let forwarded = AtomicU64::new(0); + // Set once the decoder's EOF has been forwarded. The decoder only + // sends EOF after the whole blob passed its size and digest checks + // (and `buf_channel` reports a sender dropped without EOF as an error + // rather than as EOF), so this is a positive signal that the download + // completed and was verified. + let download_complete = AtomicBool::new(false); let pump_fut = async { loop { let chunk = decoded_rx @@ -1398,9 +1407,11 @@ impl GrpcStore { .await .err_tip(|| "in GrpcStore::get_part_compressed()")?; if chunk.is_empty() { - return writer + writer .send_eof() - .err_tip(|| "in GrpcStore::get_part_compressed()"); + .err_tip(|| "in GrpcStore::get_part_compressed()")?; + download_complete.store(true, Ordering::Relaxed); + return Ok(()); } forwarded.fetch_add(chunk.len() as u64, Ordering::Relaxed); writer @@ -1430,10 +1441,31 @@ impl GrpcStore { match result { Ok(((), (), ())) => Ok(None), + // The blob was fully delivered and verified before this error + // happened (for example a transport failure while the response + // trailer was being finalized). Falling back would re-read a + // range that has already been written, into a closed writer. + Err((stage, err)) if download_complete.load(Ordering::Relaxed) => { + debug!( + ?stage, + ?err, + "Compressed read completed and verified before a late stage error" + ); + Ok(None) + } Err((CompressedReadStage::Decode, err)) if err.code == Code::InvalidArgument => { Err(err.append("in GrpcStore::get_part_compressed()")) } - Err((CompressedReadStage::Pump, err)) if !is_retryable_code(err.code) => Err(err), + // Pump failures are failures of our own writer: either the + // downstream consumer went away, or the decoder aborted and + // dropped its sender (in which case the decoder's verdict is + // reported by the arm above, since `try_join!` polls it first). + // An identity re-read helps in neither case, so never fall back — + // `buf_channel` reports a broken receiver as `Internal`, which + // `is_retryable_code` would otherwise classify as retryable. + Err((CompressedReadStage::Pump, err)) => { + Err(err.append("in GrpcStore::get_part_compressed()")) + } Err((CompressedReadStage::Feed, err)) if !is_retryable_code(err.code) => { Err(err.append("in GrpcStore::get_part_compressed()")) } @@ -1536,12 +1568,17 @@ impl StoreDriver for GrpcStore { bytes_received: i64, } + let is_digest_key = matches!(key, StoreKey::Digest(_)); let digest = key.into_digest(); if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) { return self.update_action_result_from_bytes(digest, reader).await; } + // Only real digest keys may take the compressed path: for a string key + // `into_digest()` hashes the key itself, so the remote's mandatory + // uncompressed-digest verification would reject the upload. if self.remote_cache_compression_enabled + && is_digest_key && digest.size_bytes() >= WIRE_COMPRESSION_MIN_SIZE_BYTES { return self.update_compressed(digest, reader).await; @@ -1696,6 +1733,7 @@ impl StoreDriver for GrpcStore { } let mut offset = offset; + let mut length = length; if self.remote_cache_compression_enabled && is_digest_key && offset == 0 @@ -1706,7 +1744,21 @@ impl StoreDriver for GrpcStore { None => return Ok(()), // Resume via the identity path from where the compressed // read left off. - Some(forwarded) => offset = forwarded, + Some(forwarded) => { + // Ask only for what the compressed read did not deliver. + // The guard above means the caller's range covers the + // whole blob, so the remainder is exactly + // `size - forwarded`; keeping the caller's `length` would + // request a range that runs past the end of the blob. + let remaining = digest.size_bytes().saturating_sub(forwarded); + if remaining == 0 { + return writer + .send_eof() + .err_tip(|| "Failed to send EOF in GrpcStore::get_part()"); + } + offset = forwarded; + length = Some(remaining); + } } } From 8df90caba47925e22a8e2633b559cfc1d636c0e4 Mon Sep 17 00:00:00 2001 From: Walter Gray Date: Wed, 29 Jul 2026 14:39:59 -0700 Subject: [PATCH 2/2] test(store): cover the resume path with an explicit read length too The `length: None` form of a resumed read was already correct before the range fix, because `read_limit: 0` means "to the end" on the wire. The form that actually over-requested is an explicit `length` that covers the blob, so run the same resume assertions for both. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/grpc_wire_compression_test.rs | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/nativelink-service/tests/grpc_wire_compression_test.rs b/nativelink-service/tests/grpc_wire_compression_test.rs index 112b174d5..981bc6ce7 100644 --- a/nativelink-service/tests/grpc_wire_compression_test.rs +++ b/nativelink-service/tests/grpc_wire_compression_test.rs @@ -609,9 +609,14 @@ async fn compressed_upload_returns_after_early_completion() // download must keep the already-decoded prefix and ask the identity path for // exactly the undelivered tail — not re-fetch from zero, and not request a // range that runs past the end of the blob. -#[nativelink_test] -async fn interrupted_compressed_download_resumes_via_identity() --> Result<(), Box> { +// +// `requested_length` is the caller's `length` argument. Both accepted forms +// matter: `None` (read to end) and an explicit limit that covers the blob, +// which is the form that made the resumed request over-request past the end +// and rely on the server clamping it. +async fn assert_interrupted_download_resumes( + requested_length: Option, +) -> Result<(), Box> { use nativelink_proto::build::bazel::remote::execution::v2::compressor; let (content, digest) = make_mixed_content(0x44, 8 * 1024 * 1024); @@ -642,7 +647,7 @@ async fn interrupted_compressed_download_resumes_via_identity() let fetched = tokio::time::timeout( core::time::Duration::from_secs(10), - grpc_store.get_part_unchunked(digest, 0, None), + grpc_store.get_part_unchunked(digest, 0, requested_length), ) .await .expect("interrupted compressed read must not hang")?; @@ -670,6 +675,18 @@ async fn interrupted_compressed_download_resumes_via_identity() Ok(()) } +#[nativelink_test] +async fn interrupted_compressed_download_resumes_via_identity() +-> Result<(), Box> { + assert_interrupted_download_resumes(None).await +} + +#[nativelink_test] +async fn interrupted_compressed_download_resumes_with_explicit_length() +-> Result<(), Box> { + assert_interrupted_download_resumes(Some(8 * 1024 * 1024)).await +} + // W5: a download whose consumer goes away must fail, not trigger a pointless // identity re-read of the remainder into a channel nobody is reading. #[nativelink_test]