kgo: turn DecompressBufferPool into an optional field (#6) (#7)

Signed-off-by: Miguel Ángel Ortuño <[email protected]>
grafana · Oct 8, 2024 · b76e947 · b76e947
1 parent f8f76a4
commit b76e947
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 16 deletions.
diff --git a/pkg/kgo/compression.go b/pkg/kgo/compression.go
@@ -12,8 +12,7 @@ import (
 	"github.com/klauspost/compress/s2"
 	"github.com/klauspost/compress/zstd"
 	"github.com/pierrec/lz4/v4"
-
-	"github.com/twmb/franz-go/pkg/kgo/internal/pool"
+	"github.com/twmb/franz-go/pkg/kgo/pool"
 )
 
 var byteBuffers = sync.Pool{New: func() any { return bytes.NewBuffer(make([]byte, 8<<10)) }}
@@ -274,17 +273,28 @@ func (d *decompressor) decompress(src []byte, codec byte, pool *pool.BucketedPoo
 	if compCodec == codecNone {
 		return src, nil
 	}
+	var (
+		out *bytes.Buffer
+		buf []byte
+		err error
+	)
 
-	out, buf, err := d.getDecodedBuffer(src, compCodec, pool)
-	if err != nil {
-		return nil, err
-	}
-	defer func() {
-		if compCodec == codecSnappy {
-			return
+	if pool != nil {
+		out, buf, err = d.getDecodedBuffer(src, compCodec, pool)
+		if err != nil {
+			return nil, err
 		}
-		pool.Put(buf)
-	}()
+		defer func() {
+			if compCodec == codecSnappy {
+				return
+			}
+			pool.Put(buf)
+		}()
+	} else {
+		out = byteBuffers.Get().(*bytes.Buffer)
+		out.Reset()
+		defer byteBuffers.Put(out)
+	}
 
 	switch compCodec {
 	case codecGzip:

diff --git a/pkg/kgo/compression_test.go b/pkg/kgo/compression_test.go
@@ -10,8 +10,7 @@ import (
 	"testing"
 
 	"github.com/pierrec/lz4/v4"
-
-	"github.com/twmb/franz-go/pkg/kgo/internal/pool"
+	"github.com/twmb/franz-go/pkg/kgo/pool"
 )
 
 // Regression test for #778.

diff --git a/pkg/kgo/config.go b/pkg/kgo/config.go
@@ -13,11 +13,10 @@ import (
 	"sync"
 	"time"
 
+	"github.com/twmb/franz-go/pkg/kgo/pool"
 	"github.com/twmb/franz-go/pkg/kmsg"
 	"github.com/twmb/franz-go/pkg/kversion"
 	"github.com/twmb/franz-go/pkg/sasl"
-
-	"github.com/twmb/franz-go/pkg/kgo/internal/pool"
 )
 
 // Opt is an option to configure a client.

diff --git a/pkg/kgo/internal/pool/bucketed_pool.go → pkg/kgo/pool/bucketed_pool.go b/pkg/kgo/internal/pool/bucketed_pool.go → pkg/kgo/pool/bucketed_pool.go
diff --git a/pkg/kgo/internal/pool/bucketed_pool_test.go → pkg/kgo/pool/bucketed_pool_test.go b/pkg/kgo/internal/pool/bucketed_pool_test.go → pkg/kgo/pool/bucketed_pool_test.go
diff --git a/pkg/kgo/source.go b/pkg/kgo/source.go
@@ -14,7 +14,7 @@ import (
 
 	"github.com/twmb/franz-go/pkg/kbin"
 	"github.com/twmb/franz-go/pkg/kerr"
-	"github.com/twmb/franz-go/pkg/kgo/internal/pool"
+	"github.com/twmb/franz-go/pkg/kgo/pool"
 	"github.com/twmb/franz-go/pkg/kmsg"
 )