From 9b9706a0a07445a554be70d17496369f5852bf98 Mon Sep 17 00:00:00 2001
From: Qingyang Hu <103950869+qingyang-hu@users.noreply.github.com>
Date: Wed, 26 Jun 2024 15:10:17 -0400
Subject: [PATCH] GODRIVER-3132 Use dst buf pool for zstd encoding. (#1577)

(cherry picked from commit b10dc345d022fc64ac4d9bc5796c830a7d8c71c5)
---
 x/mongo/driver/compression.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/x/mongo/driver/compression.go b/x/mongo/driver/compression.go
index d79b024b74..d9a6c68fee 100644
--- a/x/mongo/driver/compression.go
+++ b/x/mongo/driver/compression.go
@@ -30,7 +30,11 @@ type CompressionOpts struct {
 // destination writer. It panics on any errors and should only be used at
 // package initialization time.
 func mustZstdNewWriter(lvl zstd.EncoderLevel) *zstd.Encoder {
-	enc, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(lvl))
+	enc, err := zstd.NewWriter(
+		nil,
+		zstd.WithWindowSize(8<<20), // Set window size to 8MB.
+		zstd.WithEncoderLevel(lvl),
+	)
 	if err != nil {
 		panic(err)
 	}
@@ -105,6 +109,13 @@ func (e *zlibEncoder) Encode(dst, src []byte) ([]byte, error) {
 	return dst, nil
 }
 
+var zstdBufPool = sync.Pool{
+	New: func() interface{} {
+		s := make([]byte, 0)
+		return &s
+	},
+}
+
 // CompressPayload takes a byte slice and compresses it according to the options passed
 func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 	switch opts.Compressor {
@@ -123,7 +134,13 @@ func CompressPayload(in []byte, opts CompressionOpts) ([]byte, error) {
 		if err != nil {
 			return nil, err
 		}
-		return encoder.EncodeAll(in, nil), nil
+		ptr := zstdBufPool.Get().(*[]byte)
+		b := encoder.EncodeAll(in, *ptr)
+		dst := make([]byte, len(b))
+		copy(dst, b)
+		*ptr = b[:0]
+		zstdBufPool.Put(ptr)
+		return dst, nil
 	default:
 		return nil, fmt.Errorf("unknown compressor ID %v", opts.Compressor)
 	}