optimize allocations in redactStartedInformationCmd

x/mongo/driver/operation.go

@@ -158,13 +158,43 @@ type ResponseInfo struct {
 }
 
 func redactStartedInformationCmd(info startedInformation) bson.Raw {
+	intLen := func(n int) int {
+		if n == 0 {
+			return 1 // Special case: 0 has one digit
+		}
+		count := 0
+		for n > 0 {
+			n /= 10
+			count++
+		}
+		return count
+	}
+
 	var cmdCopy bson.Raw
 
 	// Make a copy of the command. Redact if the command is security
 	// sensitive and cannot be monitored. If there was a type 1 payload for
 	// the current batch, convert it to a BSON array
 	if !info.redacted {
-		cmdCopy = make([]byte, 0, len(info.cmd))
+		cmdLen := len(info.cmd)
+		for _, seq := range info.documentSequences {
+			cmdLen += 7 // 2 (header) + 4 (array length) + 1 (array end)
+			cmdLen += len(seq.identifier)
+			data := seq.data
+			i := 0
+			for {
+				doc, rest, ok := bsoncore.ReadDocument(data)
+				if !ok {
+					break
+				}
+				data = rest
+				cmdLen += len(doc)
+				cmdLen += intLen(i)
+				i++
+			}
+		}
+
+		cmdCopy = make([]byte, 0, cmdLen)
 		cmdCopy = append(cmdCopy, info.cmd...)
 
 		if len(info.documentSequences) > 0 {
@@ -1808,7 +1838,6 @@ func (op Operation) createReadPref(desc description.SelectedServer, isOpQuery bo
 	// TODO if supplied readPreference was "overwritten" with primary in description.selectForReplicaSet.
 	if desc.Server.Kind == description.ServerKindStandalone || (isOpQuery &&
 		desc.Server.Kind != description.ServerKindMongos) ||
-
 		op.Type == Write || (op.IsOutputAggregate && desc.Server.WireVersion.Max < 13) {
 		// Don't send read preference for:
 		// 1. all standalones

x/mongo/driver/operation_test.go

@@ -10,6 +10,7 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"strconv"
 	"testing"
 	"time"
 
@@ -1063,3 +1064,30 @@ func TestMarshalBSONWriteConcern(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkRedactStartedInformationCmd(b *testing.B) {
+	for _, size := range []int{0, 1, 5, 10, 100, 1000} {
+		info := startedInformation{
+			cmd: make([]byte, 100),
+			documentSequences: make([]struct {
+				identifier string
+				data       []byte
+			}, size),
+		}
+		for i := 0; i < size; i++ {
+			info.documentSequences[i] = struct {
+				identifier string
+				data       []byte
+			}{
+				identifier: strconv.Itoa(i),
+				data:       make([]byte, 100),
+			}
+		}
+		b.Run(strconv.Itoa(size), func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				redactStartedInformationCmd(info)
+			}
+		})
+	}
+}