Pool decode buffers when encoding messages using compression

encoding_bench_test.go

@@ -43,7 +43,7 @@ func BenchmarkCborDecoding(b *testing.B) {
 		for pb.Next() {
 			var got PartialGMessage
 			require.NoError(b, encoder.Decode(data, &got))
-			require.Equal(b, msg, &got)
+			requireEqualPartialMessages(b, msg, &got)
 		}
 	})
 }
@@ -79,11 +79,29 @@ func BenchmarkZstdDecoding(b *testing.B) {
 		for pb.Next() {
 			var got PartialGMessage
 			require.NoError(b, encoder.Decode(data, &got))
-			require.Equal(b, msg, &got)
+			requireEqualPartialMessages(b, msg, &got)
 		}
 	})
 }
 
+func requireEqualPartialMessages(b *testing.B, expected, actual *PartialGMessage) {
+	// Because empty ECChain gets marshaled as null, we need to use ECChain.Eq for
+	// checking equality. Hence, the custom equality check.
+	require.Equal(b, expected.Sender, actual.Sender)
+	require.Equal(b, expected.Signature, actual.Signature)
+	require.Equal(b, expected.VoteValueKey, actual.VoteValueKey)
+	require.Equal(b, expected.Ticket, actual.Ticket)
+	require.True(b, expected.Vote.Eq(&actual.Vote))
+	if expected.Justification == nil {
+		require.Nil(b, actual.Justification)
+	} else {
+		require.NotNil(b, actual.Justification)
+		require.Equal(b, expected.Justification.Signature, actual.Justification.Signature)
+		require.Equal(b, expected.Justification.Signers, actual.Justification.Signers)
+		require.True(b, expected.Justification.Vote.Eq(&actual.Justification.Vote))
+	}
+}
+
 func generateRandomPartialGMessage(b *testing.B, rng *rand.Rand) *PartialGMessage {
 	var pgmsg PartialGMessage
 	pgmsg.GMessage = generateRandomGMessage(b, rng)

internal/encoding/encoding.go

@@ -3,6 +3,7 @@ package encoding
 import (
 	"bytes"
 	"fmt"
+	"sync"
 
 	"github.com/klauspost/compress/zstd"
 	cbg "github.com/whyrusleeping/cbor-gen"
@@ -13,6 +14,13 @@ import (
 // size in GossipSub.
 const maxDecompressedSize = 1 << 20
 
+var bufferPool = sync.Pool{
+	New: func() any {
+		buf := make([]byte, maxDecompressedSize)
+		return &buf
+	},
+}
+
 type CBORMarshalUnmarshaler interface {
 	cbg.CBORMarshaler
 	cbg.CBORUnmarshaler
@@ -30,11 +38,11 @@ func NewCBOR[T CBORMarshalUnmarshaler]() *CBOR[T] {
 }
 
 func (c *CBOR[T]) Encode(m T) ([]byte, error) {
-	var buf bytes.Buffer
-	if err := m.MarshalCBOR(&buf); err != nil {
+	var out bytes.Buffer
+	if err := m.MarshalCBOR(&out); err != nil {
 		return nil, err
 	}
-	return buf.Bytes(), nil
+	return out.Bytes(), nil
 }
 
 func (c *CBOR[T]) Decode(v []byte, t T) error {
@@ -53,7 +61,9 @@ func NewZSTD[T CBORMarshalUnmarshaler]() (*ZSTD[T], error) {
 	if err != nil {
 		return nil, err
 	}
-	reader, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecompressedSize))
+	reader, err := zstd.NewReader(nil,
+		zstd.WithDecoderMaxMemory(maxDecompressedSize),
+		zstd.WithDecodeAllCapLimit(true))
 	if err != nil {
 		return nil, err
 	}
@@ -78,7 +88,10 @@ func (c *ZSTD[T]) Encode(m T) ([]byte, error) {
 }
 
 func (c *ZSTD[T]) Decode(v []byte, t T) error {
-	cborEncoded, err := c.decompressor.DecodeAll(v, make([]byte, 0, len(v)))
+	buf := bufferPool.Get().(*[]byte)
+	defer bufferPool.Put(buf)
+
+	cborEncoded, err := c.decompressor.DecodeAll(v, (*buf)[:0])
 	if err != nil {
 		return err
 	}