diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index 24af1afe0..9e0d82299 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -89,6 +89,7 @@ API APIs Backpressure Dedup +Deduplicate XCode CMake Gradle diff --git a/nativelink-config/src/backcompat.rs b/nativelink-config/src/backcompat.rs index 00b9e9b3e..f618c41d7 100644 --- a/nativelink-config/src/backcompat.rs +++ b/nativelink-config/src/backcompat.rs @@ -102,6 +102,7 @@ where max_bytes_per_stream: old_config.max_bytes_per_stream, persist_stream_on_disconnect_timeout_s: old_config .persist_stream_on_disconnect_timeout_s, + experimental_write_dedup: false, }, }) .collect(); diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 07a3d233b..9c9615cad 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -349,6 +349,27 @@ pub struct ByteStreamConfig { alias = "persist_stream_on_disconnect_timeout" )] pub persist_stream_on_disconnect_timeout_s: usize, + + /// Deduplicate concurrent uploads of the same digest ("join the flight"). + /// When multiple clients upload a blob with the same digest at the same + /// time, only the first upload is streamed to the store; the others wait + /// for it to durably commit and then complete early without transferring + /// their payload, as permitted by the REAPI specification. Uploads of + /// blobs that already exist in the store also complete early after a + /// single existence check. + /// + /// If the leading upload fails, waiting uploads receive a retryable + /// ABORTED error and one of the retrying clients becomes the new leader. + /// + /// This saves upload bandwidth and store work when many actions produce + /// identical outputs, at the cost of one existence check per new upload. + /// The check uses the digest function in the upload resource name, and + /// proxy-backed `GrpcStore` instances participate in the same local + /// single-flight behavior. + /// + /// Default: false (disabled) + #[serde(default, skip_serializing_if = "is_default")] + pub experimental_write_dedup: bool, } // Older bytestream config. All fields are as per the newer docs, but this requires diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index c678eed59..77fa6e4a1 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -55,6 +55,7 @@ use nativelink_util::store_trait::{Store, StoreLike, StoreOptimizations, UploadS use nativelink_util::task::JoinHandleDropGuard; use opentelemetry::context::FutureExt; use parking_lot::Mutex; +use tokio::sync::watch; use tokio::time::sleep; use tonic::{Request, Response, Status, Streaming}; use tracing::{Instrument, Level, debug, error, error_span, info, instrument, trace, warn}; @@ -101,6 +102,17 @@ pub struct ByteStreamMetrics { pub resumed_uploads: AtomicU64, /// Number of idle streams that timed out pub idle_stream_timeouts: AtomicU64, + /// Number of uploads that joined an in-flight upload of the same digest + pub write_dedup_flights_joined: AtomicU64, + /// Number of uploads completed early because the digest was already durable + pub write_dedup_early_completes: AtomicU64, + /// Number of leading uploads that failed or were cancelled while waiters were joined + pub write_dedup_leader_failures: AtomicU64, + /// Approximate upper bound of upload bytes not transferred thanks to + /// write dedup (payload bytes already carried by a request's first + /// message cannot be subtracted, and compressed uploads are counted at + /// their uncompressed size) + pub write_dedup_bytes_saved: AtomicU64, } impl MetricsComponent for ByteStreamMetrics { @@ -201,11 +213,55 @@ impl MetricsComponent for ByteStreamMetrics { MetricKind::Counter, "Number of idle streams that timed out" ); + { + let _write_dedup_enter = group!("write_dedup").entered(); + publish!( + "flights_joined", + &self.write_dedup_flights_joined, + MetricKind::Counter, + "Number of uploads that joined an in-flight upload of the same digest" + ); + publish!( + "early_completes", + &self.write_dedup_early_completes, + MetricKind::Counter, + "Number of uploads completed early because the digest was already durable" + ); + publish!( + "leader_failures", + &self.write_dedup_leader_failures, + MetricKind::Counter, + "Number of leading uploads that failed or were cancelled" + ); + publish!( + "bytes_saved", + &self.write_dedup_bytes_saved, + MetricKind::Counter, + "Approximate upper bound of upload bytes not transferred thanks to write dedup" + ); + } Ok(MetricPublishKnownKindData::Component) } } +impl ByteStreamMetrics { + /// Records the shared write-request epilogue: duration plus the + /// success/failure counter. `bytes_written_total` is recorded separately + /// by paths that actually stored bytes. + fn record_write_result(&self, start_time: Instant, success: bool) { + #[allow(clippy::cast_possible_truncation)] + let elapsed_ns = start_time.elapsed().as_nanos() as u64; + self.write_duration_ns + .fetch_add(elapsed_ns, Ordering::Relaxed); + if success { + self.write_requests_success.fetch_add(1, Ordering::Relaxed); + } else { + self.write_requests_failure.fetch_add(1, Ordering::Relaxed); + } + } +} + type BytesWrittenAndIdleStream = (Arc, Option); /// Type alias for the UUID key used in `active_uploads` `HashMap`. @@ -259,6 +315,11 @@ pub struct InstanceInfo { _sweeper_handle: Arc>, /// Whether this instance supports Bazel remote cache compression. remote_cache_compression_enabled: bool, + /// Whether join-the-flight write dedup is enabled for this instance. + write_dedup_enabled: bool, + /// In-flight upload flights keyed by digest and digest function. Only used when + /// `write_dedup_enabled` is set. + write_flights: Arc>>>, } impl Debug for InstanceInfo { @@ -478,6 +539,100 @@ impl IdleStream { } } +/// A single-flight registration for an in-progress upload of a digest, used +/// when `experimental_write_dedup` is enabled. Waiters subscribe to the watch +/// channel and are released when the leading upload durably commits (`Ok`) or +/// fails/cancels (`Err`). +#[derive(Debug)] +struct WriteFlight { + result_tx: watch::Sender>>, +} + +/// Owned by the leading upload of a write flight. Exactly one terminal +/// transition fires: `succeed()`, `fail()`, or (if the leader is cancelled at +/// any await point) `Drop` — each removes the flight from the map first and +/// then broadcasts the outcome, so waiters can never be stranded and a late +/// arrival always starts a fresh flight instead of joining a settled one. +struct WriteFlightGuard { + digest: DigestInfo, + flight: Option>, + digest_function: DigestHasherFunc, + write_flights: Arc>>>, + metrics: Arc, +} + +impl WriteFlightGuard { + /// The leading upload durably committed; waiters may complete early. + fn succeed(mut self) { + if let Some(flight) = self.flight.take() { + self.complete(&flight, Ok(())); + } + } + + /// The leading upload failed; waiters receive a retryable error so one + /// of their retries becomes the new leading upload. + fn fail(mut self) { + if let Some(flight) = self.flight.take() { + self.complete( + &flight, + Err(make_err!( + Code::Aborted, + "The leading upload of this digest failed; retrying may succeed" + )), + ); + } + } + + fn complete(&self, flight: &Arc, result: Result<(), Error>) { + self.write_flights + .lock() + .remove(&(self.digest, self.digest_function)); + // Only count a leader failure when waiters were actually joined; + // otherwise every routine failed upload would inflate the metric. + if result.is_err() && flight.result_tx.receiver_count() > 0 { + self.metrics + .write_dedup_leader_failures + .fetch_add(1, Ordering::Relaxed); + } + // Ignore send errors: every waiter may already be gone. + drop(flight.result_tx.send(Some(result))); + } +} + +impl Drop for WriteFlightGuard { + fn drop(&mut self) { + if let Some(flight) = self.flight.take() { + self.complete( + &flight, + Err(make_err!( + Code::Aborted, + "The leading upload of this digest was cancelled; retrying may succeed" + )), + ); + } + } +} + +/// Outcome of the write-dedup check for a new upload. +enum WriteDedupOutcome { + /// Proceed with the upload. Carries the flight guard when this request + /// was elected the leading upload for its digest; `None` when dedup does + /// not apply (resumed stream). + Proceed(Option>), + /// Respond immediately without receiving the payload: the digest is + /// already durable, per the REAPI duplicate-upload contract. + EarlyComplete { committed_size: i64 }, + /// Another upload of the same digest is in flight. The caller must drain + /// the client stream while awaiting the flight result (unread request + /// bytes would otherwise pin the HTTP/2 connection flow-control window + /// and could stall the leading upload sharing the connection), then + /// respond with `committed_size` on success. + Waiter { + result_rx: watch::Receiver>>, + committed_size: i64, + }, +} + #[derive(Debug)] pub struct ByteStreamServer { instance_infos: HashMap, @@ -600,6 +755,8 @@ impl ByteStreamServer { metrics, _sweeper_handle: Arc::new(sweeper_handle), remote_cache_compression_enabled, + write_dedup_enabled: config.experimental_write_dedup, + write_flights: Arc::new(Mutex::new(HashMap::new())), }) } @@ -607,6 +764,169 @@ impl ByteStreamServer { Server::new(self) } + /// Runs the join-the-flight write-dedup check for a new upload. + /// + /// Exactly one concurrent upload per digest is elected the leading upload + /// (`Proceed(Some(guard))`); the caller must consume the guard with + /// `succeed()`/`fail()` once the upload settles. Concurrent uploads of + /// the same digest wait for the leading upload instead of transferring + /// their payload: on durable commit they get `EarlyComplete`, on leader + /// failure a retryable `Aborted` error so one retry becomes the new + /// leader. Uploads whose digest is already durable get `EarlyComplete` + /// after a single existence check. + /// + /// Resumed streams (UUID already tracked in `active_uploads`) bypass + /// dedup entirely (`Proceed(None)`) to preserve resume semantics. + async fn check_write_dedup( + instance: &InstanceInfo, + digest: DigestInfo, + digest_function: DigestHasherFunc, + uuid_str: &str, + is_compressed: bool, + ) -> Result { + let uuid_key = parse_uuid_to_key(uuid_str); + if instance.active_uploads.lock().contains_key(&uuid_key) { + return Ok(WriteDedupOutcome::Proceed(None)); + } + let committed_size = if is_compressed { + // The REAPI compressed-blobs contract reports -1 as the + // committed_size of an upload completed by another client. + -1 + } else { + i64::try_from(digest.size_bytes()) + .err_tip(|| "Digest size not convertible to i64 in check_write_dedup")? + }; + let claim = { + let mut write_flights = instance.write_flights.lock(); + match write_flights.entry((digest, digest_function)) { + Entry::Occupied(entry) => Err(entry.get().result_tx.subscribe()), + Entry::Vacant(entry) => { + let (result_tx, _result_rx) = watch::channel(None); + let flight = Arc::new(WriteFlight { result_tx }); + entry.insert(flight.clone()); + Ok(flight) + } + } + }; + match claim { + Err(result_rx) => { + instance + .metrics + .write_dedup_flights_joined + .fetch_add(1, Ordering::Relaxed); + Ok(WriteDedupOutcome::Waiter { + result_rx, + committed_size, + }) + } + Ok(flight) => { + let guard = WriteFlightGuard { + digest, + digest_function, + flight: Some(flight), + write_flights: instance.write_flights.clone(), + metrics: instance.metrics.clone(), + }; + // One existence check: uploads of already-durable digests + // complete early per the REAPI duplicate-upload contract. + let has_result = instance + .store + .has(digest) + .with_context(make_ctx_for_hash_func(digest_function)?) + .await; + match has_result { + Ok(Some(_size)) => { + // Broadcast success so any waiter that joined while + // the existence check ran completes early too. + guard.succeed(); + instance + .metrics + .write_dedup_early_completes + .fetch_add(1, Ordering::Relaxed); + instance + .metrics + .write_dedup_bytes_saved + .fetch_add(digest.size_bytes(), Ordering::Relaxed); + Ok(WriteDedupOutcome::EarlyComplete { committed_size }) + } + Ok(None) => Ok(WriteDedupOutcome::Proceed(Some(Box::new(guard)))), + Err(err) => { + // An existence-check failure must not fail the + // upload; proceed as the leading upload. + debug!( + ?err, + "Existence check failed in write dedup; proceeding with upload" + ); + Ok(WriteDedupOutcome::Proceed(Some(Box::new(guard)))) + } + } + } + } + } + + /// Waits for the leading upload of `digest` to settle while draining the + /// waiter's client stream. + /// + /// Draining matters: the server must keep consuming request bytes so + /// they cannot accumulate against the HTTP/2 connection flow-control + /// window — with an unread waiter stream, a leading upload sharing the + /// same connection would stall once the window is exhausted, deadlocking + /// both. Drained bytes are discarded without touching the store. + async fn wait_on_write_flight( + instance: &InstanceInfo, + mut stream: WriteRequestStreamWrapper< + impl Stream> + Unpin, + >, + mut result_rx: watch::Receiver>>, + digest: DigestInfo, + ) -> Result<(), Error> { + async fn settle( + result_rx: &mut watch::Receiver>>, + ) -> Result<(), Error> { + fn closed_flight_error() -> Error { + make_err!( + Code::Aborted, + "The leading upload of this digest went away before completing; retrying may succeed" + ) + } + match result_rx.wait_for(Option::is_some).await { + Ok(value) => value.clone().unwrap_or_else(|| Err(closed_flight_error())), + Err(_closed) => Err(closed_flight_error()), + } + } + let mut drained_bytes: u64 = 0; + let mut client_stream_done = false; + let flight_result = loop { + if client_stream_done { + break settle(&mut result_rx).await; + } + tokio::select! { + settled = settle(&mut result_rx) => break settled, + message = stream.next() => match message { + Some(Ok(write_request)) => { + drained_bytes += write_request.data.len() as u64; + if write_request.finish_write { + client_stream_done = true; + } + } + // The client ended or went away early; the flight outcome + // still decides our response. + Some(Err(_)) | None => client_stream_done = true, + }, + } + }; + match flight_result { + Ok(()) => { + instance.metrics.write_dedup_bytes_saved.fetch_add( + digest.size_bytes().saturating_sub(drained_bytes), + Ordering::Relaxed, + ); + Ok(()) + } + Err(err) => Err(err), + } + } + /// Creates or joins an upload stream for the given UUID. /// /// This function handles three scenarios: @@ -1492,12 +1812,6 @@ impl ByteStream for ByteStreamServer { ) .err_tip(|| "Invalid digest input in ByteStream::write")?; - // If we are a GrpcStore we shortcut here, as this is a special store. - if let Some(grpc_store) = store.downcast_ref::(Some(digest.into())) { - let resp = grpc_store.write(stream).await.map_err(Into::into); - return resp; - } - let digest_function = stream .resource_info .digest_function @@ -1512,10 +1826,71 @@ impl ByteStream for ByteStreamServer { stream.resource_info.compressor.as_deref(), instance.remote_cache_compression_enabled, )?; + let is_compressed_upload = wire_compressor != compressor::Value::Identity; + + // Join-the-flight write dedup: concurrent uploads of the same digest + // collapse into one leading upload and uploads of already-durable + // digests complete immediately, without receiving the payload. + let write_flight_guard = if instance.write_dedup_enabled + && let Some(uuid_str) = stream.resource_info.uuid.as_deref() + { + match Self::check_write_dedup( + instance, + digest, + digest_function, + uuid_str, + is_compressed_upload, + ) + .await + { + Ok(WriteDedupOutcome::Proceed(maybe_guard)) => maybe_guard, + Ok(WriteDedupOutcome::Waiter { + result_rx, + committed_size, + }) => { + let result = + Self::wait_on_write_flight(instance, stream, result_rx, digest).await; + instance + .metrics + .record_write_result(start_time, result.is_ok()); + return match result { + Ok(()) => Ok(Response::new(WriteResponse { committed_size })), + Err(err) => Err(err.into()), + }; + } + Ok(WriteDedupOutcome::EarlyComplete { committed_size }) => { + instance.metrics.record_write_result(start_time, true); + return Ok(Response::new(WriteResponse { committed_size })); + } + Err(err) => { + instance.metrics.record_write_result(start_time, false); + return Err(err.into()); + } + } + } else { + None + }; + + // GrpcStore is a proxy for the upstream ByteStream service. Run the + // local dedup check first so enabling experimental write dedup does + // not silently bypass the feature for proxy-backed instances. + if let Some(grpc_store) = store.downcast_ref::(Some(digest.into())) { + let result = grpc_store.write(stream).await; + if let Some(guard) = write_flight_guard { + match &result { + Ok(_) => guard.succeed(), + Err(_) => guard.fail(), + } + } + instance + .metrics + .record_write_result(start_time, result.is_ok()); + return result.map_err(Into::into); + } // For compressed uploads, stream compressed wire bytes through the // decoder and store the resulting raw bytes. - if wire_compressor != compressor::Value::Identity { + if is_compressed_upload { let result = self .inner_write_compressed(instance, digest, digest_function, wire_compressor, stream) .instrument(error_span!("bytestream_write_compressed")) @@ -1526,6 +1901,13 @@ impl ByteStream for ByteStreamServer { .await .err_tip(|| "In ByteStreamServer::write_compressed"); + if let Some(guard) = write_flight_guard { + match &result { + Ok(_) => guard.succeed(), + Err(_) => guard.fail(), + } + } + // Track metrics based on result #[allow(clippy::cast_possible_truncation)] let elapsed_ns = start_time.elapsed().as_nanos() as u64; @@ -1605,6 +1987,13 @@ impl ByteStream for ByteStreamServer { .err_tip(|| "In ByteStreamServer::write") }; + if let Some(guard) = write_flight_guard { + match &result { + Ok(_) => guard.succeed(), + Err(_) => guard.fail(), + } + } + // Track metrics based on result #[allow(clippy::cast_possible_truncation)] let elapsed_ns = start_time.elapsed().as_nanos() as u64; diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index e1c78b439..c6ed04c0d 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use bytes::Bytes; @@ -23,8 +24,10 @@ use hyper::body::Frame; use hyper_util::rt::TokioIo; use hyper_util::server::conn::auto; use hyper_util::service::TowerToHyperService; -use nativelink_config::cas_server::{ByteStreamConfig, HttpListener, WithInstanceName}; -use nativelink_config::stores::{MemorySpec, StoreSpec}; +use nativelink_config::cas_server::{ + ByteStreamConfig, CasStoreConfig, HttpListener, WithInstanceName, +}; +use nativelink_config::stores::{GrpcEndpoint, GrpcSpec, MemorySpec, Retry, StoreSpec, StoreType}; use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient; @@ -33,12 +36,14 @@ use nativelink_proto::google::bytestream::{ QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, WriteRequest, WriteResponse, }; use nativelink_service::bytestream_server::ByteStreamServer; +use nativelink_service::cas_server::CasServer; use nativelink_service::wire_compression::RemoteCacheCompressionInstances; use nativelink_store::default_store_factory::store_factory; +use nativelink_store::grpc_store::GrpcStore; use nativelink_store::store_manager::StoreManager; use nativelink_util::channel_body_for_tests::ChannelBody; use nativelink_util::common::{DigestInfo, encode_stream_proto}; -use nativelink_util::store_trait::StoreLike; +use nativelink_util::store_trait::{Store, StoreLike}; use nativelink_util::task::JoinHandleDropGuard; use nativelink_util::{background_spawn, spawn}; use pretty_assertions::assert_eq; @@ -50,7 +55,8 @@ use tokio::task::yield_now; use tokio_stream::StreamExt; use tokio_stream::wrappers::UnboundedReceiverStream; use tonic::codec::{Codec, CompressionEncoding}; -use tonic::transport::{Channel, Endpoint}; +use tonic::transport::server::TcpIncoming; +use tonic::transport::{Channel, Endpoint, Server}; use tonic::{Request, Response, Streaming}; use tonic_prost::ProstCodec; use tower::service_fn; @@ -91,6 +97,7 @@ fn make_bytestream_server_with_remote_cache_compression( cas_store: "main_cas".to_string(), persist_stream_on_disconnect_timeout_s: 0, max_bytes_per_stream: 1024, + ..Default::default() }, }] }); @@ -104,6 +111,46 @@ fn make_bytestream_server_with_remote_cache_compression( ByteStreamServer::new(&config, store_manager, &remote_cache_compression_instances) } +async fn make_tcp_proxy_backend(store_manager: &StoreManager) -> Result { + let instance_name = INSTANCE_NAME.to_string(); + let remote_cache_compression_instances = RemoteCacheCompressionInstances::default(); + let bytestream = ByteStreamServer::new( + &[WithInstanceName { + instance_name: instance_name.clone(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + persist_stream_on_disconnect_timeout_s: 0, + max_bytes_per_stream: 1024, + ..Default::default() + }, + }], + store_manager, + &remote_cache_compression_instances, + )?; + let cas = CasServer::new( + &[WithInstanceName { + instance_name, + config: CasStoreConfig { + cas_store: "main_cas".to_string(), + experimental_chunking: None, + }, + }], + store_manager, + &remote_cache_compression_instances, + )?; + let listener = TcpIncoming::bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let port = listener.local_addr().unwrap().port(); + background_spawn!("bytestream_proxy_backend", async move { + Server::builder() + .add_service(bytestream.into_service()) + .add_service(cas.into_service()) + .serve_with_incoming(listener) + .await + .expect("proxy backend server failed"); + }); + Ok(port) +} + fn make_stream( encoding: Option, ) -> (mpsc::Sender>, Streaming) { @@ -1878,3 +1925,408 @@ async fn uuid_collision_does_not_deadlock() -> Result<(), Box Result { + make_bytestream_server_with_remote_cache_compression( + store_manager, + Some(vec![WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + experimental_write_dedup: true, + ..Default::default() + }, + }]), + remote_cache_compression_enabled, + ) +} + +fn make_dedup_resource_name(uuid: &str, data_len: usize) -> String { + format!("{INSTANCE_NAME}/uploads/{uuid}/blobs/{HASH1}/{data_len}") +} + +fn make_dedup_resource_name_with_digest_function( + uuid: &str, + digest_function: &str, + data_len: usize, +) -> String { + format!("{INSTANCE_NAME}/uploads/{uuid}/blobs/{digest_function}/{HASH1}/{data_len}") +} + +/// An upload whose digest is already durable must complete immediately with +/// the full committed size, without the client sending the payload. +#[nativelink_test] +pub async fn write_dedup_completes_early_for_existing_blob() +-> Result<(), Box> { + const DATA: &[u8] = b"write dedup early complete payload"; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), false).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + let digest = DigestInfo::try_new(HASH1, DATA.len())?; + store.update_oneshot(digest, DATA.into()).await?; + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None); + // Send only the first partial message; the server must respond without + // ever receiving the rest of the payload. + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name("11111111-1389-4ab5-b188-4a59f22ceb4b", DATA.len()), + write_offset: 0, + finish_write: false, + data: DATA[..5].into(), + })?)) + .await?; + + let response = join_handle.await.expect("Failed to join")?; + assert_eq!(response.into_inner().committed_size, DATA.len() as i64); + drop(tx); + Ok(()) +} + +/// A fresh (non-duplicate) upload with dedup enabled must behave exactly as +/// before: full multi-chunk transfer, then commit. +#[nativelink_test] +pub async fn write_dedup_fresh_upload_unchanged() -> Result<(), Box> { + const DATA: &[u8] = b"write dedup fresh upload payload"; + const SPLIT: usize = 7; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), false).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None); + let resource_name = + make_dedup_resource_name("22222222-1389-4ab5-b188-4a59f22ceb4b", DATA.len()); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: resource_name.clone(), + write_offset: 0, + finish_write: false, + data: DATA[..SPLIT].into(), + })?)) + .await?; + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name, + write_offset: SPLIT as i64, + finish_write: true, + data: DATA[SPLIT..].into(), + })?)) + .await?; + + let response = join_handle.await.expect("Failed to join")?; + assert_eq!(response.into_inner().committed_size, DATA.len() as i64); + let digest = DigestInfo::try_new(HASH1, DATA.len())?; + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, DATA); + Ok(()) +} + +/// A concurrent upload of an in-flight digest must wait for the leading +/// upload and complete early when it durably commits, without transferring +/// its own payload. +#[nativelink_test] +pub async fn write_dedup_joins_inflight_upload() -> Result<(), Box> { + const DATA: &[u8] = b"write dedup join the flight payload"; + const SPLIT: usize = 9; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), false).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + // Leading upload: send the first chunk only, keeping the flight open. + let (leader_tx, leader_join) = make_stream_and_writer_spawn(bs_server.clone(), None); + let leader_resource = + make_dedup_resource_name("33333333-1389-4ab5-b188-4a59f22ceb4b", DATA.len()); + leader_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: leader_resource.clone(), + write_offset: 0, + finish_write: false, + data: DATA[..SPLIT].into(), + })?)) + .await?; + tokio::time::sleep(core::time::Duration::from_millis(50)).await; + + // Duplicate upload (different UUID, same digest): sends only its first + // partial message and then waits on the flight. + let (waiter_tx, waiter_join) = make_stream_and_writer_spawn(bs_server.clone(), None); + waiter_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name( + "44444444-1389-4ab5-b188-4a59f22ceb4b", + DATA.len(), + ), + write_offset: 0, + finish_write: false, + data: DATA[..SPLIT].into(), + })?)) + .await?; + tokio::time::sleep(core::time::Duration::from_millis(50)).await; + + // The waiter must not complete before the leading upload commits. + let mut waiter_join = waiter_join; + assert!( + matches!(poll!(&mut waiter_join), Poll::Pending), + "Waiter completed before the leading upload committed" + ); + + // Finish the leading upload. + leader_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: leader_resource, + write_offset: SPLIT as i64, + finish_write: true, + data: DATA[SPLIT..].into(), + })?)) + .await?; + + let leader_response = leader_join.await.expect("Failed to join")?; + assert_eq!( + leader_response.into_inner().committed_size, + DATA.len() as i64 + ); + let waiter_response = waiter_join.await.expect("Failed to join")?; + assert_eq!( + waiter_response.into_inner().committed_size, + DATA.len() as i64 + ); + + let digest = DigestInfo::try_new(HASH1, DATA.len())?; + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, DATA); + drop(waiter_tx); + Ok(()) +} + +/// When the leading upload fails, waiters must receive a retryable ABORTED +/// error, and a retrying client must succeed as the new leading upload. +#[nativelink_test] +pub async fn write_dedup_waiter_retries_after_leader_failure() +-> Result<(), Box> { + const DATA: &[u8] = b"write dedup leader failure payload"; + const SPLIT: usize = 9; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), false).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + let (leader_tx, leader_join) = make_stream_and_writer_spawn(bs_server.clone(), None); + leader_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name( + "55555555-1389-4ab5-b188-4a59f22ceb4b", + DATA.len(), + ), + write_offset: 0, + finish_write: false, + data: DATA[..SPLIT].into(), + })?)) + .await?; + tokio::time::sleep(core::time::Duration::from_millis(50)).await; + + let (waiter_tx, waiter_join) = make_stream_and_writer_spawn(bs_server.clone(), None); + waiter_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name( + "66666666-1389-4ab5-b188-4a59f22ceb4b", + DATA.len(), + ), + write_offset: 0, + finish_write: false, + data: DATA[..SPLIT].into(), + })?)) + .await?; + tokio::time::sleep(core::time::Duration::from_millis(50)).await; + + // Disconnect the leading upload mid-stream. + drop(leader_tx); + let leader_result = leader_join.await.expect("Failed to join"); + assert!(leader_result.is_err(), "Expected leading upload to fail"); + + let waiter_status = waiter_join + .await + .expect("Failed to join") + .expect_err("Expected waiter to receive an error"); + assert_eq!( + waiter_status.code(), + Code::Aborted, + "Waiter error must be retryable: {waiter_status:?}" + ); + drop(waiter_tx); + + // A retry becomes the new leading upload and succeeds. + let (retry_tx, retry_join) = make_stream_and_writer_spawn(bs_server, None); + retry_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name( + "77777777-1389-4ab5-b188-4a59f22ceb4b", + DATA.len(), + ), + write_offset: 0, + finish_write: true, + data: DATA.into(), + })?)) + .await?; + let retry_response = retry_join.await.expect("Failed to join")?; + assert_eq!( + retry_response.into_inner().committed_size, + DATA.len() as i64 + ); + let digest = DigestInfo::try_new(HASH1, DATA.len())?; + assert_eq!(store.get_part_unchunked(digest, 0, None).await?, DATA); + Ok(()) +} + +/// A compressed upload of an already-durable digest must complete early with +/// `committed_size` -1 per the REAPI compressed-blobs contract. +#[nativelink_test] +pub async fn write_dedup_compressed_early_complete_returns_negative_one() +-> Result<(), Box> { + const DATA: &[u8] = b"write dedup compressed payload"; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), true).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + let digest = DigestInfo::try_new(HASH1, DATA.len())?; + store.update_oneshot(digest, DATA.into()).await?; + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_compressed_resource_name( + "88888888-1389-4ab5-b188-4a59f22ceb4b", + HASH1, + DATA.len(), + ), + write_offset: 0, + finish_write: false, + data: vec![].into(), + })?)) + .await?; + + let response = join_handle.await.expect("Failed to join")?; + assert_eq!(response.into_inner().committed_size, -1); + drop(tx); + Ok(()) +} + +/// A proxy-backed instance must run local write dedup before forwarding a +/// write. Otherwise an already-present upstream blob would be overwritten by +/// the forwarded payload despite dedup being enabled locally. +#[nativelink_test] +pub async fn write_dedup_works_for_grpc_store_proxy() -> Result<(), Box> { + const ORIGINAL: &[u8] = b"proxy-original-data"; + const FORWARDED: &[u8] = b"proxy-forwarded-xxx"; + assert_eq!(ORIGINAL.len(), FORWARDED.len()); + + let upstream_manager = make_store_manager().await?; + let upstream_store = upstream_manager.get_store("main_cas").unwrap(); + let digest = DigestInfo::try_new(HASH1, ORIGINAL.len())?; + upstream_store + .update_oneshot(digest, ORIGINAL.into()) + .await?; + let port = make_tcp_proxy_backend(upstream_manager.as_ref()).await?; + + let local_manager = Arc::new(StoreManager::new()); + let grpc_store = GrpcStore::new(&GrpcSpec { + instance_name: INSTANCE_NAME.to_string(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 0, + tcp_keepalive_s: 0, + http2_keepalive_interval_s: 0, + http2_keepalive_timeout_s: 0, + }], + store_type: StoreType::Cas, + retry: Retry::default(), + max_concurrent_requests: 0, + connections_per_endpoint: 1, + rpc_timeout_s: 0, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + }) + .await?; + local_manager.add_store("main_cas", Store::new(grpc_store))?; + let bs_server = Arc::new( + make_write_dedup_server(local_manager.as_ref(), false).expect("Failed to make server"), + ); + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name( + "bbbbbbbb-1389-4ab5-b188-4a59f22ceb4b", + FORWARDED.len(), + ), + write_offset: 0, + finish_write: true, + data: FORWARDED.into(), + })?)) + .await?; + let response = join_handle.await.expect("Failed to join")?; + assert_eq!(response.into_inner().committed_size, FORWARDED.len() as i64); + assert_eq!( + upstream_store.get_part_unchunked(digest, 0, None).await?, + ORIGINAL + ); + Ok(()) +} + +/// The requested digest function is part of a write flight key. A SHA-256 +/// upload must not make a concurrent BLAKE3 upload of the same hash/size wait +/// for it. +#[nativelink_test] +pub async fn write_dedup_does_not_collide_across_digest_functions() +-> Result<(), Box> { + const DATA: &[u8] = b"write dedup digest function flight key"; + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_write_dedup_server(store_manager.as_ref(), false).expect("Failed to make server"), + ); + + let (leader_tx, leader_join) = make_stream_and_writer_spawn(bs_server.clone(), None); + leader_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name_with_digest_function( + "99999999-1389-4ab5-b188-4a59f22ceb4b", + "sha256", + DATA.len(), + ), + write_offset: 0, + finish_write: false, + data: DATA[..1].into(), + })?)) + .await?; + tokio::time::sleep(core::time::Duration::from_millis(50)).await; + + let (blake3_tx, blake3_join) = make_stream_and_writer_spawn(bs_server, None); + blake3_tx + .send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_dedup_resource_name_with_digest_function( + "aaaaaaaa-1389-4ab5-b188-4a59f22ceb4b", + "blake3", + DATA.len(), + ), + write_offset: 0, + finish_write: true, + data: DATA.into(), + })?)) + .await?; + + let response = tokio::time::timeout(core::time::Duration::from_secs(1), blake3_join) + .await + .expect("BLAKE3 upload incorrectly joined the SHA-256 flight")??; + assert_eq!(response.into_inner().committed_size, DATA.len() as i64); + + drop(blake3_tx); + drop(leader_tx); + assert!(leader_join.await.expect("Failed to join").is_err()); + Ok(()) +} diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index ef6c7fd7e..0f16764f1 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -57,7 +57,7 @@ use opentelemetry::global; use opentelemetry::propagation::Injector; use parking_lot::Mutex; use prost::Message; -use tokio::sync::{Semaphore, oneshot}; +use tokio::sync::{Mutex as AsyncMutex, Semaphore, oneshot}; use tokio::time::sleep; use tonic::metadata::{Ascii, MetadataKey, MetadataValue}; use tonic::{Code, IntoRequest, Request, Response, Status, Streaming}; @@ -1166,7 +1166,7 @@ impl StoreDriver for GrpcStore { ) -> Result { struct LocalState { resource_name: String, - reader: DropCloserReadHalf, + reader: tokio::sync::OwnedMutexGuard, did_error: bool, bytes_received: i64, } @@ -1207,9 +1207,16 @@ impl StoreDriver for GrpcStore { digest_size = digest.size_bytes(), "GrpcStore::update: starting upload for digest", ); + // The reader is shared so that, when the server completes the write + // early (REAPI duplicate-upload contract: another client already + // uploaded this digest), the remainder can be drained after the RPC + // resolves. Without the drain, an upstream sender coupled to this + // reader errors with "receiver disconnected" even though the upload + // succeeded. + let shared_reader = Arc::new(AsyncMutex::new(reader)); let local_state = LocalState { resource_name, - reader, + reader: shared_reader.clone().lock_owned().await, did_error: false, bytes_received: 0, }; @@ -1254,6 +1261,18 @@ impl StoreDriver for GrpcStore { .await .err_tip(|| "in GrpcStore::update()")?; + // Drain any bytes the server chose not to consume (early-completed + // write), but do not make the successful upstream response depend on + // the producer reaching EOF. A worker can still be producing data, or + // can be cancelled while the upstream has already durably committed. + // Keep the reader alive in a detached task so a coupled producer sees + // a live receiver; producer cancellation and drain errors are local + // to that task and must not turn the committed upload into a failure. + let mut reader = shared_reader.lock_owned().await; + background_spawn!("grpc_store_update_drain", async move { + drop(reader.drain().await); + }); + Ok(digest.size_bytes()) } diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index e120402d4..053827d84 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -245,6 +245,52 @@ async fn write_update_works_with_legacy_resource_names() -> Result<(), Error> { write_update_works_core(true, upload_pattern).await } +#[nativelink_test] +async fn early_completed_update_does_not_wait_for_stalled_producer() -> Result<(), Error> { + let (_server, port) = make_fake_bytestream_server().await; + let spec = test_spec(format!("http://localhost:{port}"), false); + let store = GrpcStore::new(&spec).await?; + let digest = DigestInfo::try_new(VALID_HASH, 9)?; + let (mut tx, rx) = make_buf_channel_pair(); + + tx.send("first".into()).await?; + let result = timeout( + Duration::from_secs(1), + store.update(digest, rx, UploadSizeInfo::ExactSize(9)), + ) + .await??; + assert_eq!(result, 9); + + // The detached drain keeps the producer coupled to a live receiver after + // the upstream has already committed, without making update() wait for + // this data or its EOF. + tx.send("tail!!!".into()).await?; + tx.send_eof()?; + Ok(()) +} + +#[nativelink_test] +async fn early_completed_update_succeeds_when_producer_is_cancelled() -> Result<(), Error> { + let (_server, port) = make_fake_bytestream_server().await; + let spec = test_spec(format!("http://localhost:{port}"), false); + let store = GrpcStore::new(&spec).await?; + let digest = DigestInfo::try_new(VALID_HASH, 9)?; + let (mut tx, rx) = make_buf_channel_pair(); + + tx.send("first".into()).await?; + let result = timeout( + Duration::from_secs(1), + store.update(digest, rx, UploadSizeInfo::ExactSize(9)), + ) + .await??; + assert_eq!(result, 9); + + // A producer that is cancelled without sending EOF must not turn an + // already successful upstream commit into an update failure. + drop(tx); + Ok(()) +} + async fn read_works_core( use_legacy_resource_names: bool, upload_pattern: &str,