diff --git a/x/mongo/driver/operation.go b/x/mongo/driver/operation.go index 50136456e4..197e1a8440 100644 --- a/x/mongo/driver/operation.go +++ b/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 diff --git a/x/mongo/driver/operation_test.go b/x/mongo/driver/operation_test.go index 911f32dbf5..8fe1dafd03 100644 --- a/x/mongo/driver/operation_test.go +++ b/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) + } + }) + } +}