GODRIVER-2778 Reduce memory usage on compressed loads. (#1200)

Author: qingyang-hu

mongo/integration/mtest/opmsg_deployment.go

@@ -57,7 +57,8 @@ func (c *connection) WriteWireMessage(_ context.Context, wm []byte) error {
 }
 
 // ReadWireMessage returns the next response in the connection's list of responses.
-func (c *connection) ReadWireMessage(_ context.Context, dst []byte) ([]byte, error) {
+func (c *connection) ReadWireMessage(_ context.Context) ([]byte, error) {
+	var dst []byte
 	if len(c.responses) == 0 {
 		return dst, errors.New("no responses remaining")
 	}

mongo/integration/mtest/wiremessage_helpers.go

@@ -54,7 +54,7 @@ func parseOpCompressed(wm []byte) (wiremessage.OpCode, []byte, error) {
 		return originalOpcode, nil, errors.New("failed to read compressor ID")
 	}
 
-	compressedMsg, wm, ok := wiremessage.ReadCompressedCompressedMessage(wm, int32(len(wm)))
+	compressedMsg, _, ok := wiremessage.ReadCompressedCompressedMessage(wm, int32(len(wm)))
 	if !ok {
 		return originalOpcode, nil, errors.New("failed to read compressed message")
 	}

x/mongo/driver/compression.go

@@ -26,21 +26,61 @@ type CompressionOpts struct {
 	UncompressedSize int32
 }
 
-var zstdEncoders = &sync.Map{}
+var zstdEncoders sync.Map // map[zstd.EncoderLevel]*zstd.Encoder
 
-func getZstdEncoder(l zstd.EncoderLevel) (*zstd.Encoder, error) {
-	v, ok := zstdEncoders.Load(l)
-	if ok {
+func getZstdEncoder(level zstd.EncoderLevel) (*zstd.Encoder, error) {
+	if v, ok := zstdEncoders.Load(level); ok {
 		return v.(*zstd.Encoder), nil
 	}
-	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(l))
+	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(level))
 	if err != nil {
 		return nil, err
 	}
-	zstdEncoders.Store(l, encoder)
+	zstdEncoders.Store(level, encoder)
 	return encoder, nil
 }
 
+var zlibEncoders sync.Map // map[int /*level*/]*zlibEncoder
+
+func getZlibEncoder(level int) (*zlibEncoder, error) {
+	if v, ok := zlibEncoders.Load(level); ok {
+		return v.(*zlibEncoder), nil
+	}
+	writer, err := zlib.NewWriterLevel(nil, level)
+	if err != nil {
+		return nil, err
+	}
+	encoder := &zlibEncoder{writer: writer, buf: new(bytes.Buffer)}
+	zlibEncoders.Store(level, encoder)
+
+	return encoder, nil
+}
+
+type zlibEncoder struct {
+	mu     sync.Mutex
+	writer *zlib.Writer
+	buf    *bytes.Buffer
+}
+
+func (e *zlibEncoder) Encode(dst, src []byte) ([]byte, error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.buf.Reset()
+	e.writer.Reset(e.buf)
+
+	_, err := e.writer.Write(src)
+	if err != nil {
+		return nil, err
+	}
+	err = e.writer.Close()
+	if err != nil {
+		return nil, err
+	}
+	dst = append(dst[:0], e.buf.Bytes()...)
+	return dst, nil
+}
+
 // CompressPayload takes a byte slice and compresses it according to the options passed
 func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 	switch opts.Compressor {
@@ -49,20 +89,11 @@ func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 	case wiremessage.CompressorSnappy:
 		return snappy.Encode(nil, in), nil
 	case wiremessage.CompressorZLib:
-		var b bytes.Buffer
-		w, err := zlib.NewWriterLevel(&b, opts.ZlibLevel)
+		encoder, err := getZlibEncoder(opts.ZlibLevel)
 		if err != nil {
 			return nil, err
 		}
-		_, err = w.Write(in)
-		if err != nil {
-			return nil, err
-		}
-		err = w.Close()
-		if err != nil {
-			return nil, err
-		}
-		return b.Bytes(), nil
+		return encoder.Encode(nil, in)
 	case wiremessage.CompressorZstd:
 		encoder, err := getZstdEncoder(zstd.EncoderLevelFromZstd(opts.ZstdLevel))
 		if err != nil {
@@ -75,20 +106,23 @@ func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 }
 
 // DecompressPayload takes a byte slice that has been compressed and undoes it according to the options passed
-func DecompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
+func DecompressPayload(in []byte, opts CompressionOpts) (uncompressed []byte, err error) {
 	switch opts.Compressor {
 	case wiremessage.CompressorNoOp:
 		return in, nil
 	case wiremessage.CompressorSnappy:
-		uncompressed := make([]byte, opts.UncompressedSize)
+		uncompressed = make([]byte, opts.UncompressedSize)
 		return snappy.Decode(uncompressed, in)
 	case wiremessage.CompressorZLib:
-		decompressor, err := zlib.NewReader(bytes.NewReader(in))
+		r, err := zlib.NewReader(bytes.NewReader(in))
 		if err != nil {
 			return nil, err
 		}
-		uncompressed := make([]byte, opts.UncompressedSize)
-		_, err = io.ReadFull(decompressor, uncompressed)
+		defer func() {
+			err = r.Close()
+		}()
+		uncompressed = make([]byte, opts.UncompressedSize)
+		_, err = io.ReadFull(r, uncompressed)
 		if err != nil {
 			return nil, err
 		}
@@ -99,7 +133,7 @@ func DecompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 			return nil, err
 		}
 		defer r.Close()
-		uncompressed := make([]byte, opts.UncompressedSize)
+		uncompressed = make([]byte, opts.UncompressedSize)
 		_, err = io.ReadFull(r, uncompressed)
 		if err != nil {
 			return nil, err

x/mongo/driver/driver.go

@@ -59,7 +59,7 @@ type Server interface {
 // Connection represents a connection to a MongoDB server.
 type Connection interface {
 	WriteWireMessage(context.Context, []byte) error
-	ReadWireMessage(ctx context.Context, dst []byte) ([]byte, error)
+	ReadWireMessage(ctx context.Context) ([]byte, error)
 	Description() description.Server
 
 	// Close closes any underlying connection and returns or frees any resources held by the

x/mongo/driver/drivertest/channel_conn.go

@@ -40,22 +40,15 @@ func (c *ChannelConn) WriteWireMessage(ctx context.Context, wm []byte) error {
 }
 
 // ReadWireMessage implements the driver.Connection interface.
-func (c *ChannelConn) ReadWireMessage(ctx context.Context, dst []byte) ([]byte, error) {
-	dst = dst[:0]
+func (c *ChannelConn) ReadWireMessage(ctx context.Context) ([]byte, error) {
 	var wm []byte
 	var err error
 	select {
 	case wm = <-c.ReadResp:
 	case err = <-c.ReadErr:
 	case <-ctx.Done():
 	}
-	if l := len(wm); l > 0 {
-		if l > cap(dst) {
-			dst = make([]byte, 0, l)
-		}
-		dst = append(dst, wm...)
-	}
-	return dst, err
+	return wm, err
 }
 
 // Description implements the driver.Connection interface.

x/mongo/driver/operation.go

@@ -865,7 +865,7 @@ func (op Operation) roundTrip(ctx context.Context, conn Connection, wm []byte) (
 }
 
 func (op Operation) readWireMessage(ctx context.Context, conn Connection) (result []byte, err error) {
-	wm, err := conn.ReadWireMessage(ctx, nil)
+	wm, err := conn.ReadWireMessage(ctx)
 	if err != nil {
 		return nil, op.networkError(err)
 	}
@@ -876,14 +876,21 @@ func (op Operation) readWireMessage(ctx context.Context, conn Connection) (resul
 		streamer.SetStreaming(wiremessage.IsMsgMoreToCome(wm))
 	}
 
-	// decompress wiremessage
-	wm, err = op.decompressWireMessage(wm)
-	if err != nil {
-		return nil, err
+	length, _, _, opcode, rem, ok := wiremessage.ReadHeader(wm)
+	if !ok || len(wm) < int(length) {
+		return nil, errors.New("malformed wire message: insufficient bytes")
+	}
+	if opcode == wiremessage.OpCompressed {
+		rawsize := length - 16 // remove header size
+		// decompress wiremessage
+		opcode, rem, err = op.decompressWireMessage(rem[:rawsize])
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	// decode
-	res, err := op.decodeResult(wm)
+	res, err := op.decodeResult(opcode, rem)
 	// Update cluster/operation time and recovery tokens before handling the error to ensure we're properly updating
 	// everything.
 	op.updateClusterTimes(res)
@@ -940,51 +947,39 @@ func (op *Operation) moreToComeRoundTrip(ctx context.Context, conn Connection, w
 	return bsoncore.BuildDocument(nil, bsoncore.AppendInt32Element(nil, "ok", 1)), err
 }
 
-// decompressWireMessage handles decompressing a wiremessage. If the wiremessage
-// is not compressed, this method will return the wiremessage.
-func (Operation) decompressWireMessage(wm []byte) ([]byte, error) {
-	// read the header and ensure this is a compressed wire message
-	length, reqid, respto, opcode, rem, ok := wiremessage.ReadHeader(wm)
-	if !ok || len(wm) < int(length) {
-		return nil, errors.New("malformed wire message: insufficient bytes")
-	}
-	if opcode != wiremessage.OpCompressed {
-		return wm, nil
-	}
+// decompressWireMessage handles decompressing a wiremessage without the header.
+func (Operation) decompressWireMessage(wm []byte) (wiremessage.OpCode, []byte, error) {
 	// get the original opcode and uncompressed size
-	opcode, rem, ok = wiremessage.ReadCompressedOriginalOpCode(rem)
+	opcode, rem, ok := wiremessage.ReadCompressedOriginalOpCode(wm)
 	if !ok {
-		return nil, errors.New("malformed OP_COMPRESSED: missing original opcode")
+		return 0, nil, errors.New("malformed OP_COMPRESSED: missing original opcode")
 	}
 	uncompressedSize, rem, ok := wiremessage.ReadCompressedUncompressedSize(rem)
 	if !ok {
-		return nil, errors.New("malformed OP_COMPRESSED: missing uncompressed size")
+		return 0, nil, errors.New("malformed OP_COMPRESSED: missing uncompressed size")
 	}
 	// get the compressor ID and decompress the message
 	compressorID, rem, ok := wiremessage.ReadCompressedCompressorID(rem)
 	if !ok {
-		return nil, errors.New("malformed OP_COMPRESSED: missing compressor ID")
+		return 0, nil, errors.New("malformed OP_COMPRESSED: missing compressor ID")
 	}
-	compressedSize := length - 25 // header (16) + original opcode (4) + uncompressed size (4) + compressor ID (1)
+	compressedSize := len(wm) - 9 // original opcode (4) + uncompressed size (4) + compressor ID (1)
 	// return the original wiremessage
-	msg, _, ok := wiremessage.ReadCompressedCompressedMessage(rem, compressedSize)
+	msg, _, ok := wiremessage.ReadCompressedCompressedMessage(rem, int32(compressedSize))
 	if !ok {
-		return nil, errors.New("malformed OP_COMPRESSED: insufficient bytes for compressed wiremessage")
+		return 0, nil, errors.New("malformed OP_COMPRESSED: insufficient bytes for compressed wiremessage")
 	}
 
-	wm = make([]byte, 0, int(uncompressedSize)+16)
-	wm = wiremessage.AppendHeader(wm, uncompressedSize+16, reqid, respto, opcode)
 	opts := CompressionOpts{
 		Compressor:       compressorID,
 		UncompressedSize: uncompressedSize,
 	}
 	uncompressed, err := DecompressPayload(msg, opts)
 	if err != nil {
-		return nil, err
+		return 0, nil, err
 	}
-	wm = append(wm, uncompressed...)
 
-	return wm, nil
+	return opcode, uncompressed, nil
 }
 
 func (op Operation) createWireMessage(
@@ -1541,28 +1536,12 @@ func (Operation) canCompress(cmd string) bool {
 }
 
 // decodeOpReply extracts the necessary information from an OP_REPLY wire message.
-// includesHeader: specifies whether or not wm includes the message header
 // Returns the decoded OP_REPLY. If the err field of the returned opReply is non-nil, an error occurred while decoding
 // or validating the response and the other fields are undefined.
-func (Operation) decodeOpReply(wm []byte, includesHeader bool) opReply {
+func (Operation) decodeOpReply(wm []byte) opReply {
 	var reply opReply
 	var ok bool
 
-	if includesHeader {
-		wmLength := len(wm)
-		var length int32
-		var opcode wiremessage.OpCode
-		length, _, _, opcode, wm, ok = wiremessage.ReadHeader(wm)
-		if !ok || int(length) > wmLength {
-			reply.err = errors.New("malformed wire message: insufficient bytes")
-			return reply
-		}
-		if opcode != wiremessage.OpReply {
-			reply.err = errors.New("malformed wire message: incorrect opcode")
-			return reply
-		}
-	}
-
 	reply.responseFlags, wm, ok = wiremessage.ReadReplyFlags(wm)
 	if !ok {
 		reply.err = errors.New("malformed OP_REPLY: missing flags")
@@ -1583,7 +1562,7 @@ func (Operation) decodeOpReply(wm []byte, includesHeader bool) opReply {
 		reply.err = errors.New("malformed OP_REPLY: missing numberReturned")
 		return reply
 	}
-	reply.documents, wm, ok = wiremessage.ReadReplyDocuments(wm)
+	reply.documents, _, ok = wiremessage.ReadReplyDocuments(wm)
 	if !ok {
 		reply.err = errors.New("malformed OP_REPLY: could not read documents from reply")
 	}
@@ -1607,18 +1586,10 @@ func (Operation) decodeOpReply(wm []byte, includesHeader bool) opReply {
 	return reply
 }
 
-func (op Operation) decodeResult(wm []byte) (bsoncore.Document, error) {
-	wmLength := len(wm)
-	length, _, _, opcode, wm, ok := wiremessage.ReadHeader(wm)
-	if !ok || int(length) > wmLength {
-		return nil, errors.New("malformed wire message: insufficient bytes")
-	}
-
-	wm = wm[:wmLength-16] // constrain to just this wiremessage, incase there are multiple in the slice
-
+func (op Operation) decodeResult(opcode wiremessage.OpCode, wm []byte) (bsoncore.Document, error) {
 	switch opcode {
 	case wiremessage.OpReply:
-		reply := op.decodeOpReply(wm, false)
+		reply := op.decodeOpReply(wm)
 		if reply.err != nil {
 			return nil, reply.err
 		}
@@ -1635,7 +1606,7 @@ func (op Operation) decodeResult(wm []byte) (bsoncore.Document, error) {
 
 		return rdr, ExtractErrorFromServerResponse(rdr)
 	case wiremessage.OpMsg:
-		_, wm, ok = wiremessage.ReadMsgFlags(wm)
+		_, wm, ok := wiremessage.ReadMsgFlags(wm)
 		if !ok {
 			return nil, errors.New("malformed wire message: missing OP_MSG flags")
 		}

x/mongo/driver/operation_test.go

@@ -740,7 +740,6 @@ func (m *mockServerSelector) SelectServer(description.Topology, []description.Se
 type mockConnection struct {
 	// parameters
 	pWriteWM []byte
-	pReadDst []byte
 
 	// returns
 	rWriteErr     error
@@ -770,8 +769,7 @@ func (m *mockConnection) WriteWireMessage(_ context.Context, wm []byte) error {
 	return m.rWriteErr
 }
 
-func (m *mockConnection) ReadWireMessage(_ context.Context, dst []byte) ([]byte, error) {
-	m.pReadDst = dst
+func (m *mockConnection) ReadWireMessage(_ context.Context) ([]byte, error) {
 	return m.rReadWM, m.rReadErr
 }
 

x/mongo/driver/session/client_session.go

@@ -82,7 +82,7 @@ func (s TransactionState) String() string {
 type LoadBalancedTransactionConnection interface {
 	// Functions copied over from driver.Connection.
 	WriteWireMessage(context.Context, []byte) error
-	ReadWireMessage(ctx context.Context, dst []byte) ([]byte, error)
+	ReadWireMessage(ctx context.Context) ([]byte, error)
 	Description() description.Server
 	Close() error
 	ID() string