GODRIVER-4096 Fix panic in Collection.insert on out-of-order write errors (#2557)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Matt Dale <9760375+matthewdale@users.noreply.github.com>

Author: 3 people

internal/integration/collection_test.go

@@ -257,6 +257,96 @@ func TestCollection(t *testing.T) {
 				})
 			}
 		})
+		// The server typically returns write errors for an unordered insert in
+		// ascending "index" order. However, in rare cases, the server may
+		// return write errors indexes not in ascending order. Use a mock
+		// deployment to create server responses where the write error indexes
+		// are not in ascending order.
+		//
+		// See GODRIVER-4096 for more background.
+		mt.RunOpts("out-of-order write errors", mtest.NewOptions().ClientType(mtest.Mock), func(mt *mtest.T) {
+			testCases := []struct {
+				name        string
+				docs        []any
+				response    bson.D
+				wantErrs    int
+				wantInserts []any
+			}{
+				{
+					// Before the fix, this response caused a panic. Note the
+					// "writeErrors.index" values are not in ascending order.
+					name: "all documents fail",
+					docs: []any{
+						bson.D{{"_id", int32(0)}},
+						bson.D{{"_id", int32(1)}},
+					},
+					response: bson.D{
+						{"ok", 1},
+						{"n", 0},
+						{"writeErrors", bson.A{
+							bson.D{{"index", 1}, {"code", errorDuplicateKey}, {"errmsg", "duplicate key error"}},
+							bson.D{{"index", 0}, {"code", errorDuplicateKey}, {"errmsg", "duplicate key error"}},
+						}},
+					},
+					wantErrs:    2,
+					wantInserts: []any{},
+				},
+				{
+					// Before the fix, this removed the wrong IDs from the
+					// result without panicking. Note the "writeErrors.index"
+					// values are not in ascending order.
+					name: "some documents fail",
+					docs: []any{
+						bson.D{{"_id", int32(0)}},
+						bson.D{{"_id", int32(1)}},
+						bson.D{{"_id", int32(2)}},
+						bson.D{{"_id", int32(3)}},
+						bson.D{{"_id", int32(4)}},
+					},
+					response: bson.D{
+						{"ok", 1},
+						{"n", 3},
+						{"writeErrors", bson.A{
+							bson.D{{"index", 3}, {"code", errorDuplicateKey}, {"errmsg", "duplicate key error"}},
+							bson.D{{"index", 1}, {"code", errorDuplicateKey}, {"errmsg", "duplicate key error"}},
+						}},
+					},
+					wantErrs:    2,
+					wantInserts: []any{int32(0), int32(2), int32(4)},
+				},
+			}
+
+			for _, tc := range testCases {
+				mt.Run(tc.name, func(mt *mtest.T) {
+					mt.AddMockResponses(tc.response)
+
+					res, err := mt.Coll.InsertMany(
+						context.Background(),
+						tc.docs,
+						options.InsertMany().SetOrdered(false),
+					)
+
+					var bwe mongo.BulkWriteException
+					require.True(
+						mt,
+						errors.As(err, &bwe),
+						"expected error to be a mongo.BulkWriteException, got %#v",
+						err,
+					)
+					assert.Len(
+						mt,
+						bwe.WriteErrors,
+						tc.wantErrs,
+						"expected %v write errors, got %v",
+						tc.wantErrs,
+						len(bwe.WriteErrors),
+					)
+
+					require.NotNil(mt, res, "expected a non-nil result")
+					assert.Equal(mt, tc.wantInserts, res.InsertedIDs, "expected inserted IDs to match")
+				})
+			}
+		})
 		mt.Run("writeError index", func(mt *mtest.T) {
 			mt.Parallel()
 

mongo/collection.go

@@ -367,19 +367,49 @@ func (coll *Collection) insert(
 		return result, err
 	}
 
-	// remove the ids that had writeErrors from result
-	for i, we := range wce.WriteErrors {
-		// i indexes have been removed before the current error, so the index is we.Index-i
-		idIndex := int(we.Index) - i
-		// if the insert is ordered, nothing after the error was inserted
-		if args.Ordered == nil || *args.Ordered {
-			result = result[:idIndex]
-			break
+	ordered := args.Ordered == nil || *args.Ordered
+	result = keepInsertedIDs(result, wce.WriteErrors, ordered)
+
+	return result, err
+}
+
+// keepInsertedIDs returns the subset of result whose corresponding documents
+// were not rejected by the server. Each WriteError.Index refers to the
+// positional index of the document in the original insert request, which is
+// also the index space of result.
+//
+// writeErrors is not required to be sorted by Index: the server does not
+// guarantee that write errors for an unordered bulk write are returned in
+// ascending Index order, and assuming otherwise previously caused
+// out-of-range slice panics (GODRIVER-4096). Neither result nor writeErrors
+// is modified.
+func keepInsertedIDs(result []any, writeErrors driver.WriteErrors, ordered bool) []any {
+	// If there are no write errors, then every document was inserted.
+	if len(writeErrors) == 0 {
+		return result
+	}
+
+	failed := make(map[int]bool, len(writeErrors))
+	for _, writeError := range writeErrors {
+		failed[int(writeError.Index)] = true
+	}
+
+	kept := result[:0:0]
+	for idx, id := range result {
+		if failed[idx] {
+			// The server stops an ordered insert at the first document that
+			// fails, so nothing at or after that index was inserted.
+			if ordered {
+				break
+			}
+
+			continue
 		}
-		result = append(result[:idIndex], result[idIndex+1:]...)
+
+		kept = append(kept, id)
 	}
 
-	return result, err
+	return kept
 }
 
 // InsertOne executes an insert command to insert a single document into the collection.

mongo/collection_test.go

@@ -21,6 +21,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
 	"go.mongodb.org/mongo-driver/v2/mongo/writeconcern"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/topology"
 )
 
@@ -307,3 +308,77 @@ func TestNewFindArgsFromFindOneArgs(t *testing.T) {
 		})
 	}
 }
