Skip to content

Commit 910eb26

Browse files
authored
optimize allocations in redactStartedInformationCmd (#2129)
1 parent 779938e commit 910eb26

File tree

2 files changed

+59
-2
lines changed

2 files changed

+59
-2
lines changed

x/mongo/driver/operation.go

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,13 +158,43 @@ type ResponseInfo struct {
158158
}
159159

160160
func redactStartedInformationCmd(info startedInformation) bson.Raw {
161+
intLen := func(n int) int {
162+
if n == 0 {
163+
return 1 // Special case: 0 has one digit
164+
}
165+
count := 0
166+
for n > 0 {
167+
n /= 10
168+
count++
169+
}
170+
return count
171+
}
172+
161173
var cmdCopy bson.Raw
162174

163175
// Make a copy of the command. Redact if the command is security
164176
// sensitive and cannot be monitored. If there was a type 1 payload for
165177
// the current batch, convert it to a BSON array
166178
if !info.redacted {
167-
cmdCopy = make([]byte, 0, len(info.cmd))
179+
cmdLen := len(info.cmd)
180+
for _, seq := range info.documentSequences {
181+
cmdLen += 7 // 2 (header) + 4 (array length) + 1 (array end)
182+
cmdLen += len(seq.identifier)
183+
data := seq.data
184+
i := 0
185+
for {
186+
doc, rest, ok := bsoncore.ReadDocument(data)
187+
if !ok {
188+
break
189+
}
190+
data = rest
191+
cmdLen += len(doc)
192+
cmdLen += intLen(i)
193+
i++
194+
}
195+
}
196+
197+
cmdCopy = make([]byte, 0, cmdLen)
168198
cmdCopy = append(cmdCopy, info.cmd...)
169199

170200
if len(info.documentSequences) > 0 {
@@ -1808,7 +1838,6 @@ func (op Operation) createReadPref(desc description.SelectedServer, isOpQuery bo
18081838
// TODO if supplied readPreference was "overwritten" with primary in description.selectForReplicaSet.
18091839
if desc.Server.Kind == description.ServerKindStandalone || (isOpQuery &&
18101840
desc.Server.Kind != description.ServerKindMongos) ||
1811-
18121841
op.Type == Write || (op.IsOutputAggregate && desc.Server.WireVersion.Max < 13) {
18131842
// Don't send read preference for:
18141843
// 1. all standalones

x/mongo/driver/operation_test.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"bytes"
1111
"context"
1212
"errors"
13+
"strconv"
1314
"testing"
1415
"time"
1516

@@ -1063,3 +1064,30 @@ func TestMarshalBSONWriteConcern(t *testing.T) {
10631064
})
10641065
}
10651066
}
1067+
1068+
func BenchmarkRedactStartedInformationCmd(b *testing.B) {
1069+
for _, size := range []int{0, 1, 5, 10, 100, 1000} {
1070+
info := startedInformation{
1071+
cmd: make([]byte, 100),
1072+
documentSequences: make([]struct {
1073+
identifier string
1074+
data []byte
1075+
}, size),
1076+
}
1077+
for i := 0; i < size; i++ {
1078+
info.documentSequences[i] = struct {
1079+
identifier string
1080+
data []byte
1081+
}{
1082+
identifier: strconv.Itoa(i),
1083+
data: make([]byte, 100),
1084+
}
1085+
}
1086+
b.Run(strconv.Itoa(size), func(b *testing.B) {
1087+
b.ReportAllocs()
1088+
for i := 0; i < b.N; i++ {
1089+
redactStartedInformationCmd(info)
1090+
}
1091+
})
1092+
}
1093+
}

0 commit comments

Comments
 (0)