From 8547c95ea6f619a1050ff2d150ab6f8284d53511 Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Tue, 21 Jul 2026 15:58:40 -0700 Subject: [PATCH 1/2] Batch small worker output uploads into BatchUpdateBlobs Worker output upload opens one ByteStream Write stream per output file, paying a fixed per-stream cost for every small blob. This adds StoreDriver::update_many (default: loop over update_oneshot, zero behavior change), an opt-in GrpcSpec.experimental_write_batching that packs small blobs into BatchUpdateBlobs RPCs (3MiB budget with a 256-byte per-entry overhead charge, digest dedup within a call, per-entry status isolation: retryable entry errors fall back to the streaming path, non-retryable propagate), a FastSlowStore::update_many that registers in-flight slow writes and forwards the batch to both tiers, and a running_actions_manager seam that queues output files at or below 128KiB and publishes them with batched update_many calls. The worker seam only activates when the store chain advertises the new StoreOptimizations::SubscribesToUpdateMany (FastSlowStore forwards its slow tier's advertisement); with the flag unset every store keeps the existing per-file streaming path. Measured (real gRPC over TCP, byte-verified): BatchUpdateBlobs vs stream-per-blob 15.5x at 4KiB x2048 (30us -> 1us/blob), 3.6x at 32KiB, 1.4x at 256KiB, on a zero-RTT loopback floor; real network RTT widens the gap. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UXVtatcR9YMecBiu9RjwpC --- .../vocabularies/TraceMachina/accept.txt | 1 + nativelink-config/src/stores.rs | 55 ++ nativelink-service/tests/cas_server_test.rs | 1 + nativelink-store/BUILD.bazel | 1 + nativelink-store/src/cache_metrics_store.rs | 26 + nativelink-store/src/fast_slow_store.rs | 107 ++- nativelink-store/src/grpc_store.rs | 204 +++++- .../tests/grpc_read_batching_test.rs | 2 + nativelink-store/tests/grpc_store_test.rs | 1 + .../tests/grpc_write_batching_test.rs | 650 ++++++++++++++++++ nativelink-util/src/store_trait.rs | 38 + .../src/running_actions_manager.rs | 124 +++- .../tests/directory_cache_test.rs | 1 + .../tests/running_actions_manager_test.rs | 255 ++++++- 14 files changed, 1446 insertions(+), 20 deletions(-) create mode 100644 nativelink-store/tests/grpc_write_batching_test.rs diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index b3da2fbc2..1e59fb0bd 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -96,6 +96,7 @@ Pantsbuild [Ss]andboxing [Cc]onfig bytestream +[Bb]atcher [Ff]ailover proto quantiles diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 8aca04747..b46cb97bb 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1465,6 +1465,61 @@ pub struct GrpcSpec { /// Default: unset (disabled). When unset there is zero behavior change. #[serde(default)] pub experimental_read_batching: Option, + + /// Experimental: batch small-blob uploads issued through `update_many` + /// into `BatchUpdateBlobs` RPCs instead of one `ByteStream` `Write` + /// stream per blob. This amortizes the per-stream fixed cost when many + /// small objects are published at once (e.g. worker output files). + /// + /// Only blobs at or below `max_blob_size_bytes` are batched; larger + /// blobs and callers using the streaming `update` path are unaffected. + /// Unlike `experimental_read_batching` this does not coalesce across + /// callers, so it is compatible with `forward_headers`. + /// + /// Default: unset (disabled). When unset there is zero behavior change. + #[serde(default)] + pub experimental_write_batching: Option, +} + +/// Configuration for experimental small-blob upload batching in a gRPC +/// store. See [`GrpcSpec::experimental_write_batching`]. +#[derive(Serialize, Deserialize, Debug, Clone, Copy)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct GrpcWriteBatchingConfig { + /// Only blobs at or below this size (in bytes) are eligible for + /// batching. Larger blobs always use the `ByteStream` `Write` path. + /// + /// Default: 131072 (128 KiB). + #[serde( + default = "default_write_batching_max_blob_size_bytes", + deserialize_with = "convert_data_size_with_shellexpand" + )] + pub max_blob_size_bytes: u64, + + /// Maximum total payload bytes packed into a single `BatchUpdateBlobs` + /// request. This should leave headroom under the 4 MiB default gRPC + /// message limit for protobuf framing overhead. + /// + /// Default: 3145728 (3 MiB). + #[serde( + default = "default_write_batching_max_batch_bytes", + deserialize_with = "convert_data_size_with_shellexpand" + )] + pub max_batch_bytes: u64, +} + +/// Default value of [`GrpcWriteBatchingConfig::max_blob_size_bytes`]. +/// Exported so call sites that pre-filter blobs for batching (e.g. the +/// worker's small-output batcher) cannot drift from the config default. +pub const DEFAULT_WRITE_BATCHING_MAX_BLOB_SIZE_BYTES: u64 = 128 * 1024; + +const fn default_write_batching_max_blob_size_bytes() -> u64 { + DEFAULT_WRITE_BATCHING_MAX_BLOB_SIZE_BYTES +} + +const fn default_write_batching_max_batch_bytes() -> u64 { + 3 * 1024 * 1024 // 3 MiB. } /// Configuration for experimental small-blob read coalescing in a gRPC diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index c786a7730..fd712cbc0 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, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + // Batch metrics are all-or-nothing: a partially successful batch + // counts every item as a write error because per-entry outcomes are + // not visible at this layer (accepted approximation). + let start = Instant::now(); + let num_items = items.len() as u64; + let total_bytes: u64 = items.iter().map(|(_, data)| data.len() as u64).sum(); + let result = self.backend.update_many(items).await; + if result.is_ok() { + CACHE_METRICS + .cache_operations + .add(num_items, self.attrs.write_success()); + self.record_write_io(Some(total_bytes)); + self.record_duration(start, self.attrs.write_success()); + } else { + CACHE_METRICS + .cache_operations + .add(num_items, self.attrs.write_error()); + self.record_duration(start, self.attrs.write_error()); + } + result + } + fn optimized_for(&self, optimization: StoreOptimizations) -> bool { self.backend.optimized_for(optimization) } diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index d1f76db7c..02b588d46 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -23,6 +23,7 @@ use std::ffi::OsString; use std::sync::{Arc, Weak}; use async_trait::async_trait; +use bytes::Bytes; use futures::stream::{FuturesUnordered, StreamExt}; use futures::{FutureExt, join, try_join}; use nativelink_config::stores::{FastSlowSpec, StoreDirection}; @@ -154,6 +155,9 @@ impl Drop for LoaderGuard<'_> { } } +/// Maximum concurrent fast-tier writes for one `update_many` batch. +const FAST_TIER_UPDATE_CONCURRENCY: usize = 16; + impl FastSlowStore { pub fn new(spec: &FastSlowSpec, fast_store: Store, slow_store: Store) -> Arc { Arc::new_cyclic(|weak_self| Self { @@ -634,9 +638,108 @@ impl StoreDriver for FastSlowStore { data_stream_res.merge(fast_res).merge(slow_res) } - /// `FastSlowStore` has optimizations for dealing with files. + /// `FastSlowStore` has optimizations for dealing with files, and + /// advertises batched uploads when its slow tier can batch them. fn optimized_for(&self, optimization: StoreOptimizations) -> bool { - optimization == StoreOptimizations::FileUpdates + match optimization { + StoreOptimizations::FileUpdates => true, + StoreOptimizations::SubscribesToUpdateMany => { + self.slow_store + .optimized_for(StoreOptimizations::SubscribesToUpdateMany) + && self.slow_direction != StoreDirection::ReadOnly + && self.slow_direction != StoreDirection::Get + } + _ => false, + } + } + + /// Batched variant of `update()`: publishes every item to both tiers, + /// letting the slow tier amortize per-object wire costs via its own + /// `update_many` implementation. Only used on the batched path when the + /// slow tier subscribes to it; otherwise the default per-item loop + /// preserves `update()` semantics exactly. + /// + /// Note: in-flight slow-write guards are held until the entire batch + /// settles, so concurrent `has()` calls on any batch member block until + /// the whole batch completes (coarser than `update()`'s per-key window). + async fn update_many( + self: Pin<&Self>, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + let ignore_slow = self + .slow_store + .inner_store(None::) + .optimized_for(StoreOptimizations::NoopUpdates) + || self.slow_direction == StoreDirection::ReadOnly + || self.slow_direction == StoreDirection::Get; + let ignore_fast = self + .fast_store + .inner_store(None::) + .optimized_for(StoreOptimizations::NoopUpdates) + || self.fast_direction == StoreDirection::ReadOnly + || self.fast_direction == StoreDirection::Get; + let slow_can_batch = !ignore_slow + && self + .slow_store + .optimized_for(StoreOptimizations::SubscribesToUpdateMany); + if !slow_can_batch { + // Per-item updates through `update()` keep every existing + // semantic (noop/direction handling, in-flight registration). + for (key, data) in items { + self.as_store_driver_pin() + .update_oneshot(key, data) + .await + .err_tip(|| "In FastSlowStore::update_many fallback")?; + } + return Ok(()); + } + + // Make the in-flight slow writes visible to concurrent has() calls, + // exactly like update() does per key. + let guards: Vec<(InFlightSlowWriteGuard, u64)> = items + .iter() + .map(|(key, data)| { + ( + self.register_in_flight_slow_write(key.borrow()), + data.len() as u64, + ) + }) + .collect(); + + let slow_items: Vec<(StoreKey<'static>, Bytes)> = items + .iter() + .map(|(key, data)| (key.clone(), data.clone())) + .collect(); + let slow_fut = self.slow_store.update_many(slow_items); + let fast_fut = async { + if ignore_fast { + return Ok(()); + } + // The fast tier has no batched implementation; write with + // bounded concurrency instead of the serial default loop. + // Plain loop instead of an iterator closure: async closures + // over borrowed keys hit HRTB inference errors. + let fast_store = &self.fast_store; + let mut write_futures = Vec::with_capacity(items.len()); + for (key, data) in items { + write_futures.push(async move { fast_store.update_oneshot(key, data).await }); + } + let mut writes = + futures::stream::iter(write_futures).buffer_unordered(FAST_TIER_UPDATE_CONCURRENCY); + while let Some(result) = writes.next().await { + result.err_tip(|| "In FastSlowStore::update_many fast store item")?; + } + Ok::<(), Error>(()) + }; + let (slow_res, fast_res) = tokio::join!(slow_fut, fast_fut); + if slow_res.is_ok() { + for (guard, size) in guards { + guard.complete(Some(size)); + } + } + slow_res + .err_tip(|| "In FastSlowStore::update_many slow store") + .merge(fast_res.err_tip(|| "In FastSlowStore::update_many fast store")) } /// Optimized variation to consume the file if one of the stores is a diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index ef6c7fd7e..ddf0a15ae 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -16,14 +16,14 @@ use core::pin::Pin; use core::sync::atomic::{AtomicU64, Ordering}; use core::time::Duration; use std::borrow::Cow; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::{Arc, Weak}; use async_trait::async_trait; use bytes::{Bytes, BytesMut}; use futures::stream::{FuturesUnordered, unfold}; use futures::{Future, Stream, StreamExt, TryFutureExt, TryStreamExt, future}; -use nativelink_config::stores::{GrpcReadBatchingConfig, GrpcSpec}; +use nativelink_config::stores::{GrpcReadBatchingConfig, GrpcSpec, GrpcWriteBatchingConfig}; use nativelink_error::{Error, ResultExt, error_if, make_err}; use nativelink_metric::MetricsComponent; use nativelink_proto::build::bazel::remote::execution::v2::action_cache_client::ActionCacheClient; @@ -32,7 +32,7 @@ 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, batch_update_blobs_request, }; use nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient; use nativelink_proto::google::bytestream::{ @@ -49,7 +49,9 @@ use nativelink_util::proto_stream_utils::{ }; use nativelink_util::resource_info::ResourceInfo; use nativelink_util::retry::{Retrier, RetryResult}; -use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; +use nativelink_util::store_trait::{ + RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo, +}; use nativelink_util::telemetry::ClientHeaders; use nativelink_util::{background_spawn, default_health_status_indicator, tls_utils}; use opentelemetry::context::Context; @@ -61,7 +63,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,10 +110,18 @@ fn enrich_request( request } -/// 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. -const BATCH_READ_PER_ENTRY_OVERHEAD_BYTES: u64 = 256; +/// Estimated per-entry protobuf and framing overhead charged against a +/// batch's byte budget, so that batches of many tiny blobs cannot push a +/// `BatchReadBlobs`/`BatchUpdateBlobs` message over the gRPC size limit. +const BATCH_PER_ENTRY_OVERHEAD_BYTES: u64 = 256; + +/// Number of `BatchUpdateBlobs` RPCs `update_many` dispatches concurrently. +const BATCH_WRITE_DISPATCH_CONCURRENCY: usize = 4; + +/// Maximum concurrent per-blob streaming uploads used by `update_many`'s +/// non-batched paths (disabled batching and fallback entries), restoring +/// the concurrency callers previously got from per-file upload streams. +const STREAMING_UPLOAD_CONCURRENCY: usize = 32; /// A small-blob read waiting to be coalesced into a `BatchReadBlobs` RPC. #[derive(Debug)] @@ -211,6 +221,10 @@ pub struct GrpcStore { /// RPCs. `None` means reads always use the `ByteStream` `Read` path. #[metric(group = "read_batcher")] read_batcher: Option, + /// When configured, `update_many` packs small blobs into + /// `BatchUpdateBlobs` RPCs. `None` means every blob uses the + /// `ByteStream` `Write` path. + write_batching: Option, /// Used by the read coalescer to hand a strong reference of this store /// to detached dispatcher tasks. weak_self: Weak, @@ -258,6 +272,19 @@ impl GrpcStore { None => None, }; + if let Some(config) = &spec.experimental_write_batching { + // A blob at the batching threshold must fit in one batch, + // otherwise an over-budget request bypasses the per-entry + // streaming fallback and fails the whole call. + error_if!( + config + .max_blob_size_bytes + .saturating_add(BATCH_PER_ENTRY_OVERHEAD_BYTES) + > config.max_batch_bytes, + "experimental_write_batching.max_batch_bytes must be at least max_blob_size_bytes plus the {BATCH_PER_ENTRY_OVERHEAD_BYTES}-byte per-entry overhead" + ); + } + 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 +319,7 @@ impl GrpcStore { rpc_timeout, use_legacy_resource_names: spec.use_legacy_resource_names, read_batcher, + write_batching: spec.experimental_write_batching, headers, // We lowercase keys as HTTP headers are case-insensitive so we should match all cases forward_headers: spec @@ -392,6 +420,27 @@ impl GrpcStore { .await } + /// Uploads items through the regular streaming path with bounded + /// concurrency. Used by `update_many` when write batching is disabled + /// and as the fallback for entries that could not be batched. + async fn update_items_streaming( + self: Pin<&Self>, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + // Plain loop instead of an iterator closure: async closures over + // borrowed keys hit HRTB inference errors. + let mut upload_futures = Vec::with_capacity(items.len()); + for (key, data) in items { + upload_futures.push(async move { self.update_oneshot(key, data).await }); + } + let mut uploads = + futures::stream::iter(upload_futures).buffer_unordered(STREAMING_UPLOAD_CONCURRENCY); + while let Some(result) = uploads.next().await { + result.err_tip(|| "In GrpcStore::update_items_streaming")?; + } + Ok(()) + } + pub async fn batch_read_blobs( &self, grpc_request: Request, @@ -523,7 +572,7 @@ impl GrpcStore { let item_cost = item .digest .size_bytes() - .saturating_add(BATCH_READ_PER_ENTRY_OVERHEAD_BYTES); + .saturating_add(BATCH_PER_ENTRY_OVERHEAD_BYTES); if item.digest_function == digest_function && (batch.is_empty() || batch_bytes.saturating_add(item_cost) <= batcher.max_batch_bytes) @@ -1091,6 +1140,141 @@ impl StoreDriver for GrpcStore { Ok(()) } + fn optimized_for(&self, optimization: StoreOptimizations) -> bool { + optimization == StoreOptimizations::SubscribesToUpdateMany + && self.write_batching.is_some() + && !matches!(self.store_type, nativelink_config::stores::StoreType::Ac) + } + + async fn update_many( + self: Pin<&Self>, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + let Some(config) = self.write_batching else { + return self.update_items_streaming(items).await; + }; + if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) { + return self.update_items_streaming(items).await; + } + + // Partition into batchable small blobs and streaming fallbacks. + // Duplicate digests are uploaded once: CAS content is identical by + // definition and servers may reject duplicates within one request. + let mut seen_digests: HashSet = HashSet::with_capacity(items.len()); + let mut batchable: Vec<(DigestInfo, Bytes)> = Vec::with_capacity(items.len()); + let mut fallback: Vec<(StoreKey<'static>, Bytes)> = Vec::new(); + for (key, data) in items { + let digest = key.borrow().into_digest(); + if data.len() as u64 > config.max_blob_size_bytes { + fallback.push((key, data)); + } else if seen_digests.insert(digest) { + batchable.push((digest, data)); + } + } + + let digest_function: i32 = Context::current() + .get::() + .map_or_else(default_digest_hasher_func, |v| *v) + .proto_digest_func() + .into(); + + // Pack under the batch byte budget. + let mut chunks: Vec> = Vec::new(); + let mut chunk: Vec<(DigestInfo, Bytes)> = Vec::new(); + let mut chunk_bytes = 0u64; + for (digest, data) in batchable { + let entry_cost = data.len() as u64 + BATCH_PER_ENTRY_OVERHEAD_BYTES; + if !chunk.is_empty() && chunk_bytes + entry_cost > config.max_batch_bytes { + chunks.push(core::mem::take(&mut chunk)); + chunk_bytes = 0; + } + chunk.push((digest, data)); + chunk_bytes += entry_cost; + } + if !chunk.is_empty() { + chunks.push(chunk); + } + + // Dispatch chunk RPCs concurrently; each task returns the items of + // that chunk that must fall back to the streaming path. + let mut chunk_tasks = futures::stream::iter(chunks.into_iter().map(|chunk| async move { + let request = BatchUpdateBlobsRequest { + // batch_update_blobs() overwrites the instance name. + instance_name: String::new(), + requests: chunk + .iter() + .map(|(digest, data)| batch_update_blobs_request::Request { + digest: Some((*digest).into()), + data: data.clone(), + compressor: 0, + }) + .collect(), + digest_function, + }; + let response = match self.batch_update_blobs(Request::new(request)).await { + Ok(response) => response.into_inner(), + // A whole-RPC failure (already through the retrier) falls + // back to per-blob streams: an intermediary's message-size + // limit can reject the batch while individual streams + // would succeed. + Err(err) => { + debug!( + ?err, + "BatchUpdateBlobs failed as a whole; falling back to streaming uploads for this chunk", + ); + return Ok(chunk + .into_iter() + .map(|(digest, data)| (StoreKey::from(digest), data)) + .collect::>()); + } + }; + + let mut error_by_digest: HashMap> = + HashMap::with_capacity(response.responses.len()); + for entry in response.responses { + let Some(Ok(digest)) = entry.digest.map(DigestInfo::try_from) else { + continue; + }; + let entry_error = entry + .status + .filter(|status| status.code != 0) + .map(Error::from); + error_by_digest.insert(digest, entry_error); + } + let mut chunk_fallback = Vec::new(); + for (digest, data) in chunk { + match error_by_digest.remove(&digest) { + // Entry succeeded. + Some(None) => {} + Some(Some(err)) if is_retryable_code(err.code) => { + trace!( + ?digest, + ?err, + "Batched upload entry failed with retryable error, falling back to ByteStream write", + ); + chunk_fallback.push((digest.into(), data)); + } + Some(Some(err)) => { + return Err(err.append(format!( + "in BatchUpdateBlobs response for {digest} in GrpcStore::update_many" + ))); + } + // Server omitted the entry; fall back to the streaming + // path rather than guessing at its state. + None => chunk_fallback.push((digest.into(), data)), + } + } + Ok::<_, Error>(chunk_fallback) + })) + .buffer_unordered(BATCH_WRITE_DISPATCH_CONCURRENCY); + while let Some(chunk_fallback) = chunk_tasks.next().await { + fallback.extend(chunk_fallback?); + } + drop(chunk_tasks); + + self.update_items_streaming(fallback).await + } + // NOTE: This function can only be safely used on CAS stores. AC stores may return a size that // is incorrect. async fn has_with_results( diff --git a/nativelink-store/tests/grpc_read_batching_test.rs b/nativelink-store/tests/grpc_read_batching_test.rs index 86d419c3b..a431888d8 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_write_batching: None, }; let err = GrpcStore::new(&spec) .await diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index e120402d4..d5c492352 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_write_batching: None, } } diff --git a/nativelink-store/tests/grpc_write_batching_test.rs b/nativelink-store/tests/grpc_write_batching_test.rs new file mode 100644 index 000000000..964db9a22 --- /dev/null +++ b/nativelink-store/tests/grpc_write_batching_test.rs @@ -0,0 +1,650 @@ +// 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 std::collections::HashMap; +use std::sync::Arc; + +use async_lock::Mutex; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use nativelink_config::stores::{ + CacheMetricsSpec, FastSlowSpec, GrpcEndpoint, GrpcSpec, GrpcWriteBatchingConfig, MemorySpec, + NoopSpec, Retry, StoreSpec, StoreType, +}; +use nativelink_error::{Code, Error}; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{ + ContentAddressableStorage, ContentAddressableStorageServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::{ + BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, + BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse, GetTreeRequest, + GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse, SplitBlobRequest, SplitBlobResponse, + 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_proto::google::rpc::Status as RpcStatus; +use nativelink_store::cache_metrics_store::CacheMetricsStore; +use nativelink_store::fast_slow_store::FastSlowStore; +use nativelink_store::grpc_store::GrpcStore; +use nativelink_store::memory_store::MemoryStore; +use nativelink_util::background_spawn; +use nativelink_util::common::DigestInfo; +use nativelink_util::store_trait::{Store, StoreKey, StoreLike, StoreOptimizations}; +use tonic::transport::Server; +use tonic::transport::server::TcpIncoming; +use tonic::{Request, Response, Status, Streaming}; + +/// A fake CAS server that records every `BatchUpdateBlobs` request and +/// stores its blobs. Digests listed in `error_hashes` are answered with the +/// configured per-entry status code instead of being stored. +#[derive(Debug, Clone)] +struct FakeCasServer { + blobs: Arc>>, + error_hashes: Arc>>, + batch_update_requests: Arc>>, + /// When set, every `BatchUpdateBlobs` RPC fails as a whole. + fail_batch_rpc: Arc>, +} + +impl FakeCasServer { + fn new() -> Self { + Self { + blobs: Arc::new(Mutex::new(HashMap::new())), + error_hashes: Arc::new(Mutex::new(HashMap::new())), + batch_update_requests: Arc::new(Mutex::new(vec![])), + fail_batch_rpc: Arc::new(Mutex::new(false)), + } + } +} + +type GetTreeStream = Pin> + Send + 'static>>; + +#[tonic::async_trait] +impl ContentAddressableStorage for FakeCasServer { + type GetTreeStream = GetTreeStream; + + #[allow(clippy::unimplemented)] + async fn find_missing_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + async fn batch_update_blobs( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + self.batch_update_requests + .lock() + .await + .push(request.clone()); + if *self.fail_batch_rpc.lock().await { + return Err(Status::resource_exhausted( + "Injected whole-RPC failure (e.g. message-size limit)", + )); + } + + let mut blobs = self.blobs.lock().await; + let error_hashes = self.error_hashes.lock().await; + let mut responses = Vec::with_capacity(request.requests.len()); + for entry in request.requests { + let Some(digest) = entry.digest else { + return Err(Status::invalid_argument("Missing digest in request")); + }; + let status = if let Some(&code) = error_hashes.get(&digest.hash) { + RpcStatus { + code, + message: format!("Injected error for {}", digest.hash), + details: vec![], + } + } else { + blobs.insert(digest.hash.clone(), entry.data); + RpcStatus { + code: Code::Ok as i32, + message: String::new(), + details: vec![], + } + }; + responses.push(batch_update_blobs_response::Response { + digest: Some(digest), + status: Some(status), + }); + } + Ok(Response::new(BatchUpdateBlobsResponse { responses })) + } + + #[allow(clippy::unimplemented)] + async fn batch_read_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn get_tree( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn split_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn splice_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +type ReadStream = Pin> + Send + 'static>>; + +/// A fake `ByteStream` server that accepts uploads into the same blob map as +/// [`FakeCasServer`] and records the resource name of every `Write` stream. +#[derive(Debug, Clone)] +struct FakeByteStreamServer { + blobs: Arc>>, + write_resource_names: Arc>>, +} + +impl FakeByteStreamServer { + fn new(blobs: Arc>>) -> Self { + Self { + blobs, + write_resource_names: Arc::new(Mutex::new(vec![])), + } + } +} + +#[tonic::async_trait] +impl ByteStream for FakeByteStreamServer { + type ReadStream = ReadStream; + + #[allow(clippy::unimplemented)] + async fn read( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + async fn write( + &self, + grpc_request: Request>, + ) -> Result, Status> { + let mut stream = grpc_request.into_inner(); + let mut resource_name = String::new(); + let mut data = Vec::new(); + while let Some(request) = stream.next().await { + let request = request?; + if resource_name.is_empty() && !request.resource_name.is_empty() { + resource_name.clone_from(&request.resource_name); + } + data.extend_from_slice(&request.data); + if request.finish_write { + break; + } + } + // Resource names look like `{instance}/uploads/{uuid}/blobs/{hash}/{size}`. + let mut components = resource_name.rsplit('/'); + let _size = components.next(); + let hash = components.next().unwrap_or_default().to_string(); + self.write_resource_names + .lock() + .await + .push(resource_name.clone()); + let committed_size = + i64::try_from(data.len()).map_err(|_| Status::invalid_argument("Upload too large"))?; + self.blobs.lock().await.insert(hash, Bytes::from(data)); + Ok(Response::new(WriteResponse { committed_size })) + } + + #[allow(clippy::unimplemented)] + async fn query_write_status( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +struct TestFixture { + cas_server: FakeCasServer, + bytestream_server: FakeByteStreamServer, + store: Arc, +} + +async fn make_fixture( + write_batching: Option, +) -> Result { + let cas_server = FakeCasServer::new(); + let bytestream_server = FakeByteStreamServer::new(cas_server.blobs.clone()); + let listener = TcpIncoming::bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let port = listener.local_addr().unwrap().port(); + + let cas_service = ContentAddressableStorageServer::new(cas_server.clone()); + let bytestream_service = ByteStreamServer::new(bytestream_server.clone()); + background_spawn!("grpc_write_batching_test_server", async move { + Server::builder() + .add_service(cas_service) + .add_service(bytestream_service) + .serve_with_incoming(listener) + .await + .unwrap(); + }); + + let spec = GrpcSpec { + instance_name: String::new(), + endpoints: vec![GrpcEndpoint { + address: format!("http://localhost:{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: 0, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_write_batching: write_batching, + }; + let store = GrpcStore::new(&spec).await?; + Ok(TestFixture { + cas_server, + bytestream_server, + store, + }) +} + +const fn batching_config() -> GrpcWriteBatchingConfig { + GrpcWriteBatchingConfig { + max_blob_size_bytes: 128 * 1024, + max_batch_bytes: 3 * 1024 * 1024, + } +} + +/// Creates a unique digest and 64-byte content for blob index `i`. +fn make_blob(i: usize) -> (DigestInfo, Bytes) { + let hash = format!("{:064x}", i + 1); + let content = Bytes::from(format!("{i:0>64}")); + let digest = DigestInfo::try_new(&hash, content.len()).unwrap(); + (digest, content) +} + +fn make_items(blobs: &[(DigestInfo, Bytes)]) -> Vec<(StoreKey<'static>, Bytes)> { + blobs + .iter() + .map(|(digest, content)| (StoreKey::from(*digest), content.clone())) + .collect() +} + +// Many small uploads through update_many must be packed into BatchUpdateBlobs +// RPCs and never open a ByteStream Write stream. +#[nativelink_test] +async fn small_blob_updates_are_batched() -> Result<(), Error> { + const NUM_BLOBS: usize = 200; + + let fixture = make_fixture(Some(batching_config())).await?; + let blobs: Vec<_> = (0..NUM_BLOBS).map(make_blob).collect(); + + fixture.store.update_many(make_items(&blobs)).await?; + + for (digest, content) in &blobs { + let stored = fixture + .cas_server + .blobs + .lock() + .await + .get(&digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored.as_ref(), Some(content), "for {digest}"); + } + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!(batch_requests.len(), 1, "expected one batched RPC"); + assert_eq!(batch_requests[0].requests.len(), NUM_BLOBS); + let stream_writes = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!(stream_writes.len(), 0, "no ByteStream writes expected"); + Ok(()) +} + +// Duplicate digests within one update_many call must be uploaded exactly once. +#[nativelink_test] +async fn duplicate_digests_are_uploaded_once() -> Result<(), Error> { + let fixture = make_fixture(Some(batching_config())).await?; + let (digest, content) = make_blob(7); + let items = vec![ + (StoreKey::from(digest), content.clone()), + (StoreKey::from(digest), content.clone()), + (StoreKey::from(digest), content.clone()), + ]; + + fixture.store.update_many(items).await?; + + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!(batch_requests.len(), 1); + assert_eq!( + batch_requests[0].requests.len(), + 1, + "digest must be deduped" + ); + Ok(()) +} + +// Blobs above max_blob_size_bytes must use the ByteStream Write path even +// when batching is enabled. +#[nativelink_test] +async fn large_blobs_fall_back_to_streaming() -> Result<(), Error> { + let config = GrpcWriteBatchingConfig { + max_blob_size_bytes: 64, + max_batch_bytes: 3 * 1024 * 1024, + }; + let fixture = make_fixture(Some(config)).await?; + + let (small_digest, small_content) = make_blob(1); + let large_content = Bytes::from(vec![0xabu8; 1024]); + let large_hash = format!("{:064x}", 0xdead_beefu64); + let large_digest = DigestInfo::try_new(&large_hash, large_content.len()).unwrap(); + + fixture + .store + .update_many(vec![ + (StoreKey::from(small_digest), small_content.clone()), + (StoreKey::from(large_digest), large_content.clone()), + ]) + .await?; + + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!(batch_requests.len(), 1); + assert_eq!(batch_requests[0].requests.len(), 1, "only the small blob"); + let stream_writes = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!(stream_writes.len(), 1, "large blob must stream"); + let stored_large = fixture + .cas_server + .blobs + .lock() + .await + .get(&large_digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored_large, Some(large_content)); + Ok(()) +} + +// A retryable per-entry error must fall back to the ByteStream Write path +// for that entry while its batch peers succeed. +#[nativelink_test] +async fn retryable_entry_error_falls_back_to_streaming() -> Result<(), Error> { + let fixture = make_fixture(Some(batching_config())).await?; + let blobs: Vec<_> = (0..3).map(make_blob).collect(); + fixture.cas_server.error_hashes.lock().await.insert( + blobs[1].0.packed_hash().to_string(), + Code::Unavailable as i32, + ); + + fixture.store.update_many(make_items(&blobs)).await?; + + let stream_writes = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!( + stream_writes.len(), + 1, + "only the failed entry falls back to streaming" + ); + assert!( + stream_writes[0].contains(&blobs[1].0.packed_hash().to_string()), + "fallback must be for the failed digest" + ); + // All blobs must be durable in the end. + for (digest, content) in &blobs { + let stored = fixture + .cas_server + .blobs + .lock() + .await + .get(&digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored.as_ref(), Some(content), "for {digest}"); + } + Ok(()) +} + +// A non-retryable per-entry error must fail the whole update_many call. +#[nativelink_test] +async fn non_retryable_entry_error_propagates() -> Result<(), Error> { + let fixture = make_fixture(Some(batching_config())).await?; + let blobs: Vec<_> = (0..3).map(make_blob).collect(); + fixture.cas_server.error_hashes.lock().await.insert( + blobs[1].0.packed_hash().to_string(), + Code::InvalidArgument as i32, + ); + + let result = fixture.store.update_many(make_items(&blobs)).await; + let err = result.expect_err("expected update_many to fail"); + assert_eq!(err.code, Code::InvalidArgument, "{err:?}"); + Ok(()) +} + +// Without the config flag, update_many must use the ByteStream Write path +// for every blob (default loop behavior) and advertise no optimization. +#[nativelink_test] +async fn disabled_config_uses_streaming() -> Result<(), Error> { + const NUM_BLOBS: usize = 5; + + let fixture = make_fixture(None).await?; + assert!( + !fixture + .store + .optimized_for(StoreOptimizations::SubscribesToUpdateMany), + "must not advertise batching when disabled" + ); + let blobs: Vec<_> = (0..NUM_BLOBS).map(make_blob).collect(); + + fixture.store.update_many(make_items(&blobs)).await?; + + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!(batch_requests.len(), 0, "no batched RPCs expected"); + let stream_writes = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!(stream_writes.len(), NUM_BLOBS); + Ok(()) +} + +// The batch byte budget must split very many blobs into multiple RPCs. +#[nativelink_test] +async fn batch_byte_budget_splits_requests() -> Result<(), Error> { + // 64-byte blobs + 256-byte overhead = 320 bytes/entry. A 1600-byte + // budget fits exactly 5 entries per request. The blob threshold must + // stay within the budget to satisfy the construction-time invariant. + let config = GrpcWriteBatchingConfig { + max_blob_size_bytes: 1024, + max_batch_bytes: 1600, + }; + let fixture = make_fixture(Some(config)).await?; + let blobs: Vec<_> = (0..12).map(make_blob).collect(); + + fixture.store.update_many(make_items(&blobs)).await?; + + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + // Chunks are dispatched concurrently, so assert sizes order-free. + let mut sizes: Vec = batch_requests + .iter() + .map(|request| request.requests.len()) + .collect(); + sizes.sort_unstable(); + assert_eq!(sizes, vec![2, 5, 5], "expected 5+5+2 split"); + Ok(()) +} + +// The worker's real store shape: FastSlowStore with a batching gRPC slow +// tier must advertise SubscribesToUpdateMany, publish to both tiers, and +// send one batched RPC. +#[nativelink_test] +async fn fast_slow_chain_batches_to_slow_tier() -> Result<(), Error> { + const NUM_BLOBS: usize = 20; + + let fixture = make_fixture(Some(batching_config())).await?; + let fast_store = MemoryStore::new(&MemorySpec::default()); + let fast_slow_store = FastSlowStore::new( + &FastSlowSpec { + // The inner specs are unused by new(); the stores are passed in. + fast: StoreSpec::Noop(NoopSpec::default()), + fast_direction: nativelink_config::stores::StoreDirection::default(), + slow: StoreSpec::Noop(NoopSpec::default()), + slow_direction: nativelink_config::stores::StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(fast_store.clone()), + Store::new(fixture.store.clone()), + ); + assert!( + fast_slow_store.optimized_for(StoreOptimizations::SubscribesToUpdateMany), + "fast_slow must advertise batching when its slow tier batches" + ); + + let blobs: Vec<_> = (0..NUM_BLOBS).map(make_blob).collect(); + fast_slow_store.update_many(make_items(&blobs)).await?; + + for (digest, content) in &blobs { + // Durable on the (fake) remote CAS. + let stored = fixture + .cas_server + .blobs + .lock() + .await + .get(&digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored.as_ref(), Some(content), "slow tier for {digest}"); + // And present in the fast tier. + let fast_data = fast_store.get_part_unchunked(*digest, 0, None).await?; + assert_eq!(&fast_data, content, "fast tier for {digest}"); + } + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!(batch_requests.len(), 1, "expected one batched RPC"); + Ok(()) +} + +// Regression: composite wrappers that forward optimized_for() must also +// forward update_many(), or the advertisement routes call sites into a +// batched path that silently dissolves into per-blob streams. The store +// factory wraps every backend in CacheMetricsStore, so this is the default +// production composition. +#[nativelink_test] +async fn cache_metrics_wrapper_preserves_batching() -> Result<(), Error> { + const NUM_BLOBS: usize = 20; + + let fixture = make_fixture(Some(batching_config())).await?; + let wrapped = CacheMetricsStore::new( + &CacheMetricsSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + cache_type: "test_cas".to_string(), + }, + Store::new(fixture.store.clone()), + ); + assert!( + wrapped.optimized_for(StoreOptimizations::SubscribesToUpdateMany), + "wrapper must forward the advertisement" + ); + + let blobs: Vec<_> = (0..NUM_BLOBS).map(make_blob).collect(); + wrapped.update_many(make_items(&blobs)).await?; + + for (digest, content) in &blobs { + let stored = fixture + .cas_server + .blobs + .lock() + .await + .get(&digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored.as_ref(), Some(content), "blob for {digest}"); + } + let batch_requests = fixture.cas_server.batch_update_requests.lock().await; + assert_eq!( + batch_requests.len(), + 1, + "the wrapper must dispatch one batched RPC, not per-blob streams" + ); + Ok(()) +} + +// Regression: a blob at the batching threshold must fit in one batch; +// otherwise an over-budget request bypasses the per-entry streaming +// fallback. Reject the misconfiguration at startup. +#[nativelink_test] +async fn rejects_blob_threshold_larger_than_batch_budget() -> Result<(), Error> { + let result = make_fixture(Some(GrpcWriteBatchingConfig { + max_blob_size_bytes: 3 * 1024 * 1024, + max_batch_bytes: 3 * 1024 * 1024, + })) + .await; + let err = result.err().expect("expected construction to fail"); + assert!( + err.to_string().contains("max_batch_bytes"), + "unexpected error: {err}" + ); + Ok(()) +} + +// A whole-RPC batch failure (e.g. an intermediary's gRPC message-size +// limit rejecting the batch) must fall back to per-blob streaming uploads +// instead of failing the call. +#[nativelink_test] +async fn whole_rpc_failure_falls_back_to_streaming() -> Result<(), Error> { + const NUM_BLOBS: usize = 6; + + let fixture = make_fixture(Some(batching_config())).await?; + *fixture.cas_server.fail_batch_rpc.lock().await = true; + + let blobs: Vec<_> = (0..NUM_BLOBS).map(make_blob).collect(); + fixture.store.update_many(make_items(&blobs)).await?; + + for (digest, content) in &blobs { + let stored = fixture + .cas_server + .blobs + .lock() + .await + .get(&digest.packed_hash().to_string()) + .cloned(); + assert_eq!(stored.as_ref(), Some(content), "blob for {digest}"); + } + let streams = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!( + streams.len(), + NUM_BLOBS, + "every blob must arrive via the streaming fallback" + ); + Ok(()) +} diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index 5fb650e2b..91ad288f3 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -148,6 +148,17 @@ pub enum StoreOptimizations { /// channel overhead for direct Bytes writes. Stores with this optimization can /// accept complete data directly without going through the MPSC channel. SubscribesToUpdateOneshot, + + /// The store provides an `update_many` implementation that amortizes + /// per-object fixed costs (e.g. one RPC per batch instead of one stream + /// per object). Callers with many small, fully-in-memory objects should + /// prefer `update_many` when this optimization is present. + /// + /// Contract for wrapper stores: a wrapper that forwards this + /// optimization MUST also forward `update_many`, otherwise call sites + /// pay batching overhead for a dispatch that dissolves into the serial + /// default implementation. + SubscribesToUpdateMany, } /// A wrapper struct for [`StoreKey`] to work around @@ -560,6 +571,20 @@ pub trait StoreLike: Send + Sync + Sized + Unpin + 'static { .update_oneshot(digest.into(), data) } + /// Uploads many small, fully-in-memory objects in one operation. Stores + /// with [`StoreOptimizations::SubscribesToUpdateMany`] amortize + /// per-object fixed costs (e.g. one RPC per batch instead of one stream + /// per object); the default behavior is a loop of `update_oneshot` + /// calls. Callers should route large blobs through the streaming + /// `update` path and keep the total bytes of one call bounded. + #[inline] + fn update_many<'a>( + &'a self, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> impl Future> + Send + 'a { + self.as_store_driver_pin().update_many(items) + } + /// Retrieves part of the data from the store and writes it to the given writer. #[inline] fn get_part<'a>( @@ -726,6 +751,19 @@ pub trait StoreDriver: Ok(()) } + /// See: [`StoreLike::update_many`] for details. + async fn update_many( + self: Pin<&Self>, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + for (key, data) in items { + self.update_oneshot(key, data) + .await + .err_tip(|| "In default update_many implementation")?; + } + Ok(()) + } + /// See: [`StoreLike::get_part`] for details. async fn get_part( self: Pin<&Self>, diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 6e4b20573..a9c7ac575 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -66,7 +66,9 @@ use nativelink_util::action_messages::{ use nativelink_util::common::{DigestInfo, fs}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime}; -use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo}; +use nativelink_util::store_trait::{ + Store, StoreKey, StoreLike, StoreOptimizations, UploadSizeInfo, +}; use nativelink_util::{background_spawn, spawn, spawn_blocking}; use parking_lot::Mutex; use prost::Message; @@ -604,12 +606,64 @@ fn is_executable(metadata: &std::fs::Metadata, _full_path: &impl AsRef) -> type DigestUploader = Arc>; +/// Output files at or below this size are queued for a batched publish +/// instead of one upload stream each. Matches the default small-blob +/// threshold of the gRPC store's write batching; operators tuning +/// `max_blob_size_bytes` should keep it at or above this value so queued +/// files stay batchable. +const SMALL_OUTPUT_FILE_MAX_SIZE: u64 = + nativelink_config::stores::DEFAULT_WRITE_BATCHING_MAX_BLOB_SIZE_BYTES; + +/// Byte budget for queued small output files; when exceeded, the queue is +/// flushed in place so memory stays bounded for actions with very many +/// small outputs. +const SMALL_OUTPUT_FLUSH_BYTES: u64 = 32 * 1024 * 1024; + +/// Accumulates small output files for one action so they are published with +/// batched `update_many` calls instead of one upload stream per file. Only +/// created when the CAS store advertises +/// `StoreOptimizations::SubscribesToUpdateMany`; otherwise every file keeps +/// the streaming path and behavior is unchanged. +struct SmallFileBatcher { + /// Queued items plus their total payload bytes. + items: Mutex<(Vec<(StoreKey<'static>, Bytes)>, u64)>, +} + +impl SmallFileBatcher { + fn new() -> Arc { + Arc::new(Self { + items: Mutex::new((Vec::new(), 0)), + }) + } + + /// Queues an item. Returns a drained batch when the queue exceeds its + /// byte budget; the caller must publish that batch. + fn push(&self, key: StoreKey<'static>, data: Bytes) -> Option, Bytes)>> { + let mut inner = self.items.lock(); + inner.1 += data.len() as u64; + inner.0.push((key, data)); + if inner.1 >= SMALL_OUTPUT_FLUSH_BYTES { + inner.1 = 0; + Some(core::mem::take(&mut inner.0)) + } else { + None + } + } + + fn take_all(&self) -> Vec<(StoreKey<'static>, Bytes)> { + let mut inner = self.items.lock(); + inner.1 = 0; + core::mem::take(&mut inner.0) + } +} + async fn upload_file( cas_store: Pin<&impl StoreLike>, full_path: impl AsRef + Debug + Send + Sync, hasher: DigestHasherFunc, metadata: std::fs::Metadata, digest_uploaders: Arc>>, + small_file_batcher: Option>, ) -> Result { let is_executable = is_executable(&metadata, &full_path); let file_size = metadata.len(); @@ -637,7 +691,7 @@ async fn upload_file( // Only upload if the digest doesn't already exist, this should be // a much cheaper operation than an upload. let cas_store = cas_store.as_store_driver_pin(); - let store_key: nativelink_util::store_trait::StoreKey<'_> = digest.into(); + let store_key: StoreKey<'_> = digest.into(); let has_start = std::time::Instant::now(); if cas_store .has(store_key.borrow()) @@ -658,6 +712,34 @@ async fn upload_file( "upload_file: digest not in CAS, starting upload", ); + // Small files are queued for a batched publish at the end of the + // action's output upload instead of one upload stream each. + if let Some(batcher) = &small_file_batcher + && digest.size_bytes() <= SMALL_OUTPUT_FILE_MAX_SIZE + { + file.rewind().await.err_tip(|| "Could not rewind file")?; + let expected_len = usize::try_from(digest.size_bytes()) + .err_tip(|| "Digest size too large for memory buffer")?; + let mut data = Vec::with_capacity(expected_len); + file.read_to_end(&mut data) + .await + .err_tip(|| format!("Reading small output file {full_path:?}"))?; + if data.len() != expected_len { + return Err(make_err!( + Code::Internal, + "Small output file {full_path:?} changed size during upload ({} vs {expected_len})", + data.len(), + )); + } + if let Some(flush_items) = batcher.push(digest.into(), Bytes::from(data)) { + cas_store + .update_many(flush_items) + .await + .err_tip(|| "Flushing batched small output files")?; + } + return Ok(()); + } + file.rewind().await.err_tip(|| "Could not rewind file")?; // Note: For unknown reasons we appear to be hitting: @@ -784,6 +866,7 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( full_work_directory: &'a str, hasher: DigestHasherFunc, digest_uploaders: Arc>>, + small_file_batcher: Option>, ) -> BoxFuture<'a, Result<(Directory, VecDeque), Error>> { Box::pin(async move { let file_futures = FuturesUnordered::new(); @@ -815,6 +898,7 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( full_work_directory, hasher, digest_uploaders.clone(), + small_file_batcher.clone(), ) .and_then(|(dir, all_dirs)| async move { let directory_name = full_path @@ -849,13 +933,21 @@ fn upload_directory<'a, P: AsRef + Debug + Send + Sync + Clone + 'a>( ); } else if file_type.is_file() { let digest_uploaders = digest_uploaders.clone(); + let small_file_batcher = small_file_batcher.clone(); file_futures.push(async move { let metadata = fs::metadata(&full_path) .await .err_tip(|| format!("Could not open file {}", full_path.display()))?; - upload_file(cas_store, &full_path, hasher, metadata, digest_uploaders) - .map_ok(TryInto::try_into) - .await? + upload_file( + cas_store, + &full_path, + hasher, + metadata, + digest_uploaders, + small_file_batcher, + ) + .map_ok(TryInto::try_into) + .await? }); } else if file_type.is_symlink() { symlink_futures.push( @@ -1743,6 +1835,12 @@ impl RunningActionImpl { output_paths.append(&mut command_proto.output_directories); } let digest_uploaders = Arc::new(Mutex::new(HashMap::new())); + // Batch small output files into update_many calls when the store + // chain can amortize them (e.g. gRPC write batching); otherwise the + // streaming per-file path is kept exactly as before. + let small_file_batcher = cas_store + .optimized_for(StoreOptimizations::SubscribesToUpdateMany) + .then(SmallFileBatcher::new); for entry in output_paths { let full_path = OsString::from(if command_proto.working_directory.is_empty() { format!("{}/{}", self.work_directory, entry) @@ -1754,6 +1852,7 @@ impl RunningActionImpl { }); let work_directory = &self.work_directory; let digest_uploaders = digest_uploaders.clone(); + let small_file_batcher = small_file_batcher.clone(); output_path_futures.push(async move { let metadata = { let metadata = match fs::symlink_metadata(&full_path).await { @@ -1778,6 +1877,7 @@ impl RunningActionImpl { hasher, metadata, digest_uploaders, + small_file_batcher, ) .await .map(|mut file_info| { @@ -1797,6 +1897,7 @@ impl RunningActionImpl { work_directory, hasher, digest_uploaders, + small_file_batcher, ) .and_then(|(root_dir, children)| async move { let tree = ProtoTree { @@ -1843,6 +1944,7 @@ impl RunningActionImpl { work_directory, hasher, digest_uploaders, + small_file_batcher, ) .and_then(|(root_dir, children)| async move { let tree = ProtoTree { @@ -1878,6 +1980,7 @@ impl RunningActionImpl { hasher, resolved_meta, digest_uploaders, + small_file_batcher, ) .await .map(|mut file_info| { @@ -2034,6 +2137,17 @@ impl RunningActionImpl { Err(e) => return Err(e).err_tip(|| "Error while uploading results"), }; + // Publish all queued small output files with one batched call. + if let Some(batcher) = small_file_batcher { + let items = batcher.take_all(); + if !items.is_empty() { + cas_store + .update_many(items) + .await + .err_tip(|| "Flushing batched small output files")?; + } + } + execution_metadata.output_upload_completed_timestamp = (self.running_actions_manager.callbacks.now_fn)(); output_files.sort_unstable_by(|a, b| a.name_or_path.cmp(&b.name_or_path)); diff --git a/nativelink-worker/tests/directory_cache_test.rs b/nativelink-worker/tests/directory_cache_test.rs index 681127791..000597b9a 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_write_batching: None, }; let fast_spec = FilesystemSpec { content_path: make_temp_path("paginated_get_tree_cas_content"), diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index 8768af642..592356aca 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -16,6 +16,7 @@ use serial_test::serial; #[serial] mod tests { + use core::pin::Pin; use core::str::from_utf8; use core::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; #[cfg(target_family = "unix")] @@ -67,7 +68,10 @@ mod tests { }; use nativelink_util::common::{DigestInfo, fs, make_temp_path}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; - use nativelink_util::store_trait::{Store, StoreLike}; + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + use nativelink_util::store_trait::{ + Store, StoreDriver, StoreKey, StoreLike, StoreOptimizations, UploadSizeInfo, + }; #[cfg(target_os = "linux")] use nativelink_worker::namespace_utils; use nativelink_worker::running_actions_manager::{ @@ -1562,7 +1566,7 @@ mod tests { b"hello-from-outside".len() ); // Verify the blob actually landed in CAS by re-reading it. - let key: nativelink_util::store_trait::StoreKey<'_> = uploaded.digest.into(); + let key: StoreKey<'_> = uploaded.digest.into(); let blob = slow_store.as_ref().get_part_unchunked(key, 0, None).await?; assert_eq!(blob.as_ref(), b"hello-from-outside"); Ok(()) @@ -1723,7 +1727,7 @@ mod tests { .clone() .expect("inner.txt must have a digest") .try_into()?; - let key: nativelink_util::store_trait::StoreKey<'_> = inner_digest.into(); + let key: StoreKey<'_> = inner_digest.into(); let blob = slow_store.as_ref().get_part_unchunked(key, 0, None).await?; assert_eq!(blob.as_ref(), b"inner-payload"); Ok(()) @@ -4998,4 +5002,249 @@ done assert_eq!(parse_pgid_from_stat("123 (only) S"), None); // too few fields assert_eq!(parse_pgid_from_stat(""), None); } + + /// Slow-store wrapper that advertises `SubscribesToUpdateMany` and + /// counts `update_many` calls, standing in for a gRPC store with write + /// batching enabled. + #[derive(Debug)] + struct BatchingRecorderStore { + inner: Store, + update_many_calls: AtomicU64, + update_many_items: AtomicU64, + } + + impl nativelink_metric::MetricsComponent for BatchingRecorderStore { + fn publish( + &self, + _kind: nativelink_metric::MetricKind, + _field_metadata: nativelink_metric::MetricFieldData, + ) -> Result + { + Ok(nativelink_metric::MetricPublishKnownKindData::Component) + } + } + + #[async_trait::async_trait] + impl StoreDriver for BatchingRecorderStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + fn optimized_for(&self, optimization: StoreOptimizations) -> bool { + optimization == StoreOptimizations::SubscribesToUpdateMany + } + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .has_with_results(keys, results) + .await + } + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: nativelink_util::buf_channel::DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + self.inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await + } + async fn update_many( + self: Pin<&Self>, + items: Vec<(StoreKey<'static>, Bytes)>, + ) -> Result<(), Error> { + self.update_many_calls.fetch_add(1, Ordering::Relaxed); + self.update_many_items + .fetch_add(items.len() as u64, Ordering::Relaxed); + for (key, data) in items { + self.inner + .as_store_driver_pin() + .update_oneshot(key, data) + .await?; + } + Ok(()) + } + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut nativelink_util::buf_channel::DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .get_part(key, writer, offset, length) + .await + } + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + fn as_any_arc(self: Arc) -> Arc { + self + } + fn register_remove_callback( + self: Arc, + callback: nativelink_util::store_trait::RemoveCallback, + ) -> Result<(), Error> { + self.inner.clone().register_remove_callback(callback) + } + } + + default_health_status_indicator!(BatchingRecorderStore); + + // Small output files must be published through one batched update_many + // call when the store chain advertises batching, with identical results. + #[nativelink_test] + async fn small_outputs_batched_when_store_subscribes() -> Result<(), Box> + { + const WORKER_ID: &str = "batching_worker_id"; + + fn test_monotonic_clock() -> SystemTime { + static CLOCK: AtomicU64 = AtomicU64::new(0); + monotonic_clock(&CLOCK) + } + + let fast_config = FilesystemSpec { + content_path: make_temp_path("batching_content_path"), + temp_path: make_temp_path("batching_temp_path"), + eviction_policy: None, + ..Default::default() + }; + let fast_store: Arc = FilesystemStore::new(&fast_config).await?; + let recorder = Arc::new(BatchingRecorderStore { + inner: Store::new(MemoryStore::new(&MemorySpec::default())), + update_many_calls: AtomicU64::new(0), + update_many_items: AtomicU64::new(0), + }); + let ac_store = MemoryStore::new(&MemorySpec::default()); + let cas_store = FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_config), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(fast_store), + Store::new(recorder.clone()), + ); + assert!( + StoreDriver::optimized_for( + cas_store.as_ref(), + StoreOptimizations::SubscribesToUpdateMany + ), + "store chain must advertise batching" + ); + let root_action_directory = make_temp_path("batching_root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = Arc::new(RunningActionsManagerImpl::new_with_callbacks( + RunningActionsManagerArgs { + root_action_directory, + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + }, + Callbacks { + now_fn: test_monotonic_clock, + sleep_fn: |_duration| Box::pin(future::pending()), + }, + )?); + let action_result = { + let command = Command { + arguments: vec![ + "sh".to_string(), + "-c".to_string(), + "echo foo > out_a && echo barbar > out_b && echo foo > out_dup".to_string(), + ], + output_paths: vec![ + "out_a".to_string(), + "out_b".to_string(), + "out_dup".to_string(), + ], + working_directory: ".".to_string(), + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let execute_request = ExecuteRequest { + action_digest: Some(action_digest.into()), + ..Default::default() + }; + let running_action_impl = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + execute_request: Some(execute_request), + operation_id: OperationId::default().to_string(), + queued_timestamp: Some(make_system_time(1000).into()), + platform: action.platform.clone(), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + run_action(running_action_impl.clone()).await? + }; + + assert_eq!(action_result.output_files.len(), 3); + // "foo\n" and "barbar\n" with the duplicate deduped by digest: + // exactly one update_many flush carrying the two unique small blobs. + assert_eq!(recorder.update_many_calls.load(Ordering::Relaxed), 1); + assert_eq!(recorder.update_many_items.load(Ordering::Relaxed), 2); + // Every output must be durable in the slow tier with correct bytes. + for file_info in &action_result.output_files { + let data = cas_store + .get_part_unchunked(file_info.digest, 0, None) + .await?; + assert!(!data.is_empty() || file_info.digest.size_bytes() == 0); + } + Ok(()) + } } From 9d6fc151ffd076cd88f2d555b639ea284687a7cb Mon Sep 17 00:00:00 2001 From: Ernesto Cambuston Date: Wed, 22 Jul 2026 15:56:03 -0700 Subject: [PATCH 2/2] Harden batched output uploads and bound worker reads --- nativelink-store/src/grpc_store.rs | 91 ++++++++++--- .../tests/grpc_write_batching_test.rs | 126 ++++++++++++++++-- .../src/running_actions_manager.rs | 71 ++++++++-- 3 files changed, 245 insertions(+), 43 deletions(-) diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index ddf0a15ae..fc747abc1 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -402,20 +402,34 @@ impl GrpcStore { let mut request = grpc_request.into_inner(); request.instance_name.clone_from(&self.instance_name); + let rpc_timeout = self.rpc_timeout; self.perform_request(request, |request| async move { - let channel = self + let rpc_fut = self .connection_manager .connection("batch_update_blobs".into()) - .await - .err_tip(|| "in batch_update_blobs")?; - ContentAddressableStorageClient::new(channel) - .batch_update_blobs(enrich_request( - Request::new(request), - &self.headers, - &self.forward_headers, - )) - .await - .err_tip(|| "in GrpcStore::batch_update_blobs") + .and_then(|channel| async move { + ContentAddressableStorageClient::new(channel) + .batch_update_blobs(enrich_request( + Request::new(request), + &self.headers, + &self.forward_headers, + )) + .await + .err_tip(|| "in GrpcStore::batch_update_blobs") + }); + if rpc_timeout > Duration::ZERO { + tokio::time::timeout(rpc_timeout, rpc_fut) + .await + .map_err(|_| { + make_err!( + Code::DeadlineExceeded, + "GrpcStore::batch_update_blobs RPC timed out after {}s", + rpc_timeout.as_secs() + ) + })? + } else { + rpc_fut.await + } }) .await } @@ -1229,24 +1243,52 @@ impl StoreDriver for GrpcStore { } }; - let mut error_by_digest: HashMap> = + enum BatchEntryStatus { + Success, + Error(Error), + Malformed, + } + + let mut status_by_digest: HashMap = HashMap::with_capacity(response.responses.len()); for entry in response.responses { let Some(Ok(digest)) = entry.digest.map(DigestInfo::try_from) else { continue; }; - let entry_error = entry - .status - .filter(|status| status.code != 0) - .map(Error::from); - error_by_digest.insert(digest, entry_error); + let entry_status = match entry.status { + None => BatchEntryStatus::Malformed, + Some(status) if !(0..=16).contains(&status.code) => { + BatchEntryStatus::Malformed + } + Some(status) if status.code == Code::Ok as i32 => BatchEntryStatus::Success, + Some(status) => BatchEntryStatus::Error(Error::from(status)), + }; + + // A malformed or non-success duplicate must not be hidden by + // a later success response for the same digest. + if let Some(existing) = status_by_digest.get_mut(&digest) { + let replace = match existing { + BatchEntryStatus::Malformed => false, + BatchEntryStatus::Error(_) => { + matches!(entry_status, BatchEntryStatus::Malformed) + } + BatchEntryStatus::Success => { + !matches!(entry_status, BatchEntryStatus::Success) + } + }; + if replace { + *existing = entry_status; + } + } else { + status_by_digest.insert(digest, entry_status); + } } let mut chunk_fallback = Vec::new(); for (digest, data) in chunk { - match error_by_digest.remove(&digest) { + match status_by_digest.remove(&digest) { // Entry succeeded. - Some(None) => {} - Some(Some(err)) if is_retryable_code(err.code) => { + Some(BatchEntryStatus::Success) => {} + Some(BatchEntryStatus::Error(err)) if is_retryable_code(err.code) => { trace!( ?digest, ?err, @@ -1254,11 +1296,18 @@ impl StoreDriver for GrpcStore { ); chunk_fallback.push((digest.into(), data)); } - Some(Some(err)) => { + Some(BatchEntryStatus::Error(err)) => { return Err(err.append(format!( "in BatchUpdateBlobs response for {digest} in GrpcStore::update_many" ))); } + Some(BatchEntryStatus::Malformed) => { + trace!( + ?digest, + "Batched upload entry had a missing or malformed status, falling back to ByteStream write", + ); + chunk_fallback.push((digest.into(), data)); + } // Server omitted the entry; fall back to the streaming // path rather than guessing at its state. None => chunk_fallback.push((digest.into(), data)), diff --git a/nativelink-store/tests/grpc_write_batching_test.rs b/nativelink-store/tests/grpc_write_batching_test.rs index 964db9a22..49b0ca312 100644 --- a/nativelink-store/tests/grpc_write_batching_test.rs +++ b/nativelink-store/tests/grpc_write_batching_test.rs @@ -13,7 +13,8 @@ // limitations under the License. use core::pin::Pin; -use std::collections::HashMap; +use core::time::Duration; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use async_lock::Mutex; @@ -47,6 +48,7 @@ use nativelink_store::memory_store::MemoryStore; use nativelink_util::background_spawn; use nativelink_util::common::DigestInfo; use nativelink_util::store_trait::{Store, StoreKey, StoreLike, StoreOptimizations}; +use tokio::time::sleep; use tonic::transport::Server; use tonic::transport::server::TcpIncoming; use tonic::{Request, Response, Status, Streaming}; @@ -58,9 +60,13 @@ use tonic::{Request, Response, Status, Streaming}; struct FakeCasServer { blobs: Arc>>, error_hashes: Arc>>, + missing_status_hashes: Arc>>, + malformed_status_hashes: Arc>>, batch_update_requests: Arc>>, /// When set, every `BatchUpdateBlobs` RPC fails as a whole. fail_batch_rpc: Arc>, + /// When set, every `BatchUpdateBlobs` RPC waits before responding. + batch_update_delay: Arc>>, } impl FakeCasServer { @@ -68,8 +74,11 @@ impl FakeCasServer { Self { blobs: Arc::new(Mutex::new(HashMap::new())), error_hashes: Arc::new(Mutex::new(HashMap::new())), + missing_status_hashes: Arc::new(Mutex::new(HashSet::new())), + malformed_status_hashes: Arc::new(Mutex::new(HashSet::new())), batch_update_requests: Arc::new(Mutex::new(vec![])), fail_batch_rpc: Arc::new(Mutex::new(false)), + batch_update_delay: Arc::new(Mutex::new(None)), } } } @@ -102,31 +111,45 @@ impl ContentAddressableStorage for FakeCasServer { "Injected whole-RPC failure (e.g. message-size limit)", )); } + let batch_update_delay = *self.batch_update_delay.lock().await; + if let Some(delay) = batch_update_delay { + sleep(delay).await; + } let mut blobs = self.blobs.lock().await; let error_hashes = self.error_hashes.lock().await; + let missing_status_hashes = self.missing_status_hashes.lock().await; + let malformed_status_hashes = self.malformed_status_hashes.lock().await; let mut responses = Vec::with_capacity(request.requests.len()); for entry in request.requests { let Some(digest) = entry.digest else { return Err(Status::invalid_argument("Missing digest in request")); }; let status = if let Some(&code) = error_hashes.get(&digest.hash) { - RpcStatus { + Some(RpcStatus { code, message: format!("Injected error for {}", digest.hash), details: vec![], - } + }) } else { blobs.insert(digest.hash.clone(), entry.data); - RpcStatus { - code: Code::Ok as i32, - message: String::new(), - details: vec![], + if missing_status_hashes.contains(&digest.hash) { + None + } else { + Some(RpcStatus { + code: if malformed_status_hashes.contains(&digest.hash) { + 99 + } else { + Code::Ok as i32 + }, + message: String::new(), + details: vec![], + }) } }; responses.push(batch_update_blobs_response::Response { digest: Some(digest), - status: Some(status), + status, }); } Ok(Response::new(BatchUpdateBlobsResponse { responses })) @@ -244,6 +267,14 @@ struct TestFixture { async fn make_fixture( write_batching: Option, +) -> Result { + make_fixture_with_options(write_batching, 0, Retry::default()).await +} + +async fn make_fixture_with_options( + write_batching: Option, + rpc_timeout_s: u64, + retry: Retry, ) -> Result { let cas_server = FakeCasServer::new(); let bytestream_server = FakeByteStreamServer::new(cas_server.blobs.clone()); @@ -273,10 +304,10 @@ async fn make_fixture( http2_keepalive_timeout_s: 0, }], store_type: StoreType::Cas, - retry: Retry::default(), + retry, max_concurrent_requests: 0, connections_per_endpoint: 0, - rpc_timeout_s: 0, + rpc_timeout_s, use_legacy_resource_names: false, headers: HashMap::new(), forward_headers: vec![], @@ -441,6 +472,43 @@ async fn retryable_entry_error_falls_back_to_streaming() -> Result<(), Error> { Ok(()) } +// Missing or invalid per-entry statuses are ambiguous: the upstream may have +// committed the blob, so retry through the idempotent streaming path instead +// of treating either response as a successful batch entry. +#[nativelink_test] +async fn missing_or_malformed_status_falls_back_per_entry() -> Result<(), Error> { + let fixture = make_fixture(Some(batching_config())).await?; + let blobs: Vec<_> = (0..3).map(make_blob).collect(); + fixture + .cas_server + .missing_status_hashes + .lock() + .await + .insert(blobs[0].0.packed_hash().to_string()); + fixture + .cas_server + .malformed_status_hashes + .lock() + .await + .insert(blobs[1].0.packed_hash().to_string()); + + fixture.store.update_many(make_items(&blobs)).await?; + + let streams = fixture.bytestream_server.write_resource_names.lock().await; + assert_eq!(streams.len(), 2, "only ambiguous entries should fall back"); + assert!( + streams + .iter() + .any(|name| name.contains(&blobs[0].0.packed_hash().to_string())) + ); + assert!( + streams + .iter() + .any(|name| name.contains(&blobs[1].0.packed_hash().to_string())) + ); + Ok(()) +} + // A non-retryable per-entry error must fail the whole update_many call. #[nativelink_test] async fn non_retryable_entry_error_propagates() -> Result<(), Error> { @@ -648,3 +716,41 @@ async fn whole_rpc_failure_falls_back_to_streaming() -> Result<(), Error> { ); Ok(()) } + +// The configured batch RPC deadline must participate in the normal retrier; +// after retries are exhausted, update_many falls back to idempotent streams. +#[nativelink_test] +async fn timed_out_batch_retries_then_falls_back_to_streaming() -> Result<(), Error> { + let fixture = make_fixture_with_options( + Some(batching_config()), + 1, + Retry { + max_retries: 1, + delay: 0.0, + jitter: 0.0, + retry_on_errors: None, + }, + ) + .await?; + *fixture.cas_server.batch_update_delay.lock().await = Some(Duration::from_secs(2)); + let blobs: Vec<_> = (0..3).map(make_blob).collect(); + + fixture.store.update_many(make_items(&blobs)).await?; + + assert_eq!( + fixture.cas_server.batch_update_requests.lock().await.len(), + 2, + "the timeout should be retried once before fallback" + ); + assert_eq!( + fixture + .bytestream_server + .write_resource_names + .lock() + .await + .len(), + blobs.len(), + "all entries should use streaming fallback after batch timeout" + ); + Ok(()) +} diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index a9c7ac575..fa0fd26c0 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -657,6 +657,37 @@ impl SmallFileBatcher { } } +async fn read_small_output( + file: R, + expected_len: usize, + full_path: impl Debug, +) -> Result +where + R: tokio::io::AsyncRead + Unpin, +{ + // Read at most one byte past the expected digest size. This bounds the + // allocation even if the file grows after hashing, while retaining an + // explicit overflow signal below. + let read_limit = expected_len + .checked_add(1) + .ok_or_else(|| make_err!(Code::Internal, "Small output size is too large"))?; + let read_limit = u64::try_from(read_limit) + .err_tip(|| "Small output size cannot be represented as a read limit")?; + let mut data = Vec::with_capacity(expected_len); + file.take(read_limit) + .read_to_end(&mut data) + .await + .err_tip(|| format!("Reading small output file {full_path:?}"))?; + if data.len() != expected_len { + return Err(make_err!( + Code::Internal, + "Small output file {full_path:?} changed size during upload ({} vs {expected_len})", + data.len(), + )); + } + Ok(Bytes::from(data)) +} + async fn upload_file( cas_store: Pin<&impl StoreLike>, full_path: impl AsRef + Debug + Send + Sync, @@ -720,18 +751,8 @@ async fn upload_file( file.rewind().await.err_tip(|| "Could not rewind file")?; let expected_len = usize::try_from(digest.size_bytes()) .err_tip(|| "Digest size too large for memory buffer")?; - let mut data = Vec::with_capacity(expected_len); - file.read_to_end(&mut data) - .await - .err_tip(|| format!("Reading small output file {full_path:?}"))?; - if data.len() != expected_len { - return Err(make_err!( - Code::Internal, - "Small output file {full_path:?} changed size during upload ({} vs {expected_len})", - data.len(), - )); - } - if let Some(flush_items) = batcher.push(digest.into(), Bytes::from(data)) { + let data = read_small_output(file, expected_len, &full_path).await?; + if let Some(flush_items) = batcher.push(digest.into(), data) { cas_store .update_many(flush_items) .await @@ -1066,6 +1087,32 @@ async fn do_cleanup( } } +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use nativelink_macro::nativelink_test; + + use super::*; + + #[nativelink_test] + async fn small_output_read_caps_at_expected_size_plus_one() -> Result<(), Error> { + let data = read_small_output(Cursor::new(vec![1, 2, 3, 4]), 4, "exact") + .await + .expect("exact-sized output should be accepted"); + assert_eq!(&data[..], &[1, 2, 3, 4]); + + let err = read_small_output(Cursor::new(vec![1, 2, 3, 4, 5, 6]), 4, "overflow") + .await + .expect_err("oversized output should be rejected"); + assert!( + err.to_string().contains("5 vs 4"), + "unexpected error: {err}" + ); + Ok(()) + } +} + pub trait RunningAction: Sync + Send + Sized + Unpin + 'static { /// Returns the action id of the action. fn get_operation_id(&self) -> &OperationId;