trie: make hasher parallel when number of changes are large

trie/pure_hasher.go

@@ -26,23 +26,23 @@ import (
 type pureHasher struct {
 	sha keccakState
 
-	tmp    sliceBuffer
-	tmpKey []byte
+	tmp      sliceBuffer
+	parallel bool
 }
 
 // hashers live in a global db.
 var pureHasherPool = sync.Pool{
 	New: func() interface{} {
 		return &pureHasher{
-			tmp:    make(sliceBuffer, 0, 550), // cap is as large as a full fullNode.
-			tmpKey: make([]byte, 64),          // space for an packed key
-			sha:    sha3.NewLegacyKeccak256().(keccakState),
+			tmp: make(sliceBuffer, 0, 550), // cap is as large as a full fullNode.
+			sha: sha3.NewLegacyKeccak256().(keccakState),
 		}
 	},
 }
 
-func newPureHasher() *pureHasher {
+func newPureHasher(parallel bool) *pureHasher {
 	h := pureHasherPool.Get().(*pureHasher)
+	h.parallel = parallel
 	return h
 }
 
@@ -107,14 +107,31 @@ func (h *pureHasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cac
 	// Hash the full node's children, caching the newly hashed subtrees
 	cached = n.copy()
 	collapsed = n.copy()
-	for i := 0; i < 16; i++ {
-		if child := n.Children[i]; child != nil {
-			collapsed.Children[i], cached.Children[i] = h.hash(child, false)
-		} else {
-			collapsed.Children[i] = nilValueNode
+	if h.parallel {
+		var wg sync.WaitGroup
+		wg.Add(16)
+		for i := 0; i < 16; i++ {
+			go func(i int) {
+				hasher := newPureHasher(false)
+				if child := n.Children[i]; child != nil {
+					collapsed.Children[i], cached.Children[i] = hasher.hash(child, false)
+				} else {
+					collapsed.Children[i] = nilValueNode
+				}
+				wg.Done()
+				defer returnPureHasherToPool(hasher)
+			}(i)
+		}
+		wg.Wait()
+	} else {
+		for i := 0; i < 16; i++ {
+			if child := n.Children[i]; child != nil {
+				collapsed.Children[i], cached.Children[i] = h.hash(child, false)
+			} else {
+				collapsed.Children[i] = nilValueNode
+			}
 		}
 	}
-	cached.Children[16] = n.Children[16]
 	return collapsed, cached
 }
 
@@ -139,7 +156,7 @@ func (h *pureHasher) shortnodeToHash(n *shortNode, force bool) node {
 func (h *pureHasher) fullnodeToHash(n *fullNode, force bool) node {
 	h.tmp.Reset()
 	// Generate the RLP encoding of the node
-	if err := rlp.Encode(&h.tmp, n); err != nil {
+	if err := n.EncodeRLP(&h.tmp); err != nil {
 		panic("encode error: " + err.Error())
 	}
 

trie/trie.go

@@ -48,6 +48,10 @@ type LeafCallback func(leaf []byte, parent common.Hash) error
 type Trie struct {
 	db   *Database
 	root node
+	// Keep a rough track of the number of leafs to commit
+	dirtyCount int
+	// And leafs to hash
+	unhashedCount int
 }
 
 // newFlag returns the cache flag value for a newly created node.
@@ -163,6 +167,8 @@ func (t *Trie) Update(key, value []byte) {
 //
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *Trie) TryUpdate(key, value []byte) error {
+	t.unhashedCount++
+	t.dirtyCount++
 	k := keybytesToHex(key)
 	if len(value) != 0 {
 		_, n, err := t.insert(t.root, nil, k, valueNode(value))
@@ -259,6 +265,8 @@ func (t *Trie) Delete(key []byte) {
 // TryDelete removes any existing value for key from the trie.
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *Trie) TryDelete(key []byte) error {
+	t.unhashedCount++
+	t.dirtyCount++
 	k := keybytesToHex(key)
 	_, n, err := t.delete(t.root, nil, k)
 	if err != nil {
@@ -405,7 +413,7 @@ func (t *Trie) resolveHash(n hashNode, prefix []byte) (node, 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 {
-	hash, cached, _ := t.hashRoot(nil, nil)
+	hash, cached, _ := t.hashRoot(nil)
 	t.root = cached
 	return common.BytesToHash(hash.(hashNode))
 }
@@ -454,6 +462,7 @@ func (t *Trie) Commit(onleaf LeafCallback) (root common.Hash, err error) {
 	if err != nil {
 		return common.Hash{}, err
 	}
+	t.dirtyCount = 0
 	return rootHash, nil
 }
 
@@ -468,12 +477,14 @@ func (t *Trie) oldHashRoot(db *Database, onleaf LeafCallback) (node, node, error
 }
 
 // hashRoot calculates the root hash of the given trie
-func (t *Trie) hashRoot(db *Database, onleaf LeafCallback) (node, node, error) {
+func (t *Trie) hashRoot(db *Database) (node, node, error) {
 	if t.root == nil {
 		return hashNode(emptyRoot.Bytes()), nil, nil
 	}
-	h := newPureHasher()
+	// If the number of changes is below 100, we let one thread handle it
+	h := newPureHasher(t.unhashedCount >= 100)
 	defer returnPureHasherToPool(h)
 	hashed, cached := h.hash(t.root, true)
+	t.unhashedCount = 0
 	return hashed, cached, nil
 }