Merge pull request #2445 from CortexFoundation/dev

reduce the memory allocation in trie hashing

Author: ucwong

trie/hasher.go

@@ -1,5 +1,5 @@
-// Copyright 2019 The go-ethereum Authors
-// This file is part of The go-ethereum library.
+// Copyright 2016 The go-ethereum Authors
+// This file is part of the go-ethereum library.
 //
 // The go-ethereum library is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Lesser General Public License as published by
@@ -12,11 +12,13 @@
 // GNU Lesser General Public License for more details.
 //
 // You should have received a copy of the GNU Lesser General Public License
-// along with The go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
 
 package trie
 
 import (
+	"bytes"
+	"fmt"
 	"sync"
 
 	"github.com/CortexFoundation/CortexTheseus/crypto"
@@ -29,7 +31,7 @@ type hasher struct {
 	sha      crypto.KeccakState
 	tmp      []byte
 	encbuf   rlp.EncoderBuffer
-	parallel bool // Whether to use paralallel threads when hashing
+	parallel bool // Whether to use parallel threads when hashing
 }
 
 // hasherPool holds pureHashers
@@ -53,119 +55,119 @@ func returnHasherToPool(h *hasher) {
 	hasherPool.Put(h)
 }
 
-// hash collapses a node down into a hash node, also returning a copy of the
-// original node initialized with the computed hash to replace the original one.
-func (h *hasher) hash(n node, force bool) node {
+// hash collapses a node down into a hash node.
+func (h *hasher) hash(n node, force bool) []byte {
 	// Return the cached hash if it's available
 	if hash, _ := n.cache(); hash != nil {
 		return hash
 	}
 	// Trie not processed yet, walk the children
 	switch n := n.(type) {
 	case *shortNode:
-		collapsed := h.hashShortNodeChildren(n)
-		hashed := h.shortnodeToHash(collapsed, force)
-		// We need to retain the possibly _not_ hashed node, in case it was too
-		// small to be hashed
-		if hn, ok := hashed.(hashNode); ok {
-			n.flags.hash = hn
-		} else {
-			n.flags.hash = nil
+		enc := h.encodeShortNode(n)
+		if len(enc) < 32 && !force {
+			// Nodes smaller than 32 bytes are embedded directly in their parent.
+			// In such cases, return the raw encoded blob instead of the node hash.
+			// It's essential to deep-copy the node blob, as the underlying buffer
+			// of enc will be reused later.
+			buf := make([]byte, len(enc))
+			copy(buf, enc)
+			return buf
 		}
-		return hashed
+		hash := h.hashData(enc)
+		n.flags.hash = hash
+		return hash
+
 	case *fullNode:
-		collapsed := h.hashFullNodeChildren(n)
-		hashed := h.fullnodeToHash(collapsed, force)
-		if hn, ok := hashed.(hashNode); ok {
-			n.flags.hash = hn
-		} else {
-			n.flags.hash = nil
+		enc := h.encodeFullNode(n)
+		if len(enc) < 32 && !force {
+			// Nodes smaller than 32 bytes are embedded directly in their parent.
+			// In such cases, return the raw encoded blob instead of the node hash.
+			// It's essential to deep-copy the node blob, as the underlying buffer
+			// of enc will be reused later.
+			buf := make([]byte, len(enc))
+			copy(buf, enc)
+			return buf
 		}
-		return hashed
-	default:
-		// Value and hash nodes don't have children, so they're left as were
+		hash := h.hashData(enc)
+		n.flags.hash = hash
+		return hash
+
+	case hashNode:
+		// hash nodes don't have children, so they're left as were
 		return n
+
+	default:
+		panic(fmt.Errorf("unexpected node type, %T", n))
 	}
 }
 
-// hashShortNodeChildren collapses the short node. The returned collapsed node
-// holds a live reference to the Key, and must not be modified.
-func (h *hasher) hashShortNodeChildren(n *shortNode) *shortNode {
-	// Hash the short node's child, caching the newly hashed subtree
-	//collapsed, cached = n.copy(), n.copy()
-	var collapsed shortNode
-	// Previously, we did copy this one. We don't seem to need to actually
-	// do that, since we don't overwrite/reuse keys
-	// cached.Key = common.CopyBytes(n.Key)
-	collapsed.Key = hexToCompact(n.Key)
-	// Unless the child is a valuenode or hashnode, hash it
-	switch n.Val.(type) {
-	case *fullNode, *shortNode:
-		collapsed.Val = h.hash(n.Val, false)
-	default:
-		collapsed.Val = n.Val
+// encodeShortNode encodes the provided shortNode into the bytes. Notably, the
+// return slice must be deep-copied explicitly, otherwise the underlying slice
+// will be reused later.
+func (h *hasher) encodeShortNode(n *shortNode) []byte {
+	// Encode leaf node
+	if hasTerm(n.Key) {
+		var ln leafNodeEncoder
+		ln.Key = hexToCompact(n.Key)
+		ln.Val = n.Val.(valueNode)
+		ln.encode(h.encbuf)
+		return h.encodedBytes()
 	}
-	return &collapsed
+	// Encode extension node
+	var en extNodeEncoder
+	en.Key = hexToCompact(n.Key)
+	en.Val = h.hash(n.Val, false)
+	en.encode(h.encbuf)
+	return h.encodedBytes()
+}
+
+// fnEncoderPool is the pool for storing shared fullNode encoder to mitigate
+// the significant memory allocation overhead.
+var fnEncoderPool = sync.Pool{
+	New: func() interface{} {
+		var enc fullnodeEncoder
+		return &enc
+	},
 }
 
-func (h *hasher) hashFullNodeChildren(n *fullNode) *fullNode {
-	// Hash the full node's children, caching the newly hashed subtrees
-	var children [17]node
+// encodeFullNode encodes the provided fullNode into the bytes. Notably, the
+// return slice must be deep-copied explicitly, otherwise the underlying slice
+// will be reused later.
+func (h *hasher) encodeFullNode(n *fullNode) []byte {
+	fn := fnEncoderPool.Get().(*fullnodeEncoder)
+	fn.reset()
+
 	if h.parallel {
 		var wg sync.WaitGroup
 		for i := 0; i < 16; i++ {
-			if child := n.Children[i]; child != nil {
-				wg.Add(1)
-				go func(i int) {
-					hasher := newHasher(false)
-					children[i] = hasher.hash(child, false)
-					returnHasherToPool(hasher)
-					wg.Done()
-				}(i)
-			} else {
-				children[i] = nilValueNode
+			if n.Children[i] == nil {
+				continue
 			}
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+
+				h := newHasher(false)
+				fn.Children[i] = h.hash(n.Children[i], false)
+				returnHasherToPool(h)
+			}(i)
 		}
 		wg.Wait()
 	} else {
 		for i := 0; i < 16; i++ {
 			if child := n.Children[i]; child != nil {
-				children[i] = h.hash(child, false)
-			} else {
-				children[i] = nilValueNode
+				fn.Children[i] = h.hash(child, false)
 			}
 		}
 	}
 	if n.Children[16] != nil {
-		children[16] = n.Children[16]
-	}
-	return &fullNode{flags: nodeFlag{}, Children: children}
-}
-
-// shortnodeToHash creates a hashNode from a shortNode. The supplied shortnode
-// should have hex-type Key, which will be converted (without modification)
-// into compact form for RLP encoding.
-// If the rlp data is smaller than 32 bytes, `nil` is returned.
-func (h *hasher) shortnodeToHash(n *shortNode, force bool) node {
-	n.encode(h.encbuf)
-	enc := h.encodedBytes()
-
-	if len(enc) < 32 && !force {
-		return n // Nodes smaller than 32 bytes are stored inside their parent
+		fn.Children[16] = n.Children[16].(valueNode)
 	}
-	return h.hashData(enc)
-}
-
-// fullnodeToHash is used to create a hashNode from a fullNode, (which
-// may contain nil values)
-func (h *hasher) fullnodeToHash(n *fullNode, force bool) node {
-	n.encode(h.encbuf)
-	enc := h.encodedBytes()
+	fn.encode(h.encbuf)
+	fnEncoderPool.Put(fn)
 
-	if len(enc) < 32 && !force {
-		return n // Nodes smaller than 32 bytes are stored inside their parent
-	}
-	return h.hashData(enc)
+	return h.encodedBytes()
 }
 
 // encodedBytes returns the result of the last encoding operation on h.encbuf.
@@ -184,9 +186,10 @@ func (h *hasher) encodedBytes() []byte {
 	return h.tmp
 }
 
-// hashData hashes the provided data
-func (h *hasher) hashData(data []byte) hashNode {
-	n := make(hashNode, 32)
+// hashData hashes the provided data. It is safe to modify the returned slice after
+// the function returns.
+func (h *hasher) hashData(data []byte) []byte {
+	n := make([]byte, 32)
 	h.sha.Reset()
 	h.sha.Write(data)
 	h.sha.Read(n)
@@ -201,20 +204,17 @@ func (h *hasher) hashDataTo(dst, data []byte) {
 	h.sha.Read(dst)
 }
 
-// proofHash is used to construct trie proofs, and returns the 'collapsed'
-// node (for later RLP encoding) aswell as the hashed node -- unless the
-// node is smaller than 32 bytes, in which case it will be returned as is.
-// This method does not do anything on value- or hash-nodes.
-func (h *hasher) proofHash(original node) (collapsed, hashed node) {
+// proofHash is used to construct trie proofs, returning the rlp-encoded node blobs.
+// Note, only resolved node (shortNode or fullNode) is expected for proofing.
+//
+// It is safe to modify the returned slice after the function returns.
+func (h *hasher) proofHash(original node) []byte {
 	switch n := original.(type) {
 	case *shortNode:
-		sn := h.hashShortNodeChildren(n)
-		return sn, h.shortnodeToHash(sn, false)
+		return bytes.Clone(h.encodeShortNode(n))
 	case *fullNode:
-		fn := h.hashFullNodeChildren(n)
-		return fn, h.fullnodeToHash(fn, false)
+		return bytes.Clone(h.encodeFullNode(n))
 	default:
-		// Value and hash nodes don't have children, so they're left as were
-		return n, n
+		panic(fmt.Errorf("unexpected node type, %T", original))
 	}
 }

trie/iterator.go

@@ -198,9 +198,9 @@ func (it *nodeIterator) LeafProof() [][]byte {
 
 			for i, item := range it.stack[:len(it.stack)-1] {
 				// Gather nodes that end up as hash nodes (or the root)
-				node, hashed := hasher.proofHash(item.node)
-				if _, ok := hashed.(hashNode); ok || i == 0 {
-					proofs = append(proofs, nodeToBytes(node))
+				enc := hasher.proofHash(item.node)
+				if len(enc) >= 32 || i == 0 {
+					proofs = append(proofs, enc)
 				}
 			}
 			return proofs

trie/node.go

@@ -68,10 +68,6 @@ type (
 	}
 )
 
-// nilValueNode is used when collapsing internal trie nodes for hashing, since
-// unset children need to serialize correctly.
-var nilValueNode = valueNode(nil)
-
 // EncodeRLP encodes a full node into the consensus RLP format.
 func (n *fullNode) EncodeRLP(w io.Writer) error {
 	eb := rlp.NewEncoderBuffer(w)

trie/node_enc.go

@@ -42,18 +42,29 @@ func (n *fullNode) encode(w rlp.EncoderBuffer) {
 
 func (n *fullnodeEncoder) encode(w rlp.EncoderBuffer) {
 	offset := w.List()
-	for _, c := range n.Children {
-		if c == nil {
+	for i, c := range n.Children {
+		if len(c) == 0 {
 			w.Write(rlp.EmptyString)
-		} else if len(c) < 32 {
-			w.Write(c) // rawNode
 		} else {
-			w.WriteBytes(c) // hashNode
+			// valueNode or hashNode
+			if i == 16 || len(c) >= 32 {
+				w.WriteBytes(c)
+			} else {
+				w.Write(c) // rawNode
+			}
 		}
 	}
 	w.ListEnd(offset)
 }
 
+func (n *fullnodeEncoder) reset() {
+	for i, c := range n.Children {
+		if len(c) != 0 {
+			n.Children[i] = n.Children[i][:0]
+		}
+	}
+}
+
 func (n *shortNode) encode(w rlp.EncoderBuffer) {
 	offset := w.List()
 	w.WriteBytes(n.Key)
@@ -70,7 +81,7 @@ func (n *extNodeEncoder) encode(w rlp.EncoderBuffer) {
 	w.WriteBytes(n.Key)
 
 	if n.Val == nil {
-		w.Write(rlp.EmptyString)
+		w.Write(rlp.EmptyString) // theoretically impossible to happen
 	} else if len(n.Val) < 32 {
 		w.Write(n.Val) // rawNode
 	} else {

trie/proof.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 
 	"github.com/CortexFoundation/CortexTheseus/common"
+	"github.com/CortexFoundation/CortexTheseus/crypto"
 	"github.com/CortexFoundation/CortexTheseus/ctxcdb"
 	"github.com/CortexFoundation/CortexTheseus/ctxcdb/memorydb"
 	"github.com/CortexFoundation/CortexTheseus/log"
@@ -69,20 +70,9 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ctxcdb.KeyValueWriter)
 	defer returnHasherToPool(hasher)
 
 	for i, n := range nodes {
-		if fromLevel > 0 {
-			fromLevel--
-			continue
-		}
-		var hn node
-		n, hn = hasher.proofHash(n)
-		if hash, ok := hn.(hashNode); ok || i == 0 {
-			// If the node's database encoding is a hash (or is the
-			// root node), it becomes a proof element.
-			enc := nodeToBytes(n)
-			if !ok {
-				hash = hasher.hashData(enc)
-			}
-			proofDb.Put(hash, enc)
+		enc := hasher.proofHash(n)
+		if len(enc) >= 32 || i == 0 {
+			proofDb.Put(crypto.Keccak256(enc), enc)
 		}
 	}
 	return nil

trie/trie.go

@@ -615,7 +615,7 @@ func (t *Trie) resolveBlob(n hashNode, prefix []byte) ([]byte, error) {
 // Hash returns the root hash of the trie. It does not write to the
 // database and can be used even if the trie doesn't have one.
 func (t *Trie) Hash() common.Hash {
-	return common.BytesToHash(t.hashRoot().(hashNode))
+	return common.BytesToHash(t.hashRoot())
 }
 
 // Commit writes all nodes to the trie's memory database, tracking the internal
@@ -674,9 +674,9 @@ func (t *Trie) Commit(onleaf LeafCallback) (root common.Hash, err error) {
 }
 
 // hashRoot calculates the root hash of the given trie
-func (t *Trie) hashRoot() node {
+func (t *Trie) hashRoot() []byte {
 	if t.root == nil {
-		return hashNode(types.EmptyRootHash.Bytes())
+		return types.EmptyRootHash.Bytes()
 	}
 	// If the number of changes is below 100, we let one thread handle it
 	h := newHasher(t.unhashed >= 100)