From 77690da3a8a35d154111faa4dda86d7609e67cb5 Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Tue, 21 Jul 2026 18:54:55 -0700 Subject: [PATCH 1/4] Add opt-in wire compression to GrpcStore's own transfers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NativeLink's internal hops never used REAPI compressed-blobs: GrpcStore built plain blobs/ resource names for every upload and read, so worker output uploads and input fetches shipped raw bytes in both directions regardless of any compression setting (the zstd support from the server work is accept/advertise only). New opt-in GrpcSpec.experimental_remote_cache_compression: uploads and full-blob downloads of blobs at or above 64KiB stream through zstd (level 1; higher levels measured slower with no meaningful wire reduction) using compressed-blobs/zstd resource names. Ranged reads and small blobs keep the identity path. Downloads decode with size and digest verification; a retryable mid-stream failure resumes via the identity path at the already-forwarded uncompressed offset. The shared streaming codecs move from nativelink-service to nativelink-util (the service re-exports them), since nativelink-store cannot depend on the service crate. Also fixes a latent decoder bug found by the new tests: a blob whose decompressed size is an exact multiple of the 128KiB decoder output buffer made the loop poll the finished decoder once more, and the new-frame input hint made the EOF check misreport a fully decoded stream as "ended in the middle of a zstd frame" — spuriously failing compressed uploads of exact-multiple-size blobs from any client. Measured (token-bucket throttled real TCP, byte-verified, off -> on): 96MiB 3.7:1-compressible blob @1Gbps: upload 712ms -> 203ms and download 706ms -> 204ms (3.5x), wire 101MB -> 27MB; @100Mbps 24MiB: 1975ms -> 494ms (4.0x). Incompressible payloads: identical wall time, no size penalty. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UXVtatcR9YMecBiu9RjwpC --- Cargo.lock | 1 + nativelink-config/src/stores.rs | 24 + nativelink-service/src/bytestream_server.rs | 1 + nativelink-service/src/wire_compression.rs | 304 +--------- nativelink-service/tests/cas_server_test.rs | 1 + .../tests/grpc_wire_compression_test.rs | 530 ++++++++++++++++++ .../tests/wire_compression_test.rs | 3 +- nativelink-store/src/grpc_store.rs | 356 +++++++++++- .../tests/grpc_read_batching_test.rs | 2 + nativelink-store/tests/grpc_store_test.rs | 1 + nativelink-util/BUILD.bazel | 4 + nativelink-util/Cargo.toml | 1 + nativelink-util/src/lib.rs | 1 + nativelink-util/src/proto_stream_utils.rs | 30 +- nativelink-util/src/wire_compression.rs | 331 +++++++++++ .../tests/wire_compression_test.rs | 132 +++++ .../tests/directory_cache_test.rs | 1 + 17 files changed, 1415 insertions(+), 308 deletions(-) create mode 100644 nativelink-service/tests/grpc_wire_compression_test.rs create mode 100644 nativelink-util/src/wire_compression.rs create mode 100644 nativelink-util/tests/wire_compression_test.rs diff --git a/Cargo.lock b/Cargo.lock index a3f60eb46..102972474 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3308,6 +3308,7 @@ dependencies = [ "uuid", "walkdir", "wincode", + "zstd", ] [[package]] diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 8aca04747..b0cf292af 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1465,6 +1465,30 @@ pub struct GrpcSpec { /// Default: unset (disabled). When unset there is zero behavior change. #[serde(default)] pub experimental_read_batching: Option, + + /// Compress this store's own blob transfers on the wire with REAPI + /// `compressed-blobs/zstd`. Uploads and full-blob downloads of blobs at + /// or above 64 KiB are zstd-compressed; smaller blobs and ranged reads + /// keep the identity path. + /// + /// The upstream instance must have + /// `capabilities.remote_cache_compression` enabled, otherwise compressed + /// requests fail with `InvalidArgument`. This setting is what makes + /// NativeLink-to-NativeLink hops (for example worker to CAS) benefit + /// from wire compression; it is independent of what external clients + /// such as Bazel negotiate for themselves. + /// + /// Compressed uploads do not resume mid-stream (mirroring the REAPI + /// server contract): a transport failure part-way through a compressed + /// upload surfaces immediately to the caller instead of retrying, and + /// outer callers retry the whole upload. + /// + /// When combined with `experimental_chunked_uploads`, chunked uploads + /// take precedence for blobs at or above the chunking threshold. + /// + /// Default: false (disabled). + #[serde(default)] + pub experimental_remote_cache_compression: bool, } /// Configuration for experimental small-blob read coalescing in a gRPC diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index c678eed59..d02689bc6 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -1212,6 +1212,7 @@ impl ByteStreamServer { let encode_fut = Box::pin(crate::wire_compression::stream_encode_compressed_download( raw_rx, wire_compressor, + crate::wire_compression::ZSTD_COMPRESSION_LEVEL, compressed_tx, )); diff --git a/nativelink-service/src/wire_compression.rs b/nativelink-service/src/wire_compression.rs index dc4de4a88..12ccd79f3 100644 --- a/nativelink-service/src/wire_compression.rs +++ b/nativelink-service/src/wire_compression.rs @@ -19,33 +19,20 @@ //! This is orthogonal to at-rest compression (`CompressionStore` with LZ4). use std::collections::HashSet; -use std::io::Read; use bytes::Bytes; use nativelink_config::cas_server::{CapabilitiesConfig, InstanceName, WithInstanceName}; -use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; +use nativelink_error::{Code, Error, make_err, make_input_err}; use nativelink_proto::build::bazel::remote::execution::v2::compressor; -use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; -use nativelink_util::common::DigestInfo; -use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; use nativelink_util::spawn_blocking; +// The codecs are shared with client-side GrpcStore transfers; re-export so +// existing service callers keep their import paths. +pub use nativelink_util::wire_compression::{ + ZSTD_COMPRESSION_LEVEL, compress, decompress, stream_decode_compressed_upload, + stream_encode_compressed_download, +}; use tracing::warn; -/// Zstd compression level for wire compression. -/// Level 0 in the zstd crate means "use default" (currently 3). -/// We use an explicit level for clarity. -pub const ZSTD_COMPRESSION_LEVEL: i32 = 3; - -/// Upper bound on the buffer `decompress` reserves up front for a zstd blob. -/// `expected_size` comes from the client's claimed digest, so it must never be -/// used as an allocation hint directly: a small payload claiming a huge size -/// would otherwise force a large pre-emptive allocation before the real -/// decompressed length is ever known. We reserve `min(expected_size, this)` -/// so common payloads never reallocate while a hostile claim allocates at most -/// this. Sized comfortably above any honest `BatchUpdateBlobs` payload (the -/// only caller of this bulk path; large blobs stream through `ByteStream`). -const ZSTD_DECOMPRESS_PREALLOC_CAP: usize = 1024 * 1024; - /// Which instances accept and advertise REAPI compressed-blobs (zstd). /// /// Derived from `CapabilitiesConfig.remote_cache_compression` so that @@ -104,88 +91,6 @@ pub fn resolve_wire_compressor( } } -/// Compress data using the specified wire compressor. -/// -/// `data` is the raw (uncompressed) bytes from the store. -/// Returns the compressed bytes suitable for sending on the wire. -pub fn compress(data: Bytes, compressor_value: compressor::Value) -> Result { - match compressor_value { - compressor::Value::Identity => Ok(data), - compressor::Value::Zstd => { - let compressed = zstd::bulk::compress(&data, ZSTD_COMPRESSION_LEVEL) - .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; - Ok(Bytes::from(compressed)) - } - _ => Err(make_input_err!( - "Unsupported wire compressor for compression: {:?}", - compressor_value - )), - } -} - -/// Decompress data using the specified wire compressor. -/// -/// `data` is the compressed bytes received from the wire. -/// `expected_size` is the uncompressed size (from the client's digest). It is -/// the hard cap on the decompressed output, but never a direct allocation -/// hint: the buffer grows with the real decoded bytes so a small payload -/// claiming a huge size cannot force a large up-front allocation. -/// Returns the decompressed bytes suitable for storing. -pub fn decompress( - data: &[u8], - compressor_value: compressor::Value, - expected_size: usize, -) -> Result { - match compressor_value { - compressor::Value::Identity => { - if data.len() != expected_size { - return Err(make_err!( - Code::InvalidArgument, - "Identity data size {} does not match expected size {}", - data.len(), - expected_size - )); - } - Ok(Bytes::copy_from_slice(data)) - } - compressor::Value::Zstd => { - // Decode incrementally so `expected_size` (which is attacker - // controlled — it is the client's claimed digest size) can bound - // the output without being trusted as an allocation size. We - // reserve only `min(expected_size, ZSTD_DECOMPRESS_PREALLOC_CAP)`, - // then `take(expected_size + 1)` hard-caps the decoder so a - // decompression bomb is rejected as soon as it overshoots. This - // mirrors the real-byte-count validation the identity arm and the - // streaming upload path already perform. - let decoder = zstd::stream::read::Decoder::new(data) - .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {e}"))?; - let mut output = Vec::with_capacity(expected_size.min(ZSTD_DECOMPRESS_PREALLOC_CAP)); - // `+ 1` lets an oversized stream produce one byte past the cap so - // the size check below rejects it rather than silently truncating. - let cap = u64::try_from(expected_size) - .err_tip(|| "expected_size did not fit in u64")? - .saturating_add(1); - decoder - .take(cap) - .read_to_end(&mut output) - .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {e}"))?; - if output.len() != expected_size { - return Err(make_err!( - Code::InvalidArgument, - "Decompressed size {} does not match expected size {}", - output.len(), - expected_size - )); - } - Ok(Bytes::from(output)) - } - _ => Err(make_input_err!( - "Unsupported wire compressor for decompression: {:?}", - compressor_value - )), - } -} - #[must_use] pub fn compress_for_batch_read(data: Bytes) -> (Bytes, compressor::Value) { match compress(data.clone(), compressor::Value::Zstd) { @@ -241,198 +146,3 @@ pub async fn decompress_batch_update( other => Err(make_input_err!("Unsupported wire compressor: {:?}", other)), } } - -/// Decode a client's zstd wire stream into raw bytes on `tx`, asynchronously. -/// -/// Like [`stream_encode_compressed_download`], this must not occupy a tokio -/// blocking-pool thread for the stream's lifetime: the input arrives at the -/// client's upload pace and `tx` drains at the store's write pace, so a -/// blocking implementation parks a pool thread on whichever side is slower -/// for as long as the upload lasts. The zstd frame is consumed incrementally -/// with the raw streaming API instead; per-chunk decode cost is bounded by -/// the channel chunk size, so it runs inline on the async runtime with -/// channel-native backpressure on both sides. -/// -/// Validation semantics match the REAPI compressed-blobs contract: the -/// decoded byte count may never exceed the digest size (checked per chunk so -/// a decompression bomb is rejected as soon as it overshoots), the final -/// count must equal it exactly, and the decoded bytes must hash to `digest`. -pub async fn stream_decode_compressed_upload( - mut compressed_rx: DropCloserReadHalf, - wire_compressor: compressor::Value, - digest: DigestInfo, - digest_function: DigestHasherFunc, - mut tx: DropCloserWriteHalf, -) -> Result<(), Error> { - use zstd::stream::raw::{Decoder, InBuffer, Operation, OutBuffer}; - - if wire_compressor != compressor::Value::Zstd { - return Err(make_input_err!( - "Streaming upload decompression only supports zstd, got {:?}", - wire_compressor - )); - } - - let expected_size = digest.size_bytes(); - let mut hasher = digest_function.hasher(); - let mut decoded_size = 0u64; - let mut decoder = Decoder::new() - .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {}", e))?; - // `DCtx::out_size()` guarantees a full decompressed block always fits, so - // the decoder never stalls for lack of output space within one `run`. - let mut out_buf = vec![0u8; zstd::zstd_safe::DCtx::out_size()]; - // Last input-size hint from the decoder: nonzero at input EOF means the - // stream ended in the middle of a frame and must be rejected. - let mut frame_input_hint = 0usize; - loop { - let chunk = compressed_rx - .recv() - .await - .err_tip(|| "Failed to receive compressed data in stream_decode_compressed_upload")?; - if chunk.is_empty() { - break; // EOF. - } - let mut in_buffer = InBuffer::around(&chunk); - loop { - let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); - frame_input_hint = decoder.run(&mut in_buffer, &mut out_buffer).map_err(|e| { - make_err!(Code::InvalidArgument, "Zstd decompression failed: {}", e) - })?; - let produced = Bytes::copy_from_slice(out_buffer.as_slice()); - // A completely full output buffer means the decoder may still - // have buffered output to flush, even with no input left. - let output_was_full = produced.len() == out_buf.len(); - if !produced.is_empty() { - let produced_u64 = u64::try_from(produced.len()) - .err_tip(|| "Decoded chunk size was not convertible to u64")?; - decoded_size = decoded_size.checked_add(produced_u64).ok_or_else(|| { - make_err!( - Code::InvalidArgument, - "Decoded compressed upload size overflow" - ) - })?; - if decoded_size > expected_size { - return Err(make_err!( - Code::InvalidArgument, - "Decoded compressed upload size {} bytes exceeds digest size {} bytes", - decoded_size, - expected_size - )); - } - hasher.update(&produced); - tx.send(produced).await?; - } - if in_buffer.pos() == in_buffer.src.len() && !output_was_full { - break; - } - } - } - if frame_input_hint != 0 { - return Err(make_err!( - Code::InvalidArgument, - "Compressed upload stream ended in the middle of a zstd frame" - )); - } - - if decoded_size != expected_size { - return Err(make_err!( - Code::InvalidArgument, - "Decompressed size {} does not match expected size {}", - decoded_size, - expected_size - )); - } - let actual_digest = hasher.finalize_digest(); - if actual_digest != digest { - return Err(make_err!( - Code::InvalidArgument, - "Decompressed digest {} does not match expected digest {}", - actual_digest, - digest - )); - } - - tx.send_eof() - .err_tip(|| "Failed to send decompressed upload EOF")?; - Ok(()) -} - -/// Encode a raw byte stream into a single zstd frame on `tx`, asynchronously. -/// -/// This runs entirely on the async runtime and must never occupy a tokio -/// blocking-pool thread for the stream's lifetime: `tx` drains at the gRPC -/// client's pace, so a blocking implementation (blocking reads from `raw_rx` -/// plus `blocking_send` into `tx`) parks one pool thread per concurrent -/// compressed download until the client finishes. Enough concurrent downloads -/// with slow consumers then exhaust the blocking pool and starve every other -/// `spawn_blocking` user (filesystem store I/O, upload decode, credential -/// resolution). The zstd frame is instead produced incrementally with the raw -/// streaming API: the CPU cost per iteration is bounded by the channel chunk -/// size (small — micro/milliseconds), so it is acceptable inline on a worker -/// thread, and `tx.send(...).await` gives backpressure without a parked -/// thread. -/// -/// The output is one well-formed zstd frame, identical in wire format to what -/// `zstd::stream::read::Encoder` produces (both drive `ZSTD_compressStream` -/// on a fresh `CCtx`). -pub async fn stream_encode_compressed_download( - mut raw_rx: DropCloserReadHalf, - wire_compressor: compressor::Value, - mut tx: DropCloserWriteHalf, -) -> Result<(), Error> { - use zstd::stream::raw::{Encoder, InBuffer, Operation, OutBuffer}; - - if wire_compressor != compressor::Value::Zstd { - return Err(make_input_err!( - "Streaming download compression only supports zstd, got {:?}", - wire_compressor - )); - } - - let mut encoder = Encoder::new(ZSTD_COMPRESSION_LEVEL) - .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; - // `CCtx::out_size()` guarantees a full compressed block always fits, so - // the encoder never stalls for lack of output space within one `run`. - let mut out_buf = vec![0u8; zstd::zstd_safe::CCtx::out_size()]; - loop { - let chunk = raw_rx - .recv() - .await - .err_tip(|| "Failed to receive raw data in stream_encode_compressed_download")?; - if chunk.is_empty() { - break; // EOF. - } - let mut in_buffer = InBuffer::around(&chunk); - while in_buffer.pos() < in_buffer.src.len() { - let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); - encoder - .run(&mut in_buffer, &mut out_buffer) - .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; - let produced = out_buffer.as_slice(); - if !produced.is_empty() { - tx.send(Bytes::copy_from_slice(produced)).await?; - } - } - } - - // Finish the frame: flush any internally buffered compressed data plus - // the frame epilogue. `finish` reports the bytes still pending, so loop - // until it reports none. - loop { - let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); - let remaining = encoder - .finish(&mut out_buffer, true) - .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; - let produced = out_buffer.as_slice(); - if !produced.is_empty() { - tx.send(Bytes::copy_from_slice(produced)).await?; - } - if remaining == 0 { - break; - } - } - - tx.send_eof() - .err_tip(|| "Failed to send compressed download EOF")?; - Ok(()) -} diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index c786a7730..35fd53762 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -1562,6 +1562,7 @@ async fn chunking_on_grpc_store_forbids_index_store() -> Result<(), Box hyper::rt::Executor for Executor +where + F: Future + Send + 'static, + F::Output: Send + 'static, +{ + fn execute(&self, fut: F) { + background_spawn!("test_executor", fut); + } +} + +async fn start_real_cas_server( + memory_store: Arc, + server_compression_enabled: bool, +) -> Result { + let store_manager = Arc::new(StoreManager::new()); + store_manager.add_store("main_cas", Store::new(memory_store))?; + let compression_instances = if server_compression_enabled { + RemoteCacheCompressionInstances::from_enabled_instance_names([INSTANCE_NAME.to_string()]) + } else { + RemoteCacheCompressionInstances::default() + }; + + let bs_server = ByteStreamServer::new( + &[WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + ..Default::default() + }, + }], + &store_manager, + &compression_instances, + )?; + let cas_server = CasServer::new( + &[WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: CasStoreConfig { + cas_store: "main_cas".to_string(), + experimental_chunking: None, + }, + }], + &store_manager, + &compression_instances, + )?; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + + let routes = tonic::service::Routes::new(bs_server.into_service()) + .add_service(cas_server.into_service()); + let adapted_service = tower::ServiceBuilder::new() + .map_request(|req: hyper::Request| { + let (parts, body) = req.into_parts(); + let body = body + .map_err(|e| tonic::Status::internal(e.to_string())) + .boxed_unsync(); + hyper::Request::from_parts(parts, body) + }) + .service(routes); + let hyper_service = TowerToHyperService::new(adapted_service); + + background_spawn!("test_server_accept", async move { + loop { + let Ok((stream, _addr)) = listener.accept().await else { + break; + }; + stream.set_nodelay(true).unwrap(); + let hyper_service = hyper_service.clone(); + background_spawn!("test_server_conn", async move { + drop( + auto::Builder::new(Executor) + .serve_connection_with_upgrades(TokioIo::new(stream), hyper_service) + .await, + ); + }); + } + }); + Ok(port) +} + +fn grpc_spec(port: u16, compression: bool) -> GrpcSpec { + GrpcSpec { + instance_name: INSTANCE_NAME.to_string(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 0, + tcp_keepalive_s: 0, + http2_keepalive_interval_s: 0, + http2_keepalive_timeout_s: 0, + }], + store_type: StoreType::Cas, + retry: Retry::default(), + max_concurrent_requests: 0, + connections_per_endpoint: 0, + rpc_timeout_s: 120, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_remote_cache_compression: compression, + } +} + +/// Structured, ~3:1-compressible content with a real sha256 digest (both +/// wire directions digest-verify). +fn make_content(tag: u8, len: usize) -> (Bytes, DigestInfo) { + let mut data = vec![0u8; len]; + for (i, b) in data.iter_mut().enumerate() { + *b = match i % 8 { + 0 | 1 | 2 | 3 | 4 => tag, + 5 => (i / 8) as u8, + _ => (i / 4096) as u8, + }; + } + 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] +async fn compressed_upload_round_trip() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let port = start_real_cas_server(memory_store.clone(), true).await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (content, digest) = make_content(0xA1, 1024 * 1024); + grpc_store.update_oneshot(digest, content.clone()).await?; + + let stored = memory_store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(stored, content, "decoded upload must be byte-identical"); + Ok(()) +} + +// A full-blob read above the threshold downloads via compressed-blobs; the +// client decodes and digest-verifies. +#[nativelink_test] +async fn compressed_download_round_trip() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let port = start_real_cas_server(memory_store.clone(), true).await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (content, digest) = make_content(0xB2, 1024 * 1024); + memory_store.update_oneshot(digest, content.clone()).await?; + + let fetched = grpc_store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(fetched, content, "decoded download must be byte-identical"); + Ok(()) +} + +// Blobs below the threshold must use the identity path: they round-trip +// even against a server with compression disabled. +#[nativelink_test] +async fn small_blob_uses_identity_path() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let port = start_real_cas_server(memory_store.clone(), false).await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (content, digest) = make_content(0xC3, 4 * 1024); + grpc_store.update_oneshot(digest, content.clone()).await?; + let fetched = grpc_store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(fetched, content); + Ok(()) +} + +// Ranged reads must use the identity path even above the threshold: they +// round-trip against a compression-disabled server. +#[nativelink_test] +async fn ranged_read_uses_identity_path() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let compressed_port = start_real_cas_server(memory_store.clone(), true).await?; + let grpc_store = GrpcStore::new(&grpc_spec(compressed_port, true)).await?; + + let (content, digest) = make_content(0xD4, 1024 * 1024); + memory_store.update_oneshot(digest, content.clone()).await?; + + let fetched = grpc_store + .get_part_unchunked(digest, 100, Some(1000)) + .await?; + assert_eq!(fetched, content.slice(100..1100)); + Ok(()) +} + +// A compressed upload to an upstream without compression enabled fails with +// InvalidArgument (config-trust contract, documented on the field). +#[nativelink_test] +async fn upload_to_non_compressed_upstream_fails() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let port = start_real_cas_server(memory_store.clone(), false).await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (content, digest) = make_content(0xE5, 1024 * 1024); + let err = grpc_store + .update_oneshot(digest, content) + .await + .expect_err("expected compressed upload to fail against plain upstream"); + assert_eq!(err.code, Code::InvalidArgument, "unexpected error: {err:?}"); + Ok(()) +} + +// Flag off: byte-identical legacy behavior against a plain upstream. +#[nativelink_test] +async fn flag_off_plain_round_trip() -> Result<(), Box> { + let memory_store = MemoryStore::new(&MemorySpec::default()); + let port = start_real_cas_server(memory_store.clone(), false).await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, false)).await?; + + let (content, digest) = make_content(0xF6, 1024 * 1024); + grpc_store.update_oneshot(digest, content.clone()).await?; + let fetched = grpc_store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(fetched, content); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Review-round regression tests: W1 (no resume storm), W2 (corruption is +// terminal), W3 (no hang on early completion). These need a controllable +// fake ByteStream server rather than the real one. +// --------------------------------------------------------------------------- + +#[derive(Debug, Default)] +struct FakeByteStream { + write_calls: Arc, + identity_reads: Arc, + /// Write behavior: consume this many messages, then abort the stream + /// with UNAVAILABLE. Zero means accept-one-then-early-complete (W3). + abort_write_after_messages: Option, + 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, + /// Payload served for identity reads (fallback detector). + identity_read_payload: Option, +} + +type ReadStreamT = + core::pin::Pin> + Send>>; + +#[tonic::async_trait] +impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for FakeByteStream { + type ReadStream = ReadStreamT; + + async fn read( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let resource_name = request.into_inner().resource_name; + if resource_name.contains("compressed-blobs") { + let payload = self + .compressed_read_payload + .clone() + .ok_or_else(|| tonic::Status::not_found("no compressed payload configured"))?; + let chunks: Vec = payload + .chunks(64 * 1024) + .map(Bytes::copy_from_slice) + .collect(); + let pause_after = self.pause_compressed_read_after_chunks; + let stream = + futures::stream::unfold((chunks, 0usize), move |(chunks, index)| async move { + if index >= chunks.len() { + return None; + } + 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; + } + let data = chunks[index].clone(); + Some((Ok(ReadResponse { data }), (chunks, index + 1))) + }); + let boxed: Self::ReadStream = Box::pin(stream); + return Ok(tonic::Response::new(boxed)); + } + self.identity_reads.fetch_add(1, Ordering::Relaxed); + let payload = self + .identity_read_payload + .clone() + .ok_or_else(|| tonic::Status::not_found("no identity payload configured"))?; + let chunks: Vec> = payload + .chunks(64 * 1024) + .map(|c| { + Ok(ReadResponse { + data: Bytes::copy_from_slice(c), + }) + }) + .collect(); + let boxed: Self::ReadStream = Box::pin(futures::stream::iter(chunks)); + Ok(tonic::Response::new(boxed)) + } + + async fn write( + &self, + request: tonic::Request>, + ) -> Result, tonic::Status> { + self.write_calls.fetch_add(1, Ordering::Relaxed); + let mut stream = request.into_inner(); + if self.early_complete_write { + // Read exactly one message, then settle the write early per the + // REAPI duplicate-upload contract without consuming the rest. + let _first = stream.message().await?; + return Ok(tonic::Response::new(WriteResponse { committed_size: -1 })); + } + if let Some(abort_after) = self.abort_write_after_messages { + let mut seen = 0u64; + while seen < abort_after { + match stream.message().await? { + Some(_msg) => seen += 1, + None => break, + } + } + return Err(tonic::Status::unavailable("injected mid-stream failure")); + } + Err(tonic::Status::unimplemented("no write behavior configured")) + } + + async fn query_write_status( + &self, + _request: tonic::Request, + ) -> Result, tonic::Status> { + Err(tonic::Status::unimplemented("not used in these tests")) + } +} + +async fn start_fake_bytestream_server(fake: FakeByteStream) -> Result { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let routes = tonic::service::Routes::new( + nativelink_proto::google::bytestream::byte_stream_server::ByteStreamServer::new(fake), + ); + let adapted_service = tower::ServiceBuilder::new() + .map_request(|req: hyper::Request| { + let (parts, body) = req.into_parts(); + let body = body + .map_err(|e| tonic::Status::internal(e.to_string())) + .boxed_unsync(); + hyper::Request::from_parts(parts, body) + }) + .service(routes); + let hyper_service = TowerToHyperService::new(adapted_service); + background_spawn!("fake_server_accept", async move { + loop { + let Ok((stream, _addr)) = listener.accept().await else { + break; + }; + stream.set_nodelay(true).unwrap(); + let hyper_service = hyper_service.clone(); + background_spawn!("fake_server_conn", async move { + drop( + auto::Builder::new(Executor) + .serve_connection_with_upgrades(TokioIo::new(stream), hyper_service) + .await, + ); + }); + } + }); + Ok(port) +} + +// W1: a mid-stream failure of a compressed upload must fail fast with the +// transport error — exactly one write attempt, no resume storm ending in +// InvalidArgument. +#[nativelink_test] +async fn compressed_upload_mid_stream_failure_fails_fast() -> Result<(), Box> +{ + let write_calls = Arc::new(AtomicU64::new(0)); + let port = start_fake_bytestream_server(FakeByteStream { + write_calls: write_calls.clone(), + abort_write_after_messages: Some(2), + ..Default::default() + }) + .await?; + let mut spec = grpc_spec(port, true); + spec.retry = Retry { + max_retries: 3, + delay: 0.01, + ..Default::default() + }; + let grpc_store = GrpcStore::new(&spec).await?; + + let (content, digest) = make_content(0x11, 4 * 1024 * 1024); + let result = grpc_store.update_oneshot(digest, content).await; + let err = result.expect_err("expected mid-stream failure to surface"); + assert_eq!( + err.code, + Code::Unavailable, + "must surface the transport error, not a resume-replay InvalidArgument: {err:?}" + ); + assert_eq!( + write_calls.load(Ordering::Relaxed), + 1, + "compressed uploads must not replay-resume after partial consumption" + ); + Ok(()) +} + +// W2: decoder-detected corruption mid-stream must be terminal +// InvalidArgument deterministically — never silently retried via the +// unverified identity path. +#[nativelink_test] +async fn corrupt_compressed_download_is_terminal() -> Result<(), Box> { + // The decoder must abort MID-STREAM while the feed still has chunks in + // flight (the old classification then saw a cascade "receiver + // disconnected" feed error and silently fell back to the unverified + // identity path). Bit-flips inside a zstd block decode to garbage and + // only fail the digest check at EOF, so instead append a second valid + // frame: its first decoded block overshoots the expected size, which + // the decoder rejects immediately and deterministically. + use nativelink_proto::build::bazel::remote::execution::v2::compressor; + let (content, digest) = make_content(0x22, 8 * 1024 * 1024); + let valid_frame = + nativelink_util::wire_compression::compress(content.clone(), compressor::Value::Zstd)?; + let (extra, _) = make_content(0x99, 8 * 1024 * 1024); + let overflow_frame = + nativelink_util::wire_compression::compress(extra, compressor::Value::Zstd)?; + let valid_frame_chunks = valid_frame.len().div_ceil(64 * 1024); + let mut corrupt = valid_frame.to_vec(); + corrupt.extend_from_slice(&overflow_frame); + let identity_reads = Arc::new(AtomicU64::new(0)); + let port = start_fake_bytestream_server(FakeByteStream { + identity_reads: identity_reads.clone(), + // Pause just after the first frame ends: the decoder hits the + // size overshoot while the feed is still blocked on the wire, so + // the feed's next send hits the dropped receiver. + pause_compressed_read_after_chunks: Some(valid_frame_chunks + 1), + compressed_read_payload: Some(Bytes::from(corrupt)), + identity_read_payload: Some(content), + ..Default::default() + }) + .await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let result = grpc_store.get_part_unchunked(digest, 0, None).await; + let err = result.expect_err("corrupt compressed data must be terminal"); + assert_eq!( + err.code, + Code::InvalidArgument, + "decoder verdict must win over cascade transport classification: {err:?}" + ); + assert_eq!( + identity_reads.load(Ordering::Relaxed), + 0, + "corrupt data must never fall back to the unverified identity path" + ); + Ok(()) +} + +// W3: a server that early-completes a compressed upload while the producer +// is stalled must not hang the update. +#[nativelink_test] +async fn compressed_upload_returns_after_early_completion() +-> Result<(), Box> { + use nativelink_util::buf_channel::make_buf_channel_pair; + use nativelink_util::store_trait::UploadSizeInfo; + + let port = start_fake_bytestream_server(FakeByteStream { + early_complete_write: true, + ..Default::default() + }) + .await?; + let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; + + let (content, digest) = make_content(0x33, 1024 * 1024); + let (mut tx, rx) = make_buf_channel_pair(); + // Send half the payload, then stall (no EOF) while keeping tx alive. + tx.send(content.slice(0..content.len() / 2)).await?; + let update_fut = grpc_store.update(digest, rx, UploadSizeInfo::ExactSize(content.len() as u64)); + let result = tokio::time::timeout(core::time::Duration::from_secs(10), update_fut) + .await + .expect("early-completed compressed upload must not hang on a stalled producer"); + result?; + drop(tx); + Ok(()) +} diff --git a/nativelink-service/tests/wire_compression_test.rs b/nativelink-service/tests/wire_compression_test.rs index 592f0c565..8463943cf 100644 --- a/nativelink-service/tests/wire_compression_test.rs +++ b/nativelink-service/tests/wire_compression_test.rs @@ -24,7 +24,7 @@ use nativelink_error::{Code, Error}; use nativelink_macro::nativelink_test; use nativelink_proto::build::bazel::remote::execution::v2::compressor; use nativelink_service::wire_compression::{ - compress, decompress, decompress_batch_update, resolve_wire_compressor, + ZSTD_COMPRESSION_LEVEL, compress, decompress, decompress_batch_update, resolve_wire_compressor, stream_encode_compressed_download, }; use nativelink_util::buf_channel::make_buf_channel_pair; @@ -240,6 +240,7 @@ fn held_compressed_download_streams_must_not_starve_blocking_pool() { stream_encode_compressed_download( raw_rx, compressor::Value::Zstd, + ZSTD_COMPRESSION_LEVEL, compressed_tx, ) ); diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index ef6c7fd7e..6b2c9a14c 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -32,14 +32,16 @@ use nativelink_proto::build::bazel::remote::execution::v2::{ ActionResult, BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse, GetActionResultRequest, GetTreeRequest, GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse, - SplitBlobRequest, SplitBlobResponse, UpdateActionResultRequest, + SplitBlobRequest, SplitBlobResponse, UpdateActionResultRequest, compressor, }; use nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient; use nativelink_proto::google::bytestream::{ QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest, WriteResponse, }; -use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; +use nativelink_util::buf_channel::{ + DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, +}; use nativelink_util::common::DigestInfo; use nativelink_util::connection_manager::ConnectionManager; use nativelink_util::digest_hasher::{DigestHasherFunc, default_digest_hasher_func}; @@ -51,6 +53,9 @@ use nativelink_util::resource_info::ResourceInfo; use nativelink_util::retry::{Retrier, RetryResult}; use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; use nativelink_util::telemetry::ClientHeaders; +use nativelink_util::wire_compression::{ + stream_decode_compressed_upload, stream_encode_compressed_download, +}; use nativelink_util::{background_spawn, default_health_status_indicator, tls_utils}; use opentelemetry::context::Context; use opentelemetry::global; @@ -61,7 +66,7 @@ use tokio::sync::{Semaphore, oneshot}; use tokio::time::sleep; use tonic::metadata::{Ascii, MetadataKey, MetadataValue}; use tonic::{Code, IntoRequest, Request, Response, Status, Streaming}; -use tracing::{error, trace, warn}; +use tracing::{debug, error, trace, warn}; use uuid::Uuid; struct TonicMetadataInjector<'a>(&'a mut tonic::metadata::MetadataMap); @@ -108,6 +113,17 @@ fn enrich_request( request } +/// Minimum blob size for wire compression when +/// `experimental_remote_cache_compression` is enabled. Below this, zstd +/// framing overhead and per-blob CPU outweigh the wire savings (small blobs +/// are the batching paths' domain). +const WIRE_COMPRESSION_MIN_SIZE_BYTES: u64 = 64 * 1024; + +/// Zstd level for `GrpcStore`'s own compressed transfers. Level 1: measured on +/// artifact-shaped corpora, higher levels bought no meaningful wire reduction +/// while encoding slower, and internal hops are throughput-sensitive. +const WIRE_COMPRESSION_ZSTD_LEVEL: i32 = 1; + /// Estimated per-entry protobuf and framing overhead charged against /// `max_batch_bytes`, so that batches of many tiny blobs cannot push a /// `BatchReadBlobs` response over the gRPC message size limit. @@ -211,6 +227,7 @@ pub struct GrpcStore { /// RPCs. `None` means reads always use the `ByteStream` `Read` path. #[metric(group = "read_batcher")] read_batcher: Option, + remote_cache_compression_enabled: bool, /// Used by the read coalescer to hand a strong reference of this store /// to detached dispatcher tasks. weak_self: Weak, @@ -292,6 +309,7 @@ impl GrpcStore { rpc_timeout, use_legacy_resource_names: spec.use_legacy_resource_names, read_batcher, + remote_cache_compression_enabled: spec.experimental_remote_cache_compression, headers, // We lowercase keys as HTTP headers are case-insensitive so we should match all cases forward_headers: spec @@ -777,6 +795,19 @@ impl GrpcStore { &self, stream: WriteRequestStreamWrapper, ) -> Result, Error> + where + T: Stream> + Unpin + Send + 'static, + E: Into + 'static, + { + const RESUMABLE: bool = true; + self.write_internal(stream, RESUMABLE).await + } + + async fn write_internal( + &self, + stream: WriteRequestStreamWrapper, + resumable: bool, + ) -> Result, Error> where T: Stream> + Unpin + Send + 'static, E: Into + 'static, @@ -786,10 +817,11 @@ impl GrpcStore { "CAS operation on AC store" ); - let local_state = Arc::new(Mutex::new(WriteState::new( - self.instance_name.clone(), - stream, - ))); + let mut write_state = WriteState::new(self.instance_name.clone(), stream); + if !resumable { + write_state.set_non_resumable(); + } + let local_state = Arc::new(Mutex::new(write_state)); let write_start = std::time::Instant::now(); let instance_name = self.instance_name.clone(); @@ -1083,6 +1115,295 @@ impl GrpcStore { .await .map(|_| len) } + + /// Uploads `digest` as a REAPI `compressed-blobs/zstd` write: the raw + /// bytes from `reader` are zstd-encoded on the fly and streamed with + /// compressed write offsets. Used when + /// `experimental_remote_cache_compression` is enabled and the blob meets + /// the size threshold. + async fn update_compressed( + self: Pin<&Self>, + digest: DigestInfo, + reader: DropCloserReadHalf, + ) -> Result { + // Compressed writes are NON-resumable: the server-side protocol + // rejects replays from a nonzero compressed offset, so a mid-stream + // failure must surface immediately instead of burning the retry + // budget on guaranteed-rejected resumes. + const NON_RESUMABLE: bool = false; + + struct LocalState { + resource_name: String, + compressed_rx: DropCloserReadHalf, + did_error: bool, + bytes_received: i64, + } + + let mut buf = Uuid::encode_buffer(); + let uuid = Uuid::new_v4().hyphenated().encode_lower(&mut buf); + let resource_name = if self.use_legacy_resource_names { + format!( + "{}/uploads/{}/compressed-blobs/zstd/{}/{}", + &self.instance_name, + uuid, + digest.packed_hash(), + digest.size_bytes(), + ) + } else { + let digest_function = Context::current() + .get::() + .map_or_else(default_digest_hasher_func, |v| *v) + .proto_digest_func() + .as_str_name() + .to_ascii_lowercase(); + format!( + "{}/uploads/{}/compressed-blobs/zstd/{}/{}/{}", + &self.instance_name, + uuid, + digest_function, + digest.packed_hash(), + digest.size_bytes(), + ) + }; + + let (compressed_tx, compressed_rx) = make_buf_channel_pair(); + let encode_fut = stream_encode_compressed_download( + reader, + compressor::Value::Zstd, + WIRE_COMPRESSION_ZSTD_LEVEL, + compressed_tx, + ); + + let local_state = LocalState { + resource_name, + compressed_rx, + did_error: false, + bytes_received: 0, + }; + let stream = Box::pin(unfold(local_state, |mut local_state| async move { + if local_state.did_error { + error!("GrpcStore::update_compressed() polled stream after error was returned"); + return None; + } + let data = match local_state + .compressed_rx + .recv() + .await + .err_tip(|| "In GrpcStore::update_compressed()") + { + Ok(data) => data, + Err(err) => { + local_state.did_error = true; + return Some((Err(err), local_state)); + } + }; + let write_offset = local_state.bytes_received; + local_state.bytes_received += data.len() as i64; + Some(( + Ok(WriteRequest { + resource_name: local_state.resource_name.clone(), + write_offset, + finish_write: data.is_empty(), // EOF is when no data was polled. + data, + }), + local_state, + )) + })); + + // The encoder must be driven concurrently with the RPC: the request + // stream's first message is the encoder's first output chunk. + let write_fut = async { + self.write_internal( + WriteRequestStreamWrapper::from(stream) + .await + .err_tip(|| "in GrpcStore::update_compressed()")?, + NON_RESUMABLE, + ) + .await + .map(|_| ()) + .err_tip(|| "in GrpcStore::update_compressed()") + }; + match future::select(core::pin::pin!(write_fut), core::pin::pin!(encode_fut)).await { + future::Either::Left((write_result, _encode_fut)) => { + write_result?; + // The server settled the write before consuming the whole + // stream (REAPI early completion of a duplicate upload). Do + // NOT await the encoder: it may be blocked on a stalled + // producer. Returning drops the pinned encoder, releasing + // the raw reader so the producer is notified, mirroring the + // plain path. Full reader-drain parity arrives with the + // write-dedup branch's drain pattern on rebase. + } + future::Either::Right((encode_result, write_fut)) => { + write_fut.await?; + // An encode error with a successful write means the server + // finished without consuming the whole stream; the upload + // itself succeeded. + if let Err(err) = encode_result { + debug!( + ?err, + "Compressed upload encoder ended early after successful write" + ); + } + } + } + Ok(digest.size_bytes()) + } + + /// Reads all of `digest` as a REAPI `compressed-blobs/zstd` read, + /// streaming decode into `writer` with size and digest verification at + /// EOF. Returns `Ok(None)` on success. On a retryable transport failure + /// it returns `Ok(Some(n))` where `n` is the count of decoded bytes + /// already forwarded, so the caller can resume via the identity path at + /// uncompressed offset `n`. Terminal errors (including decode/digest + /// mismatches) propagate as `Err`. + async fn get_part_compressed( + self: Pin<&Self>, + digest: DigestInfo, + writer: &mut DropCloserWriteHalf, + ) -> Result, Error> { + let resource_name = if self.use_legacy_resource_names { + format!( + "{}/compressed-blobs/zstd/{}/{}", + &self.instance_name, + digest.packed_hash(), + digest.size_bytes(), + ) + } else { + let digest_function = Context::current() + .get::() + .map_or_else(default_digest_hasher_func, |v| *v) + .proto_digest_func() + .as_str_name() + .to_ascii_lowercase(); + format!( + "{}/compressed-blobs/zstd/{}/{}/{}", + &self.instance_name, + digest_function, + digest.packed_hash(), + digest.size_bytes(), + ) + }; + + let mut stream = match self + .read_internal(ReadRequest { + resource_name, + read_offset: 0, + read_limit: 0, + }) + .await + { + Ok(stream) => stream, + Err(err) if is_retryable_code(err.code) => { + warn!( + ?err, + "Compressed read failed to start, falling back to identity read" + ); + return Ok(Some(0)); + } + Err(err) => return Err(err.append("in GrpcStore::get_part_compressed()")), + }; + + let digest_function = Context::current() + .get::() + .map_or_else(default_digest_hasher_func, |v| *v); + let (mut compressed_tx, compressed_rx) = make_buf_channel_pair(); + let (decoded_tx, mut decoded_rx) = make_buf_channel_pair(); + let decode_fut = stream_decode_compressed_upload( + compressed_rx, + compressor::Value::Zstd, + digest, + digest_function, + decoded_tx, + ); + let feed_fut = async { + loop { + match stream.next().await { + None => { + // A send_eof failure means the decoder already + // settled and dropped its receiver; its result is + // authoritative, so this is not a feed error. + drop(compressed_tx.send_eof()); + return Ok(()); + } + Some(Ok(message)) => { + // Empty chunks are legal on the wire but are the EOF + // marker in buf_channel; skip them. + if !message.data.is_empty() + && compressed_tx.send(message.data).await.is_err() + { + // The decoder stopped consuming (it settled or + // aborted on bad data). Its result decides the + // outcome; reporting a feed error here would + // misclassify a decoder-detected data error as + // retryable transport fallout. + return Ok(()); + } + } + Some(Err(status)) => return Err(Into::::into(status)), + } + } + }; + let forwarded = AtomicU64::new(0); + let pump_fut = async { + loop { + let chunk = decoded_rx + .recv() + .await + .err_tip(|| "in GrpcStore::get_part_compressed()")?; + if chunk.is_empty() { + return writer + .send_eof() + .err_tip(|| "in GrpcStore::get_part_compressed()"); + } + forwarded.fetch_add(chunk.len() as u64, Ordering::Relaxed); + writer + .send(chunk) + .await + .err_tip(|| "in GrpcStore::get_part_compressed()")?; + } + }; + let (feed_result, decode_result, pump_result) = + tokio::join!(feed_fut, decode_fut, pump_fut); + + if feed_result.is_ok() && decode_result.is_ok() && pump_result.is_ok() { + return Ok(None); + } + // A local writer failure is terminal regardless of transport state. + if let Err(err) = pump_result + && !is_retryable_code(err.code) + { + return Err(err); + } + // Transport interruptions (and the decoder errors they cascade + // into) are retryable: hand back how much was already forwarded so + // the caller resumes via the identity path. + let ((Err(transport_err), _) | (Ok(()), Err(transport_err))) = + (&feed_result, &decode_result) + else { + // Only the pump failed, and retryably. + return Ok(Some(forwarded.load(Ordering::Relaxed))); + }; + if matches!(&feed_result, Err(err) if !is_retryable_code(err.code)) { + return Err(transport_err + .clone() + .append("in GrpcStore::get_part_compressed()")); + } + // Digest/size mismatches from a COMPLETE transfer are terminal data + // errors; decode failures after a broken feed are retryable fallout. + if feed_result.is_ok() + && let Err(err) = &decode_result + && err.code == Code::InvalidArgument + { + return Err(err.clone().append("in GrpcStore::get_part_compressed()")); + } + warn!( + ?transport_err, + forwarded = forwarded.load(Ordering::Relaxed), + "Compressed read interrupted, falling back to identity read" + ); + Ok(Some(forwarded.load(Ordering::Relaxed))) + } } #[async_trait] @@ -1176,6 +1497,12 @@ impl StoreDriver for GrpcStore { return self.update_action_result_from_bytes(digest, reader).await; } + if self.remote_cache_compression_enabled + && digest.size_bytes() >= WIRE_COMPRESSION_MIN_SIZE_BYTES + { + return self.update_compressed(digest, reader).await; + } + let mut buf = Uuid::encode_buffer(); let resource_name = if self.use_legacy_resource_names { format!( @@ -1324,6 +1651,21 @@ impl StoreDriver for GrpcStore { } } + let mut offset = offset; + if self.remote_cache_compression_enabled + && is_digest_key + && offset == 0 + && length.is_none_or(|len| len >= digest.size_bytes()) + && digest.size_bytes() >= WIRE_COMPRESSION_MIN_SIZE_BYTES + { + match self.get_part_compressed(digest, writer).await? { + None => return Ok(()), + // Resume via the identity path from where the compressed + // read left off. + Some(forwarded) => offset = forwarded, + } + } + let resource_name = if self.use_legacy_resource_names { format!( "{}/blobs/{}/{}", diff --git a/nativelink-store/tests/grpc_read_batching_test.rs b/nativelink-store/tests/grpc_read_batching_test.rs index 86d419c3b..a31e63feb 100644 --- a/nativelink-store/tests/grpc_read_batching_test.rs +++ b/nativelink-store/tests/grpc_read_batching_test.rs @@ -277,6 +277,7 @@ async fn make_fixture(read_batching: Option) -> Result Result<(), Error> { headers: HashMap::new(), forward_headers: vec!["authorization".to_string()], experimental_read_batching: Some(batching_config()), + experimental_remote_cache_compression: false, }; let err = GrpcStore::new(&spec) .await diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index e120402d4..fac51cba7 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -62,6 +62,7 @@ fn test_spec>(endpoint: T, use_legacy_resource_names: bool) -> G headers: HashMap::new(), forward_headers: vec![], experimental_read_batching: None, + experimental_remote_cache_compression: false, } } diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index f9e0dcd0a..a7c5f499a 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -41,6 +41,7 @@ rust_library( "src/task.rs", "src/telemetry.rs", "src/tls_utils.rs", + "src/wire_compression.rs", ], proc_macro_deps = [ "@crates//:async-trait", @@ -90,6 +91,7 @@ rust_library( "@crates//:uuid", "@crates//:walkdir", "@crates//:wincode", + "@crates//:zstd", ], ) @@ -115,6 +117,7 @@ rust_test_suite( "tests/store_trait_test.rs", "tests/telemetry_test.rs", "tests/tls_utils_test.rs", + "tests/wire_compression_test.rs", ], compile_data = [ "tests/data/SekienAkashita.jpg", @@ -163,6 +166,7 @@ rust_test_suite( "@crates//:tracing", "@crates//:tracing-test", "@crates//:uuid", + "@crates//:zstd", ], ) diff --git a/nativelink-util/Cargo.toml b/nativelink-util/Cargo.toml index 89fa102d3..131614b36 100644 --- a/nativelink-util/Cargo.toml +++ b/nativelink-util/Cargo.toml @@ -96,6 +96,7 @@ uuid = { version = "1.16.0", default-features = false, features = [ ] } walkdir = { version = "2.5.0", default-features = false } wincode = { version = "0.5.4", default-features = false, features = ["derive"] } +zstd = { version = "0.13.3", default-features = false } [dev-dependencies] anyhow = { version = "1.0.103", default-features = false } diff --git a/nativelink-util/src/lib.rs b/nativelink-util/src/lib.rs index 932c21bb0..d082f3c3a 100644 --- a/nativelink-util/src/lib.rs +++ b/nativelink-util/src/lib.rs @@ -39,6 +39,7 @@ pub mod store_trait; pub mod task; pub mod telemetry; pub mod tls_utils; +pub mod wire_compression; // Re-export tracing mostly for use in macros. pub use tracing as __tracing; diff --git a/nativelink-util/src/proto_stream_utils.rs b/nativelink-util/src/proto_stream_utils.rs index 08362005d..39878c48d 100644 --- a/nativelink-util/src/proto_stream_utils.rs +++ b/nativelink-util/src/proto_stream_utils.rs @@ -227,6 +227,12 @@ where resume_queue: [Option; 2], // An optimisation to avoid having to manage resume_queue when it's empty. is_resumed: bool, + // When false, a partially-consumed stream never reports `can_resume()`: + // uploads whose server-side protocol cannot accept a replay from a + // nonzero offset (REAPI compressed-blobs writes) must fail fast instead + // of burning retries on guaranteed-rejected resumes. A stream that has + // not yet been consumed can always be retried from the start. + resumable: bool, } impl WriteState @@ -242,9 +248,23 @@ where cached_messages: [None, None], resume_queue: [None, None], is_resumed: false, + resumable: true, } } + /// Marks this write as non-resumable: once the stream has been partially + /// consumed, `can_resume()` reports false so the caller fails fast with + /// the original error instead of replaying messages the server-side + /// protocol is guaranteed to reject. Retrying an unconsumed stream from + /// the start remains allowed. + pub const fn set_non_resumable(&mut self) { + self.resumable = false; + } + + pub(crate) const fn is_resumable(&self) -> bool { + self.resumable + } + fn push_message(&mut self, message: WriteRequest) { self.cached_messages.swap(0, 1); self.cached_messages[0] = Some(message); @@ -267,7 +287,8 @@ where pub const fn can_resume(&self) -> bool { self.read_stream_error.is_none() - && (self.cached_messages[0].is_some() || self.read_stream.is_first_msg()) + && ((self.resumable && self.cached_messages[0].is_some()) + || self.read_stream.is_first_msg()) } pub fn resume(&mut self) { @@ -345,8 +366,11 @@ where } } // Cache the last request in case there is an error to allow - // the upload to be resumed. - local_state.push_message(message.clone()); + // the upload to be resumed. Non-resumable writes skip the + // clone: cached messages would never be replayed. + if local_state.is_resumable() { + local_state.push_message(message.clone()); + } Some(message) } Some(Err(err)) => { diff --git a/nativelink-util/src/wire_compression.rs b/nativelink-util/src/wire_compression.rs new file mode 100644 index 000000000..9ffd1668b --- /dev/null +++ b/nativelink-util/src/wire_compression.rs @@ -0,0 +1,331 @@ +// 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. + +//! Zstd wire-compression codecs for REAPI compressed-blobs. +//! +//! Shared by the server-side accept/serve paths (nativelink-service) and the +//! client-side `GrpcStore` transfers (nativelink-store). This is orthogonal +//! to at-rest compression (`CompressionStore` with LZ4). + +use std::io::Read; + +use bytes::Bytes; +use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; +use nativelink_proto::build::bazel::remote::execution::v2::compressor; + +use crate::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; +use crate::common::DigestInfo; +use crate::digest_hasher::{DigestHasher, DigestHasherFunc}; + +/// Zstd compression level for wire compression. +/// Level 0 in the zstd crate means "use default" (currently 3). +/// We use an explicit level for clarity. +pub const ZSTD_COMPRESSION_LEVEL: i32 = 3; + +/// Upper bound on the buffer `decompress` reserves up front for a zstd blob. +/// `expected_size` comes from the client's claimed digest, so it must never be +/// used as an allocation hint directly: a small payload claiming a huge size +/// would otherwise force a large pre-emptive allocation before the real +/// decompressed length is ever known. We reserve `min(expected_size, this)` +/// so common payloads never reallocate while a hostile claim allocates at most +/// this. Sized comfortably above any honest `BatchUpdateBlobs` payload (the +/// only caller of this bulk path; large blobs stream through `ByteStream`). +const ZSTD_DECOMPRESS_PREALLOC_CAP: usize = 1024 * 1024; + +/// Compress data using the specified wire compressor. +/// +/// `data` is the raw (uncompressed) bytes from the store. +/// Returns the compressed bytes suitable for sending on the wire. +pub fn compress(data: Bytes, compressor_value: compressor::Value) -> Result { + match compressor_value { + compressor::Value::Identity => Ok(data), + compressor::Value::Zstd => { + let compressed = zstd::bulk::compress(&data, ZSTD_COMPRESSION_LEVEL) + .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; + Ok(Bytes::from(compressed)) + } + _ => Err(make_input_err!( + "Unsupported wire compressor for compression: {:?}", + compressor_value + )), + } +} + +/// Decompress data using the specified wire compressor. +/// +/// `data` is the compressed bytes received from the wire. +/// `expected_size` is the uncompressed size (from the client's digest). It is +/// the hard cap on the decompressed output, but never a direct allocation +/// hint: the buffer grows with the real decoded bytes so a small payload +/// claiming a huge size cannot force a large up-front allocation. +/// Returns the decompressed bytes suitable for storing. +pub fn decompress( + data: &[u8], + compressor_value: compressor::Value, + expected_size: usize, +) -> Result { + match compressor_value { + compressor::Value::Identity => { + if data.len() != expected_size { + return Err(make_err!( + Code::InvalidArgument, + "Identity data size {} does not match expected size {}", + data.len(), + expected_size + )); + } + Ok(Bytes::copy_from_slice(data)) + } + compressor::Value::Zstd => { + // Decode incrementally so `expected_size` (which is attacker + // controlled — it is the client's claimed digest size) can bound + // the output without being trusted as an allocation size. We + // reserve only `min(expected_size, ZSTD_DECOMPRESS_PREALLOC_CAP)`, + // then `take(expected_size + 1)` hard-caps the decoder so a + // decompression bomb is rejected as soon as it overshoots. This + // mirrors the real-byte-count validation the identity arm and the + // streaming upload path already perform. + let decoder = zstd::stream::read::Decoder::new(data) + .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {e}"))?; + let mut output = Vec::with_capacity(expected_size.min(ZSTD_DECOMPRESS_PREALLOC_CAP)); + // `+ 1` lets an oversized stream produce one byte past the cap so + // the size check below rejects it rather than silently truncating. + let cap = u64::try_from(expected_size) + .err_tip(|| "expected_size did not fit in u64")? + .saturating_add(1); + decoder + .take(cap) + .read_to_end(&mut output) + .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {e}"))?; + if output.len() != expected_size { + return Err(make_err!( + Code::InvalidArgument, + "Decompressed size {} does not match expected size {}", + output.len(), + expected_size + )); + } + Ok(Bytes::from(output)) + } + _ => Err(make_input_err!( + "Unsupported wire compressor for decompression: {:?}", + compressor_value + )), + } +} + +/// Decode a client's zstd wire stream into raw bytes on `tx`, asynchronously. +/// +/// Like [`stream_encode_compressed_download`], this must not occupy a tokio +/// blocking-pool thread for the stream's lifetime: the input arrives at the +/// client's upload pace and `tx` drains at the store's write pace, so a +/// blocking implementation parks a pool thread on whichever side is slower +/// for as long as the upload lasts. The zstd frame is consumed incrementally +/// with the raw streaming API instead; per-chunk decode cost is bounded by +/// the channel chunk size, so it runs inline on the async runtime with +/// channel-native backpressure on both sides. +/// +/// Validation semantics match the REAPI compressed-blobs contract: the +/// decoded byte count may never exceed the digest size (checked per chunk so +/// a decompression bomb is rejected as soon as it overshoots), the final +/// count must equal it exactly, and the decoded bytes must hash to `digest`. +pub async fn stream_decode_compressed_upload( + mut compressed_rx: DropCloserReadHalf, + wire_compressor: compressor::Value, + digest: DigestInfo, + digest_function: DigestHasherFunc, + mut tx: DropCloserWriteHalf, +) -> Result<(), Error> { + use zstd::stream::raw::{Decoder, InBuffer, Operation, OutBuffer}; + + if wire_compressor != compressor::Value::Zstd { + return Err(make_input_err!( + "Streaming upload decompression only supports zstd, got {:?}", + wire_compressor + )); + } + + let expected_size = digest.size_bytes(); + let mut hasher = digest_function.hasher(); + let mut decoded_size = 0u64; + let mut decoder = Decoder::new() + .map_err(|e| make_err!(Code::InvalidArgument, "Zstd decompression failed: {}", e))?; + // `DCtx::out_size()` guarantees a full decompressed block always fits, so + // the decoder never stalls for lack of output space within one `run`. + let mut out_buf = vec![0u8; zstd::zstd_safe::DCtx::out_size()]; + // Last input-size hint from the decoder: nonzero at input EOF means the + // stream ended in the middle of a frame and must be rejected. + let mut frame_input_hint = 0usize; + loop { + let chunk = compressed_rx + .recv() + .await + .err_tip(|| "Failed to receive compressed data in stream_decode_compressed_upload")?; + if chunk.is_empty() { + break; // EOF. + } + let mut in_buffer = InBuffer::around(&chunk); + loop { + let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); + frame_input_hint = decoder.run(&mut in_buffer, &mut out_buffer).map_err(|e| { + make_err!(Code::InvalidArgument, "Zstd decompression failed: {}", e) + })?; + let produced = Bytes::copy_from_slice(out_buffer.as_slice()); + // A completely full output buffer means the decoder may still + // have buffered output to flush, even with no input left. + let output_was_full = produced.len() == out_buf.len(); + if !produced.is_empty() { + let produced_u64 = u64::try_from(produced.len()) + .err_tip(|| "Decoded chunk size was not convertible to u64")?; + decoded_size = decoded_size.checked_add(produced_u64).ok_or_else(|| { + make_err!( + Code::InvalidArgument, + "Decoded compressed upload size overflow" + ) + })?; + if decoded_size > expected_size { + return Err(make_err!( + Code::InvalidArgument, + "Decoded compressed upload size {} bytes exceeds digest size {} bytes", + decoded_size, + expected_size + )); + } + hasher.update(&produced); + tx.send(produced).await?; + } + // `hint == 0` means the frame is completely decoded AND fully + // flushed. It must terminate the loop even when the output + // buffer was filled exactly: polling the decoder again after + // frame end would return the input-size hint for a NEW frame + // header, and the post-EOF `frame_input_hint != 0` check would + // then misreport a fully-decoded stream as truncated. This is + // deterministic for blobs whose decompressed size is an exact + // multiple of the decoder output buffer size. + if in_buffer.pos() == in_buffer.src.len() && (frame_input_hint == 0 || !output_was_full) + { + break; + } + } + } + if frame_input_hint != 0 { + return Err(make_err!( + Code::InvalidArgument, + "Compressed upload stream ended in the middle of a zstd frame" + )); + } + + if decoded_size != expected_size { + return Err(make_err!( + Code::InvalidArgument, + "Decompressed size {} does not match expected size {}", + decoded_size, + expected_size + )); + } + let actual_digest = hasher.finalize_digest(); + if actual_digest != digest { + return Err(make_err!( + Code::InvalidArgument, + "Decompressed digest {} does not match expected digest {}", + actual_digest, + digest + )); + } + + tx.send_eof() + .err_tip(|| "Failed to send decompressed upload EOF")?; + Ok(()) +} + +/// Encode a raw byte stream into a single zstd frame on `tx`, asynchronously. +/// +/// This runs entirely on the async runtime and must never occupy a tokio +/// blocking-pool thread for the stream's lifetime: `tx` drains at the gRPC +/// client's pace, so a blocking implementation (blocking reads from `raw_rx` +/// plus `blocking_send` into `tx`) parks one pool thread per concurrent +/// compressed download until the client finishes. Enough concurrent downloads +/// with slow consumers then exhaust the blocking pool and starve every other +/// `spawn_blocking` user (filesystem store I/O, upload decode, credential +/// resolution). The zstd frame is instead produced incrementally with the raw +/// streaming API: the CPU cost per iteration is bounded by the channel chunk +/// size (small — micro/milliseconds), so it is acceptable inline on a worker +/// thread, and `tx.send(...).await` gives backpressure without a parked +/// thread. +/// +/// The output is one well-formed zstd frame, identical in wire format to what +/// `zstd::stream::read::Encoder` produces (both drive `ZSTD_compressStream` +/// on a fresh `CCtx`). +pub async fn stream_encode_compressed_download( + mut raw_rx: DropCloserReadHalf, + wire_compressor: compressor::Value, + compression_level: i32, + mut tx: DropCloserWriteHalf, +) -> Result<(), Error> { + use zstd::stream::raw::{Encoder, InBuffer, Operation, OutBuffer}; + + if wire_compressor != compressor::Value::Zstd { + return Err(make_input_err!( + "Streaming download compression only supports zstd, got {:?}", + wire_compressor + )); + } + + let mut encoder = Encoder::new(compression_level) + .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; + // `CCtx::out_size()` guarantees a full compressed block always fits, so + // the encoder never stalls for lack of output space within one `run`. + let mut out_buf = vec![0u8; zstd::zstd_safe::CCtx::out_size()]; + loop { + let chunk = raw_rx + .recv() + .await + .err_tip(|| "Failed to receive raw data in stream_encode_compressed_download")?; + if chunk.is_empty() { + break; // EOF. + } + let mut in_buffer = InBuffer::around(&chunk); + while in_buffer.pos() < in_buffer.src.len() { + let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); + encoder + .run(&mut in_buffer, &mut out_buffer) + .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; + let produced = out_buffer.as_slice(); + if !produced.is_empty() { + tx.send(Bytes::copy_from_slice(produced)).await?; + } + } + } + + // Finish the frame: flush any internally buffered compressed data plus + // the frame epilogue. `finish` reports the bytes still pending, so loop + // until it reports none. + loop { + let mut out_buffer = OutBuffer::around(out_buf.as_mut_slice()); + let remaining = encoder + .finish(&mut out_buffer, true) + .map_err(|e| make_err!(Code::Internal, "Zstd compression failed: {}", e))?; + let produced = out_buffer.as_slice(); + if !produced.is_empty() { + tx.send(Bytes::copy_from_slice(produced)).await?; + } + if remaining == 0 { + break; + } + } + + tx.send_eof() + .err_tip(|| "Failed to send compressed download EOF")?; + Ok(()) +} diff --git a/nativelink-util/tests/wire_compression_test.rs b/nativelink-util/tests/wire_compression_test.rs new file mode 100644 index 000000000..935bb2119 --- /dev/null +++ b/nativelink-util/tests/wire_compression_test.rs @@ -0,0 +1,132 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use bytes::Bytes; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::compressor; +use nativelink_util::buf_channel::make_buf_channel_pair; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; +use nativelink_util::wire_compression::{ + ZSTD_COMPRESSION_LEVEL, stream_decode_compressed_upload, stream_encode_compressed_download, +}; +use pretty_assertions::assert_eq; + +fn make_content_with_digest(tag: u8, len: usize) -> (Bytes, DigestInfo) { + let mut data = vec![0u8; len]; + for (i, b) in data.iter_mut().enumerate() { + #[allow(clippy::cast_possible_truncation)] + let byte = match i % 8 { + 0..=4 => tag, + 5 => (i / 8) as u8, + _ => (i / 4096) as u8, + }; + *b = byte; + } + let mut hasher = DigestHasherFunc::Sha256.hasher(); + hasher.update(&data); + let digest = hasher.finalize_digest(); + (Bytes::from(data), digest) +} + +async fn decode_chunks( + frame: Bytes, + split_at: usize, + digest: DigestInfo, +) -> (Result<(), Error>, Result) { + let (mut compressed_tx, compressed_rx) = make_buf_channel_pair(); + let (decoded_tx, mut decoded_rx) = make_buf_channel_pair(); + let decode_fut = stream_decode_compressed_upload( + compressed_rx, + compressor::Value::Zstd, + digest, + DigestHasherFunc::Sha256, + decoded_tx, + ); + let feed_fut = async move { + compressed_tx.send(frame.slice(0..split_at)).await?; + if split_at < frame.len() { + compressed_tx.send(frame.slice(split_at..)).await?; + } + compressed_tx.send_eof() + }; + let pump_fut = async move { + let mut total = 0usize; + loop { + let chunk = decoded_rx.recv().await?; + if chunk.is_empty() { + return Ok(total); + } + total += chunk.len(); + } + }; + let (feed_result, decode_result, pump_result) = tokio::join!(feed_fut, decode_fut, pump_fut); + feed_result.expect("feed must succeed"); + (decode_result, pump_result) +} + +// Regression: a blob whose decompressed size is an exact multiple of the +// decoder output buffer (128KiB) made the final `run` fill the output +// exactly with `hint == 0` (frame done); the loop then polled the finished +// decoder once more, received the new-frame-header hint, and the EOF check +// misreported the complete stream as truncated. +#[nativelink_test] +async fn decode_accepts_exact_output_buffer_multiple() -> Result<(), Error> { + let (content, digest) = make_content_with_digest(0xB2, 1024 * 1024); + let frame = + Bytes::from(zstd::bulk::compress(&content, ZSTD_COMPRESSION_LEVEL).expect("compress")); + let (decode_result, pump_result) = decode_chunks(frame, 65536, digest).await; + decode_result.expect("exact-multiple frame must decode"); + assert_eq!(pump_result.expect("pump"), content.len()); + Ok(()) +} + +// Encode -> decode round trip through the streaming codecs. +#[nativelink_test] +async fn stream_encode_decode_round_trip() -> Result<(), Error> { + let (content, digest) = make_content_with_digest(0xC4, 300_000); + let (mut raw_tx, raw_rx) = make_buf_channel_pair(); + let (compressed_tx, mut compressed_rx) = make_buf_channel_pair(); + let encode_fut = stream_encode_compressed_download( + raw_rx, + compressor::Value::Zstd, + ZSTD_COMPRESSION_LEVEL, + compressed_tx, + ); + let content_for_send = content.clone(); + let send_fut = async move { + raw_tx.send(content_for_send).await?; + raw_tx.send_eof() + }; + let collect_fut = async move { + let mut frame = Vec::new(); + loop { + let chunk = compressed_rx.recv().await?; + if chunk.is_empty() { + return Ok::<_, Error>(Bytes::from(frame)); + } + frame.extend_from_slice(&chunk); + } + }; + let (send_result, encode_result, frame) = tokio::join!(send_fut, encode_fut, collect_fut); + send_result.expect("send"); + encode_result.expect("encode"); + let frame = frame.expect("collect"); + + let (decode_result, pump_result) = decode_chunks(frame, 1, digest).await; + decode_result.expect("decode"); + assert_eq!(pump_result.expect("pump"), content.len()); + Ok(()) +} diff --git a/nativelink-worker/tests/directory_cache_test.rs b/nativelink-worker/tests/directory_cache_test.rs index 681127791..beb85f99c 100644 --- a/nativelink-worker/tests/directory_cache_test.rs +++ b/nativelink-worker/tests/directory_cache_test.rs @@ -1602,6 +1602,7 @@ async fn get_tree_prefetch_follows_server_pagination() -> Result<(), Error> { headers: HashMap::new(), forward_headers: vec![], experimental_read_batching: None, + experimental_remote_cache_compression: false, }; let fast_spec = FilesystemSpec { content_path: make_temp_path("paginated_get_tree_cas_content"), From 2c786cd4eab002345084e59dcd00fb5733175196 Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Wed, 22 Jul 2026 15:56:03 -0700 Subject: [PATCH 2/4] Make wire compression streaming and cancellation safe --- nativelink-config/src/stores.rs | 2 +- .../tests/deserialization_test.rs | 28 ++++ nativelink-service/BUILD.bazel | 1 + .../tests/grpc_wire_compression_test.rs | 50 ++++-- nativelink-store/src/grpc_store.rs | 142 ++++++++++++------ nativelink-util/src/wire_compression.rs | 18 +++ 6 files changed, 176 insertions(+), 65 deletions(-) diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index b0cf292af..3182e7f75 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1487,7 +1487,7 @@ pub struct GrpcSpec { /// take precedence for blobs at or above the chunking threshold. /// /// Default: false (disabled). - #[serde(default)] + #[serde(default, deserialize_with = "convert_boolean_with_shellexpand")] pub experimental_remote_cache_compression: bool, } diff --git a/nativelink-config/tests/deserialization_test.rs b/nativelink-config/tests/deserialization_test.rs index 8d476ae24..48a79d308 100644 --- a/nativelink-config/tests/deserialization_test.rs +++ b/nativelink-config/tests/deserialization_test.rs @@ -382,6 +382,8 @@ mod optional_values_tests { } mod boolean_tests { + use nativelink_config::stores::GrpcSpec; + use crate::BoolEntity; #[test] @@ -404,6 +406,32 @@ mod boolean_tests { assert_eq!(deserialized.value, expected, "{input}"); } } + + #[test] + fn grpc_compression_accepts_string_and_environment_boolean() { + let string_spec: GrpcSpec = serde_json5::from_str( + r#"{ + endpoints: [{ address: "http://localhost:1234" }], + store_type: "cas", + experimental_remote_cache_compression: "true" + }"#, + ) + .expect("string boolean should deserialize"); + assert!(string_spec.experimental_remote_cache_compression); + + // Safety: this test uses a unique variable and does not run code that + // reads mutable environment state concurrently. + unsafe { std::env::set_var("TEST_GRPC_REMOTE_CACHE_COMPRESSION", "false") }; + let environment_spec: GrpcSpec = serde_json5::from_str( + r#"{ + endpoints: [{ address: "http://localhost:1234" }], + store_type: "cas", + experimental_remote_cache_compression: "${TEST_GRPC_REMOTE_CACHE_COMPRESSION}" + }"#, + ) + .expect("environment boolean should deserialize"); + assert!(!environment_spec.experimental_remote_cache_compression); + } } mod shellexpand_tests { diff --git a/nativelink-service/BUILD.bazel b/nativelink-service/BUILD.bazel index 43aa6d168..ff4b92439 100644 --- a/nativelink-service/BUILD.bazel +++ b/nativelink-service/BUILD.bazel @@ -71,6 +71,7 @@ rust_test_suite( "tests/fetch_server_test.rs", "tests/health_server_test.rs", "tests/push_server_test.rs", + "tests/grpc_wire_compression_test.rs", "tests/wire_compression_test.rs", "tests/worker_api_server_test.rs", ], diff --git a/nativelink-service/tests/grpc_wire_compression_test.rs b/nativelink-service/tests/grpc_wire_compression_test.rs index d7a5e81ff..8ded2f18f 100644 --- a/nativelink-service/tests/grpc_wire_compression_test.rs +++ b/nativelink-service/tests/grpc_wire_compression_test.rs @@ -16,9 +16,9 @@ //! (`GrpcSpec.experimental_remote_cache_compression`) against a real //! compression-enabled `ByteStream`/CAS server. +use core::sync::atomic::{AtomicU64, Ordering}; use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; use bytes::Bytes; use futures::Future; @@ -160,9 +160,9 @@ fn make_content(tag: u8, len: usize) -> (Bytes, DigestInfo) { let mut data = vec![0u8; len]; for (i, b) in data.iter_mut().enumerate() { *b = match i % 8 { - 0 | 1 | 2 | 3 | 4 => tag, - 5 => (i / 8) as u8, - _ => (i / 4096) as u8, + 0..=4 => tag, + 5 => u8::try_from((i / 8) % 256).expect("value is reduced modulo 256"), + _ => u8::try_from((i / 4096) % 256).expect("value is reduced modulo 256"), }; } let mut hasher = Sha256::new(); @@ -286,6 +286,10 @@ struct FakeByteStream { /// 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). identity_read_payload: Option, } @@ -312,11 +316,15 @@ impl nativelink_proto::google::bytestream::byte_stream_server::ByteStream for Fa .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 stream = futures::stream::unfold((chunks, 0usize), move |(chunks, index)| async move { if index >= chunks.len() { return None; } + 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. @@ -478,7 +486,7 @@ async fn corrupt_compressed_download_is_terminal() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { @@ -516,15 +530,23 @@ async fn compressed_upload_returns_after_early_completion() .await?; let grpc_store = GrpcStore::new(&grpc_spec(port, true)).await?; - let (content, digest) = make_content(0x33, 1024 * 1024); + let (content, digest) = make_content(0x33, 4 * 1024 * 1024); let (mut tx, rx) = make_buf_channel_pair(); - // Send half the payload, then stall (no EOF) while keeping tx alive. - tx.send(content.slice(0..content.len() / 2)).await?; + let producer_content = content.clone(); + let producer = async move { + for chunk in producer_content.chunks(64 * 1024) { + tx.send(Bytes::copy_from_slice(chunk)).await?; + } + tx.send_eof() + }; let update_fut = grpc_store.update(digest, rx, UploadSizeInfo::ExactSize(content.len() as u64)); - let result = tokio::time::timeout(core::time::Duration::from_secs(10), update_fut) + let (result, producer_result) = + tokio::time::timeout(core::time::Duration::from_secs(10), async { + tokio::join!(update_fut, producer) + }) .await - .expect("early-completed compressed upload must not hang on a stalled producer"); + .expect("early-completed compressed upload must not hang the producer"); result?; - drop(tx); + producer_result?; Ok(()) } diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 6b2c9a14c..5ac8f9fff 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -54,7 +54,7 @@ use nativelink_util::retry::{Retrier, RetryResult}; use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; use nativelink_util::telemetry::ClientHeaders; use nativelink_util::wire_compression::{ - stream_decode_compressed_upload, stream_encode_compressed_download, + stream_decode_compressed_upload, stream_encode_compressed_download_from_reader, }; use nativelink_util::{background_spawn, default_health_status_indicator, tls_utils}; use opentelemetry::context::Context; @@ -1167,8 +1167,9 @@ impl GrpcStore { }; let (compressed_tx, compressed_rx) = make_buf_channel_pair(); - let encode_fut = stream_encode_compressed_download( - reader, + let mut reader = reader; + let encode_fut = stream_encode_compressed_download_from_reader( + &mut reader, compressor::Value::Zstd, WIRE_COMPRESSION_ZSTD_LEVEL, compressed_tx, @@ -1223,19 +1224,44 @@ impl GrpcStore { .map(|_| ()) .err_tip(|| "in GrpcStore::update_compressed()") }; - match future::select(core::pin::pin!(write_fut), core::pin::pin!(encode_fut)).await { - future::Either::Left((write_result, _encode_fut)) => { + enum UploadCompletion { + Write(Result<(), Error>), + Encode(Result<(), Error>, Result<(), Error>), + } + let completion = async { + let write_fut = Box::pin(write_fut); + let encode_fut = Box::pin(encode_fut); + match future::select(write_fut, encode_fut).await { + future::Either::Left((write_result, encode_fut)) => { + drop(encode_fut); + UploadCompletion::Write(write_result) + } + future::Either::Right((encode_result, write_fut)) => { + UploadCompletion::Encode(encode_result, write_fut.await) + } + } + } + .await; + match completion { + UploadCompletion::Write(write_result) => { write_result?; // The server settled the write before consuming the whole // stream (REAPI early completion of a duplicate upload). Do - // NOT await the encoder: it may be blocked on a stalled - // producer. Returning drops the pinned encoder, releasing - // the raw reader so the producer is notified, mirroring the - // plain path. Full reader-drain parity arrives with the - // write-dedup branch's drain pattern on rebase. + // not await the encoder: it may be blocked on a stalled + // producer. Cancel it, then drain the raw reader in the + // background so the producer can finish without observing a + // broken pipe from a successful upload. + 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" + ); + } + }); } - future::Either::Right((encode_result, write_fut)) => { - write_fut.await?; + UploadCompletion::Encode(encode_result, write_result) => { + write_result?; // An encode error with a successful write means the server // finished without consuming the whole stream; the upload // itself succeeded. @@ -1245,6 +1271,18 @@ impl GrpcStore { "Compressed upload encoder ended early after successful write" ); } + // 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()) @@ -1363,46 +1401,50 @@ impl GrpcStore { .err_tip(|| "in GrpcStore::get_part_compressed()")?; } }; - let (feed_result, decode_result, pump_result) = - tokio::join!(feed_fut, decode_fut, pump_fut); - - if feed_result.is_ok() && decode_result.is_ok() && pump_result.is_ok() { - return Ok(None); - } - // A local writer failure is terminal regardless of transport state. - if let Err(err) = pump_result - && !is_retryable_code(err.code) - { - return Err(err); - } - // Transport interruptions (and the decoder errors they cascade - // into) are retryable: hand back how much was already forwarded so - // the caller resumes via the identity path. - let ((Err(transport_err), _) | (Ok(()), Err(transport_err))) = - (&feed_result, &decode_result) - else { - // Only the pump failed, and retryably. - return Ok(Some(forwarded.load(Ordering::Relaxed))); - }; - if matches!(&feed_result, Err(err) if !is_retryable_code(err.code)) { - return Err(transport_err - .clone() - .append("in GrpcStore::get_part_compressed()")); + #[derive(Debug)] + enum CompressedReadStage { + Feed, + Decode, + Pump, } - // Digest/size mismatches from a COMPLETE transfer are terminal data - // errors; decode failures after a broken feed are retryable fallout. - if feed_result.is_ok() - && let Err(err) = &decode_result - && err.code == Code::InvalidArgument - { - return Err(err.clone().append("in GrpcStore::get_part_compressed()")); - } - warn!( - ?transport_err, - forwarded = forwarded.load(Ordering::Relaxed), - "Compressed read interrupted, falling back to identity read" + + let result = tokio::try_join!( + async { + feed_fut + .await + .map_err(|err| (CompressedReadStage::Feed, err)) + }, + async { + decode_fut + .await + .map_err(|err| (CompressedReadStage::Decode, err)) + }, + async { + pump_fut + .await + .map_err(|err| (CompressedReadStage::Pump, err)) + }, ); - Ok(Some(forwarded.load(Ordering::Relaxed))) + + match result { + Ok(((), (), ())) => 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), + Err((CompressedReadStage::Feed, err)) if !is_retryable_code(err.code) => { + Err(err.append("in GrpcStore::get_part_compressed()")) + } + Err((stage, err)) => { + debug!(?stage, ?err, "Compressed read interrupted"); + warn!( + ?err, + forwarded = forwarded.load(Ordering::Relaxed), + "Compressed read interrupted, falling back to identity read" + ); + Ok(Some(forwarded.load(Ordering::Relaxed))) + } + } } } diff --git a/nativelink-util/src/wire_compression.rs b/nativelink-util/src/wire_compression.rs index 9ffd1668b..c9760c7ea 100644 --- a/nativelink-util/src/wire_compression.rs +++ b/nativelink-util/src/wire_compression.rs @@ -271,6 +271,24 @@ pub async fn stream_encode_compressed_download( mut raw_rx: DropCloserReadHalf, wire_compressor: compressor::Value, compression_level: i32, + tx: DropCloserWriteHalf, +) -> Result<(), Error> { + stream_encode_compressed_download_from_reader( + &mut raw_rx, + wire_compressor, + compression_level, + tx, + ) + .await +} + +/// Encode a raw byte stream into a single zstd frame, borrowing the input +/// reader so a caller can continue draining it if the downstream consumer +/// finishes before the encoder does. +pub async fn stream_encode_compressed_download_from_reader( + raw_rx: &mut DropCloserReadHalf, + wire_compressor: compressor::Value, + compression_level: i32, mut tx: DropCloserWriteHalf, ) -> Result<(), Error> { use zstd::stream::raw::{Encoder, InBuffer, Operation, OutBuffer}; From 0698228958d6355db5e35a850c3430c736983e96 Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Fri, 24 Jul 2026 19:43:46 -0700 Subject: [PATCH 3/4] Fix clippy items-after-statements and buildifier ordering Hoist the two function-local enums above the first statements of their functions and restore alphabetical order in the service test suite list. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UXVtatcR9YMecBiu9RjwpC --- nativelink-service/BUILD.bazel | 2 +- nativelink-store/src/grpc_store.rs | 22 ++++++++++++---------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/nativelink-service/BUILD.bazel b/nativelink-service/BUILD.bazel index ff4b92439..72b9cbd8f 100644 --- a/nativelink-service/BUILD.bazel +++ b/nativelink-service/BUILD.bazel @@ -69,9 +69,9 @@ rust_test_suite( "tests/execution_server_test.rs", "tests/fastcdc_conformance_test.rs", "tests/fetch_server_test.rs", + "tests/grpc_wire_compression_test.rs", "tests/health_server_test.rs", "tests/push_server_test.rs", - "tests/grpc_wire_compression_test.rs", "tests/wire_compression_test.rs", "tests/worker_api_server_test.rs", ], diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 5ac8f9fff..98cfa29a6 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -1126,6 +1126,11 @@ impl GrpcStore { digest: DigestInfo, reader: DropCloserReadHalf, ) -> Result { + enum UploadCompletion { + Write(Result<(), Error>), + Encode(Result<(), Error>, Result<(), Error>), + } + // Compressed writes are NON-resumable: the server-side protocol // rejects replays from a nonzero compressed offset, so a mid-stream // failure must surface immediately instead of burning the retry @@ -1224,10 +1229,6 @@ impl GrpcStore { .map(|_| ()) .err_tip(|| "in GrpcStore::update_compressed()") }; - enum UploadCompletion { - Write(Result<(), Error>), - Encode(Result<(), Error>, Result<(), Error>), - } let completion = async { let write_fut = Box::pin(write_fut); let encode_fut = Box::pin(encode_fut); @@ -1300,6 +1301,13 @@ impl GrpcStore { digest: DigestInfo, writer: &mut DropCloserWriteHalf, ) -> Result, Error> { + #[derive(Debug)] + enum CompressedReadStage { + Feed, + Decode, + Pump, + } + let resource_name = if self.use_legacy_resource_names { format!( "{}/compressed-blobs/zstd/{}/{}", @@ -1401,12 +1409,6 @@ impl GrpcStore { .err_tip(|| "in GrpcStore::get_part_compressed()")?; } }; - #[derive(Debug)] - enum CompressedReadStage { - Feed, - Decode, - Pump, - } let result = tokio::try_join!( async { From ef052da1f66b4c8ca6f7142b39e703d4559c5a4d Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Sun, 26 Jul 2026 17:13:06 -0700 Subject: [PATCH 4/4] Fix clippy cast_possible_wrap in compressed upload accounting Match the identity path's idiom from the cast_possible_wrap cleanup on main: saturating try_into instead of a bare usize-to-i64 cast. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UXVtatcR9YMecBiu9RjwpC --- nativelink-store/src/grpc_store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index f1c07d82d..a40afdb85 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -1204,7 +1204,7 @@ impl GrpcStore { } }; let write_offset = local_state.bytes_received; - local_state.bytes_received += data.len() as i64; + local_state.bytes_received += data.len().try_into().unwrap_or(i64::MAX); Some(( Ok(WriteRequest { resource_name: local_state.resource_name.clone(),