fix(store): two ExperimentalMongoStore update fixes (truncated stream, version-CAS)

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
             {

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(())
+}