From 5d1e46295e63e5afa619a97bef28b3e3081bd4ef Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:54:06 -0400 Subject: [PATCH 1/2] fix(store): propagate mid-stream errors in ExperimentalMongoStore::update update read the upload stream with while let Ok(chunk) = reader.recv().await, which treats a stream Err identically to the empty-chunk EOF. An upload that fails partway (client disconnect, or an upstream store erroring mid-copy) was upserted as whatever bytes arrived and returned Ok(partial_len) -- a truncated blob stored under the full digest key, served as success on every later read. Propagate the recv error instead, and (belt-and-suspenders) reject a completed read whose size disagrees with UploadSizeInfo::ExactSize. Regression update_rejects_truncated_stream_instead_of_committing_it drops the writer mid-stream without send_eof and asserts update returns Err and the digest is absent; fails without the fix (commits the truncation as Ok). Co-authored-by: Claude Opus 4.8 --- nativelink-store/src/mongo_store.rs | 23 +++++++++++- nativelink-store/tests/mongo_store_test.rs | 43 ++++++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index c59b4b1764c..b6a0fc853cd 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -484,15 +484,34 @@ impl StoreDriver for ExperimentalMongoStore { } } - // Read all data into memory with proper EOF handling + // Read all data into memory. A mid-stream error must be propagated, + // not treated as EOF: the old `while let Ok(chunk)` swallowed a stream + // `Err` identically to the empty-chunk EOF, committing a truncated blob + // under the digest's key (a torn write served as success). let mut data = Vec::new(); - while let Ok(chunk) = reader.recv().await { + loop { + let chunk = reader.recv().await.map_err(|e| { + make_err!( + Code::Internal, + "Failed to read stream in ExperimentalMongoStore::update: {e}" + ) + })?; if chunk.is_empty() { break; // Empty chunk signals EOF } data.extend_from_slice(&chunk); } + if let UploadSizeInfo::ExactSize(expected) = upload_size + && data.len() as u64 != expected + { + return Err(make_err!( + Code::Internal, + "ExperimentalMongoStore::update received {} bytes but expected {expected} (ExactSize)", + data.len() + )); + } + let size = data.len().try_into().unwrap_or(i64::MAX); // Create document diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs index dac6689d4f2..3c047d78e83 100644 --- a/nativelink-store/tests/mongo_store_test.rs +++ b/nativelink-store/tests/mongo_store_test.rs @@ -1027,3 +1027,46 @@ async fn single_request_permit() -> Result<(), Error> { Ok(()) } + +#[nativelink_test] +async fn update_rejects_truncated_stream_instead_of_committing_it() -> Result<(), Error> { + // A mid-stream failure (sender dropped before send_eof) must surface as an + // Err, not be committed as a successful upload. Before the fix, the + // `while let Ok(chunk)` loop treated the stream Err identically to the + // empty-chunk EOF and upserted the partial bytes as Ok(partial_len), + // storing a truncated blob under the full digest's key. + let data = Bytes::from(vec![b'A'; 256]); + let digest = DigestInfo::try_new(VALID_HASH1, data.len())?; + + let helper = TestMongoHelper::new(None).await?; + let store = helper.store.clone(); + + let (mut tx, rx) = make_buf_channel_pair(); + + let (update_result, send_result) = tokio::join!( + async { + store + .update(digest, rx, UploadSizeInfo::ExactSize(data.len() as u64)) + .await + }, + async { + // Send one non-empty chunk, then drop the writer WITHOUT send_eof, + // so the reader sees the chunk followed by an Err. + tx.send(data.clone()).await?; + drop(tx); + Ok::<_, Error>(()) + }, + ); + send_result?; + assert!( + update_result.is_err(), + "update must return Err on a mid-stream failure, not commit a truncated blob as success", + ); + + assert!( + store.has(digest).await?.is_none(), + "a torn upload must not leave the digest present", + ); + + Ok(()) +} From c36fa2e3b01b4ebb49f842f8d816defec5abc717 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 11:28:56 -0400 Subject: [PATCH 2/2] fix(store): mongo update_data version-CAS must return retryable Ok(None), not Err ExperimentalMongoStore::update_data used find_one_and_update(...).upsert(true). On a version-CAS conflict (a concurrent writer already bumped the version) the {_id, version} filter matches nothing, so upsert tries to INSERT a doc with the existing _id and MongoDB returns a DuplicateKeyError (E11000), surfaced as Err(Internal). The version-CAS contract (and the Redis backend, which returns Ok(None)) require a conflict to be retryable; instead concurrent updates on a Mongo-backed record spuriously failed. Fix: .upsert(current_version == 0) so only a genuine first-create upserts; a non-matching filter on an existing record now returns Ok(None) (retryable). Adds a deterministic fails-without regression and corrects an existing (ignored) scheduler test that had pinned the buggy Err behavior. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/mongo_store.rs | 2 +- nativelink-store/tests/mongo_store_test.rs | 52 +++++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index b6a0fc853cd..e4ee3ed9381 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -992,7 +992,7 @@ impl SchedulerStore for ExperimentalMongoStore { let result = match self .scheduler_collection .find_one_and_update(filter, update_doc) - .upsert(true) + .upsert(current_version == 0) .return_document(ReturnDocument::After) .await { diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs index 3c047d78e83..f8e1519ca2e 100644 --- a/nativelink-store/tests/mongo_store_test.rs +++ b/nativelink-store/tests/mongo_store_test.rs @@ -839,7 +839,10 @@ async fn test_scheduler_store_operations() -> Result<(), Error> { data.version = version1; // Using old version let result = helper.store.update_data(data.clone(), None).await; - assert!(result.is_err(), "Update with old version should fail"); + assert!( + matches!(result, Ok(None)), + "Stale-version update must return retryable Ok(None), got {result:?}" + ); eprintln!("Correctly rejected update with stale version"); } @@ -1070,3 +1073,50 @@ async fn update_rejects_truncated_stream_instead_of_committing_it() -> Result<() Ok(()) } + +#[nativelink_test] +async fn update_data_stale_version_returns_retryable_none_not_err() -> Result<(), Error> { + // N12 regression: ExperimentalMongoStore::update_data used upsert(true), so a + // version-CAS miss (a concurrent writer already bumped the version) hit the + // unique _id and surfaced as a DuplicateKeyError (Err) instead of the + // retryable Ok(None) the scheduler contract (and the Redis backend) require. + let helper = TestMongoHelper::new(None).await?; + + let mut data = TestSchedulerData { + key: "test:n12_version_cas".to_string(), + content: "v1".to_string(), + version: 0, + }; + + // Seed at version 0 -> creates version 1 (genuine first create still upserts). + let v1 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "create should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on create"))?; + + // One writer wins with the correct version -> version 2. + data.content = "v2".to_string(); + data.version = v1; + let v2 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "correct-version update should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on update"))?; + assert!(v2 > v1, "version should increment ({v1} -> {v2})"); + + // The losing concurrent writer carries the now-stale version. It MUST get a + // retryable Ok(None), not Err. Before the fix this returned Err(Internal) + // (duplicate _id), destroying the retry signal. + data.content = "stale".to_string(); + data.version = v1; + let stale = helper.store.update_data(data.clone(), None).await; + assert!( + matches!(stale, Ok(None)), + "stale-version update must return retryable Ok(None), got {stale:?}" + ); + + Ok(()) +}