Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
@qingyang-hu
qingyang-hu committed Dec 4, 2023
1 parent d9a1e55 commit 60e5756
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 170 deletions.
112 changes: 0 additions & 112 deletions bson/marshal.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ package bson
import (
"bytes"
"encoding/json"
"sync"

"go.mongodb.org/mongo-driver/bson/bsoncodec"
"go.mongodb.org/mongo-driver/bson/bsonrw"
Expand Down Expand Up @@ -64,117 +63,6 @@ func Marshal(val interface{}) ([]byte, error) {
return buf.Bytes(), nil
}

// MarshalWithContext returns the BSON encoding of val as a BSON document using EncodeContext ec. If val is not a type
// that can be transformed into a document, MarshalValueWithContext should be used instead.
//
// Deprecated: Use [NewEncoder] and use the Encoder configuration methods to set the desired marshal
// behavior instead:
//
// buf := bytes.NewBuffer(dst)
// vw, err := bsonrw.NewBSONValueWriter(buf)
// if err != nil {
// panic(err)
// }
// enc, err := bson.NewEncoder(vw)
// if err != nil {
// panic(err)
// }
// enc.IntMinSize()
//
// See [Encoder] for more examples.
func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) {
dst := make([]byte, 0)
return MarshalAppendWithContext(ec, dst, val)
}

// MarshalAppendWithRegistry will encode val as a BSON document using Registry r and append the bytes to dst. If dst is
// not large enough to hold the bytes, it will be grown. If val is not a type that can be transformed into a document,
// MarshalValueAppendWithRegistry should be used instead.
//
// Deprecated: Use [NewEncoder], and pass the dst byte slice (wrapped by a bytes.Buffer) into
// [bsonrw.NewBSONValueWriter], and specify the Registry by calling [Encoder.SetRegistry] instead:
//
// buf := bytes.NewBuffer(dst)
// vw, err := bsonrw.NewBSONValueWriter(buf)
// if err != nil {
// panic(err)
// }
// enc, err := bson.NewEncoder(vw)
// if err != nil {
// panic(err)
// }
// enc.SetRegistry(reg)
//
// See [Encoder] for more examples.
func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) {
return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val)
}

// Pool of buffers for marshalling BSON.
var bufPool = sync.Pool{
New: func() interface{} {
return new(bytes.Buffer)
},
}

// MarshalAppendWithContext will encode val as a BSON document using Registry r and EncodeContext ec and append the
// bytes to dst. If dst is not large enough to hold the bytes, it will be grown. If val is not a type that can be
// transformed into a document, MarshalValueAppendWithContext should be used instead.
//
// Deprecated: Use [NewEncoder], pass the dst byte slice (wrapped by a bytes.Buffer) into
// [bsonrw.NewBSONValueWriter], and use the Encoder configuration methods to set the desired marshal
// behavior instead:
//
// buf := bytes.NewBuffer(dst)
// vw, err := bsonrw.NewBSONValueWriter(buf)
// if err != nil {
// panic(err)
// }
// enc, err := bson.NewEncoder(vw)
// if err != nil {
// panic(err)
// }
// enc.IntMinSize()
//
// See [Encoder] for more examples.
func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) {
sw := bufPool.Get().(*bytes.Buffer)
defer func() {
// Proper usage of a sync.Pool requires each entry to have approximately
// the same memory cost. To obtain this property when the stored type
// contains a variably-sized buffer, we add a hard limit on the maximum
// buffer to place back in the pool. We limit the size to 16MiB because
// that's the maximum wire message size supported by any current MongoDB
// server.
//
// Comment based on
// https://cs.opensource.google/go/go/+/refs/tags/go1.19:src/fmt/print.go;l=147
//
// Recycle byte slices that are smaller than 16MiB and at least half
// occupied.
if sw.Cap() < 16*1024*1024 && sw.Cap()/2 < sw.Len() {
bufPool.Put(sw)
}
}()

sw.Reset()
vw := bvwPool.Get(sw)
defer bvwPool.Put(vw)

enc := encPool.Get().(*Encoder)
defer encPool.Put(enc)

enc.Reset(vw)
enc.ec = ec

err := enc.Encode(val)
if err != nil {
return nil, err
}

return append(dst, sw.Bytes()...), nil
}

