core/state/snapshot, tests: sync snap gen + snaps in consensus tests

core/blockchain.go

@@ -121,6 +121,8 @@ type CacheConfig struct {
 	TrieDirtyDisabled   bool          // Whether to disable trie write caching and GC altogether (archive node)
 	TrieTimeLimit       time.Duration // Time limit after which to flush the current in-memory trie to disk
 	SnapshotLimit       int           // Memory allowance (MB) to use for caching snapshot entries in memory
+
+	SnapshotWait bool // Wait for snapshot construction on startup. TODO(karalabe): This is a dirty hack for testing, nuke it
 }
 
 // BlockChain represents the canonical chain given a database with a genesis
@@ -303,7 +305,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par
 	}
 	// Load any existing snapshot, regenerating it if loading failed
 	if bc.cacheConfig.SnapshotLimit > 0 {
-		bc.snaps = snapshot.New(bc.db, bc.stateCache.TrieDB(), bc.cacheConfig.SnapshotLimit, bc.CurrentBlock().Root())
+		bc.snaps = snapshot.New(bc.db, bc.stateCache.TrieDB(), bc.cacheConfig.SnapshotLimit, bc.CurrentBlock().Root(), !bc.cacheConfig.SnapshotWait)
 	}
 	// Take ownership of this particular state
 	go bc.update()

core/state/snapshot/disklayer.go

@@ -37,8 +37,9 @@ type diskLayer struct {
 	root  common.Hash // Root hash of the base snapshot
 	stale bool        // Signals that the layer became stale (state progressed)
 
-	genMarker []byte                    // Marker for the state that's indexed during initial layer generation
-	genAbort  chan chan *generatorStats // Notification channel to abort generating the snapshot in this layer
+	genMarker  []byte                    // Marker for the state that's indexed during initial layer generation
+	genPending chan struct{}             // Notification channel when generation is done (test synchronicity)
+	genAbort   chan chan *generatorStats // Notification channel to abort generating the snapshot in this layer
 
 	lock sync.RWMutex
 }

core/state/snapshot/generate.go

@@ -101,12 +101,13 @@ func generateSnapshot(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache i
 	rawdb.WriteSnapshotRoot(diskdb, root)
 
 	base := &diskLayer{
-		diskdb:    diskdb,
-		triedb:    triedb,
-		root:      root,
-		cache:     fastcache.New(cache * 1024 * 1024),
-		genMarker: []byte{}, // Initialized but empty!
-		genAbort:  make(chan chan *generatorStats),
+		diskdb:     diskdb,
+		triedb:     triedb,
+		root:       root,
+		cache:      fastcache.New(cache * 1024 * 1024),
+		genMarker:  []byte{}, // Initialized but empty!
+		genPending: make(chan struct{}),
+		genAbort:   make(chan chan *generatorStats),
 	}
 	go base.generate(&generatorStats{wiping: wiper, start: time.Now()})
 	return base
@@ -252,6 +253,7 @@ func (dl *diskLayer) generate(stats *generatorStats) {
 
 	dl.lock.Lock()
 	dl.genMarker = nil
+	close(dl.genPending)
 	dl.lock.Unlock()
 
 	// Someone will be looking for us, wait it out

core/state/snapshot/journal.go

@@ -108,6 +108,7 @@ func loadSnapshot(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int,
 		if base.genMarker == nil {
 			base.genMarker = []byte{}
 		}
+		base.genPending = make(chan struct{})
 		base.genAbort = make(chan chan *generatorStats)
 
 		var origin uint64

core/state/snapshot/snapshot.go

@@ -164,7 +164,7 @@ type Tree struct {
 // If the snapshot is missing or inconsistent, the entirety is deleted and will
 // be reconstructed from scratch based on the tries in the key-value store, on a
 // background thread.
-func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root common.Hash) *Tree {
+func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root common.Hash, async bool) *Tree {
 	// Create a new, empty snapshot tree
 	snap := &Tree{
 		diskdb: diskdb,
@@ -172,6 +172,9 @@ func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root comm
 		cache:  cache,
 		layers: make(map[common.Hash]snapshot),
 	}
+	if !async {
+		defer snap.waitBuild()
+	}
 	// Attempt to load a previously persisted snapshot and rebuild one if failed
 	head, err := loadSnapshot(diskdb, triedb, cache, root)
 	if err != nil {
@@ -187,6 +190,27 @@ func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root comm
 	return snap
 }
 
+// waitBuild blocks until the snapshot finishes rebuilding. This method is meant
+// to  be used by tests to ensure we're testing what we believe we are.
+func (t *Tree) waitBuild() {
+	// Find the rebuild termination channel
+	var done chan struct{}
+
+	t.lock.RLock()
+	for _, layer := range t.layers {
+		if layer, ok := layer.(*diskLayer); ok {
+			done = layer.genPending
+			break
+		}
+	}
+	t.lock.RUnlock()
+
+	// Wait until the snapshot is generated
+	if done != nil {
+		<-done
+	}
+}
+
 // Snapshot retrieves a snapshot belonging to the given block root, or nil if no
 // snapshot is maintained for that block.
 func (t *Tree) Snapshot(blockRoot common.Hash) Snapshot {
@@ -477,11 +501,12 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 		log.Crit("Failed to write leftover snapshot", "err", err)
 	}
 	res := &diskLayer{
-		root:      bottom.root,
-		cache:     base.cache,
-		diskdb:    base.diskdb,
-		triedb:    base.triedb,
-		genMarker: base.genMarker,
+		root:       bottom.root,
+		cache:      base.cache,
+		diskdb:     base.diskdb,
+		triedb:     base.triedb,
+		genMarker:  base.genMarker,
+		genPending: base.genPending,
 	}
 	// If snapshot generation hasn't finished yet, port over all the starts and
 	// continue where the previous round left off.