core/state, eth/downloader, trie: construct direct cache during fast

cmd/geth/chaincmd.go

@@ -135,7 +135,10 @@ func importChain(ctx *cli.Context) error {
 	}
 	fmt.Println(stats)
 	fmt.Printf("Trie cache misses:   %d\n", trie.CacheMisses())
-	fmt.Printf("Trie cache unloads: %d\n\n", trie.CacheUnloads())
+	fmt.Printf("Trie cache unloads:  %d\n", trie.CacheUnloads())
+	fmt.Printf("Direct cache reads:  %d\n", trie.DirectCacheReads())
+	fmt.Printf("Direct cache writes: %d\n", trie.DirectCacheWrites())
+	fmt.Printf("Direct cache misses: %d\n\n", trie.DirectCacheMisses())
 
 	// Print the memory statistics used by the importing
 	mem := new(runtime.MemStats)

core/state/sync.go

@@ -32,10 +32,11 @@ import (
 type StateSync trie.TrieSync
 
 // NewStateSync create a new state trie download scheduler.
-func NewStateSync(root common.Hash, database ethdb.Database) *StateSync {
+func NewStateSync(number uint64, hash common.Hash, root common.Hash, database ethdb.Database) *StateSync {
 	var syncer *trie.TrieSync
 
-	callback := func(leaf []byte, parent common.Hash) error {
+	callback := func(keys [][]byte, leaf []byte, parent common.Hash) error {
+		// Try to decode any leaf nodes as accounts
 		var obj struct {
 			Nonce    uint64
 			Balance  *big.Int
@@ -45,12 +46,19 @@ func NewStateSync(root common.Hash, database ethdb.Database) *StateSync {
 		if err := rlp.Decode(bytes.NewReader(leaf), &obj); err != nil {
 			return err
 		}
-		syncer.AddSubTrie(obj.Root, 64, parent, nil)
+		// Populate the direct account caches in the database
+		for _, key := range keys {
+			if err := trie.WriteDirectCache(CachePrefix, key, leaf, number, hash, database); err != nil {
+				return err
+			}
+		}
+		// Schedule downloading all the dependencies of the account
+		syncer.AddSubTrie(obj.Root, 64, parent, nil, nil)
 		syncer.AddRawEntry(common.BytesToHash(obj.CodeHash), 64, parent)
 
 		return nil
 	}
-	syncer = trie.NewTrieSync(root, database, callback)
+	syncer = trie.NewTrieSync(root, database, callback, CachePrefix)
 	return (*StateSync)(syncer)
 }
 

core/state/sync_test.go

@@ -47,8 +47,8 @@ func makeTestState() (ethdb.Database, common.Hash, []*testAccount) {
 		obj := state.GetOrNewStateObject(common.BytesToAddress([]byte{i}))
 		acc := &testAccount{address: common.BytesToAddress([]byte{i})}
 
-		obj.AddBalance(big.NewInt(int64(11 * i)))
+		obj.AddBalance(big.NewInt(11 * int64(i)))
-		acc.balance = big.NewInt(int64(11 * i))
+		acc.balance = big.NewInt(11 * int64(i))
 
 		obj.SetNonce(uint64(42 * i))
 		acc.nonce = uint64(42 * i)
@@ -110,7 +110,7 @@ func checkStateConsistency(db ethdb.Database, root common.Hash) error {
 func TestEmptyStateSync(t *testing.T) {
 	empty := common.HexToHash("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421")
 	db, _ := ethdb.NewMemDatabase()
-	if req := NewStateSync(empty, db).Missing(1); len(req) != 0 {
+	if req := NewStateSync(0, common.Hash{}, empty, db).Missing(1); len(req) != 0 {
 		t.Errorf("content requested for empty state: %v", req)
 	}
 }
@@ -126,7 +126,7 @@ func testIterativeStateSync(t *testing.T, batch int) {
 
 	// Create a destination state and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewStateSync(srcRoot, dstDb)
+	sched := NewStateSync(0, common.Hash{}, srcRoot, dstDb)
 
 	queue := append([]common.Hash{}, sched.Missing(batch)...)
 	for len(queue) > 0 {
@@ -155,7 +155,7 @@ func TestIterativeDelayedStateSync(t *testing.T) {
 
 	// Create a destination state and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewStateSync(srcRoot, dstDb)
+	sched := NewStateSync(0, common.Hash{}, srcRoot, dstDb)
 
 	queue := append([]common.Hash{}, sched.Missing(0)...)
 	for len(queue) > 0 {
@@ -189,7 +189,7 @@ func testIterativeRandomStateSync(t *testing.T, batch int) {
 
 	// Create a destination state and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewStateSync(srcRoot, dstDb)
+	sched := NewStateSync(0, common.Hash{}, srcRoot, dstDb)
 
 	queue := make(map[common.Hash]struct{})
 	for _, hash := range sched.Missing(batch) {
@@ -226,7 +226,7 @@ func TestIterativeRandomDelayedStateSync(t *testing.T) {
 
 	// Create a destination state and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewStateSync(srcRoot, dstDb)
+	sched := NewStateSync(0, common.Hash{}, srcRoot, dstDb)
 
 	queue := make(map[common.Hash]struct{})
 	for _, hash := range sched.Missing(0) {
@@ -268,7 +268,7 @@ func TestIncompleteStateSync(t *testing.T) {
 
 	// Create a destination state and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewStateSync(srcRoot, dstDb)
+	sched := NewStateSync(0, common.Hash{}, srcRoot, dstDb)
 
 	added := []common.Hash{}
 	queue := append([]common.Hash{}, sched.Missing(1)...)

eth/downloader/queue.go

@@ -399,7 +399,7 @@ func (q *queue) Schedule(headers []*types.Header, from uint64) []*types.Header {
 			}
 
 			q.stateSchedLock.Lock()
-			q.stateScheduler = state.NewStateSync(header.Root, q.stateDatabase)
+			q.stateScheduler = state.NewStateSync(header.Number.Uint64(), header.Hash(), header.Root, q.stateDatabase)
 			q.stateSchedLock.Unlock()
 		}
 		inserts = append(inserts, header)
@@ -1152,6 +1152,6 @@ func (q *queue) Prepare(offset uint64, mode SyncMode, pivot uint64, head *types.
 
 	// If long running fast sync, also start up a head stateretrieval immediately
 	if mode == FastSync && pivot > 0 {
-		q.stateScheduler = state.NewStateSync(head.Root, q.stateDatabase)
+		q.stateScheduler = state.NewStateSync(head.Number.Uint64(), head.Hash(), head.Root, q.stateDatabase)
 	}
 }

trie/directcache.go

@@ -30,6 +30,27 @@ var directCacheWrites = metrics.NewCounter("directcache/writes")
 var directCacheHitTimer = metrics.NewTimer("directcache/timer/hits")
 var directCacheMissTimer = metrics.NewTimer("directcache/timer/misses")
 
+// DirectCacheReads retrieves a global counter measuring the number of direct
+// cache reads from the disk since process startup. This isn't useful for anything
+// apart from trie debugging purposes.
+func DirectCacheReads() int64 {
+	return directCacheHitTimer.Count() + directCacheMissTimer.Count()
+}
+
+// DirectCacheWrites retrieves a global counter measuring the number of direct
+// cache writes from the disk since process startup. This isn't useful for anything
+// apart from trie debugging purposes.
+func DirectCacheWrites() int64 {
+	return directCacheWrites.Count()
+}
+
+// DirectCacheMisses retrieves a global counter measuring the number of direct
+// cache writes from the disk since process startup. This isn't useful for anything
+// apart from trie debugging purposes.
+func DirectCacheMisses() int64 {
+	return directCacheMissTimer.Count()
+}
+
 type cachedValue struct {
 	Value     []byte
 	BlockNum  uint64
@@ -157,7 +178,6 @@ func (dc *DirectCache) TryDelete(key []byte) error {
 }
 
 func (dc *DirectCache) CommitTo(dbw DatabaseWriter) (root common.Hash, err error) {
-	directCacheWrites.Inc(int64(len(dc.dirty)))
 	for k, _ := range dc.dirty {
 		v, err := dc.data.TryGet([]byte(k))
 		if err, ok := err.(*MissingNodeError); err != nil && !ok {
@@ -172,9 +192,26 @@ func (dc *DirectCache) CommitTo(dbw DatabaseWriter) (root common.Hash, err error
 }
 
 func (dc *DirectCache) putCache(dbw DatabaseWriter, key, value []byte) error {
-	enc, _ := rlp.EncodeToBytes(cachedValue{value, dc.blockNum, dc.blockHash})
+	return WriteDirectCache(dc.keyPrefix, key, value, dc.blockNum, dc.blockHash, dbw)
-	if err := dbw.Put(append(dc.keyPrefix, key...), enc); err != nil {
-		return err
 }
-	return nil
+
+// WriteDirectCache places a value node directly into the database along with
+// block metadata to validate its relevancy.
+//
+// The method is meant to be used by code that circumvents the state database
+// and its integrated cache, namely during fast sync and database upgrades.
+func WriteDirectCache(prefix, key, value []byte, number uint64, hash common.Hash, dbw DatabaseWriter) error {
+	directCacheWrites.Inc(1)
+	enc, _ := rlp.EncodeToBytes(cachedValue{value, number, hash})
+	return dbw.Put(append(prefix, key...), enc)
+}
+
+// GetDirectCache retrieves a value node directly from the database along with
+// block metadata to validate its relevancy.
+//
+// The method is meant to be used by code that circumvents the state database
+// and its integrated cache, namely during fast sync and database upgrades.
+func GetDirectCache(prefix, key []byte, db Database) ([]byte, error) {
+	defer func(start time.Time) { directCacheHitTimer.UpdateSince(start) }(time.Now())
+	return db.Get(append(prefix, key...))
 }

trie/iterator.go

@@ -77,6 +77,7 @@ type nodeIteratorState struct {
 	hash    common.Hash // Hash of the node being iterated (nil if not standalone)
 	node    node        // Trie node being iterated
 	parent  common.Hash // Hash of the first full ancestor node (nil if current is the root)
+	nibbles []byte      // Key nibbles leading up to this node for key reconstruction
 	child   int         // Child to be processed next
 }
 
@@ -88,6 +89,7 @@ type NodeIterator struct {
 	Hash     common.Hash // Hash of the current node being iterated (nil if not standalone)
 	Node     node        // Current node being iterated (internal representation)
 	Parent   common.Hash // Hash of the first full ancestor node (nil if current is the root)
+	Nibbles  []byte      // Key nibbles leading up to this node for key reconstruction
 	Leaf     bool        // Flag whether the current node is a value (data) node
 	LeafBlob []byte      // Data blob contained within a leaf (otherwise nil)
 
@@ -155,10 +157,15 @@ func (it *NodeIterator) step() error {
 			}
 			for parent.child++; parent.child < len(node.Children); parent.child++ {
 				if current := node.Children[parent.child]; current != nil {
+					nibbles := parent.nibbles
+					if parent.child < 16 {
+						nibbles = append(nibbles, byte(parent.child))
+					}
 					it.stack = append(it.stack, &nodeIteratorState{
 						hash:    common.BytesToHash(node.flags.hash),
 						node:    current,
 						parent:  ancestor,
+						nibbles: nibbles,
 						child:   -1,
 					})
 					break
@@ -174,6 +181,7 @@ func (it *NodeIterator) step() error {
 				hash:    common.BytesToHash(node.flags.hash),
 				node:    node.Val,
 				parent:  ancestor,
+				nibbles: append(parent.nibbles, node.Key...),
 				child:   -1,
 			})
 		} else if hash, ok := parent.node.(hashNode); ok {
@@ -191,6 +199,7 @@ func (it *NodeIterator) step() error {
 				hash:    common.BytesToHash(hash),
 				node:    node,
 				parent:  ancestor,
+				nibbles: parent.nibbles,
 				child:   -1,
 			})
 		} else {
@@ -207,7 +216,7 @@ func (it *NodeIterator) step() error {
 // The method returns whether there are any more data left for inspection.
 func (it *NodeIterator) retrieve() bool {
 	// Clear out any previously set values
-	it.Hash, it.Node, it.Parent, it.Leaf, it.LeafBlob = common.Hash{}, nil, common.Hash{}, false, nil
+	it.Hash, it.Node, it.Parent, it.Nibbles, it.Leaf, it.LeafBlob = common.Hash{}, nil, common.Hash{}, nil, false, nil
 
 	// If the iteration's done, return no available data
 	if it.trie == nil {
@@ -216,7 +225,7 @@ func (it *NodeIterator) retrieve() bool {
 	// Otherwise retrieve the current node and resolve leaf accessors
 	state := it.stack[len(it.stack)-1]
 
-	it.Hash, it.Node, it.Parent = state.hash, state.node, state.parent
+	it.Hash, it.Node, it.Parent, it.Nibbles = state.hash, state.node, state.parent, state.nibbles
 	if value, ok := it.Node.(valueNode); ok {
 		it.Leaf, it.LeafBlob = true, []byte(value)
 	}

trie/sync.go

@@ -36,10 +36,39 @@ type request struct {
 	raw  bool        // Whether this is a raw entry (code) or a trie node
 
 	parents []*request // Parent state nodes referencing this entry (notify all upon completion)
+	nibbles [][]byte   // Trie path nibbles that accumulate up to the key at the leaf (paired with parents)
 	depth   int        // Depth level within the trie the node is located to prioritise DFS
 	deps    int        // Number of dependencies before allowed to commit this node
 
 	callback TrieSyncLeafCallback // Callback to invoke if a leaf node it reached on this branch
+	dcPrefix []byte               // Database prefix to use for leaf direct caching
+}
+
+// keys traverses the request ancestry tree, reconstructing the key paths leading
+// to the current request through the nibbles represented by each node.
+func (r *request) keys(suffix []byte) [][]byte {
+	keys := r.paths(suffix)
+	for i := 0; i < len(keys); i++ {
+		keys[i] = compactHexEncode(keys[i])
+	}
+	return keys
+}
+
+// paths traverses the request ancestry tree, reconstructing the key paths leading
+// to the current request through the nibbles represented by each node.
+func (r *request) paths(suffix []byte) [][]byte {
+	// A root level request just returns the trailing suffix
+	if len(r.parents) == 0 {
+		return [][]byte{suffix}
+	}
+	// Otherwise collect all the ancestor paths, and append the current one
+	var keys [][]byte
+	for i, parent := range r.parents {
+		for _, key := range parent.paths(r.nibbles[i]) {
+			keys = append(keys, append(key, suffix...))
+		}
+	}
+	return keys
 }
 
 // SyncResult is a simple list to return missing nodes along with their request
@@ -52,7 +81,7 @@ type SyncResult struct {
 // TrieSyncLeafCallback is a callback type invoked when a trie sync reaches a
 // leaf node. It's used by state syncing to check if the leaf node requires some
 // further data syncing.
-type TrieSyncLeafCallback func(leaf []byte, parent common.Hash) error
+type TrieSyncLeafCallback func(keys [][]byte, leaf []byte, parent common.Hash) error
 
 // TrieSync is the main state trie synchronisation scheduler, which provides yet
 // unknown trie hashes to retrieve, accepts node data associated with said hashes
@@ -64,18 +93,18 @@ type TrieSync struct {
 }
 
 // NewTrieSync creates a new trie data download scheduler.
-func NewTrieSync(root common.Hash, database ethdb.Database, callback TrieSyncLeafCallback) *TrieSync {
+func NewTrieSync(root common.Hash, database ethdb.Database, callback TrieSyncLeafCallback, dcPrefix []byte) *TrieSync {
 	ts := &TrieSync{
 		database: database,
 		requests: make(map[common.Hash]*request),
 		queue:    prque.New(),
 	}
-	ts.AddSubTrie(root, 0, common.Hash{}, callback)
+	ts.AddSubTrie(root, 0, common.Hash{}, callback, dcPrefix)
 	return ts
 }
 
 // AddSubTrie registers a new trie to the sync code, rooted at the designated parent.
-func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback) {
+func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback, dcPrefix []byte) {
 	// Short circuit if the trie is empty or already known
 	if root == emptyRoot {
 		return
@@ -90,6 +119,7 @@ func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, c
 		hash:     root,
 		depth:    depth,
 		callback: callback,
+		dcPrefix: dcPrefix,
 	}
 	// If this sub-trie has a designated parent, link them together
 	if parent != (common.Hash{}) {
@@ -198,6 +228,7 @@ func (s *TrieSync) schedule(req *request) {
 	// If we're already requesting this node, add a new reference and stop
 	if old, ok := s.requests[req.hash]; ok {
 		old.parents = append(old.parents, req.parents...)
+		old.nibbles = append(old.nibbles, req.nibbles...)
 		return
 	}
 	// Schedule the request for future retrieval
@@ -210,6 +241,7 @@ func (s *TrieSync) schedule(req *request) {
 func (s *TrieSync) children(req *request, object node) ([]*request, error) {
 	// Gather all the children of the node, irrelevant whether known or not
 	type child struct {
+		path  []byte
 		node  node
 		depth int
 	}
@@ -218,13 +250,19 @@ func (s *TrieSync) children(req *request, object node) ([]*request, error) {
 	switch node := (object).(type) {
 	case *shortNode:
 		children = []child{{
+			path:  node.Key,
 			node:  node.Val,
 			depth: req.depth + len(node.Key),
 		}}
 	case *fullNode:
 		for i := 0; i < 17; i++ {
 			if node.Children[i] != nil {
+				var path []byte
+				if i < 16 {
+					path = []byte{byte(i)}
+				}
 				children = append(children, child{
+					path:  path,
 					node:  node.Children[i],
 					depth: req.depth + 1,
 				})
@@ -239,7 +277,7 @@ func (s *TrieSync) children(req *request, object node) ([]*request, error) {
 		// Notify any external watcher of a new key/value node
 		if req.callback != nil {
 			if node, ok := (child.node).(valueNode); ok {
-				if err := req.callback(node, req.hash); err != nil {
+				if err := req.callback(req.keys(child.path), node, req.hash); err != nil {
 					return nil, err
 				}
 			}
@@ -249,14 +287,40 @@ func (s *TrieSync) children(req *request, object node) ([]*request, error) {
 			// Try to resolve the node from the local database
 			blob, _ := s.database.Get(node)
 			if local, err := decodeNode(node[:], blob, 0); local != nil && err == nil {
+				// Local node found, iterate and invoke the leaf callback for all subleaves
+				if req.callback != nil && req.dcPrefix != nil {
+					trie, _ := New(common.BytesToHash(node[:]), s.database, 0)
+					unknown := false
+
+					it := NewNodeIterator(trie)
+					for it.Next() {
+						if it.Leaf {
+							// If this branch was already visited, abort iteration
+							keys := req.keys(append(child.path, it.Nibbles...))
+							if !unknown {
+								if val, err := GetDirectCache(req.dcPrefix, keys[0], s.database); val != nil && err == nil {
+									break
+								}
+								// Branch is unknown, mark as so to avoid repeated db lookups
+								unknown = true
+							}
+							// Otherwise write out all accounts ending in this leaf
+							if err := req.callback(keys, nil, req.hash); err != nil {
+								return nil, err
+							}
+						}
+					}
+				}
 				continue
 			}
 			// Locally unknown node, schedule for retrieval
 			requests = append(requests, &request{
 				hash:     common.BytesToHash(node),
 				parents:  []*request{req},
+				nibbles:  [][]byte{child.path},
 				depth:    child.depth,
 				callback: req.callback,
+				dcPrefix: req.dcPrefix,
 			})
 		}
 	}

trie/sync_test.go

@@ -55,6 +55,28 @@ func makeTestTrie() (ethdb.Database, *Trie, map[string][]byte) {
 	return db, trie, content
 }
 
+// makeRepetitiveTestTrie creates a test trie to test node-wise reconstruction,
+// where all values are the same, forcing large parts of the trie to share nodes.
+func makeRepetitiveTestTrie() (ethdb.Database, *Trie, map[string][]byte) {
+	// Create an empty trie
+	db, _ := ethdb.NewMemDatabase()
+	trie, _ := New(common.Hash{}, db, 0)
+
+	// Fill it with some arbitrary data
+	content := make(map[string][]byte)
+	for i := byte(0); i < 255; i++ {
+		for j := byte(0); j < 255; j++ {
+			key, val := common.LeftPadBytes([]byte{j, i}, 32), common.LeftPadBytes(nil, 32)
+			content[string(key)] = val
+			trie.Update(key, val)
+		}
+	}
+	trie.Commit()
+
+	// Return the generated trie
+	return db, trie, content
+}
+
 // checkTrieContents cross references a reconstructed trie with an expected data
 // content map.
 func checkTrieContents(t *testing.T, db Database, root []byte, content map[string][]byte) {
@@ -93,7 +115,7 @@ func TestEmptyTrieSync(t *testing.T) {
 
 	for i, trie := range []*Trie{emptyA, emptyB} {
 		db, _ := ethdb.NewMemDatabase()
-		if req := NewTrieSync(common.BytesToHash(trie.Root()), db, nil).Missing(1); len(req) != 0 {
+		if req := NewTrieSync(common.BytesToHash(trie.Root()), db, nil, nil).Missing(1); len(req) != 0 {
 			t.Errorf("test %d: content requested for empty trie: %v", i, req)
 		}
 	}
@@ -110,7 +132,7 @@ func testIterativeTrieSync(t *testing.T, batch int) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	queue := append([]common.Hash{}, sched.Missing(batch)...)
 	for len(queue) > 0 {
@@ -139,7 +161,7 @@ func TestIterativeDelayedTrieSync(t *testing.T) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	queue := append([]common.Hash{}, sched.Missing(10000)...)
 	for len(queue) > 0 {
@@ -173,7 +195,7 @@ func testIterativeRandomTrieSync(t *testing.T, batch int) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	queue := make(map[common.Hash]struct{})
 	for _, hash := range sched.Missing(batch) {
@@ -210,7 +232,7 @@ func TestIterativeRandomDelayedTrieSync(t *testing.T) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	queue := make(map[common.Hash]struct{})
 	for _, hash := range sched.Missing(10000) {
@@ -253,7 +275,7 @@ func TestDuplicateAvoidanceTrieSync(t *testing.T) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	queue := append([]common.Hash{}, sched.Missing(0)...)
 	requested := make(map[common.Hash]struct{})
@@ -289,7 +311,7 @@ func TestIncompleteTrieSync(t *testing.T) {
 
 	// Create a destination trie and sync with the scheduler
 	dstDb, _ := ethdb.NewMemDatabase()
-	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil)
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, nil, nil)
 
 	added := []common.Hash{}
 	queue := append([]common.Hash{}, sched.Missing(1)...)
@@ -331,3 +353,49 @@ func TestIncompleteTrieSync(t *testing.T) {
 		dstDb.Put(key, value)
 	}
 }
+
+// Tests that if a leaf callback is specified, all leaf nodes are invoked with
+// it, even if sync short circuits with existing local nodes.
+func TestAllLeafCallbackTrieSync(t *testing.T) {
+	// Create a random trie to copy
+	srcDb, srcTrie, srcData := makeRepetitiveTestTrie()
+
+	// Create a callback to accumulate all keys for later verification
+	leaves, count := make(map[string]struct{}), 0
+	callback := func(keys [][]byte, leaf []byte, parent common.Hash) error {
+		for _, key := range keys {
+			leaves[string(key)] = struct{}{}
+			count++
+		}
+		return nil
+	}
+	// Create a destination trie and sync with the scheduler
+	dstDb, _ := ethdb.NewMemDatabase()
+	sched := NewTrieSync(common.BytesToHash(srcTrie.Root()), dstDb, callback, []byte("leaf:"))
+
+	queue := append([]common.Hash{}, sched.Missing(1)...)
+	for len(queue) > 0 {
+		results := make([]SyncResult, len(queue))
+		for i, hash := range queue {
+			data, err := srcDb.Get(hash.Bytes())
+			if err != nil {
+				t.Fatalf("failed to retrieve node data for %x: %v", hash, err)
+			}
+			results[i] = SyncResult{hash, data}
+		}
+		if _, index, err := sched.Process(results); err != nil {
+			t.Fatalf("failed to process result #%d: %v", index, err)
+		}
+		queue = append(queue[:0], sched.Missing(1)...)
+	}
+	// Cross check that the two tries are in sync
+	checkTrieContents(t, dstDb, srcTrie.Root(), srcData)
+
+	// Ensure the leaf callback was invoked for all leaves, including cached ones
+	if len(leaves) != len(srcData) {
+		t.Errorf("Leaf callback count mismatch: have %v, want %v", len(leaves), len(srcData))
+	}
+	if len(leaves) != count {
+		t.Errorf("Duplicate leaf callbacks: have %v, want %v", count, len(leaves))
+	}
+}