// MarshalValue returns the BSON encoding of val.
//
// MarshalValue will use bson.DefaultRegistry to transform val into a BSON value. If val is a struct, this function will
Expand Down
55 changes: 9 additions & 46 deletions bson/marshal_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,49 +25,6 @@ import (

var tInt32 = reflect.TypeOf(int32(0))

func TestMarshalAppendWithRegistry(t *testing.T) {
for _, tc := range marshalingTestCases {
t.Run(tc.name, func(t *testing.T) {
dst := make([]byte, 0, 1024)
var reg *bsoncodec.Registry
if tc.reg != nil {
reg = tc.reg
} else {
reg = DefaultRegistry
}
got, err := MarshalAppendWithRegistry(reg, dst, tc.val)
noerr(t, err)

if !bytes.Equal(got, tc.want) {
t.Errorf("Bytes are not equal. got %v; want %v", got, tc.want)
t.Errorf("Bytes:\n%v\n%v", got, tc.want)
}
})
}
}

func TestMarshalAppendWithContext(t *testing.T) {
for _, tc := range marshalingTestCases {
t.Run(tc.name, func(t *testing.T) {
dst := make([]byte, 0, 1024)
var reg *bsoncodec.Registry
if tc.reg != nil {
reg = tc.reg
} else {
reg = DefaultRegistry
}
ec := bsoncodec.EncodeContext{Registry: reg}
got, err := MarshalAppendWithContext(ec, dst, tc.val)
noerr(t, err)

if !bytes.Equal(got, tc.want) {
t.Errorf("Bytes are not equal. got %v; want %v", got, tc.want)
t.Errorf("Bytes:\n%v\n%v", got, tc.want)
}
})
}
}

func TestMarshalWithRegistry(t *testing.T) {
buf := new(bytes.Buffer)
for _, tc := range marshalingTestCases {
Expand Down Expand Up @@ -95,6 +52,7 @@ func TestMarshalWithRegistry(t *testing.T) {
}

func TestMarshalWithContext(t *testing.T) {
buf := new(bytes.Buffer)
for _, tc := range marshalingTestCases {
t.Run(tc.name, func(t *testing.T) {
var reg *bsoncodec.Registry
Expand All @@ -103,11 +61,16 @@ func TestMarshalWithContext(t *testing.T) {
} else {
reg = DefaultRegistry
}
ec := bsoncodec.EncodeContext{Registry: reg}
got, err := MarshalWithContext(ec, tc.val)
buf.Reset()
vw, err := bsonrw.NewBSONValueWriter(buf)
noerr(t, err)
enc := NewEncoder(vw)
enc.IntMinSize()
enc.SetRegistry(reg)
err = enc.Encode(tc.val)
noerr(t, err)

if !bytes.Equal(got, tc.want) {
if got := buf.Bytes(); !bytes.Equal(got, tc.want) {
t.Errorf("Bytes are not equal. got %v; want %v", got, tc.want)
t.Errorf("Bytes:\n%v\n%v", got, tc.want)
}
Expand Down
24 changes: 18 additions & 6 deletions bson/truncation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,11 @@
package bson

import (
"bytes"
"testing"

"go.mongodb.org/mongo-driver/bson/bsoncodec"
"go.mongodb.org/mongo-driver/bson/bsonrw"
"go.mongodb.org/mongo-driver/internal/assert"
)

Expand All @@ -29,9 +31,14 @@ func TestTruncation(t *testing.T) {
inputVal := 4.7892

input := inputArgs{Name: inputName, Val: &inputVal}
ec := bsoncodec.EncodeContext{Registry: DefaultRegistry}

doc, err := MarshalWithContext(ec, &input)
buf := new(bytes.Buffer)
vw, err := bsonrw.NewBSONValueWriter(buf)
assert.Nil(t, err)
enc := NewEncoder(vw)
enc.IntMinSize()
enc.SetRegistry(DefaultRegistry)
err = enc.Encode(&input)
assert.Nil(t, err)

var output outputArgs
Expand All @@ -40,7 +47,7 @@ func TestTruncation(t *testing.T) {
Truncate: true,
}

err = UnmarshalWithContext(dc, doc, &output)
err = UnmarshalWithContext(dc, buf.Bytes(), &output)
assert.Nil(t, err)

assert.Equal(t, inputName, output.Name)
Expand All @@ -51,9 +58,14 @@ func TestTruncation(t *testing.T) {
inputVal := 7.382

input := inputArgs{Name: inputName, Val: &inputVal}
ec := bsoncodec.EncodeContext{Registry: DefaultRegistry}

doc, err := MarshalWithContext(ec, &input)
buf := new(bytes.Buffer)
vw, err := bsonrw.NewBSONValueWriter(buf)
assert.Nil(t, err)
enc := NewEncoder(vw)
enc.IntMinSize()
enc.SetRegistry(DefaultRegistry)
err = enc.Encode(&input)
assert.Nil(t, err)

var output outputArgs
Expand All @@ -63,7 +75,7 @@ func TestTruncation(t *testing.T) {
}

// case throws an error when truncation is disabled
err = UnmarshalWithContext(dc, doc, &output)
err = UnmarshalWithContext(dc, buf.Bytes(), &output)
assert.NotNil(t, err)
})
}
20 changes: 14 additions & 6 deletions mongo/cursor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
package mongo

import (
"bytes"
"context"
"errors"
"fmt"
Expand Down Expand Up @@ -91,7 +92,7 @@ func NewCursorFromDocuments(documents []interface{}, err error, registry *bsonco
}

// Convert documents slice to a sequence-style byte array.
var docsBytes []byte
buf := new(bytes.Buffer)
for _, doc := range documents {
switch t := doc.(type) {
case nil:
Expand All @@ -100,15 +101,22 @@ func NewCursorFromDocuments(documents []interface{}, err error, registry *bsonco
// Slight optimization so we'll just use MarshalBSON and not go through the codec machinery.
doc = bson.Raw(t)
}
var marshalErr error
docsBytes, marshalErr = bson.MarshalAppendWithRegistry(registry, docsBytes, doc)
if marshalErr != nil {
return nil, marshalErr
vw, err := bsonrw.NewBSONValueWriter(buf)
if err != nil {
return nil, err
}
enc := bson.NewEncoder(vw)
enc.SetRegistry(registry)
err = enc.Encode(doc)
// var marshalErr error
// docsBytes, marshalErr = bson.MarshalAppendWithRegistry(registry, docsBytes, doc)
if err != nil {
return nil, err
}
}

c := &Cursor{
bc: driver.NewBatchCursorFromDocuments(docsBytes),
bc: driver.NewBatchCursorFromDocuments(buf.Bytes()),
registry: registry,
err: err,
}
Expand Down

0 comments on commit 60e5756

Please sign in to comment.