Optimize dockey extraction by avoiding bson.Marshal.

This code gets called repeatedly, so it’s worth optimizing.
This also optimizes the field-names uniqueness check.

Author: FGasper

dockey/agg_test.go

@@ -4,13 +4,14 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/10gen/migration-verifier/mslices"
 	"github.com/pkg/errors"
-	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
-// This extracts the document key from a document gets its field names.
+// ExtractTrueDocKeyFromDoc extracts the document key from a document
+// given its field names.
 //
 // NB: This avoids the problem documented in SERVER-109340; as a result,
 // the returned key may not always match the change stream’s `documentKey`
@@ -21,7 +22,8 @@ func ExtractTrueDocKeyFromDoc(
 ) (bson.Raw, error) {
 	assertFieldNameUniqueness(fieldNames)
 
-	var dk bson.D
+	docBuilder := bsoncore.NewDocumentBuilder()
+
 	for _, field := range fieldNames {
 		var val bson.RawValue
 
@@ -38,19 +40,20 @@ func ExtractTrueDocKeyFromDoc(
 			return nil, errors.Wrapf(err, "extracting doc key field %#q from doc %+v", field, doc)
 		}
 
-		dk = append(dk, bson.E{field, val})
-	}
-
-	docKey, err := bson.Marshal(dk)
-	if err != nil {
-		return nil, errors.Wrapf(err, "marshaling doc key %v from doc %v", dk, docKey)
+		docBuilder.AppendValue(
+			field,
+			bsoncore.Value{
+				Type: val.Type,
+				Data: val.Value,
+			},
+		)
 	}
 
-	return docKey, nil
+	return bson.Raw(docBuilder.Build()), nil
 }
 
 func assertFieldNameUniqueness(fieldNames []string) {
-	if len(lo.Uniq(fieldNames)) != len(fieldNames) {
+	if mslices.FindFirstDupe(fieldNames).IsSome() {
 		panic(fmt.Sprintf("Duplicate field names: %v", fieldNames))
 	}
 }

dockey/raw.go

@@ -46,15 +46,4 @@ func TestExtractTrueDocKeyFromDoc(t *testing.T) {
 			)
 		}
 	}
-
-	assert.Panics(
-		t,
-		func() {
-			_, _ = ExtractTrueDocKeyFromDoc(
-				[]string{"foo", "bar", "foo"},
-				bson.Raw{0},
-			)
-		},
-		"duplicate field name should cause panic",
-	)
 }

dockey/raw_test.go

@@ -3,6 +3,8 @@ package mslices
 import (
 	"reflect"
 	"slices"
+
+	"github.com/10gen/migration-verifier/option"
 )
 
 // This package complements the Go standard library’s package of the
@@ -43,3 +45,23 @@ func Compact[T any, S ~[]T](slc S) S {
 func isZero[T any](val T) bool {
 	return reflect.ValueOf(&val).Elem().IsZero()
 }
+
+// FindFirstDupe returns the first item in the slice that has at
+// least 1 duplicate, or none if all slice members are unique.
+func FindFirstDupe[T comparable](items []T) option.Option[T] {
+	for i := range items {
+		j := 1 + i
+
+		if j == len(items) {
+			break
+		}
+
+		for ; j < len(items); j++ {
+			if items[i] == items[j] {
+				return option.Some(items[i])
+			}
+		}
+	}
+
+	return option.None[T]()
+}