trie, core: use unsafe pool internally in stacktrie

core/types/hashing_test.go

@@ -93,10 +93,12 @@ func BenchmarkDeriveSha200(b *testing.B) {
 		}
 	})
 
+	st := trie.NewStackTrie(nil)
 	b.Run("stack_trie", func(b *testing.B) {
 		b.ReportAllocs()
 		for b.Loop() {
-			have = types.DeriveSha(txs, trie.NewStackTrie(nil))
+			st.Reset()
+			have = types.DeriveSha(txs, st)
 		}
 		if have != want {
 			b.Errorf("have %x want %x", have, want)

trie/bytepool.go

@@ -32,8 +32,8 @@ func newBytesPool(sliceCap, nitems int) *bytesPool {
 	}
 }
 
-// Get returns a slice. Safe for concurrent use.
-func (bp *bytesPool) Get() []byte {
+// get returns a slice. Safe for concurrent use.
+func (bp *bytesPool) get() []byte {
 	select {
 	case b := <-bp.c:
 		return b
@@ -42,18 +42,18 @@ func (bp *bytesPool) Get() []byte {
 	}
 }
 
-// GetWithSize returns a slice with specified byte slice size.
-func (bp *bytesPool) GetWithSize(s int) []byte {
-	b := bp.Get()
+// getWithSize returns a slice with specified byte slice size.
+func (bp *bytesPool) getWithSize(s int) []byte {
+	b := bp.get()
 	if cap(b) < s {
 		return make([]byte, s)
 	}
 	return b[:s]
 }
 
-// Put returns a slice to the pool. Safe for concurrent use. This method
+// put returns a slice to the pool. Safe for concurrent use. This method
 // will ignore slices that are too small or too large (>3x the cap)
-func (bp *bytesPool) Put(b []byte) {
+func (bp *bytesPool) put(b []byte) {
 	if c := cap(b); c < bp.w || c > 3*bp.w {
 		return
 	}
@@ -62,3 +62,40 @@ func (bp *bytesPool) Put(b []byte) {
 	default:
 	}
 }
+
+// unsafeBytesPool is a pool for byte slices. It is not safe for concurrent use.
+type unsafeBytesPool struct {
+	items [][]byte
+	w     int
+}
+
+// newUnsafeBytesPool creates a new unsafeBytesPool. The sliceCap sets the
+// capacity of newly allocated slices, and the nitems determines how many
+// items the pool will hold, at maximum.
+func newUnsafeBytesPool(sliceCap, nitems int) *unsafeBytesPool {
+	return &unsafeBytesPool{
+		items: make([][]byte, 0, nitems),
+		w:     sliceCap,
+	}
+}
+
+// Get returns a slice.
+func (bp *unsafeBytesPool) get() []byte {
+	if len(bp.items) > 0 {
+		last := bp.items[len(bp.items)-1]
+		bp.items = bp.items[:len(bp.items)-1]
+		return last
+	}
+	return make([]byte, 0, bp.w)
+}
+
+// put returns a slice to the pool. This method
+// will ignore slices that are too small or too large (>3x the cap)
+func (bp *unsafeBytesPool) put(b []byte) {
+	if c := cap(b); c < bp.w || c > 3*bp.w {
+		return
+	}
+	if len(bp.items) < cap(bp.items) {
+		bp.items = append(bp.items, b)
+	}
+}

trie/stacktrie.go

@@ -50,7 +50,7 @@ type StackTrie struct {
 	onTrieNode OnTrieNode
 	kBuf       []byte // buf space used for hex-key during insertions
 	pBuf       []byte // buf space used for path during insertions
-	vPool      *bytesPool
+	vPool      *unsafeBytesPool
 }
 
 // NewStackTrie allocates and initializes an empty trie. The committed nodes
@@ -62,7 +62,7 @@ func NewStackTrie(onTrieNode OnTrieNode) *StackTrie {
 		onTrieNode: onTrieNode,
 		kBuf:       make([]byte, 64),
 		pBuf:       make([]byte, 64),
-		vPool:      newBytesPool(300, 20),
+		vPool:      newUnsafeBytesPool(300, 20),
 	}
 }
 
@@ -99,7 +99,7 @@ func (t *StackTrie) Update(key, value []byte) error {
 	} else {
 		t.last = append(t.last[:0], k...) // reuse key slice
 	}
-	vBuf := t.vPool.Get()
+	vBuf := t.vPool.get()
 	if cap(vBuf) < len(value) {
 		vBuf = common.CopyBytes(value)
 	} else {
@@ -114,6 +114,7 @@ func (t *StackTrie) Update(key, value []byte) error {
 func (t *StackTrie) Reset() {
 	t.root = stPool.Get().(*stNode)
 	t.last = nil
+	t.onTrieNode = nil
 }
 
 // TrieKey returns the internal key representation for the given user key.
@@ -166,7 +167,7 @@ func (n *stNode) reset() *stNode {
 		// On hashnodes, we 'own' the val: it is guaranteed to be not held
 		// by external caller. Hence, when we arrive here, we can put it back
 		// into the pool
-		bPool.Put(n.val)
+		bPool.put(n.val)
 	}
 	n.key = n.key[:0]
 	n.val = nil
@@ -413,20 +414,20 @@ func (t *StackTrie) hash(st *stNode, path []byte) {
 
 	// Release reference to (potentially externally held) value-slice.
 	if cap(st.val) > 0 && t.vPool != nil {
-		t.vPool.Put(st.val)
+		t.vPool.put(st.val)
 	}
 	st.val = nil
 
 	// Skip committing the non-root node if the size is smaller than 32 bytes
 	// as tiny nodes are always embedded in their parent except root node.
 	if len(blob) < 32 && len(path) > 0 {
-		st.val = bPool.GetWithSize(len(blob))
+		st.val = bPool.getWithSize(len(blob))
 		copy(st.val, blob)
 		return
 	}
 	// Write the hash to the 'val'. We allocate a new val here to not mutate
 	// input values.
-	st.val = bPool.GetWithSize(32)
+	st.val = bPool.getWithSize(32)
 	t.h.hashDataTo(st.val, blob)
 
 	// Invoke the callback it's provided. Notably, the path and blob slices are