Skip to content

Commit f28d5b4

Browse files
committed
Pool decode buffers when encoding messages using compression
Pool the decode buffers and set strict max cap allocation for zstd decompressor. See Benchamark results below. Before: ``` BenchmarkCborEncoding BenchmarkCborEncoding-12 467352 2376 ns/op 14680 B/op 10 allocs/op BenchmarkCborDecoding BenchmarkCborDecoding-12 104410 11347 ns/op 14944 B/op 27 allocs/op BenchmarkZstdEncoding BenchmarkZstdEncoding-12 286735 3897 ns/op 46748 B/op 12 allocs/op BenchmarkZstdDecoding BenchmarkZstdDecoding-12 110794 10783 ns/op 28512 B/op 28 allocs/op ``` After: ``` BenchmarkCborEncoding BenchmarkCborEncoding-12 436754 2383 ns/op 14680 B/op 10 allocs/op BenchmarkCborDecoding BenchmarkCborDecoding-12 106809 11280 ns/op 14944 B/op 27 allocs/op BenchmarkZstdEncoding BenchmarkZstdEncoding-12 294043 3918 ns/op 46746 B/op 12 allocs/op BenchmarkZstdDecoding BenchmarkZstdDecoding-12 114854 10747 ns/op 18314 B/op 27 allocs/op ``` Fixes: #850 #849
1 parent bae93e2 commit f28d5b4

File tree

2 files changed

+38
-7
lines changed

2 files changed

+38
-7
lines changed

encoding_bench_test.go

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ func BenchmarkCborDecoding(b *testing.B) {
4343
for pb.Next() {
4444
var got PartialGMessage
4545
require.NoError(b, encoder.Decode(data, &got))
46-
require.Equal(b, msg, &got)
46+
requireEqualPartialMessages(b, msg, &got)
4747
}
4848
})
4949
}
@@ -79,11 +79,29 @@ func BenchmarkZstdDecoding(b *testing.B) {
7979
for pb.Next() {
8080
var got PartialGMessage
8181
require.NoError(b, encoder.Decode(data, &got))
82-
require.Equal(b, msg, &got)
82+
requireEqualPartialMessages(b, msg, &got)
8383
}
8484
})
8585
}
8686

87+
func requireEqualPartialMessages(b *testing.B, expected, actual *PartialGMessage) {
88+
// Because empty ECChain gets marshaled as null, we need to use ECChain.Eq for
89+
// checking equality. Hence, the custom equality check.
90+
require.Equal(b, expected.Sender, actual.Sender)
91+
require.Equal(b, expected.Signature, actual.Signature)
92+
require.Equal(b, expected.VoteValueKey, actual.VoteValueKey)
93+
require.Equal(b, expected.Ticket, actual.Ticket)
94+
require.True(b, expected.Vote.Eq(&actual.Vote))
95+
if expected.Justification == nil {
96+
require.Nil(b, actual.Justification)
97+
} else {
98+
require.NotNil(b, actual.Justification)
99+
require.Equal(b, expected.Justification.Signature, actual.Justification.Signature)
100+
require.Equal(b, expected.Justification.Signers, actual.Justification.Signers)
101+
require.True(b, expected.Justification.Vote.Eq(&actual.Justification.Vote))
102+
}
103+
}
104+
87105
func generateRandomPartialGMessage(b *testing.B, rng *rand.Rand) *PartialGMessage {
88106
var pgmsg PartialGMessage
89107
pgmsg.GMessage = generateRandomGMessage(b, rng)

internal/encoding/encoding.go

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package encoding
33
import (
44
"bytes"
55
"fmt"
6+
"sync"
67

78
"github.com/klauspost/compress/zstd"
89
cbg "github.com/whyrusleeping/cbor-gen"
@@ -13,6 +14,13 @@ import (
1314
// size in GossipSub.
1415
const maxDecompressedSize = 1 << 20
1516

17+
var bufferPool = sync.Pool{
18+
New: func() any {
19+
buf := make([]byte, maxDecompressedSize)
20+
return &buf
21+
},
22+
}
23+
1624
type CBORMarshalUnmarshaler interface {
1725
cbg.CBORMarshaler
1826
cbg.CBORUnmarshaler
@@ -30,11 +38,11 @@ func NewCBOR[T CBORMarshalUnmarshaler]() *CBOR[T] {
3038
}
3139

3240
func (c *CBOR[T]) Encode(m T) ([]byte, error) {
33-
var buf bytes.Buffer
34-
if err := m.MarshalCBOR(&buf); err != nil {
41+
var out bytes.Buffer
42+
if err := m.MarshalCBOR(&out); err != nil {
3543
return nil, err
3644
}
37-
return buf.Bytes(), nil
45+
return out.Bytes(), nil
3846
}
3947

4048
func (c *CBOR[T]) Decode(v []byte, t T) error {
@@ -53,7 +61,9 @@ func NewZSTD[T CBORMarshalUnmarshaler]() (*ZSTD[T], error) {
5361
if err != nil {
5462
return nil, err
5563
}
56-
reader, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecompressedSize))
64+
reader, err := zstd.NewReader(nil,
65+
zstd.WithDecoderMaxMemory(maxDecompressedSize),
66+
zstd.WithDecodeAllCapLimit(true))
5767
if err != nil {
5868
return nil, err
5969
}
@@ -78,7 +88,10 @@ func (c *ZSTD[T]) Encode(m T) ([]byte, error) {
7888
}
7989

8090
func (c *ZSTD[T]) Decode(v []byte, t T) error {
81-
cborEncoded, err := c.decompressor.DecodeAll(v, make([]byte, 0, len(v)))
91+
buf := bufferPool.Get().(*[]byte)
92+
defer bufferPool.Put(buf)
93+
94+
cborEncoded, err := c.decompressor.DecodeAll(v, (*buf)[:0])
8295
if err != nil {
8396
return err
8497
}

0 commit comments

Comments
 (0)