+
+func TestKeepInsertedIDs(t *testing.T) {
+	newResult := func(n int) []any {
+		result := make([]any, n)
+		for i := range result {
+			result[i] = i
+		}
+		return result
+	}
+
+	tests := []struct {
+		name        string
+		result      []any
+		writeErrors driver.WriteErrors
+		ordered     bool
+		want        []any
+	}{
+		{
+			name:    "no write errors",
+			result:  newResult(3),
+			ordered: false,
+			want:    newResult(3),
+		},
+		{
+			name:        "unordered, ascending write errors",
+			result:      newResult(5),
+			writeErrors: driver.WriteErrors{{Index: 1}, {Index: 3}},
+			ordered:     false,
+			want:        []any{0, 2, 4},
+		},
+		{
+			// The server does not guarantee ascending Index order for an
+			// unordered bulk write. Before the fix, this silently removed
+			// the wrong document.
+			name:        "unordered, out-of-order write errors",
+			result:      newResult(5),
+			writeErrors: driver.WriteErrors{{Index: 3}, {Index: 1}},
+			ordered:     false,
+			want:        []any{0, 2, 4},
+		},
+		{
+			// Regression test for GODRIVER-4096: before the fix, this
+			// input caused "slice bounds out of range [:-1]".
+			name:        "unordered, all documents fail, out of order",
+			result:      newResult(2),
+			writeErrors: driver.WriteErrors{{Index: 1}, {Index: 0}},
+			ordered:     false,
+			want:        []any{},
+		},
+		{
+			name:        "ordered, single write error",
+			result:      newResult(5),
+			writeErrors: driver.WriteErrors{{Index: 2}},
+			ordered:     true,
+			want:        []any{0, 1},
+		},
+		{
+			// Sanity check for the ordered path: stop at the earliest
+			// failed index, not at whichever error is listed first.
+			name:        "ordered, out-of-order write errors",
+			result:      newResult(2),
+			writeErrors: driver.WriteErrors{{Index: 1}, {Index: 0}},
+			ordered:     true,
+			want:        []any{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := keepInsertedIDs(test.result, test.writeErrors, test.ordered)
+			require.Equal(t, test.want, got)
+		})
+	}
+}