diff --git a/bson/bson.go b/bson/bson.go index ac1c02c7f..6bfad20b0 100644 --- a/bson/bson.go +++ b/bson/bson.go @@ -489,8 +489,15 @@ func handleErr(err *error) { // } // func Marshal(in interface{}) (out []byte, err error) { + return MarshalBuffer(in, make([]byte, 0, initialBufferSize)) +} + +// MarshalBuffer behaves the same way as Marshal, except that instead of +// allocating a new byte slice it tries to use the received byte slice and +// only allocates more memory if necessary to fit the marshaled value. +func MarshalBuffer(in interface{}, buf []byte) (out []byte, err error) { defer handleErr(&err) - e := &encoder{make([]byte, 0, initialBufferSize)} + e := &encoder{buf} e.addDoc(reflect.ValueOf(in)) return e.out, nil } diff --git a/bson/bson_test.go b/bson/bson_test.go index b77ec877f..5866539b6 100644 --- a/bson/bson_test.go +++ b/bson/bson_test.go @@ -249,6 +249,13 @@ func (s *S) TestUnmarshalNonNilInterface(c *C) { c.Assert(m, DeepEquals, bson.M{"a": 1}) } +func (s *S) TestMarshalBuffer(c *C) { + buf := make([]byte, 0, 256) + data, err := bson.MarshalBuffer(bson.M{"a": 1}, buf) + c.Assert(err, IsNil) + c.Assert(data, DeepEquals, buf[:len(data)]) +} + // -------------------------------------------------------------------------- // Some one way marshaling operations which would unmarshal differently. diff --git a/socket.go b/socket.go index a2343354d..8b85c7256 100644 --- a/socket.go +++ b/socket.go @@ -372,13 +372,22 @@ func (socket *mongoSocket) SimpleQuery(op *queryOp) (data []byte, err error) { return data, err } +var bytesBufferPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 0, 256) + }, +} + func (socket *mongoSocket) Query(ops ...interface{}) (err error) { if lops := socket.flushLogout(); len(lops) > 0 { ops = append(lops, ops...) } - buf := make([]byte, 0, 256) + buf := bytesBufferPool.Get().([]byte) + defer func() { + bytesBufferPool.Put(buf[:0]) + }() // Serialize operations synchronously to avoid interrupting // other goroutines while we can't really be sending data. @@ -674,7 +683,9 @@ func addBSON(b []byte, doc interface{}) ([]byte, error) { if doc == nil { return append(b, 5, 0, 0, 0, 0), nil } - data, err := bson.Marshal(doc) + buf := bytesBufferPool.Get().([]byte) + data, err := bson.MarshalBuffer(doc, buf) + defer bytesBufferPool.Put(data[:0]) if err != nil { return b, err }