Ensure insertedIds contain ids from all batches (#850)

The previous code was incorrect because it was comparing absolute write indexes
with indexes that are relative to the current batch. This patch avoids that
by using the insertedId map from SplittablePayload directly, which already
contains absolute write indexes.

JAVA-4436

Author: jyemin

driver-core/src/main/com/mongodb/internal/operation/BulkWriteBatch.java

@@ -53,12 +53,11 @@
 import org.bson.codecs.configuration.CodecRegistry;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 import static com.mongodb.internal.bulk.WriteRequest.Type.DELETE;
 import static com.mongodb.internal.bulk.WriteRequest.Type.INSERT;
@@ -282,21 +281,11 @@ private BulkWriteResult getBulkWriteResult(final BsonDocument result) {
     }
 
     private List<BulkWriteInsert> getInsertedItems(final BsonDocument result) {
-        if (payload.getPayloadType() == SplittablePayload.Type.INSERT) {
-
-            Stream<WriteRequestWithIndex> writeRequests = payload.getWriteRequestWithIndexes().stream();
-            List<Integer> writeErrors = getWriteErrors(result).stream().map(BulkWriteError::getIndex).collect(Collectors.toList());
-            if (!writeErrors.isEmpty()) {
-                writeRequests = writeRequests.filter(wr -> !writeErrors.contains(wr.getIndex()));
-            }
-            if (payload.getPosition() < payload.size()) {
-                writeRequests = writeRequests.filter(wr -> wr.getIndex() < payload.getPosition());
-            }
-            return writeRequests
-                    .map(wr -> new BulkWriteInsert(wr.getIndex(), payload.getInsertedIds().get(wr.getIndex())))
-                    .collect(Collectors.toList());
-        }
-        return Collections.emptyList();
+        Set<Integer> writeErrors = getWriteErrors(result).stream().map(BulkWriteError::getIndex).collect(Collectors.toSet());
+        return payload.getInsertedIds().entrySet().stream()
+                .filter(entry -> !writeErrors.contains(entry.getKey()))
+                .map(entry -> new BulkWriteInsert(entry.getKey(), entry.getValue()))
+                .collect(Collectors.toList());
     }
 
 

driver-core/src/test/unit/com/mongodb/internal/operation/BulkWriteBatchSpecification.groovy

@@ -16,6 +16,7 @@
 
 package com.mongodb.internal.operation
 
+import com.mongodb.MongoBulkWriteException
 import com.mongodb.MongoNamespace
 import com.mongodb.ReadConcern
 import com.mongodb.ServerAddress
@@ -228,23 +229,71 @@ class BulkWriteBatchSpecification extends Specification {
         !bulkWriteBatch.hasAnotherBatch()
     }
 
-    def 'should only map inserts up to the payload position'() {
+    def 'should map all inserted ids'() {
         when:
         def bulkWriteBatch = BulkWriteBatch.createBulkWriteBatch(namespace, serverDescription, connectionDescription, false,
-                WriteConcern.ACKNOWLEDGED, null, false, getWriteRequests()[3..4], sessionContext)
+                WriteConcern.ACKNOWLEDGED, null, false,
+                [new InsertRequest(toBsonDocument('{_id: 0}')),
+                 new InsertRequest(toBsonDocument('{_id: 1}')),
+                 new InsertRequest(toBsonDocument('{_id: 2}'))
+                ],
+                sessionContext)
         def payload = bulkWriteBatch.getPayload()
         payload.setPosition(1)
+        payload.insertedIds.put(0, new BsonInt32(0))
         bulkWriteBatch.addResult(BsonDocument.parse('{"n": 1, "ok": 1.0}'))
 
         then:
-        bulkWriteBatch.getResult().inserts == [new BulkWriteInsert(0, null)]
+        bulkWriteBatch.getResult().inserts == [new BulkWriteInsert(0, new BsonInt32(0))]
 
         when:
-        payload.setPosition(2)
+        bulkWriteBatch = bulkWriteBatch.getNextBatch()
+        payload = bulkWriteBatch.getPayload()
+        payload.setPosition(1)
+        payload.insertedIds.put(1, new BsonInt32(1))
+        bulkWriteBatch.addResult(BsonDocument.parse('{"n": 1, "ok": 1.0}'))
+
+        then:
+        bulkWriteBatch.getResult().inserts == [new BulkWriteInsert(0, new BsonInt32(0)),
+                                               new BulkWriteInsert(1, new BsonInt32(1))]
+
+        when:
+        bulkWriteBatch = bulkWriteBatch.getNextBatch()
+        payload = bulkWriteBatch.getPayload()
+        payload.setPosition(1)
+        payload.insertedIds.put(2, new BsonInt32(2))
         bulkWriteBatch.addResult(BsonDocument.parse('{"n": 1, "ok": 1.0}'))
 
         then:
-        bulkWriteBatch.getResult().inserts == [new BulkWriteInsert(0, null), new BulkWriteInsert(1, null)]
+        bulkWriteBatch.getResult().inserts == [new BulkWriteInsert(0, new BsonInt32(0)),
+                                               new BulkWriteInsert(1, new BsonInt32(1)),
+                                               new BulkWriteInsert(2, new BsonInt32(2))]
+    }
+
+    def 'should not map inserted id with a write error'() {
+        given:
+        def bulkWriteBatch = BulkWriteBatch.createBulkWriteBatch(namespace, serverDescription, connectionDescription, false,
+                WriteConcern.ACKNOWLEDGED, null, false,
+                [new InsertRequest(toBsonDocument('{_id: 0}')),
+                 new InsertRequest(toBsonDocument('{_id: 1}')),
+                 new InsertRequest(toBsonDocument('{_id: 2}'))
+                ],
+                sessionContext)
+        def payload = bulkWriteBatch.getPayload()
+        payload.setPosition(3)
+        payload.insertedIds.put(0, new BsonInt32(0))
+        payload.insertedIds.put(1, new BsonInt32(1))
+        payload.insertedIds.put(2, new BsonInt32(2))
+
+        when:
+        bulkWriteBatch.addResult(toBsonDocument('''{"ok": 1, "n": 2,
+            "writeErrors": [{ "index" : 1, "code" : 11000, "errmsg": "duplicate key error"}] }'''))
+        bulkWriteBatch.getResult()
+
+        then:
+        def ex = thrown(MongoBulkWriteException)
+        ex.getWriteResult().inserts == [new BulkWriteInsert(0, new BsonInt32(0)),
+                                        new BulkWriteInsert(2, new BsonInt32(2))]
     }
 
     def 'should not retry when at least one write is not retryable'() {