Avoid BSON unmarshal in extractTimestampFromResumeToken.

Unmarshaling is slow due to reflection. This function gets called
frequently enough to justify the optimization.

Author: FGasper

internal/verifier/change_stream.go

@@ -10,6 +10,7 @@ import (
 	"github.com/10gen/migration-verifier/internal/retry"
 	"github.com/10gen/migration-verifier/internal/types"
 	"github.com/10gen/migration-verifier/internal/util"
+	"github.com/10gen/migration-verifier/mbson"
 	"github.com/10gen/migration-verifier/msync"
 	"github.com/10gen/migration-verifier/option"
 	mapset "github.com/deckarep/golang-set/v2"
@@ -857,17 +858,19 @@ func (csr *ChangeStreamReader) persistChangeStreamResumeToken(ctx context.Contex
 }
 
 func extractTimestampFromResumeToken(resumeToken bson.Raw) (primitive.Timestamp, error) {
-	tokenStruct := struct {
-		Data string `bson:"_data"`
-	}{}
-
 	// Change stream token is always a V1 keystring in the _data field
-	err := bson.Unmarshal(resumeToken, &tokenStruct)
+	tokenDataRV, err := resumeToken.LookupErr("_data")
+
+	if err != nil {
+		return primitive.Timestamp{}, errors.Wrapf(err, "extracting %#q from resume token (%v)", "_data", resumeToken)
+	}
+
+	tokenData, err := mbson.CastRawValue[string](tokenDataRV)
 	if err != nil {
-		return primitive.Timestamp{}, errors.Wrapf(err, "failed to extract %#q from resume token (%v)", "_data", resumeToken)
+		return primitive.Timestamp{}, errors.Wrapf(err, "parsing resume token (%v)", "_data", resumeToken)
 	}
 
-	resumeTokenBson, err := keystring.KeystringToBson(keystring.V1, tokenStruct.Data)
+	resumeTokenBson, err := keystring.KeystringToBson(keystring.V1, tokenData)
 	if err != nil {
 		return primitive.Timestamp{}, err
 	}

mbson/raw_value.go

@@ -0,0 +1,50 @@
+package mbson
+
+import (
+	"fmt"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsontype"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+)
+
+type bsonType interface {
+	bson.Raw | primitive.Timestamp | string
+}
+
+type cannotCastErr struct {
+	gotBSONType bsontype.Type
+	toGoType    any
+}
+
+func (ce cannotCastErr) Error() string {
+	return fmt.Sprintf("cannot cast BSON %s to %T", ce.gotBSONType, ce.toGoType)
+}
+
+// CastRawValue is a “one-stop-shop” interface around bson.RawValue’s various
+// casting interfaces. Unlike those functions, though, this returns an error
+// if the target type doesn’t match the value.
+//
+// Augment bsonType if you find a type here that’s missing.
+func CastRawValue[T bsonType](in bson.RawValue) (T, error) {
+	retPtr := new(T)
+
+	switch any(*retPtr).(type) {
+	case bson.Raw:
+		if doc, isDoc := in.DocumentOK(); isDoc {
+			return any(doc).(T), nil
+		}
+	case primitive.Timestamp:
+		if t, i, ok := in.TimestampOK(); ok {
+			return any(primitive.Timestamp{t, i}).(T), nil
+		}
+	case string:
+		if str, ok := in.StringValueOK(); ok {
+			return any(str).(T), nil
+		}
+	default:
+		panic(fmt.Sprintf("Unrecognized Go type: %T (maybe augment bsonType?)", *retPtr))
+	}
+
+	return *retPtr, cannotCastErr{in.Type, *retPtr}
+}