Use bson2 package for wire queries and replies (#4108)

FerretDB · Feb 22, 2024 · 033dfef · 033dfef
1 parent 9de2931
commit 033dfef
Show file tree

Hide file tree

Showing 11 changed files with 758 additions and 721 deletions.
diff --git a/internal/handler/cmd_query.go b/internal/handler/cmd_query.go
@@ -41,9 +41,7 @@ func (h *Handler) CmdQuery(ctx context.Context, query *wire.OpQuery) (*wire.OpRe
 
 	if cmd == "saslStart" && strings.HasSuffix(collection, ".$cmd") {
 		var emptyPayload types.Binary
-		reply := wire.OpReply{
-			NumberReturned: 1,
-		}
+		var reply wire.OpReply
 		reply.SetDocument(must.NotFail(types.NewDocument(
 			"conversationId", int32(1),
 			"done", true,

diff --git a/internal/handler/common/ismaster.go b/internal/handler/common/ismaster.go
@@ -31,9 +31,7 @@ func IsMaster(ctx context.Context, query *types.Document, tcpHost, name string)
 		return nil, lazyerrors.Error(err)
 	}
 
-	reply := wire.OpReply{
-		NumberReturned: 1,
-	}
+	var reply wire.OpReply
 	reply.SetDocument(IsMasterDocument(tcpHost, name))
 
 	return &reply, nil

diff --git a/internal/handler/sjson/sjson_test.go b/internal/handler/sjson/sjson_test.go
@@ -216,7 +216,12 @@ func addRecordedFuzzDocs(f *testing.F, needDocument, needSchema bool) int {
 			}
 
 		case *wire.OpReply:
-			docs = append(docs, b.Documents()...)
+			doc, err := b.Document()
+			require.NoError(f, err)
+
+			if doc != nil {
+				docs = append(docs, doc)
+			}
 		}
 
 		for _, doc := range docs {

diff --git a/internal/wire/msg_body.go b/internal/wire/msg_body.go
@@ -32,8 +32,9 @@ import (
 type MsgBody interface {
 	msgbody() // seal for sumtype
 
-	readFrom(*bufio.Reader) error
-	encoding.BinaryUnmarshaler
+	// UnmarshalBinaryNocopy is a variant of [encoding.BinaryUnmarshaler] that does not have to copy the data.
+	UnmarshalBinaryNocopy([]byte) error
+
 	encoding.BinaryMarshaler
 	fmt.Stringer
 }
@@ -59,7 +60,7 @@ func ReadMessage(r *bufio.Reader) (*MsgHeader, MsgBody, error) {
 	switch header.OpCode {
 	case OpCodeReply: // not sent by clients, but we should be able to read replies from a proxy
 		var reply OpReply
-		if err := reply.UnmarshalBinary(b); err != nil {
+		if err := reply.UnmarshalBinaryNocopy(b); err != nil {
 			return nil, nil, lazyerrors.Error(err)
 		}
 
@@ -71,15 +72,15 @@ func ReadMessage(r *bufio.Reader) (*MsgHeader, MsgBody, error) {
 		}
 
 		var msg OpMsg
-		if err := msg.UnmarshalBinary(b); err != nil {
+		if err := msg.UnmarshalBinaryNocopy(b); err != nil {
 			return &header, nil, lazyerrors.Error(err)
 		}
 
 		return &header, &msg, nil
 
 	case OpCodeQuery:
 		var query OpQuery
-		if err := query.UnmarshalBinary(b); err != nil {
+		if err := query.UnmarshalBinaryNocopy(b); err != nil {
 			return nil, nil, lazyerrors.Error(err)
 		}
 

diff --git a/internal/wire/op_msg.go b/internal/wire/op_msg.go
@@ -164,7 +164,11 @@ func (msg *OpMsg) RawDocument() (bson2.RawDocument, error) {
 
 func (msg *OpMsg) msgbody() {}
 
-func (msg *OpMsg) readFrom(bufr *bufio.Reader) error {
+// UnmarshalBinaryNocopy implements [MsgBody] interface.
+func (msg *OpMsg) UnmarshalBinaryNocopy(b []byte) error {
+	br := bytes.NewReader(b)
+	bufr := bufio.NewReader(br)
+
 	if err := binary.Read(bufr, binary.LittleEndian, &msg.FlagBits); err != nil {
 		return lazyerrors.Error(err)
 	}
@@ -255,18 +259,6 @@ func (msg *OpMsg) readFrom(bufr *bufio.Reader) error {
 		return err
 	}
 
-	return nil
-}
-
-// UnmarshalBinary reads an OpMsg from a byte array.
-func (msg *OpMsg) UnmarshalBinary(b []byte) error {
-	br := bytes.NewReader(b)
-	bufr := bufio.NewReader(br)
-
-	if err := msg.readFrom(bufr); err != nil {
-		return lazyerrors.Error(err)
-	}
-
 	if _, err := bufr.Peek(1); err != io.EOF {
 		return lazyerrors.Errorf("unexpected end of the OpMsg: %v", err)
 	}