Batch small worker output uploads into BatchUpdateBlobs

.github/styles/config/vocabularies/TraceMachina/accept.txt

@@ -96,6 +96,7 @@ Pantsbuild
 [Ss]andboxing
 [Cc]onfig
 bytestream
+[Bb]atcher
 [Ff]ailover
 proto
 quantiles

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<GrpcReadBatchingConfig>,
+
+    /// 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<GrpcWriteBatchingConfig>,
+}
+
+/// 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

nativelink-service/tests/cas_server_test.rs

@@ -1562,6 +1562,7 @@ async fn chunking_on_grpc_store_forbids_index_store() -> Result<(), Box<dyn core
                 headers: std::collections::HashMap::new(),
                 forward_headers: vec![],
                 experimental_read_batching: None,
+                experimental_write_batching: None,
             }),
             &store_manager,
             None,

nativelink-store/BUILD.bazel

@@ -134,6 +134,7 @@ rust_test_suite(
         "tests/gcs_store_test.rs",
         "tests/grpc_read_batching_test.rs",
         "tests/grpc_store_test.rs",
+        "tests/grpc_write_batching_test.rs",
         "tests/memory_store_test.rs",
         "tests/mongo_store_test.rs",
         "tests/oci_store_test.rs",

nativelink-store/src/cache_metrics_store.rs

@@ -172,6 +172,32 @@ impl StoreDriver for CacheMetricsStore {
         result
     }
 
+    async fn update_many(
+        self: Pin<&Self>,
+        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)
     }

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<Self> {
         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::<StoreKey>)
+            .optimized_for(StoreOptimizations::NoopUpdates)
+            || self.slow_direction == StoreDirection::ReadOnly
+            || self.slow_direction == StoreDirection::Get;
+        let ignore_fast = self
+            .fast_store
+            .inner_store(None::<StoreKey>)
+            .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