From ed7cd8c77f4981192fd4e59979beb352274b9843 Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Fri, 19 Apr 2024 17:34:46 -0700 Subject: [PATCH 1/3] feat: shrink allocations when setting new elements `bytes.Buffer` will often over-allocate. Here, we copy to an exact-sized buffer to avoid this. We'll pay some extra copying cost, but we can at least avoid the allocation overhead with a buffer pool in most cases. --- amt.go | 7 +++---- util.go | 35 ++++++++++++++++++++++++++++++++++- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/amt.go b/amt.go index 586a088..cc438ca 100644 --- a/amt.go +++ b/amt.go @@ -1,7 +1,6 @@ package amt import ( - "bytes" "context" "fmt" "math" @@ -149,11 +148,11 @@ func (r *Root) Set(ctx context.Context, i uint64, val cbg.CBORMarshaler) error { if val == nil { d.Raw = cbg.CborNull } else { - valueBuf := new(bytes.Buffer) - if err := val.MarshalCBOR(valueBuf); err != nil { + data, err := cborToBytes(val) + if err != nil { return err } - d.Raw = valueBuf.Bytes() + d.Raw = data } // where the index is greater than the number of elements we can fit into the diff --git a/util.go b/util.go index adf86a6..577debd 100644 --- a/util.go +++ b/util.go @@ -1,6 +1,12 @@ package amt -import "math" +import ( + "bytes" + "math" + "sync" + + cbg "github.com/whyrusleeping/cbor-gen" +) // Given height 'height', how many nodes in a maximally full tree can we // build? (bitWidth^2)^height = width^height. If we pass in height+1 we can work @@ -13,3 +19,30 @@ func nodesForHeight(bitWidth uint, height int) uint64 { } return 1 << heightLogTwo } + +var bufferPool sync.Pool = sync.Pool{ + New: func() any { + return bytes.NewBuffer(nil) + }, +} + +func cborToBytes(val cbg.CBORMarshaler) ([]byte, error) { + // Temporary location to put values. We'll copy them to an exact-sized buffer when done. + valueBuf := bufferPool.Get().(*bytes.Buffer) + defer func() { + valueBuf.Reset() + bufferPool.Put(valueBuf) + }() + + if err := val.MarshalCBOR(valueBuf); err != nil { + + return nil, err + } + + // Copy to shrink the allocation. + buf := valueBuf.Bytes() + cpy := make([]byte, len(buf)) + copy(cpy, buf) + + return cpy, nil +} From 784c15d56cac2422a8a1f6a21347f189dffaa3cc Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Fri, 19 Apr 2024 20:15:48 -0700 Subject: [PATCH 2/3] remove extra line Co-authored-by: Rod Vagg --- util.go | 1 - 1 file changed, 1 deletion(-) diff --git a/util.go b/util.go index 577debd..84ccb0a 100644 --- a/util.go +++ b/util.go @@ -35,7 +35,6 @@ func cborToBytes(val cbg.CBORMarshaler) ([]byte, error) { }() if err := val.MarshalCBOR(valueBuf); err != nil { - return nil, err } From c64ccd156b18ffbd7ba6348b32f0b0787bbb551b Mon Sep 17 00:00:00 2001 From: Steven Allen Date: Sat, 20 Apr 2024 09:52:43 -0700 Subject: [PATCH 3/3] chore: remove unnecessary explicit type Co-authored-by: Masih H. Derkani --- util.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/util.go b/util.go index 84ccb0a..8d10265 100644 --- a/util.go +++ b/util.go @@ -20,7 +20,7 @@ func nodesForHeight(bitWidth uint, height int) uint64 { return 1 << heightLogTwo } -var bufferPool sync.Pool = sync.Pool{ +var bufferPool = sync.Pool{ New: func() any { return bytes.NewBuffer(nil) },