diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index b3da2fbc2..dc29391f6 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -14,6 +14,9 @@ Colab composable CPUs [Dd]eduplication +[Dd]emotes? +[Dd]emoted +[Dd]emotion eviction_policy ELB Eskandar diff --git a/Cargo.lock b/Cargo.lock index eaf8c840b..ea8111ee7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3190,6 +3190,7 @@ dependencies = [ "bytes", "const_format", "dirs", + "fastcdc", "flate2", "fs-set-times", "futures", diff --git a/nativelink-config/examples/chunking_cas.json5 b/nativelink-config/examples/chunking_cas.json5 index f2ca4b966..5341e47db 100644 --- a/nativelink-config/examples/chunking_cas.json5 +++ b/nativelink-config/examples/chunking_cas.json5 @@ -7,6 +7,11 @@ // before and does not advertise chunking support. When enabled, clients // upload and download large blobs as content-defined chunks, so small // changes to large outputs only transfer the chunks that changed. +// +// This example configures the CAS service for CDC-aware external clients. +// NativeLink worker/StoreDriver uploads use a separate opt-in: +// `experimental_chunked_uploads` on the worker's upstream grpc CAS store. +// Ordinary external ByteStream uploads remain unchanged. { stores: [ { diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index c72c6e4b3..8f1df6b08 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1475,6 +1475,105 @@ pub struct GrpcSpec { /// Default: unset (disabled). When unset there is zero behavior change. #[serde(default)] pub experimental_read_batching: Option, + + /// Experimental: upload large blobs from the worker/`StoreDriver` path as + /// content-defined chunks. Blobs at or above `min_blob_size_bytes` are + /// split locally with `FastCDC` 2020, only the chunks the backend reports + /// missing are transferred, and the blob is assembled remotely with + /// `SpliceBlob`. On incremental changes to large artifacts this transfers + /// only the changed chunks. + /// + /// The backend MUST support the REAPI chunking extension (a `NativeLink` + /// CAS with `experimental_chunking` configured, or another compatible + /// server). During startup, `NativeLink` calls `GetCapabilities` and + /// requires the configured upstream instance to advertise `SplitBlob`, + /// `SpliceBlob`, and `FastCDC` 2020 with the same average chunk size and + /// seed 0. Startup fails with a configuration error when these + /// requirements are not met. This check is validation, not negotiation: + /// `NativeLink` never changes the configured parameters to match the + /// backend. + /// + /// If a chunk is evicted from the backend between its upload and the + /// final `SpliceBlob`, the upload fails with a retryable ABORTED error. + /// This option does not change the `ByteStream.Write` proxy path: external + /// `ByteStream` uploads remain ordinary streaming writes. If a chunked + /// worker upload fails after its input stream is consumed, the failure + /// surfaces to the caller's higher-level retry (normally action retry). + /// + /// Takes precedence over `experimental_remote_cache_compression` for + /// blobs at or above `min_blob_size_bytes` (chunks are transferred + /// uncompressed; the two features do not compose on the chunked path). + /// + /// WARNING (CAS sizing): a chunked upload stores each large blob twice + /// in the backend CAS — the chunk blobs (retained for future + /// incremental deduplication) plus the assembled blob. On a + /// size-capped CAS, budget roughly 2x the large-output working set as + /// headroom, and prefer a backend with post-splice chunk demotion so + /// chunks are evicted before primary blobs. A CAS whose eviction can + /// outpace a build's working set risks evicting still-referenced blobs + /// under build-without-the-bytes regardless of chunking; chunking + /// increases that pressure. + /// + /// Default: unset (disabled). When unset there is zero behavior change. + #[serde(default)] + pub experimental_chunked_uploads: Option, +} + +/// Configuration for experimental chunked uploads in a gRPC store. +/// See [`GrpcSpec::experimental_chunked_uploads`]. +#[derive(Serialize, Deserialize, Debug, Clone, Copy)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct GrpcChunkedUploadsConfig { + /// Only blobs at or above this size (in bytes) are uploaded as chunks. + /// Smaller blobs always use the plain `ByteStream` `Write` path. + /// + /// Default: 8388608 (8 MiB). + #[serde( + default = "default_chunked_uploads_min_blob_size_bytes", + deserialize_with = "convert_data_size_with_shellexpand" + )] + pub min_blob_size_bytes: u64, + + /// The average `FastCDC` 2020 chunk size in bytes. The minimum and + /// maximum chunk sizes are derived from this value (avg / 4 and + /// avg * 4). MUST match the average chunk size the backend uses for its + /// own chunking so worker-uploaded and server-split chunks share digests. + /// The worker path currently caps this at 768 KiB so a largest possible + /// chunk fits under its `BatchUpdateBlobs` request budget. Must be between + /// 1 KiB and 768 KiB for worker uploads. + /// + /// Default: 524288 (512 KiB), the REAPI-recommended value. + #[serde( + default = "default_chunked_uploads_avg_chunk_size_bytes", + deserialize_with = "convert_data_size_with_shellexpand" + )] + pub avg_chunk_size_bytes: u64, + + /// Blobs that could produce more than this many chunks (at the minimum + /// chunk size) use the plain streaming path instead, since a chunked + /// upload cannot fall back once the stream is partially consumed. + /// Should not exceed the `max_chunk_count` of the backend. + /// + /// Default: 50000 (matches the backend default; ~25 GiB at the default + /// average chunk size). + #[serde( + default = "default_chunked_uploads_max_chunk_count", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub max_chunk_count: u64, +} + +const fn default_chunked_uploads_min_blob_size_bytes() -> u64 { + 8 * 1024 * 1024 // 8 MiB. +} + +const fn default_chunked_uploads_avg_chunk_size_bytes() -> u64 { + 512 * 1024 // 512 KiB. +} + +const fn default_chunked_uploads_max_chunk_count() -> u64 { + 50_000 } /// Configuration for experimental small-blob read coalescing in a gRPC @@ -1876,3 +1975,44 @@ impl Retry { } } } + +#[cfg(test)] +mod tests { + use super::{GrpcChunkedUploadsConfig, GrpcSpec}; + + #[test] + fn grpc_chunked_uploads_defaults_are_backward_compatible() { + let config: GrpcChunkedUploadsConfig = serde_json5::from_str("{}").unwrap(); + assert_eq!(config.min_blob_size_bytes, 8 * 1024 * 1024); + assert_eq!(config.avg_chunk_size_bytes, 512 * 1024); + assert_eq!(config.max_chunk_count, 50_000); + + let spec: GrpcSpec = serde_json5::from_str( + r#"{ + endpoints: [{ address: "http://127.0.0.1:50051" }], + store_type: "cas" + }"#, + ) + .unwrap(); + assert!(spec.experimental_chunked_uploads.is_none()); + } + + #[test] + fn grpc_chunked_uploads_parse_data_sizes_and_reject_unknown_fields() { + let config: GrpcChunkedUploadsConfig = serde_json5::from_str( + r#"{ + min_blob_size_bytes: "16MiB", + avg_chunk_size_bytes: "256KiB", + max_chunk_count: 1234 + }"#, + ) + .unwrap(); + assert_eq!(config.min_blob_size_bytes, 16 * 1024 * 1024); + assert_eq!(config.avg_chunk_size_bytes, 256 * 1024); + assert_eq!(config.max_chunk_count, 1234); + + let error = + serde_json5::from_str::(r"{ unknown: true }").unwrap_err(); + assert!(error.to_string().contains("unknown")); + } +} diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 953c2bc6e..85ea6ed2e 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -232,6 +232,15 @@ impl CasServer { "'experimental_chunking.index_store' of instance '{}' must not be set when 'cas_store' is a grpc store: SplitBlob/SpliceBlob are forwarded to the backend", config.instance_name ); + if let Some(grpc_store) = store.downcast_ref::(None) + && let Some(upload_avg) = grpc_store.chunked_upload_avg_chunk_size_bytes() + { + error_if!( + upload_avg != avg_chunk_size_bytes, + "'experimental_chunking.avg_chunk_size_bytes' of instance '{}' ({avg_chunk_size_bytes}) must match the grpc store's 'experimental_chunked_uploads.avg_chunk_size_bytes' ({upload_avg})", + config.instance_name + ); + } // No ChunkingInstance: the forwarding shortcut in the // handlers takes over before local chunking is reached. stores.insert(config.instance_name.clone(), store); diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index ecf522bbf..30f46d4d7 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -1561,6 +1561,7 @@ async fn chunking_on_grpc_store_forbids_index_store() -> Result<(), Box(&'a mut tonic::metadata::MetadataMap); @@ -113,6 +119,40 @@ fn enrich_request( /// `BatchReadBlobs` response over the gRPC message size limit. const BATCH_READ_PER_ENTRY_OVERHEAD_BYTES: u64 = 256; +/// v1 cap on the configured average chunk size for chunked uploads: the +/// largest possible chunk (avg * 4) plus per-entry overhead must fit into a +/// single `BatchUpdateBlobs` request under `CHUNKED_UPLOAD_MAX_BATCH_BYTES`. +const CHUNKED_UPLOAD_MAX_AVG_CHUNK_SIZE_BYTES: u64 = 768 * 1024; + +/// Maximum payload bytes per `BatchUpdateBlobs` request used to upload +/// missing chunks (leaves headroom under the 4 MiB default gRPC message +/// limit for protobuf framing). +const CHUNKED_UPLOAD_MAX_BATCH_BYTES: u64 = 3 * 1024 * 1024; + +/// Chunked uploads buffer at most this many payload bytes (one window) +/// before checking chunk existence and transferring only the missing ones. +const CHUNKED_UPLOAD_WINDOW_BYTES: usize = 8 * 1024 * 1024; + +/// Upper bound on chunks per window (relevant for tiny average chunk sizes). +const CHUNKED_UPLOAD_WINDOW_CHUNKS: usize = 128; + +/// Ceiling on blocking tasks chunk digests are fanned out across per +/// window; the effective width is capped at the host's available +/// parallelism so small workers are not oversubscribed. Hashing dominates +/// the chunked-upload CPU cost (~4x the boundary pass) and parallelizes +/// near-linearly since chunks are independent. Note the CPU accounting: +/// parallel hashing performs the same total work as serial hashing — it +/// narrows the burst rather than growing it — and runs in the post-action +/// upload phase alongside the already-concurrent output upload fan-out. +const CHUNKED_UPLOAD_HASH_PARALLELISM_CEILING: usize = 8; + +/// Effective chunk-hash fan-out for this host. +fn chunk_hash_parallelism() -> usize { + std::thread::available_parallelism() + .map_or(4, usize::from) + .min(CHUNKED_UPLOAD_HASH_PARALLELISM_CEILING) +} + /// A small-blob read waiting to be coalesced into a `BatchReadBlobs` RPC. #[derive(Debug)] struct PendingRead { @@ -195,6 +235,38 @@ const fn is_retryable_code(code: Code) -> bool { // This store is usually a pass-through store, but can also be used as a CAS store. Using it as an // AC store has one major side-effect... The has() function may not give the proper size of the // underlying data. This might cause issues if embedded in certain stores. + +/// State and metrics for experimental chunked uploads. Large blobs are +/// split locally with `FastCDC` 2020, only missing chunks are transferred, +/// and the blob is assembled remotely with `SpliceBlob`. +#[derive(Debug, MetricsComponent)] +struct ChunkedUploader { + config: GrpcChunkedUploadsConfig, + #[metric(help = "Number of uploads that took the chunked path")] + chunked_uploads_total: AtomicU64, + #[metric(help = "Number of chunks transferred by chunked uploads")] + chunks_sent: AtomicU64, + #[metric(help = "Number of chunks skipped because the backend already had them")] + chunks_deduped: AtomicU64, + #[metric(help = "Payload bytes transferred by chunked uploads")] + bytes_sent: AtomicU64, + #[metric(help = "Payload bytes skipped because the backend already had their chunks")] + bytes_deduped: AtomicU64, +} + +impl ChunkedUploader { + const fn new(config: GrpcChunkedUploadsConfig) -> Self { + Self { + config, + chunked_uploads_total: AtomicU64::new(0), + chunks_sent: AtomicU64::new(0), + chunks_deduped: AtomicU64::new(0), + bytes_sent: AtomicU64::new(0), + bytes_deduped: AtomicU64::new(0), + } + } +} + #[derive(Debug, MetricsComponent)] pub struct GrpcStore { #[metric(help = "Instance name for the store")] @@ -211,6 +283,11 @@ pub struct GrpcStore { /// RPCs. `None` means reads always use the `ByteStream` `Read` path. #[metric(group = "read_batcher")] read_batcher: Option, + /// When configured, uploads large blobs as content-defined chunks via + /// `FindMissingBlobs` + `BatchUpdateBlobs` + `SpliceBlob`. `None` means + /// uploads always use the plain `ByteStream` `Write` path. + #[metric(group = "chunked_uploads")] + chunked_uploader: Option, /// Used by the read coalescer to hand a strong reference of this store /// to detached dispatcher tasks. weak_self: Weak, @@ -258,6 +335,36 @@ impl GrpcStore { None => None, }; + let chunked_uploader = match &spec.experimental_chunked_uploads { + Some(config) => { + error_if!( + matches!(spec.store_type, nativelink_config::stores::StoreType::Ac), + "experimental_chunked_uploads is not supported on AC stores" + ); + // REAPI FastCDC bounds are 1KiB..=1MiB, but the largest chunk + // (avg * 4) must also fit into one BatchUpdateBlobs request + // under the batch byte budget, so v1 caps the average lower. + error_if!( + !(1024..=CHUNKED_UPLOAD_MAX_AVG_CHUNK_SIZE_BYTES) + .contains(&config.avg_chunk_size_bytes), + "experimental_chunked_uploads.avg_chunk_size_bytes is {}, must be between 1024 and {CHUNKED_UPLOAD_MAX_AVG_CHUNK_SIZE_BYTES}", + config.avg_chunk_size_bytes + ); + error_if!( + config.min_blob_size_bytes < config.avg_chunk_size_bytes * 4, + "experimental_chunked_uploads.min_blob_size_bytes ({}) must be at least avg_chunk_size_bytes * 4 ({})", + config.min_blob_size_bytes, + config.avg_chunk_size_bytes * 4 + ); + error_if!( + config.max_chunk_count == 0, + "experimental_chunked_uploads.max_chunk_count must be greater than zero" + ); + Some(ChunkedUploader::new(*config)) + } + None => None, + }; + let mut headers = Vec::with_capacity(spec.headers.len()); for (name, value) in &spec.headers { // We lowercase keys as HTTP headers are case-insensitive so we should match all cases @@ -292,6 +399,7 @@ impl GrpcStore { rpc_timeout, use_legacy_resource_names: spec.use_legacy_resource_names, read_batcher, + chunked_uploader, headers, // We lowercase keys as HTTP headers are case-insensitive so we should match all cases forward_headers: spec @@ -302,6 +410,118 @@ impl GrpcStore { })) } + /// Returns the `FastCDC` average used by worker/StoreDriver-only chunked + /// uploads, when that upload path is enabled. + #[must_use] + pub fn chunked_upload_avg_chunk_size_bytes(&self) -> Option { + self.chunked_uploader + .as_ref() + .map(|uploader| uploader.config.avg_chunk_size_bytes) + } + + /// Verifies the configured upstream instance advertises the exact + /// chunking behavior this uploader requires. This is validation, not + /// negotiation: the local average and seed are never changed to match the + /// backend. + async fn validate_chunked_upload_capabilities(&self) -> Result<(), Error> { + let uploader = self + .chunked_uploader + .as_ref() + .err_tip(|| "Chunked-upload capability validation called while disabled")?; + let request = GetCapabilitiesRequest { + instance_name: self.instance_name.clone(), + }; + let capabilities = self + .perform_request(request, |request| async move { + let channel = self + .connection_manager + .connection("validate_chunked_upload_capabilities".into()) + .await + .err_tip(|| "Connecting for chunked-upload capability validation")?; + CapabilitiesClient::new(channel) + .get_capabilities(enrich_request( + Request::new(request), + &self.headers, + &self.forward_headers, + )) + .await + .err_tip(|| "Calling GetCapabilities for experimental_chunked_uploads") + }) + .await + .map_err(|error| { + error.append(format!( + "Cannot enable experimental_chunked_uploads for upstream instance '{}': \ + the backend capability check failed. Ensure the endpoint serves REAPI \ + Capabilities for this instance and advertises SplitBlob, SpliceBlob, and \ + FastCDC 2020, or remove experimental_chunked_uploads", + self.instance_name + )) + })? + .into_inner(); + + let cache_capabilities = capabilities.cache_capabilities.ok_or_else(|| { + make_err!( + Code::FailedPrecondition, + "Cannot enable experimental_chunked_uploads for upstream instance '{}': \ + GetCapabilities returned no cache_capabilities. Configure a chunking-capable CAS \ + backend or remove experimental_chunked_uploads", + self.instance_name + ) + })?; + if !cache_capabilities.split_blob_support || !cache_capabilities.splice_blob_support { + return Err(make_err!( + Code::FailedPrecondition, + "Cannot enable experimental_chunked_uploads for upstream instance '{}': the \ + backend must advertise both SplitBlob and SpliceBlob support \ + (split_blob_support={}, splice_blob_support={}). Enable experimental_chunking on \ + the backend or remove experimental_chunked_uploads", + self.instance_name, + cache_capabilities.split_blob_support, + cache_capabilities.splice_blob_support + )); + } + + let fast_cdc_params = cache_capabilities.fast_cdc_2020_params.ok_or_else(|| { + make_err!( + Code::FailedPrecondition, + "Cannot enable experimental_chunked_uploads for upstream instance '{}': the \ + backend did not advertise FastCDC 2020 parameters. Enable FastCDC 2020 on the \ + backend or remove experimental_chunked_uploads", + self.instance_name + ) + })?; + if fast_cdc_params.avg_chunk_size_bytes != uploader.config.avg_chunk_size_bytes { + return Err(make_err!( + Code::FailedPrecondition, + "Cannot enable experimental_chunked_uploads for upstream instance '{}': \ + avg_chunk_size_bytes is {}, but the backend advertises {}. Configure both sides \ + with the same FastCDC 2020 average", + self.instance_name, + uploader.config.avg_chunk_size_bytes, + fast_cdc_params.avg_chunk_size_bytes + )); + } + if fast_cdc_params.seed != 0 { + return Err(make_err!( + Code::FailedPrecondition, + "Cannot enable experimental_chunked_uploads for upstream instance '{}': the \ + backend advertises FastCDC 2020 seed {}, but NativeLink worker uploads currently \ + require seed 0. Configure the backend with seed 0 or remove \ + experimental_chunked_uploads", + self.instance_name, + fast_cdc_params.seed + )); + } + + info!( + instance_name = %self.instance_name, + avg_chunk_size_bytes = fast_cdc_params.avg_chunk_size_bytes, + seed = fast_cdc_params.seed, + "Validated backend support for experimental chunked uploads", + ); + Ok(()) + } + async fn perform_request(&self, input: I, mut request: F) -> Result where F: FnMut(I) -> Fut + Send + Copy, @@ -1083,11 +1303,439 @@ impl GrpcStore { .await .map(|_| len) } + + /// Validates that a `BatchUpdateBlobs` response covers exactly the request + /// entries and returns the entries whose per-entry status failed. REAPI + /// permits response entries to be returned in a different order, so + /// coverage is checked by digest rather than by position. + fn batch_update_failures( + requests: &[batch_update_blobs_request::Request], + response: BatchUpdateBlobsResponse, + ) -> Result, Error> { + let mut requests_by_digest = HashMap::with_capacity(requests.len()); + for request in requests { + let digest = request + .digest + .clone() + .err_tip(|| "Missing digest in BatchUpdateBlobs request")?; + let digest_info = DigestInfo::try_from(digest) + .err_tip(|| "Invalid digest in BatchUpdateBlobs request")?; + error_if!( + requests_by_digest + .insert(digest_info, request.clone()) + .is_some(), + "Duplicate digest {digest_info} in BatchUpdateBlobs request" + ); + } + + let mut seen = HashSet::with_capacity(response.responses.len()); + let mut failures = Vec::new(); + for entry in response.responses { + let digest = entry + .digest + .err_tip(|| "Missing digest in BatchUpdateBlobs response")?; + let digest_info = DigestInfo::try_from(digest) + .err_tip(|| "Invalid digest in BatchUpdateBlobs response")?; + let request = requests_by_digest.get(&digest_info).ok_or_else(|| { + make_err!( + Code::Internal, + "BatchUpdateBlobs response contains unrequested digest {digest_info}" + ) + })?; + error_if!( + !seen.insert(digest_info), + "Duplicate digest {digest_info} in BatchUpdateBlobs response" + ); + + let status = entry.status.ok_or_else(|| { + make_err!( + Code::Internal, + "BatchUpdateBlobs response omitted status for {digest_info}" + ) + })?; + let status_code = status.code; + if status_code != 0 { + failures.push(( + request.clone(), + make_err!( + Code::from(status_code), + "BatchUpdateBlobs entry failed for {digest_info}: {status:?}" + ), + )); + } + } + + let missing_responses: Vec<_> = requests_by_digest + .keys() + .filter(|digest| !seen.contains(digest)) + .collect(); + if !missing_responses.is_empty() { + return Err(make_err!( + Code::Internal, + "BatchUpdateBlobs response omitted {} requested digest(s): {:?}", + missing_responses.len(), + missing_responses + )); + } + Ok(failures) + } + + /// Retries one failed `BatchUpdateBlobs` entry using the store's configured + /// retry policy. The first error is yielded to Retrier so its delay and + /// max-retry semantics apply between the original attempt and retries. + async fn retry_batch_update_entry( + &self, + request: batch_update_blobs_request::Request, + digest_function: i32, + initial_error: Error, + ) -> Result<(), Error> { + let mut initial_error = Some(initial_error); + self.retrier + .retry(unfold((), move |()| { + let request = request.clone(); + let initial_error = initial_error.take(); + async move { + let result = if let Some(error) = initial_error { + RetryResult::Retry(error) + } else { + match self + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: self.instance_name.clone(), + requests: vec![request.clone()], + digest_function, + })) + .await + { + Err(error) => RetryResult::Retry(error), + Ok(response) => match Self::batch_update_failures( + core::slice::from_ref(&request), + response.into_inner(), + ) { + Ok(mut failures) => { + failures.pop().map_or(RetryResult::Ok(()), |(_, error)| { + RetryResult::Retry(error) + }) + } + Err(error) => RetryResult::Err(error), + }, + } + }; + Some((result, ())) + } + })) + .await + } + /// Uploads one window of chunks: checks which are missing on the + /// backend and transfers only those, packed into `BatchUpdateBlobs` + /// requests under the batch byte budget. + async fn flush_chunk_window( + &self, + uploader: &ChunkedUploader, + window: Vec<(DigestInfo, Bytes)>, + digest_function: i32, + ) -> Result<(), Error> { + // Dedup digests within the window (identical chunks are common in + // sparse or repetitive regions). + let mut unique: HashMap = HashMap::with_capacity(window.len()); + for (chunk_digest, data) in window { + unique.entry(chunk_digest).or_insert(data); + } + let requested_digests: Vec = unique.keys().copied().collect(); + let request = FindMissingBlobsRequest { + instance_name: self.instance_name.clone(), + blob_digests: unique.keys().map(|d| Digest::from(*d)).collect(), + digest_function, + }; + let missing_digests = self + .find_missing_blobs(Request::new(request)) + .await + .err_tip(|| "In GrpcStore::flush_chunk_window")? + .into_inner() + .missing_blob_digests; + + let mut missing: Vec<(DigestInfo, Bytes)> = Vec::with_capacity(missing_digests.len()); + let mut taken: HashSet = HashSet::with_capacity(missing_digests.len()); + for digest in missing_digests { + let digest_info = DigestInfo::try_from(digest) + .err_tip(|| "Invalid missing digest in flush_chunk_window")?; + let Some(data) = unique.remove(&digest_info) else { + // REAPI does not forbid a backend repeating a digest in + // missing_blob_digests; tolerate repeats but still reject + // digests that were never requested. + error_if!( + !taken.contains(&digest_info), + "Backend reported chunk {digest_info} missing that was never requested in flush_chunk_window" + ); + continue; + }; + taken.insert(digest_info); + missing.push((digest_info, data)); + } + for (chunk_digest, data) in unique { + drop(data); + uploader.chunks_deduped.fetch_add(1, Ordering::Relaxed); + uploader + .bytes_deduped + .fetch_add(chunk_digest.size_bytes(), Ordering::Relaxed); + } + if missing.is_empty() { + return Ok(()); + } + + // Pack missing chunks into batch requests under the byte budget. + let mut batch: Vec = Vec::new(); + let mut batch_bytes: u64 = 0; + let mut batches = Vec::new(); + for (chunk_digest, data) in missing { + let entry_cost = data.len() as u64 + BATCH_READ_PER_ENTRY_OVERHEAD_BYTES; + if batch_bytes + entry_cost > CHUNKED_UPLOAD_MAX_BATCH_BYTES && !batch.is_empty() { + batches.push(core::mem::take(&mut batch)); + batch_bytes = 0; + } + uploader.chunks_sent.fetch_add(1, Ordering::Relaxed); + uploader + .bytes_sent + .fetch_add(data.len() as u64, Ordering::Relaxed); + batch.push(batch_update_blobs_request::Request { + digest: Some(chunk_digest.into()), + data, + compressor: 0, + }); + batch_bytes += entry_cost; + } + if !batch.is_empty() { + batches.push(batch); + } + for requests in batches { + let request_entries = requests.clone(); + let response = self + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: self.instance_name.clone(), + requests, + digest_function, + })) + .await + .err_tip(|| "In GrpcStore::flush_chunk_window")?; + for (request, error) in + Self::batch_update_failures(&request_entries, response.into_inner())? + { + self.retry_batch_update_entry(request, digest_function, error) + .await + .err_tip(|| "Retrying failed BatchUpdateBlobs entry")?; + } + } + + // A server may report a transient per-entry error after it has + // durably stored the chunk. Recheck by digest before SpliceBlob so a + // retry never depends on whether the failed response was received + // before or after the write committed. + for requested in requested_digests.chunks(CHUNKED_UPLOAD_WINDOW_CHUNKS) { + let missing = self + .find_missing_blobs(Request::new(FindMissingBlobsRequest { + instance_name: self.instance_name.clone(), + blob_digests: requested.iter().copied().map(Into::into).collect(), + digest_function, + })) + .await + .err_tip(|| "Rechecking chunks after BatchUpdateBlobs")? + .into_inner() + .missing_blob_digests; + if !missing.is_empty() { + return Err(make_err!( + Code::Aborted, + "{} chunk(s) remained missing after BatchUpdateBlobs: {:?}", + missing.len(), + missing + )); + } + } + Ok(()) + } + + /// Digests a window of chunks in parallel on the blocking pool, + /// preserving order. Chunk hashing is the dominant CPU cost of a + /// chunked upload; the chunks are independent so this parallelizes + /// near-linearly. Dropping the returned future merely abandons pure + /// hash work (no side effects), so cancellation stays safe. + async fn hash_chunks_parallel( + hasher_func: DigestHasherFunc, + chunks: Vec, + ) -> Result, Error> { + let num_chunks = chunks.len(); + let parallelism = chunk_hash_parallelism(); + let partition_len = num_chunks.div_ceil(parallelism).max(1); + let mut handles = Vec::with_capacity(parallelism); + for partition in chunks.chunks(partition_len) { + // Bytes clones are refcount bumps, not copies. + let partition: Vec = partition.to_vec(); + handles.push(spawn_blocking!("grpc_chunked_upload_hash", move || { + partition + .into_iter() + .map(|data| { + let mut hasher = hasher_func.hasher(); + hasher.update(&data); + (hasher.finalize_digest(), data) + }) + .collect::>() + })); + } + let mut hashed = Vec::with_capacity(num_chunks); + for handle in handles { + hashed.extend( + handle + .await + .map_err(|e| make_err!(Code::Internal, "Chunk hashing task failed: {e:?}")) + .err_tip(|| "In GrpcStore::hash_chunks_parallel")?, + ); + } + Ok(hashed) + } + + /// Uploads a large blob as content-defined chunks: `FastCDC`-splits the + /// stream, transfers only the chunks the backend is missing, and + /// assembles the blob remotely with `SpliceBlob`. + async fn chunked_update( + &self, + digest: DigestInfo, + reader: DropCloserReadHalf, + expected_size: u64, + hasher_func: DigestHasherFunc, + ) -> Result { + let uploader = self + .chunked_uploader + .as_ref() + .err_tip(|| "chunked_update called without chunked_uploader")?; + uploader + .chunked_uploads_total + .fetch_add(1, Ordering::Relaxed); + let avg_size = u32::try_from(uploader.config.avg_chunk_size_bytes) + .err_tip(|| "avg_chunk_size_bytes did not fit in u32 in chunked_update")?; + let (min_size, max_size) = (avg_size / 4, avg_size * 4); + // The hasher is resolved once in update() — the same source the + // plain path uses for its resource name — so chunk digests, the + // digest_function field, and the blob digest can never diverge. + let digest_function = hasher_func.proto_digest_func() as i32; + + let mut bytes_reader = StreamReader::new(reader); + let mut cdc = AsyncStreamCDC::with_level( + &mut bytes_reader, + min_size, + avg_size, + max_size, + Normalization::Level2, + ); + let mut cdc_stream = core::pin::pin!(cdc.as_stream()); + + let mut all_chunk_digests: Vec = Vec::new(); + let mut window: Vec = Vec::new(); + let mut window_bytes = 0usize; + let mut total_bytes = 0u64; + let mut chunk_count = 0u64; + while let Some(chunk_result) = cdc_stream.next().await { + let chunk = chunk_result + .map_err(|e| make_err!(Code::Internal, "Failed to chunk blob: {e:?}")) + .err_tip(|| "In GrpcStore::chunked_update")?; + total_bytes += chunk.data.len() as u64; + window_bytes += chunk.data.len(); + chunk_count += 1; + window.push(chunk.data.into()); + error_if!( + chunk_count > uploader.config.max_chunk_count, + "Blob {digest} produced more than max_chunk_count ({}) chunks in chunked_update", + uploader.config.max_chunk_count + ); + if window.len() >= CHUNKED_UPLOAD_WINDOW_CHUNKS + || window_bytes >= CHUNKED_UPLOAD_WINDOW_BYTES + { + let hashed = + Self::hash_chunks_parallel(hasher_func, core::mem::take(&mut window)).await?; + all_chunk_digests.extend(hashed.iter().map(|(d, _)| Digest::from(*d))); + self.flush_chunk_window(uploader, hashed, digest_function) + .await?; + window_bytes = 0; + } + } + if !window.is_empty() { + let hashed = Self::hash_chunks_parallel(hasher_func, window).await?; + all_chunk_digests.extend(hashed.iter().map(|(d, _)| Digest::from(*d))); + self.flush_chunk_window(uploader, hashed, digest_function) + .await?; + } + error_if!( + total_bytes != expected_size, + "Chunked upload of {digest} received {total_bytes} bytes, expected {expected_size}" + ); + + let response = match self + .splice_blob(Request::new(SpliceBlobRequest { + instance_name: self.instance_name.clone(), + blob_digest: Some(digest.into()), + chunk_digests: all_chunk_digests, + digest_function, + chunking_function: chunking_function::Value::FastCdc2020.into(), + })) + .await + { + Ok(response) => Some(response), + Err(err) if err.code == Code::AlreadyExists => { + // REAPI permits this when the assembled blob already exists + // and the server keeps a different chunk layout. The + // StoreDriver update has still achieved its contract. + trace!( + %digest, + "SpliceBlob reported that the target blob already exists", + ); + None + } + Err(err) if err.code == Code::NotFound => { + // A chunk was evicted between its upload and the splice. The + // stream is consumed, so this upload cannot be re-run here: + // worker output uploads surface this to their action-level + // retry. + return Err(make_err!( + Code::Aborted, + "Chunk evicted before SpliceBlob completed; re-running the upload may succeed: {err}" + )) + .err_tip(|| "In GrpcStore::chunked_update"); + } + Err(err) => return Err(err).err_tip(|| "In GrpcStore::chunked_update"), + }; + let Some(response) = response else { + return Ok(expected_size); + }; + let response_digest = response + .into_inner() + .blob_digest + .ok_or_else(|| { + make_err!( + Code::Internal, + "SpliceBlob response omitted blob_digest for {digest}" + ) + }) + .and_then(|response_digest| { + DigestInfo::try_from(response_digest).map_err(|error| { + make_err!( + Code::Internal, + "SpliceBlob returned an invalid blob_digest for {digest}: {error}" + ) + }) + })?; + if response_digest != digest { + return Err(make_err!( + Code::Internal, + "SpliceBlob returned blob_digest {response_digest}, expected {digest}" + )); + } + Ok(expected_size) + } } #[async_trait] impl StoreDriver for GrpcStore { async fn post_init(self: Arc) -> Result<(), Error> { + if self.chunked_uploader.is_some() { + self.validate_chunked_upload_capabilities().await?; + } Ok(()) } @@ -1162,7 +1810,7 @@ impl StoreDriver for GrpcStore { self: Pin<&Self>, key: StoreKey<'_>, reader: DropCloserReadHalf, - _size_info: UploadSizeInfo, + size_info: UploadSizeInfo, ) -> Result { struct LocalState { resource_name: String, @@ -1176,6 +1824,27 @@ impl StoreDriver for GrpcStore { return self.update_action_result_from_bytes(digest, reader).await; } + // Resolved once and shared by the chunked and plain paths so their + // digest-function handling can never diverge. + let hasher_func = Context::current() + .get::() + .map_or_else(default_digest_hasher_func, |v| *v); + + // Large blobs take the chunked-upload path when configured: split + // locally, transfer only missing chunks, assemble with SpliceBlob. + // Blobs that could exceed max_chunk_count use plain streaming (a + // chunked upload cannot fall back once the stream is consumed). + if let Some(uploader) = &self.chunked_uploader + && let UploadSizeInfo::ExactSize(expected_size) = size_info + && expected_size >= uploader.config.min_blob_size_bytes + && expected_size.div_ceil((uploader.config.avg_chunk_size_bytes / 4).max(1)) + <= uploader.config.max_chunk_count + { + return self + .chunked_update(digest, reader, expected_size, hasher_func) + .await; + } + let mut buf = Uuid::encode_buffer(); let resource_name = if self.use_legacy_resource_names { format!( @@ -1186,9 +1855,7 @@ impl StoreDriver for GrpcStore { digest.size_bytes(), ) } else { - let digest_function = Context::current() - .get::() - .map_or_else(default_digest_hasher_func, |v| *v) + let digest_function = hasher_func .proto_digest_func() .as_str_name() .to_ascii_lowercase(); diff --git a/nativelink-store/tests/grpc_chunked_upload_test.rs b/nativelink-store/tests/grpc_chunked_upload_test.rs new file mode 100644 index 000000000..21b5c90f1 --- /dev/null +++ b/nativelink-store/tests/grpc_chunked_upload_test.rs @@ -0,0 +1,708 @@ +// 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 core::pin::Pin; +use core::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::collections::HashMap; +use std::sync::Arc; + +use bytes::Bytes; +use futures::Stream; +use nativelink_config::stores::{ + GrpcChunkedUploadsConfig, GrpcEndpoint, GrpcSpec, Retry, StoreType, +}; +use nativelink_error::{Code, Error, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::capabilities_server::{ + Capabilities, CapabilitiesServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{ + ContentAddressableStorage, ContentAddressableStorageServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::{ + BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, + BatchUpdateBlobsResponse, CacheCapabilities, FastCdc2020Params, FindMissingBlobsRequest, + FindMissingBlobsResponse, GetCapabilitiesRequest, GetTreeRequest, GetTreeResponse, + ServerCapabilities, SpliceBlobRequest, SpliceBlobResponse, SplitBlobRequest, SplitBlobResponse, + batch_update_blobs_response, +}; +use nativelink_proto::google::bytestream::byte_stream_server::{ByteStream, ByteStreamServer}; +use nativelink_proto::google::bytestream::{ + QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest, + WriteResponse, +}; +use nativelink_store::grpc_store::GrpcStore; +use nativelink_util::buf_channel::make_buf_channel_pair; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; +use nativelink_util::store_trait::{StoreDriver, StoreLike, UploadSizeInfo}; +use pretty_assertions::assert_eq; +use tokio::sync::Mutex; +use tokio_stream::StreamExt; +use tonic::{Request, Response, Status, Streaming}; + +/// Minimal in-memory CAS implementing exactly what chunked uploads use: +/// `FindMissingBlobs`, `BatchUpdateBlobs`, `SpliceBlob` — plus a `ByteStream` `Write` +/// sink so the plain path is observable. +#[derive(Debug, Default)] +struct FakeCas { + blobs: Mutex>, + batch_update_payload_bytes: AtomicU64, + batch_update_calls: AtomicU64, + batch_update_entry_failures: AtomicU64, + omit_batch_update_response: AtomicBool, + bytestream_writes: AtomicU64, + splice_requests: AtomicU64, + fail_splice_already_exists: AtomicBool, + fail_splice_not_found: AtomicBool, + duplicate_missing_digests: AtomicBool, +} + +/// Local handle type so tonic service traits can be implemented (coherence +/// forbids implementing them for `Arc` directly). +#[derive(Debug, Clone)] +struct FakeCasHandle(Arc); + +impl core::ops::Deref for FakeCasHandle { + type Target = FakeCas; + fn deref(&self) -> &FakeCas { + &self.0 + } +} + +fn digest_key(digest: &nativelink_proto::build::bazel::remote::execution::v2::Digest) -> String { + format!("{}-{}", digest.hash, digest.size_bytes) +} + +#[tonic::async_trait] +impl ContentAddressableStorage for FakeCasHandle { + async fn find_missing_blobs( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let blobs = self.blobs.lock().await; + let mut missing_blob_digests: Vec<_> = request + .blob_digests + .into_iter() + .filter(|digest| !blobs.contains_key(&digest_key(digest))) + .collect(); + // REAPI does not forbid repeated entries; some backends emit them. + if self.duplicate_missing_digests.load(Ordering::Relaxed) { + let duplicates = missing_blob_digests.clone(); + missing_blob_digests.extend(duplicates); + } + Ok(Response::new(FindMissingBlobsResponse { + missing_blob_digests, + })) + } + + async fn batch_update_blobs( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.batch_update_calls.fetch_add(1, Ordering::Relaxed); + let mut blobs = self.blobs.lock().await; + let mut responses = Vec::with_capacity(request.requests.len()); + for entry in request.requests { + let digest = entry.digest.clone().expect("digest must be set"); + self.batch_update_payload_bytes + .fetch_add(entry.data.len() as u64, Ordering::Relaxed); + if self + .batch_update_entry_failures + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |remaining| { + if remaining > 0 { + Some(remaining - 1) + } else { + None + } + }) + .is_ok() + { + responses.push(batch_update_blobs_response::Response { + digest: Some(digest), + status: Some(nativelink_proto::google::rpc::Status { + code: Code::Unavailable as i32, + message: "injected retryable entry failure".to_string(), + details: vec![], + }), + }); + continue; + } + blobs.insert(digest_key(&digest), entry.data); + responses.push(batch_update_blobs_response::Response { + digest: Some(digest), + status: Some(nativelink_proto::google::rpc::Status::default()), + }); + } + if self + .omit_batch_update_response + .swap(false, Ordering::Relaxed) + { + responses.pop(); + } + Ok(Response::new(BatchUpdateBlobsResponse { responses })) + } + + async fn splice_blob( + &self, + request: Request, + ) -> Result, Status> { + self.splice_requests.fetch_add(1, Ordering::Relaxed); + if self.fail_splice_not_found.load(Ordering::Relaxed) { + return Err(Status::not_found("chunk evicted (injected)")); + } + if self.fail_splice_already_exists.load(Ordering::Relaxed) { + return Err(Status::already_exists("blob already exists (injected)")); + } + let request = request.into_inner(); + let blob_digest = request.blob_digest.expect("blob_digest must be set"); + let mut blobs = self.blobs.lock().await; + let mut assembled = Vec::with_capacity(usize::try_from(blob_digest.size_bytes).unwrap()); + for chunk_digest in &request.chunk_digests { + let chunk = blobs + .get(&digest_key(chunk_digest)) + .ok_or_else(|| Status::not_found(format!("chunk {chunk_digest:?} missing")))?; + assembled.extend_from_slice(chunk); + } + if assembled.len().try_into().unwrap_or(i64::MAX) != blob_digest.size_bytes { + return Err(Status::invalid_argument("assembled size mismatch")); + } + blobs.insert(digest_key(&blob_digest), assembled.into()); + Ok(Response::new(SpliceBlobResponse { + blob_digest: Some(blob_digest), + })) + } + + async fn batch_read_blobs( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not needed")) + } + + type GetTreeStream = + Pin> + Send + 'static>>; + async fn get_tree( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not needed")) + } + + async fn split_blob( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not needed")) + } +} + +#[tonic::async_trait] +impl ByteStream for FakeCasHandle { + type ReadStream = Pin> + Send + 'static>>; + async fn read( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not needed")) + } + + async fn write( + &self, + request: Request>, + ) -> Result, Status> { + self.bytestream_writes.fetch_add(1, Ordering::Relaxed); + let mut stream = request.into_inner(); + let mut first_resource_name = None; + let mut data = Vec::new(); + while let Some(message) = stream.next().await { + let message = message?; + if first_resource_name.is_none() { + first_resource_name = Some(message.resource_name.clone()); + } + data.extend_from_slice(&message.data); + if message.finish_write { + break; + } + } + let resource_name = first_resource_name.unwrap_or_default(); + // uploads/{uuid}/blobs/{hash}/{size} + let mut parts = resource_name.split('/').rev(); + let size: i64 = parts.next().unwrap().parse().unwrap(); + let hash = parts.next().unwrap().to_string(); + self.blobs + .lock() + .await + .insert(format!("{hash}-{size}"), Bytes::from(data.clone())); + Ok(Response::new(WriteResponse { + committed_size: data.len().try_into().unwrap_or(i64::MAX), + })) + } + + async fn query_write_status( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not needed")) + } +} + +#[derive(Debug, Clone)] +struct FakeCapabilities { + split_blob_support: bool, + splice_blob_support: bool, + fast_cdc_2020_params: Option, +} + +impl Default for FakeCapabilities { + fn default() -> Self { + Self { + split_blob_support: true, + splice_blob_support: true, + fast_cdc_2020_params: Some(FastCdc2020Params { + avg_chunk_size_bytes: 16 * 1024, + seed: 0, + }), + } + } +} + +#[tonic::async_trait] +impl Capabilities for FakeCapabilities { + async fn get_capabilities( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(ServerCapabilities { + cache_capabilities: Some(CacheCapabilities { + split_blob_support: self.split_blob_support, + splice_blob_support: self.splice_blob_support, + fast_cdc_2020_params: self.fast_cdc_2020_params, + ..Default::default() + }), + ..Default::default() + })) + } +} + +async fn start_fake_cas_with_capabilities( + cas: Arc, + capabilities: FakeCapabilities, +) -> u16 { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + nativelink_util::background_spawn!("fake_cas_server", async move { + tonic::transport::Server::builder() + .add_service(ContentAddressableStorageServer::new(FakeCasHandle( + cas.clone(), + ))) + .add_service(ByteStreamServer::new(FakeCasHandle(cas))) + .add_service(CapabilitiesServer::new(capabilities)) + .serve_with_incoming(incoming) + .await + .unwrap(); + }); + port +} + +async fn start_fake_cas(cas: Arc) -> u16 { + start_fake_cas_with_capabilities(cas, FakeCapabilities::default()).await +} + +fn chunked_spec(port: u16, config: GrpcChunkedUploadsConfig) -> GrpcSpec { + GrpcSpec { + instance_name: String::new(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 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: 60, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_chunked_uploads: Some(config), + } +} + +const fn test_config() -> GrpcChunkedUploadsConfig { + GrpcChunkedUploadsConfig { + min_blob_size_bytes: 1024 * 1024, + avg_chunk_size_bytes: 16 * 1024, + max_chunk_count: 50_000, + } +} + +async fn new_chunked_store( + port: u16, + config: GrpcChunkedUploadsConfig, +) -> Result, Error> { + let store = GrpcStore::new(&chunked_spec(port, config)).await?; + store.clone().post_init().await?; + Ok(store) +} + +fn make_payload(len: usize, seed: u64) -> Vec { + let mut data = vec![0u8; len]; + let mut state = seed; + for word in data.chunks_mut(8) { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + word.copy_from_slice(&state.to_le_bytes()[..word.len()]); + } + data +} + +fn digest_of(data: &[u8]) -> DigestInfo { + let mut hasher = DigestHasherFunc::Sha256.hasher(); + hasher.update(data); + hasher.finalize_digest() +} + +async fn stream_upload(store: &Arc, data: &[u8]) -> Result { + let digest = digest_of(data); + let content = Bytes::copy_from_slice(data); + let (mut tx, rx) = make_buf_channel_pair(); + let blob_len = content.len(); + let send_fut = async move { + let mut offset = 0usize; + while offset < blob_len { + let end = (offset + 64 * 1024).min(blob_len); + tx.send(content.slice(offset..end)).await?; + offset = end; + } + tx.send_eof() + }; + let update_fut = store.update(digest, rx, UploadSizeInfo::ExactSize(blob_len as u64)); + let (send_result, update_result) = tokio::join!(send_fut, update_fut); + send_result?; + update_result?; + Ok(digest) +} + +async fn stored_blob(cas: &FakeCas, digest: DigestInfo) -> Option { + cas.blobs + .lock() + .await + .get(&format!("{}-{}", digest.packed_hash(), digest.size_bytes())) + .cloned() +} + +// A blob below min_blob_size_bytes must use the plain ByteStream path. +#[nativelink_test] +async fn small_blob_uses_plain_path() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + let data = make_payload(256 * 1024, 1); + let digest = stream_upload(&store, &data).await?; + + assert_eq!(cas.bytestream_writes.load(Ordering::Relaxed), 1); + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 0); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} + +// A large blob must be uploaded as chunks and assembled with SpliceBlob, +// byte-identical. +#[nativelink_test] +async fn large_blob_chunk_uploads_and_splices() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + let data = make_payload(4 * 1024 * 1024, 2); + let digest = stream_upload(&store, &data).await?; + + assert_eq!(cas.bytestream_writes.load(Ordering::Relaxed), 0); + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 1); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} + +// Retryable per-entry statuses must use the configured Retrier rather than +// failing the whole chunked upload. +#[nativelink_test] +async fn retryable_batch_update_entry_failure_is_retried() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + cas.batch_update_entry_failures.store(1, Ordering::Relaxed); + let port = start_fake_cas(cas.clone()).await; + let mut spec = chunked_spec(port, test_config()); + spec.retry.max_retries = 1; + let store = GrpcStore::new(&spec).await?; + store.clone().post_init().await?; + + let data = make_payload(2 * 1024 * 1024, 8); + let digest = stream_upload(&store, &data).await?; + + assert!( + cas.batch_update_calls.load(Ordering::Relaxed) >= 2, + "the failed BatchUpdateBlobs entry was not retried" + ); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} + +// A successful BatchUpdateBlobs RPC with incomplete response coverage must +// fail before SpliceBlob is attempted. +#[nativelink_test] +async fn missing_batch_update_response_is_rejected() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + cas.omit_batch_update_response + .store(true, Ordering::Relaxed); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + let data = make_payload(2 * 1024 * 1024, 9); + let error = stream_upload(&store, &data) + .await + .expect_err("missing BatchUpdateBlobs response should fail"); + assert_eq!(error.code, Code::Internal); + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 0); + Ok(()) +} + +// Re-uploading a churned version must transfer only the changed chunks. +#[nativelink_test] +async fn churned_blob_transfers_only_missing_chunks() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + let v1 = make_payload(4 * 1024 * 1024, 3); + stream_upload(&store, &v1).await?; + let v1_payload_bytes = cas.batch_update_payload_bytes.load(Ordering::Relaxed); + + // v2: one contiguous 128KiB region rewritten. + let mut v2 = v1.clone(); + let mutated = make_payload(128 * 1024, 4); + v2[1024 * 1024..1024 * 1024 + 128 * 1024].copy_from_slice(&mutated); + let v2_digest = stream_upload(&store, &v2).await?; + + let v2_payload_bytes = + cas.batch_update_payload_bytes.load(Ordering::Relaxed) - v1_payload_bytes; + assert!( + v2_payload_bytes < v2.len() as u64 / 4, + "churned upload transferred {v2_payload_bytes} bytes, expected far less than {}", + v2.len() + ); + assert_eq!(stored_blob(&cas, v2_digest).await.as_deref(), Some(&v2[..])); + Ok(()) +} + +// A SpliceBlob NotFound (chunk evicted mid-upload) must surface as a +// retryable ABORTED error. +#[nativelink_test] +async fn splice_not_found_maps_to_aborted() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + cas.fail_splice_not_found.store(true, Ordering::Relaxed); + let data = make_payload(2 * 1024 * 1024, 5); + let err = stream_upload(&store, &data) + .await + .expect_err("expected splice failure"); + assert_eq!(err.code, Code::Aborted, "unexpected error: {err:?}"); + Ok(()) +} + +// Blobs that could exceed max_chunk_count must use the plain path (a chunked +// upload cannot fall back once the stream is consumed). +#[nativelink_test] +async fn oversized_chunk_count_uses_plain_path() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas(cas.clone()).await; + let mut config = test_config(); + config.max_chunk_count = 8; + let store = new_chunked_store(port, config).await?; + + let data = make_payload(4 * 1024 * 1024, 6); + let digest = stream_upload(&store, &data).await?; + + assert_eq!(cas.bytestream_writes.load(Ordering::Relaxed), 1); + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 0); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} + +// Invalid configurations must be rejected at construction. +#[nativelink_test] +async fn invalid_configs_rejected() -> Result<(), Error> { + let make = |config| chunked_spec(1, config); + + let too_big_avg = GrpcStore::new(&make(GrpcChunkedUploadsConfig { + min_blob_size_bytes: 8 * 1024 * 1024, + avg_chunk_size_bytes: 1024 * 1024, + max_chunk_count: 50_000, + })) + .await; + assert!(too_big_avg.is_err(), "avg over v1 cap must be rejected"); + + let min_below_max_chunk = GrpcStore::new(&make(GrpcChunkedUploadsConfig { + min_blob_size_bytes: 64 * 1024, + avg_chunk_size_bytes: 512 * 1024, + max_chunk_count: 50_000, + })) + .await; + assert!( + min_below_max_chunk.is_err(), + "min_blob_size below avg*4 must be rejected" + ); + + let zero_chunk_count = GrpcStore::new(&make(GrpcChunkedUploadsConfig { + min_blob_size_bytes: 8 * 1024 * 1024, + avg_chunk_size_bytes: 512 * 1024, + max_chunk_count: 0, + })) + .await; + assert!(zero_chunk_count.is_err(), "zero max_chunk_count rejected"); + drop(make_err!(Code::Ok, "unused")); // keep make_err import used + Ok(()) +} + +// A backend that repeats digests in missing_blob_digests must not fail the +// upload: repeats are tolerated, unknown digests still error. +#[nativelink_test] +async fn duplicate_missing_digests_are_tolerated() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + cas.duplicate_missing_digests.store(true, Ordering::Relaxed); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + + let data = make_payload(12 * 1024 * 1024, 7); + let digest = stream_upload(&store, &data).await?; + + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 1); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} + +// Opting in must fail during startup when the configured backend does not +// advertise both SplitBlob and SpliceBlob support. +#[nativelink_test] +async fn backend_without_chunking_is_rejected_at_startup() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas_with_capabilities( + cas, + FakeCapabilities { + split_blob_support: false, + splice_blob_support: false, + ..Default::default() + }, + ) + .await; + let store = GrpcStore::new(&chunked_spec(port, test_config())).await?; + + let error = store + .post_init() + .await + .expect_err("a backend without chunking support must fail startup"); + assert_eq!(error.code, Code::FailedPrecondition); + let message = error.to_string(); + assert!(message.contains("experimental_chunked_uploads")); + assert!(message.contains("SplitBlob")); + assert!(message.contains("SpliceBlob")); + assert!(message.contains("Enable experimental_chunking")); + Ok(()) +} + +// Startup validation is fail-closed for FastCDC parameters as well: this +// client does not silently negotiate a different average or seed. +#[nativelink_test] +async fn incompatible_fastcdc_parameters_are_rejected_at_startup() -> Result<(), Error> { + for capabilities in [ + FakeCapabilities { + fast_cdc_2020_params: Some(FastCdc2020Params { + avg_chunk_size_bytes: 32 * 1024, + seed: 0, + }), + ..Default::default() + }, + FakeCapabilities { + fast_cdc_2020_params: Some(FastCdc2020Params { + avg_chunk_size_bytes: 16 * 1024, + seed: 1, + }), + ..Default::default() + }, + ] { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas_with_capabilities(cas, capabilities).await; + let store = GrpcStore::new(&chunked_spec(port, test_config())).await?; + let error = store + .post_init() + .await + .expect_err("incompatible FastCDC parameters must fail startup"); + assert_eq!(error.code, Code::FailedPrecondition); + } + Ok(()) +} + +// Stores without the opt-in do not perform a capability check, preserving +// startup behavior for ordinary gRPC stores. +#[nativelink_test] +async fn disabled_chunked_uploads_skip_capability_validation() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + let port = start_fake_cas_with_capabilities( + cas, + FakeCapabilities { + split_blob_support: false, + splice_blob_support: false, + fast_cdc_2020_params: None, + }, + ) + .await; + let mut spec = chunked_spec(port, test_config()); + spec.experimental_chunked_uploads = None; + GrpcStore::new(&spec).await?.post_init().await +} + +// ALREADY_EXISTS is an allowed SpliceBlob result when the target blob is +// already present, so it completes the StoreDriver upload successfully. +#[nativelink_test] +async fn splice_already_exists_completes_upload() -> Result<(), Error> { + let cas = Arc::new(FakeCas::default()); + cas.fail_splice_already_exists + .store(true, Ordering::Relaxed); + let port = start_fake_cas(cas.clone()).await; + let store = new_chunked_store(port, test_config()).await?; + let data = make_payload(2 * 1024 * 1024, 10); + let digest = digest_of(&data); + cas.blobs + .lock() + .await + .insert(digest_key(&digest.into()), Bytes::copy_from_slice(&data)); + + let uploaded_digest = stream_upload(&store, &data).await?; + + assert_eq!(uploaded_digest, digest); + assert_eq!(cas.splice_requests.load(Ordering::Relaxed), 1); + assert_eq!(stored_blob(&cas, digest).await.as_deref(), Some(&data[..])); + Ok(()) +} diff --git a/nativelink-store/tests/grpc_read_batching_test.rs b/nativelink-store/tests/grpc_read_batching_test.rs index 86d419c3b..1f476a1a5 100644 --- a/nativelink-store/tests/grpc_read_batching_test.rs +++ b/nativelink-store/tests/grpc_read_batching_test.rs @@ -276,6 +276,7 @@ async fn make_fixture(read_batching: Option) -> Result Result<(), Error> { use_legacy_resource_names: false, headers: HashMap::new(), forward_headers: vec!["authorization".to_string()], + experimental_chunked_uploads: None, experimental_read_batching: Some(batching_config()), }; let err = GrpcStore::new(&spec) diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index 689307b2a..baa42f6e4 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -61,6 +61,7 @@ fn test_spec>(endpoint: T, use_legacy_resource_names: bool) -> G use_legacy_resource_names, headers: HashMap::new(), forward_headers: vec![], + experimental_chunked_uploads: None, experimental_read_batching: None, } } diff --git a/nativelink-worker/tests/directory_cache_test.rs b/nativelink-worker/tests/directory_cache_test.rs index 681127791..c2309482b 100644 --- a/nativelink-worker/tests/directory_cache_test.rs +++ b/nativelink-worker/tests/directory_cache_test.rs @@ -1601,6 +1601,7 @@ async fn get_tree_prefetch_follows_server_pagination() -> Result<(), Error> { use_legacy_resource_names: false, headers: HashMap::new(), forward_headers: vec![], + experimental_chunked_uploads: None, experimental_read_batching: None, }; let fast_spec = FilesystemSpec { diff --git a/web/apps/docs/content/docs/configuration/chunking.mdx b/web/apps/docs/content/docs/configuration/chunking.mdx index bfa28bda1..b94f82ea6 100644 --- a/web/apps/docs/content/docs/configuration/chunking.mdx +++ b/web/apps/docs/content/docs/configuration/chunking.mdx @@ -19,6 +19,17 @@ chunking blobs on demand with FastCDC 2020 when they were uploaded whole, which is what makes chunked downloads work for outputs produced by remote execution workers. +CDC supports two independent upload paths: + +- External clients such as Bazel opt in with their own chunking flag and use + the CAS `SplitBlob` and `SpliceBlob` services. +- NativeLink workers opt in on their upstream `grpc` CAS store. Their + `StoreDriver` uploads are chunked before crossing the worker-to-CAS link. + +Enabling worker uploads does not rewrite incoming `ByteStream.Write` +requests. Ordinary external ByteStream uploads through a NativeLink proxy +remain streaming writes. + ## Requirements - **Bazel 9.1.1+ or 8.7.0+** on the client, with @@ -29,54 +40,145 @@ execution workers. below, NativeLink behaves exactly as before and does not advertise chunking support, so clients fall back to regular transfers. -## Enabling it +## Supported configurations + +| Data path | Where to enable it | Required backend behavior | +| --- | --- | --- | +| Bazel or another CDC-aware client to a NativeLink CAS | CAS service `experimental_chunking`, a chunk index store, and a capabilities service | NativeLink serves `SplitBlob` and `SpliceBlob` locally | +| NativeLink worker to a remote NativeLink CAS | Worker-side `grpc.experimental_chunked_uploads`; backend CAS `experimental_chunking` and capabilities for the same instance | Backend advertises `SplitBlob`, `SpliceBlob`, FastCDC 2020, the configured average, and seed 0 | +| NativeLink gRPC proxy serving CDC-aware external clients | Frontend CAS `experimental_chunking` without `index_store`; backend chunking and capabilities enabled | Frontend forwards `SplitBlob` and `SpliceBlob`; backend owns chunk layouts | +| Ordinary external ByteStream upload | No CDC setting changes this path | Upload remains an ordinary streaming `ByteStream.Write` | + +Worker-side chunking only applies to CAS `grpc` stores and to +`StoreDriver` uploads that provide an exact size. Blobs below +`min_blob_size_bytes`, blobs whose worst-case chunk count exceeds +`max_chunk_count`, unknown-size uploads, and AC stores do not take the +chunked path. A local filesystem, object, or memory store cannot use the +worker-side option directly. + +NativeLink does not negotiate parameters. At startup, it validates the +support advertised by the backend and fails with a message naming the +missing or mismatched capability. The operator must still configure both +ends with the same average and ensure the worker's `max_chunk_count` does +not exceed the backend limit, because that limit is not advertised by REAPI. + +## Enable client-to-CAS chunking Add an `experimental_chunking` block to the CAS service and give it a small store for chunk layouts. The index store must not verify content digests and must not be the CAS store itself: ```json5 -stores: [ - { - name: "CAS_MAIN_STORE", - // ... your existing CAS store ... - }, - { - // Blob-to-chunks layouts: roughly 80-140 bytes per chunk. - name: "CHUNK_INDEX_STORE", - filesystem: { - content_path: "/tmp/nativelink/data/content_path-chunk-index", - temp_path: "/tmp/nativelink/data/tmp_path-chunk-index", - eviction_policy: { max_bytes: 100000000 }, +// Demonstrates REAPI content-defined chunking: the SplitBlob/SpliceBlob +// RPCs used by Bazel's --experimental_remote_cache_chunking flag +// (available in Bazel 8.7.0+ / 9.1.0+). +// +// Chunking is entirely optional and disabled by default: without the +// `experimental_chunking` block below, NativeLink behaves exactly as +// before and does not advertise chunking support. When enabled, clients +// upload and download large blobs as content-defined chunks, so small +// changes to large outputs only transfer the chunks that changed. +// +// This example configures the CAS service for CDC-aware external clients. +// NativeLink worker/StoreDriver uploads use a separate opt-in: +// `experimental_chunked_uploads` on the worker's upstream grpc CAS store. +// Ordinary external ByteStream uploads remain unchanged. +{ + stores: [ + { + name: "CAS_MAIN_STORE", + filesystem: { + content_path: "/tmp/nativelink/data/content_path-cas", + temp_path: "/tmp/nativelink/data/tmp_path-cas", + eviction_policy: { + // 10gb. + max_bytes: 10000000000, + }, + }, + }, + { + // Holds the blob-to-chunks layouts registered via SpliceBlob or + // created by on-demand chunking in SplitBlob. Layout entries are + // small (roughly 80-140 bytes per chunk). This store must not verify + // content digests and must not be the same store as the CAS itself. + name: "CHUNK_INDEX_STORE", + filesystem: { + content_path: "/tmp/nativelink/data/content_path-chunk-index", + temp_path: "/tmp/nativelink/data/tmp_path-chunk-index", + eviction_policy: { + // 100mb. + max_bytes: 100000000, + }, + }, }, - }, -], -servers: [ - { - // ... - services: { - cas: [ - { - cas_store: "CAS_MAIN_STORE", - experimental_chunking: { - index_store: "CHUNK_INDEX_STORE", - // Optional; the REAPI-recommended default. Must be between - // 1 KiB and 1 MiB. Blobs smaller than 4x this value are - // never chunked. - avg_chunk_size_bytes: 524288, - // Optional; blobs producing more chunks than this are served - // without chunking (~25 GiB at the default average). - max_chunk_count: 50000, + { + name: "AC_MAIN_STORE", + filesystem: { + content_path: "/tmp/nativelink/data/content_path-ac", + temp_path: "/tmp/nativelink/data/tmp_path-ac", + eviction_policy: { + // 500mb. + max_bytes: 500000000, + }, + }, + }, + ], + servers: [ + { + listener: { + http: { + socket_address: "0.0.0.0:50051", + }, + }, + services: { + cas: [ + { + instance_name: "main", + cas_store: "CAS_MAIN_STORE", + + // Optional: omit this block to disable chunking entirely. + experimental_chunking: { + // Required, unless `cas_store` is a grpc store — in that + // case the chunking RPCs are forwarded to the backend and + // `index_store` must be omitted. + index_store: "CHUNK_INDEX_STORE", + + // Optional: the average chunk size in bytes advertised to + // clients and used for server-side chunking. Must be between + // 1 KiB and 1 MiB. + // Default: 524288 (512 KiB). + avg_chunk_size_bytes: 524288, + + // Optional: blobs that would produce more chunks than this + // are served without chunking. + // Default: 50000. + max_chunk_count: 50000, + }, + }, + ], + ac: [ + { + instance_name: "main", + ac_store: "AC_MAIN_STORE", + }, + ], + + // The capabilities service advertises chunking support; Bazel only + // issues SplitBlob/SpliceBlob when it is advertised. + capabilities: [ + { + instance_name: "main", + }, + ], + bytestream: { + cas_stores: { + main: "CAS_MAIN_STORE", }, }, - ], - // The capabilities service advertises chunking support; clients - // only use it when advertised. - capabilities: [{}], - // ... + }, }, - }, -], + ], +} ``` A complete runnable example lives at @@ -92,6 +194,78 @@ bazel build //... \ --experimental_remote_cache_chunking ``` +## Enable NativeLink worker uploads + +When a worker's `cas_fast_slow_store` ultimately uploads to a remote +`grpc` CAS store, add `experimental_chunked_uploads` to that `grpc` store. +If the deployment has opted into CDC but this block is omitted, worker +outputs continue to use full-blob ByteStream uploads: + +```json5 +{ + stores: [ + { + name: "WORKER_FAST_SLOW_STORE", + fast_slow: { + // Keep the worker's normal local filesystem store here. + fast: { + filesystem: { + content_path: "/var/lib/nativelink/worker-cas", + temp_path: "/var/lib/nativelink/worker-cas-tmp", + eviction_policy: { max_bytes: 30000000000 }, + }, + }, + slow: { + grpc: { + instance_name: "main", + endpoints: [{ address: "grpc://cas.internal:50051" }], + store_type: "cas", + + // Use static headers here when the backend requires + // authentication. Worker uploads and the startup check do not + // inherit credentials from an incoming ByteStream request. + headers: { + // authorization: "Bearer ${NATIVELINK_CAS_TOKEN}", + }, + + experimental_chunked_uploads: { + // Only exact-size blobs at or above 8 MiB take the CDC path. + min_blob_size_bytes: "8MiB", + + // Must exactly match the backend's advertised FastCDC value. + avg_chunk_size_bytes: "512KiB", + + // Must not exceed the backend's configured limit. + max_chunk_count: 50000, + }, + }, + }, + }, + }, + ], + workers: [ + { + local: { + cas_fast_slow_store: "WORKER_FAST_SLOW_STORE", + // ...the rest of the worker configuration... + }, + }, + ], +} +``` + +The backend instance named by `instance_name` must expose the capabilities +service and advertise both split and splice support, FastCDC 2020 with seed +0, and the same `avg_chunk_size_bytes`. A missing capabilities service, +unsupported split or splice operation, nonzero seed, or average mismatch +stops NativeLink during startup instead of waiting for the first large +output to fail. + +For blobs that take this path, +`experimental_chunked_uploads` takes precedence over +`experimental_remote_cache_compression`; chunk batches are sent +uncompressed. + ## When it helps, and when it doesn't Chunking pays off when clients reach the cache across a real network (WAN, @@ -100,8 +274,25 @@ uncompressed archives, linked binaries, and container layers typically save 80–90% of transfer bytes per change. It does little on same-rack links — saved bytes only save time when the wire is the bottleneck — and little for compressed artifacts, where everything after the first changed byte -re-transfers. Small blobs (below 4x the average chunk size) are never -chunked, so hot small-object traffic is unaffected. +re-transfers. On the client-to-CAS path, blobs below 4x the average chunk size +are never chunked. Worker uploads instead use `min_blob_size_bytes`, so hot +small-object traffic remains on ordinary `ByteStream.Write`. + +The worker-upload benchmark from the pull request used a 256 MiB artifact, +real gRPC over counted TCP, and byte-verified output: + +| Upload | Plain wire | CDC wire | Wire saved | Plain wall | CDC wall | +| --- | ---: | ---: | ---: | ---: | ---: | +| Cold v1 | 269.2 MB | 268.7 MB | 0.2% | 125 ms | 912 ms | +| v2 with 5% clustered churn | 269.2 MB | 28.3 MB | 89.5% | 112 ms | 737 ms | +| v3 with a 64 KiB insertion | 269.3 MB | 1.07 MB | 99.6% | 139 ms | 706 ms | +| Unchanged v3 re-upload | 269.3 MB | 67 KB | 99.97% | 129 ms | 201 ms | + +The wall-time columns are intentional: on this fast local link, CDC reduced +wire bytes but added hashing, boundary detection, and RPC work. It is a +bandwidth optimization, not an unconditional latency optimization. Clustered +changes benefit because most chunk boundaries remain stable; uniformly +scattered changes can touch every chunk and save little or nothing. The server verifies every spliced blob's digest before committing it and materializes the full blob, so non-chunking clients and every existing read