Skip to content

Commit

Permalink
feat: shrink allocations when setting new elements
Browse files Browse the repository at this point in the history
`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.
  • Loading branch information
Stebalien committed Apr 20, 2024
1 parent f6a28b2 commit ed7cd8c
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 5 deletions.
7 changes: 3 additions & 4 deletions amt.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package amt

import (
"bytes"
"context"
"fmt"
"math"
Expand Down Expand Up @@ -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
Expand Down
35 changes: 34 additions & 1 deletion util.go
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
}

0 comments on commit ed7cd8c

Please sign in to comment.