diff --git a/Cargo.lock b/Cargo.lock index eaf8c840b..f1f631dd8 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 c72c6e4b3..8f95b9004 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1475,6 +1475,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, deserialize_with = "convert_boolean_with_shellexpand")] + pub experimental_remote_cache_compression: bool, } /// Configuration for experimental small-blob read coalescing in a gRPC 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..72b9cbd8f 100644 --- a/nativelink-service/BUILD.bazel +++ b/nativelink-service/BUILD.bazel @@ -69,6 +69,7 @@ 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/wire_compression_test.rs", diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index c1063106b..179abe29b 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -1242,6 +1242,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 ecf522bbf..e317a8f81 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..=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(); + 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, + /// 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, +} + +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 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. + 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. + hold_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 = tokio::time::timeout( + core::time::Duration::from_secs(2), + grpc_store.get_part_unchunked(digest, 0, None), + ) + .await + .expect("decoder failure must cancel the held compressed read stream"); + 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 worker +// producer is still sending must not turn the producer's send/drain into an +// error. +#[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, 4 * 1024 * 1024); + let (mut tx, rx) = make_buf_channel_pair(); + 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, 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 the producer"); + result?; + producer_result?; + 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 5d9c47cac..a40afdb85 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_from_reader, +}; 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,339 @@ 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 { + 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 + // 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 mut reader = reader; + let encode_fut = stream_encode_compressed_download_from_reader( + &mut 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().try_into().unwrap_or(i64::MAX); + 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()") + }; + 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. 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" + ); + } + }); + } + 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. + if let Err(err) = encode_result { + debug!( + ?err, + "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()) + } + + /// 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> { + #[derive(Debug)] + enum CompressedReadStage { + Feed, + Decode, + Pump, + } + + 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 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)) + }, + ); + + 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))) + } + } + } } #[async_trait] @@ -1176,6 +1541,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 +1695,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 689307b2a..9e735a9eb 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 26d2a9271..b665c5626 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..c9760c7ea --- /dev/null +++ b/nativelink-util/src/wire_compression.rs @@ -0,0 +1,349 @@ +// 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, + 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}; + + 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"),