trie: batch pruning better, cache unreferencing trie prefixes

trie/database.go

@@ -571,7 +571,12 @@ func (db *Database) Dereference(root common.Hash, prune bool) error {
 
 	nodes, storage, start := len(db.dirties), db.dirtiesSize, time.Now()
 	prunetime, prunenodes, prunesize := db.prunetime, db.prunenodes, db.prunesize
-	if err := db.dereference(common.Hash{}, root, common.Hash{}, common.Hash{}, prune, nil, db.newPruner(root)); err != nil {
+
+	pruner := db.newPruner(root)
+	if err := db.dereference(common.Hash{}, root, common.Hash{}, common.Hash{}, prune, nil, pruner); err != nil {
+		return err
+	}
+	if err := pruner.flush(); err != nil {
 		return err
 	}
 	db.gcnodes += uint64(nodes - len(db.dirties))
@@ -610,15 +615,9 @@ func (db *Database) dereference(childOwner common.Hash, childHash common.Hash, p
 	child, ok := db.dirties[childKey]
 	if !ok {
 		if prune {
-			batch := db.diskdb.NewBatch()
-
 			start := time.Now()
-			pruner.prune(childOwner, childHash, path, batch)
+			pruner.prune(childOwner, childHash, path)
 			db.prunetime += time.Since(start)
-
-			if err := batch.Write(); err != nil {
-				return err
-			}
 		}
 		return nil
 	}

trie/database_pruning.go

@@ -33,6 +33,7 @@ import (
 type pruner struct {
 	db    *Database    // Trie database for accessing dirty and clean data
 	tries []*traverser // Individual stateful trie traversers for fast liveness checks
+	batch ethdb.Batch  // Write batch to minimize database trashing
 }
 
 // newPruner creates a new trie pruner tied to the liveness of all the currently
@@ -44,8 +45,9 @@ func (db *Database) newPruner(skip common.Hash) *pruner {
 		if _, root := splitNodeKey(key); root != skip {
 			traversers = append(traversers, &traverser{
 				db: db,
-				state: &tranverserState{
+				state: &traverserState{
 					node: hashNode(root[:]),
+					hash: root,
 				},
 			})
 		}
@@ -54,12 +56,22 @@ func (db *Database) newPruner(skip common.Hash) *pruner {
 	return &pruner{
 		db:    db,
 		tries: traversers,
+		batch: db.diskdb.NewBatch(),
 	}
 }
 
+// flush commits any pending database writes.
+func (p *pruner) flush() error {
+	if err := p.batch.Write(); err != nil {
+		return err
+	}
+	p.batch.Reset()
+	return nil
+}
+
 // prune deletes a trie node from disk if there are no more live references to
 // it, cascading until all dangling nodes are removed.
-func (p *pruner) prune(owner common.Hash, hash common.Hash, path []byte, batch ethdb.Batch) {
+func (p *pruner) prune(owner common.Hash, hash common.Hash, path []byte) {
 	// If the node is still live in the memory cache, it's still referenced so we
 	// can abort. This case is important when and old trie being pruned references
 	// a new node (maybe that node was recreted since), since currently live nodes
@@ -73,10 +85,14 @@ func (p *pruner) prune(owner common.Hash, hash common.Hash, path []byte, batch e
 	if owner != (common.Hash{}) {
 		crosspath = append(append(keybytesToHex(owner[:]), 0xff), crosspath...)
 	}
+	unrefs := make(map[common.Hash]bool)
 	for _, trie := range p.tries {
-		if trie.live(owner, hash, crosspath) {
+		// If the node is still live, abort
+		if trie.live(owner, hash, crosspath, unrefs) {
 			return
 		}
+		// Node dead in this trie, cache the result for subsequent traversals
+		trie.unref(2, unrefs)
 	}
 	// Dead node found, delete it from the database
 	dead := []byte(makeNodeKey(owner, hash))
@@ -89,12 +105,12 @@ func (p *pruner) prune(owner common.Hash, hash common.Hash, path []byte, batch e
 
 	// Prune the node and its children if it's not a bytecode blob
 	p.db.cleans.Delete(key)
-	batch.Delete(dead)
+	p.batch.Delete(dead)
 	p.db.prunenodes++
 	p.db.prunesize += common.StorageSize(len(blob))
 
 	iterateRefs(node, path, func(path []byte, hash common.Hash) error {
-		p.prune(owner, hash, path, batch)
+		p.prune(owner, hash, path)
 		return nil
 	})
 }
@@ -104,15 +120,16 @@ func (p *pruner) prune(owner common.Hash, hash common.Hash, path []byte, batch e
 // a separate data structure is to allow reusing previous traversals to check
 // the liveness of nested nodes (i.e. entire subtried during pruning).
 type traverser struct {
-	db    *Database        // Trie database for accessing dirty and clean data
+	db    *Database       // Trie database for accessing dirty and clean data
-	state *tranverserState // Leftover state from the previous traversals
+	state *traverserState // Leftover state from the previous traversals
 }
 
-// tranverserState is the internal state of a trie traverser.
+// traverserState is the internal state of a trie traverser.
-type tranverserState struct {
+type traverserState struct {
-	parent *tranverserState // Parent traverser to allow backtracking
+	parent *traverserState // Parent traverser to allow backtracking
-	prefix []byte           // Path leading up to the root of this traverser
+	prefix []byte          // Path leading up to the root of this traverser
-	node   node             // Trie node where this traverser is currently at
+	node   node            // Trie node where this traverser is currently at
+	hash   common.Hash     // Hash of the trie node at the traversed position
 }
 
 // live checks whether the trie iterated by this traverser contains the hashnode
@@ -124,37 +141,54 @@ type tranverserState struct {
 // paths are separated by a 0xff byte (nibbles range from 0x00-0x10). This byte
 // is needed to differentiate between the leaf of the account trie and the root
 // of a storage trie (which otherwise would have the same traversal path).
-func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool {
+func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte, unrefs map[common.Hash]bool) bool {
 	// Rewind the traverser until it's prefix is actually a prefix of the path
 	for !bytes.HasPrefix(path, t.state.prefix) {
 		t.state = t.state.parent
 	}
+	// Short circuit the liveness check if we already covered this prefix (if this
+	// prefix path was not yet seen in previous tries, no parent could have been
+	// seen either, so no point in checkin upwards further than the first hash).
+	state := t.state
+	for state != nil {
+		// If we've found a hash node, check if it's an already known result
+		if state.hash != (common.Hash{}) {
+			if unrefs[state.hash] {
+				return false
+			}
+			break
+		}
+		// Not a hash node, traverse further up
+		state = state.parent
+	}
 	// Traverse downward until the prefix matches the path completely
 	path = path[len(t.state.prefix):]
 	for len(path) > 0 {
 		// If we're at a hash node, expand before continuing
 		if n, ok := t.state.node.(hashNode); ok {
+			// Short circuit if we already encountered this node
+			t.state.hash = common.BytesToHash(n)
+			if unrefs[t.state.hash] {
+				return false
+			}
 			// Generate the database key for this hash node
-			var (
+			var key string
-				key  string
-				hash = common.BytesToHash(n)
-			)
 			if len(t.state.prefix) < 2*common.HashLength {
-				key = makeNodeKey(common.Hash{}, hash)
+				key = makeNodeKey(common.Hash{}, t.state.hash)
 			} else {
-				key = makeNodeKey(owner, hash)
+				key = makeNodeKey(owner, t.state.hash)
 			}
 			// Replace the node in the traverser with the expanded one
 			if enc, err := t.db.cleans.Get(key); err == nil && enc != nil {
-				t.state.node = mustDecodeNode(hash[:], enc, 0)
+				t.state.node = mustDecodeNode(t.state.hash[:], enc, 0)
 			} else if node := t.db.dirties[key]; node != nil {
 				t.state.node = node.node
 			} else {
 				blob, err := t.db.diskdb.Get([]byte(key))
 				if blob == nil || err != nil {
-					panic(fmt.Sprintf("missing referenced node %x (searching for %x:%x at %x%x)", key, owner, hash, t.state.prefix, path))
+					panic(fmt.Sprintf("missing referenced node %x (searching for %x:%x at %x%x)", key, owner, t.state.hash, t.state.prefix, path))
 				}
-				t.state.node = mustDecodeNode(hash[:], blob, 0)
+				t.state.node = mustDecodeNode(t.state.hash[:], blob, 0)
 			}
 		}
 		// If we reached an account node, extract the storage trie root to continue on
@@ -174,7 +208,7 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 					return false
 				}
 				// Create a new nesting in the traversal and continue on that depth
-				t.state, path = &tranverserState{
+				t.state, path = &traverserState{
 					parent: t.state,
 					prefix: append(t.state.prefix, 0xff),
 					node:   hashNode(account.Root[:]),
@@ -189,7 +223,7 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 		switch n := t.state.node.(type) {
 		case *rawShortNode:
 			if prefixLen(n.Key, path) == len(n.Key) {
-				t.state, path = &tranverserState{
+				t.state, path = &traverserState{
 					parent: t.state,
 					prefix: append(t.state.prefix, path[:len(n.Key)]...),
 					node:   n.Val,
@@ -200,7 +234,7 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 
 		case *shortNode:
 			if prefixLen(n.Key, path) == len(n.Key) {
-				t.state, path = &tranverserState{
+				t.state, path = &traverserState{
 					parent: t.state,
 					prefix: append(t.state.prefix, path[:len(n.Key)]...),
 					node:   n.Val,
@@ -211,7 +245,7 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 
 		case rawFullNode:
 			if child := n[path[0]]; child != nil {
-				t.state, path = &tranverserState{
+				t.state, path = &traverserState{
 					parent: t.state,
 					prefix: append(t.state.prefix, path[0]),
 					node:   child,
@@ -222,7 +256,7 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 
 		case *fullNode:
 			if child := n.Children[path[0]]; child != nil {
-				t.state, path = &tranverserState{
+				t.state, path = &traverserState{
 					parent: t.state,
 					prefix: append(t.state.prefix, path[0]),
 					node:   child,
@@ -236,11 +270,28 @@ func (t *traverser) live(owner common.Hash, hash common.Hash, path []byte) bool
 		}
 	}
 	// The prefix should match perfectly here, check if the hashes matches
-	if have, ok := t.state.node.(hashNode); ok {
+	if t.state.hash != (common.Hash{}) { // expanded/cached hash node
-		return common.BytesToHash(have) == hash
+		return t.state.hash == hash
 	}
-	if have, _ := t.state.node.cache(); have != nil {
+	if have, ok := t.state.node.(hashNode); ok { // collapsed hash node
-		return common.BytesToHash(have) == hash
+		t.state.hash = common.BytesToHash(have)
+		return t.state.hash == hash
 	}
 	return false
 }
+
+// unref marks the current traversal nodes as *not* containing the specific trie
+// node having been searched for. It is used by searches in subsequent tries to
+// avoid reiterating the exact same sub-tries.
+func (t *traverser) unref(count int, unrefs map[common.Hash]bool) {
+	state := t.state
+	for state != nil && count > 0 {
+		// If we've found a hash node, store it as a subresult
+		if state.hash != (common.Hash{}) {
+			unrefs[state.hash] = true
+			count--
+		}
+		// Traverse further up to the next hash node
+		state = state.parent
+	}
+}