diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index c59b4b1764c..e4ee3ed9381 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 @@ -973,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 dac6689d4f2..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"); } @@ -1027,3 +1030,93 @@ 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(()) +} + +#[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(()) +}