Godriver-2259 Reduce memory consumption. (#985)

Author: qingyang-hu

x/mongo/driver/compression.go

@@ -11,6 +11,7 @@ import (
 	"compress/zlib"
 	"fmt"
 	"io"
+	"sync"
 
 	"github.com/golang/snappy"
 	"github.com/klauspost/compress/zstd"
@@ -25,6 +26,21 @@ type CompressionOpts struct {
 	UncompressedSize int32
 }
 
+var zstdEncoders = &sync.Map{}
+
+func getZstdEncoder(l zstd.EncoderLevel) (*zstd.Encoder, error) {
+	v, ok := zstdEncoders.Load(l)
+	if ok {
+		return v.(*zstd.Encoder), nil
+	}
+	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(l))
+	if err != nil {
+		return nil, err
+	}
+	zstdEncoders.Store(l, encoder)
+	return encoder, 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 {
@@ -48,21 +64,11 @@ func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 		}
 		return b.Bytes(), nil
 	case wiremessage.CompressorZstd:
-		var b bytes.Buffer
-		w, err := zstd.NewWriter(&b, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(opts.ZstdLevel)))
-		if err != nil {
-			return nil, err
-		}
-		_, err = io.Copy(w, bytes.NewBuffer(in))
+		encoder, err := getZstdEncoder(zstd.EncoderLevelFromZstd(opts.ZstdLevel))
 		if err != nil {
-			_ = w.Close()
 			return nil, err
 		}
-		err = w.Close()
-		if err != nil {
-			return nil, err
-		}
-		return b.Bytes(), nil
+		return encoder.EncodeAll(in, nil), nil
 	default:
 		return nil, fmt.Errorf("unknown compressor ID %v", opts.Compressor)
 	}

x/mongo/driver/compression_test.go

@@ -7,7 +7,7 @@
 package driver
 
 import (
-	"strconv"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -23,7 +23,7 @@ func TestCompression(t *testing.T) {
 	}
 
 	for _, compressor := range compressors {
-		t.Run(strconv.Itoa(int(compressor)), func(t *testing.T) {
+		t.Run(compressor.String(), func(t *testing.T) {
 			payload := []byte("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt")
 			opts := CompressionOpts{
 				Compressor:       compressor,
@@ -40,3 +40,41 @@ func TestCompression(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkCompressPayload(b *testing.B) {
+	payload := func() []byte {
+		buf, err := os.ReadFile("compression.go")
+		if err != nil {
+			b.Log(err)
+			b.FailNow()
+		}
+		for i := 1; i < 10; i++ {
+			buf = append(buf, buf...)
+		}
+		return buf
+	}()
+
+	compressors := []wiremessage.CompressorID{
+		wiremessage.CompressorSnappy,
+		wiremessage.CompressorZLib,
+		wiremessage.CompressorZstd,
+	}
+
+	for _, compressor := range compressors {
+		b.Run(compressor.String(), func(b *testing.B) {
+			opts := CompressionOpts{
+				Compressor: compressor,
+				ZlibLevel:  wiremessage.DefaultZlibLevel,
+				ZstdLevel:  wiremessage.DefaultZstdLevel,
+			}
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					_, err := CompressPayload(payload, opts)
+					if err != nil {
+						b.Error(err)
+					}
+				}
+			})
+		})
+	}
+}

x/mongo/driver/topology/connection.go

@@ -638,6 +638,9 @@ func (c *Connection) CompressWireMessage(src, dst []byte) ([]byte, error) {
 		return dst, ErrConnectionClosed
 	}
 	if c.connection.compressor == wiremessage.CompressorNoOp {
+		if len(dst) == 0 {
+			return src, nil
+		}
 		return append(dst, src...), nil
 	}
 	_, reqid, respto, origcode, rem, ok := wiremessage.ReadHeader(src)

x/mongo/driver/topology/connection_test.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
+	"math/rand"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -22,6 +23,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/address"
 	"go.mongodb.org/mongo-driver/mongo/description"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/wiremessage"
 )
 
 type testHandshaker struct {
@@ -926,6 +928,24 @@ func TestConnection(t *testing.T) {
 	})
 }
 
+func BenchmarkConnection(b *testing.B) {
+	b.Run("CompressWireMessage CompressorNoOp", func(b *testing.B) {
+		buf := make([]byte, 256)
+		_, err := rand.Read(buf)
+		if err != nil {
+			b.Log(err)
+			b.FailNow()
+		}
+		conn := Connection{connection: &connection{compressor: wiremessage.CompressorNoOp}}
+		for i := 0; i < b.N; i++ {
+			_, err := conn.CompressWireMessage(buf, nil)
+			if err != nil {
+				b.Error(err)
+			}
+		}
+	})
+}
+
 // cancellationTestNetConn is a net.Conn implementation that is used to test context.Cancellation during an in-progress
 // network read or write. This type has two unbuffered channels: operationStartedChan and continueChan. When Read/Write
 // starts, the type will write to operationStartedChan, which will block until the test reads from it. This signals to

x/mongo/driver/wiremessage/wiremessage.go

@@ -188,6 +188,22 @@ const (
 	CompressorZstd
 )
 
+// String implements the fmt.Stringer interface.
+func (id CompressorID) String() string {
+	switch id {
+	case CompressorNoOp:
+		return "CompressorNoOp"
+	case CompressorSnappy:
+		return "CompressorSnappy"
+	case CompressorZLib:
+		return "CompressorZLib"
+	case CompressorZstd:
+		return "CompressorZstd"
+	default:
+		return "CompressorInvalid"
+	}
+}
+
 const (
 	// DefaultZlibLevel is the default level for zlib compression
 	DefaultZlibLevel = 6