core/state/snapshot: decouple state generation from disk layer

core/state/snapshot/context.go

@@ -36,7 +36,9 @@ const (
 )
 
 // generatorStats is a collection of statistics gathered by the snapshot generator
-// for logging purposes.
+// for logging purposes. This data structure is used throughout the entire
+// lifecycle of the snapshot generation process and is shared across multiple
+// generation cycles.
 type generatorStats struct {
 	origin   uint64             // Origin prefix where generation started
 	start    time.Time          // Timestamp when generation started
@@ -46,9 +48,9 @@ type generatorStats struct {
 	storage  common.StorageSize // Total account and storage slot size(generation or recovery)
 }
 
-// Log creates a contextual log with the given message and the context pulled
+// log creates a contextual log with the given message and the context pulled
 // from the internally maintained statistics.
-func (gs *generatorStats) Log(msg string, root common.Hash, marker []byte) {
+func (gs *generatorStats) log(msg string, root common.Hash, marker []byte) {
 	var ctx []interface{}
 	if root != (common.Hash{}) {
 		ctx = append(ctx, []interface{}{"root", root}...)
@@ -85,29 +87,42 @@ func (gs *generatorStats) Log(msg string, root common.Hash, marker []byte) {
 	log.Info(msg, ctx...)
 }
 
-// generatorContext carries a few global values to be shared by all generation functions.
+// generatorContext holds several global fields that are used throughout the
+// current generation cycle.
 type generatorContext struct {
-	stats   *generatorStats     // Generation statistic collection
+	root      common.Hash         // State root of the generation target
-	db      ethdb.KeyValueStore // Key-value store containing the snapshot data
+	marker    []byte              // Generation progress marker
+	setMarker func(marker []byte) // Function to notify the generation progress
 	account   *holdableIterator   // Iterator of account snapshot data
 	storage   *holdableIterator   // Iterator of storage snapshot data
-	batch   ethdb.Batch         // Database batch for writing batch data atomically
+	db        ethdb.KeyValueStore // Key-value store containing the snapshot data
+	batch     ethdb.Batch         // Database batch for writing data atomically
 	logged    time.Time           // The timestamp when last generation progress was displayed
 }
 
 // newGeneratorContext initializes the context for generation.
-func newGeneratorContext(stats *generatorStats, db ethdb.KeyValueStore, accMarker []byte, storageMarker []byte) *generatorContext {
+func newGeneratorContext(root common.Hash, marker []byte, setMarker func(marker []byte), db ethdb.KeyValueStore) *generatorContext {
 	ctx := &generatorContext{
-		stats:  stats,
+		root:      root,
+		marker:    marker,
+		setMarker: setMarker,
 		db:        db,
 		batch:     db.NewBatch(),
 		logged:    time.Now(),
 	}
+	accMarker, storageMarker := splitMarker(marker)
 	ctx.openIterator(snapAccount, accMarker)
 	ctx.openIterator(snapStorage, storageMarker)
 	return ctx
 }
 
+// setGenMarker updates the generation progress marker locally and cascades it
+// to associated disk layer.
+func (ctx *generatorContext) setGenMarker(marker []byte) {
+	ctx.marker = marker
+	ctx.setMarker(marker)
+}
+
 // openIterator constructs global account and storage snapshot iterators
 // at the interrupted position. These iterators should be reopened from time
 // to time to avoid blocking leveldb compaction for a long time.
@@ -164,7 +179,7 @@ func (ctx *generatorContext) iterator(kind string) *holdableIterator {
 // the specified account. When the iterator touches the storage entry which
 // is located in or outside the given account, it stops and holds the current
 // iterated element locally.
-func (ctx *generatorContext) removeStorageBefore(account common.Hash) {
+func (ctx *generatorContext) removeStorageBefore(account common.Hash) uint64 {
 	var (
 		count uint64
 		start = time.Now()
@@ -183,8 +198,8 @@ func (ctx *generatorContext) removeStorageBefore(account common.Hash) {
 			ctx.batch.Reset()
 		}
 	}
-	ctx.stats.dangling += count
 	snapStorageCleanCounter.Inc(time.Since(start).Nanoseconds())
+	return count
 }
 
 // removeStorageAt deletes all storage entries which are located in the specified
@@ -221,7 +236,7 @@ func (ctx *generatorContext) removeStorageAt(account common.Hash) error {
 
 // removeStorageLeft deletes all storage entries which are located after
 // the current iterator position.
-func (ctx *generatorContext) removeStorageLeft() {
+func (ctx *generatorContext) removeStorageLeft() uint64 {
 	var (
 		count uint64
 		start = time.Now()
@@ -235,7 +250,7 @@ func (ctx *generatorContext) removeStorageLeft() {
 			ctx.batch.Reset()
 		}
 	}
-	ctx.stats.dangling += count
 	snapDanglingStorageMeter.Mark(int64(count))
 	snapStorageCleanCounter.Inc(time.Since(start).Nanoseconds())
+	return count
 }

core/state/snapshot/disklayer.go

@@ -26,23 +26,22 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/rlp"
-	"github.com/ethereum/go-ethereum/triedb"
 )
 
 // diskLayer is a low level persistent snapshot built on top of a key-value store.
 type diskLayer struct {
 	diskdb ethdb.KeyValueStore // Key-value store containing the base snapshot
-	triedb *triedb.Database    // Trie node cache for reconstruction purposes
 	cache  *fastcache.Cache    // Cache to avoid hitting the disk for direct access
 
 	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
-	genPending chan struct{}             // Notification channel when generation is done (test synchronicity)
+	lock      sync.RWMutex // Lock to protect stale and genMarker
-	genAbort   chan chan *generatorStats // Notification channel to abort generating the snapshot in this layer
 
-	lock sync.RWMutex
+	// State snapshot generator, set only if background generation is granted.
+	// Normally, a non-nil generator indicates that background snapshot generation
+	// is actively running, except for very short periods during restarts.
+	generator *generator
 }
 
 // Release releases underlying resources; specifically the fastcache requires
@@ -74,6 +73,17 @@ func (dl *diskLayer) Stale() bool {
 	return dl.stale
 }
 
+// markStale sets the stale flag as true.
+func (dl *diskLayer) markStale() {
+	dl.lock.Lock()
+	defer dl.lock.Unlock()
+
+	if dl.stale {
+		panic("disk layer is stale")
+	}
+	dl.stale = true
+}
+
 // Account directly retrieves the account associated with a particular hash in
 // the snapshot slim data format.
 func (dl *diskLayer) Account(hash common.Hash) (*types.SlimAccount, error) {
@@ -175,3 +185,11 @@ func (dl *diskLayer) Storage(accountHash, storageHash common.Hash) ([]byte, erro
 func (dl *diskLayer) Update(blockHash common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte) *diffLayer {
 	return newDiffLayer(dl, blockHash, destructs, accounts, storage)
 }
+
+// setGenMarker updates the generation progress marker with provided value.
+func (dl *diskLayer) setGenMarker(marker []byte) {
+	dl.lock.Lock()
+	defer dl.lock.Unlock()
+
+	dl.genMarker = marker
+}

core/state/snapshot/disklayer_test.go

@@ -24,7 +24,6 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/ethdb/memorydb"
-	"github.com/ethereum/go-ethereum/rlp"
 )
 
 // reverse reverses the contents of a byte slice. It's used to update random accs
@@ -426,81 +425,6 @@ func TestDiskPartialMerge(t *testing.T) {
 	}
 }
 
-// Tests that when the bottom-most diff layer is merged into the disk
-// layer whether the corresponding generator is persisted correctly.
-func TestDiskGeneratorPersistence(t *testing.T) {
-	var (
-		accOne        = randomHash()
-		accTwo        = randomHash()
-		accOneSlotOne = randomHash()
-		accOneSlotTwo = randomHash()
-
-		accThree     = randomHash()
-		accThreeSlot = randomHash()
-		baseRoot     = randomHash()
-		diffRoot     = randomHash()
-		diffTwoRoot  = randomHash()
-		genMarker    = append(randomHash().Bytes(), randomHash().Bytes()...)
-	)
-	// Testing scenario 1, the disk layer is still under the construction.
-	db := rawdb.NewMemoryDatabase()
-
-	rawdb.WriteAccountSnapshot(db, accOne, accOne[:])
-	rawdb.WriteStorageSnapshot(db, accOne, accOneSlotOne, accOneSlotOne[:])
-	rawdb.WriteStorageSnapshot(db, accOne, accOneSlotTwo, accOneSlotTwo[:])
-	rawdb.WriteSnapshotRoot(db, baseRoot)
-
-	// Create a disk layer based on all above updates
-	snaps := &Tree{
-		layers: map[common.Hash]snapshot{
-			baseRoot: &diskLayer{
-				diskdb:    db,
-				cache:     fastcache.New(500 * 1024),
-				root:      baseRoot,
-				genMarker: genMarker,
-			},
-		},
-	}
-	// Modify or delete some accounts, flatten everything onto disk
-	if err := snaps.Update(diffRoot, baseRoot, nil, map[common.Hash][]byte{
-		accTwo: accTwo[:],
-	}, nil); err != nil {
-		t.Fatalf("failed to update snapshot tree: %v", err)
-	}
-	if err := snaps.Cap(diffRoot, 0); err != nil {
-		t.Fatalf("failed to flatten snapshot tree: %v", err)
-	}
-	blob := rawdb.ReadSnapshotGenerator(db)
-	var generator journalGenerator
-	if err := rlp.DecodeBytes(blob, &generator); err != nil {
-		t.Fatalf("Failed to decode snapshot generator %v", err)
-	}
-	if !bytes.Equal(generator.Marker, genMarker) {
-		t.Fatalf("Generator marker is not matched")
-	}
-	// Test scenario 2, the disk layer is fully generated
-	// Modify or delete some accounts, flatten everything onto disk
-	if err := snaps.Update(diffTwoRoot, diffRoot, nil, map[common.Hash][]byte{
-		accThree: accThree.Bytes(),
-	}, map[common.Hash]map[common.Hash][]byte{
-		accThree: {accThreeSlot: accThreeSlot.Bytes()},
-	}); err != nil {
-		t.Fatalf("failed to update snapshot tree: %v", err)
-	}
-	diskLayer := snaps.layers[snaps.diskRoot()].(*diskLayer)
-	diskLayer.genMarker = nil // Construction finished
-	if err := snaps.Cap(diffTwoRoot, 0); err != nil {
-		t.Fatalf("failed to flatten snapshot tree: %v", err)
-	}
-	blob = rawdb.ReadSnapshotGenerator(db)
-	if err := rlp.DecodeBytes(blob, &generator); err != nil {
-		t.Fatalf("Failed to decode snapshot generator %v", err)
-	}
-	if len(generator.Marker) != 0 {
-		t.Fatalf("Failed to update snapshot generator")
-	}
-}
-
 // Tests that merging something into a disk layer persists it into the database
 // and invalidates any previously written and cached values, discarding anything
 // after the in-progress generation marker.

core/state/snapshot/generate.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"github.com/VictoriaMetrics/fastcache"
@@ -53,7 +54,64 @@ var (
 	errMissingTrie = errors.New("missing trie")
 )
 
-// generateSnapshot regenerates a brand new snapshot based on an existing state
+// generator is the struct for initial state snapshot generation.
+type generator struct {
+	active atomic.Bool         // Flag if the background generation is running
+	diskdb ethdb.KeyValueStore // Key-value store containing the base snapshot
+	triedb *triedb.Database    // Trie node cache for reconstruction purposes
+	stats  *generatorStats     // Generation statistics used throughout the entire life cycle
+	abort  chan chan struct{}  // Notification channel to abort generating the snapshot in this layer
+	done   chan struct{}       // Notification channel when generation is done (test synchronicity)
+}
+
+// newGenerator constructs the state snapshot generator.
+func newGenerator(diskdb ethdb.KeyValueStore, triedb *triedb.Database, stats *generatorStats) *generator {
+	if stats == nil {
+		stats = &generatorStats{start: time.Now()}
+	}
+	return &generator{
+		diskdb: diskdb,
+		triedb: triedb,
+		stats:  stats,
+		abort:  make(chan chan struct{}),
+		done:   make(chan struct{}),
+	}
+}
+
+// run starts the state snapshot generation in the background. It will panic
+// if the previous cycle is not terminated yet.
+func (g *generator) run(root common.Hash, marker []byte, setMarker func([]byte)) {
+	if g.active.Load() {
+		panic("the previous generation cycle is not aborted")
+	}
+	g.active.Store(true)
+
+	go g.generate(newGeneratorContext(root, marker, setMarker, g.diskdb))
+}
+
+// stop terminates the background generation if it's actively running.
+func (g *generator) stop() {
+	if !g.active.Load() {
+		log.Error("State snapshot is not generating")
+		return
+	}
+	ch := make(chan struct{})
+	g.abort <- ch
+	<-ch
+	g.active.Store(false)
+}
+
+// splitMarker is an internal helper which splits the generation progress marker
+// into two parts.
+func splitMarker(marker []byte) ([]byte, []byte) {
+	var accMarker []byte
+	if len(marker) > 0 { // []byte{} is the start, use nil for that
+		accMarker = marker[:common.HashLength]
+	}
+	return accMarker, marker
+}
+
+// generateSnapshot regenerates a brand-new snapshot based on an existing state
 // database and head block asynchronously. The snapshot is returned immediately
 // and generation is continued in the background until done.
 func generateSnapshot(diskdb ethdb.KeyValueStore, triedb *triedb.Database, cache int, root common.Hash) *diskLayer {
@@ -70,14 +128,12 @@ func generateSnapshot(diskdb ethdb.KeyValueStore, triedb *triedb.Database, cache
 	}
 	base := &diskLayer{
 		diskdb:    diskdb,
-		triedb:     triedb,
-		root:       root,
 		cache:     fastcache.New(cache * 1024 * 1024),
+		root:      root,
 		genMarker: genMarker,
-		genPending: make(chan struct{}),
+		generator: newGenerator(diskdb, triedb, stats),
-		genAbort:   make(chan chan *generatorStats),
 	}
-	go base.generate(stats)
+	base.generator.run(root, genMarker, base.setGenMarker)
 	log.Debug("Start snapshot generation", "root", root)
 	return base
 }
@@ -160,7 +216,7 @@ func (result *proofResult) forEach(callback func(key []byte, val []byte) error)
 //
 // The proof result will be returned if the range proving is finished, otherwise
 // the error will be returned to abort the entire procedure.
-func (dl *diskLayer) proveRange(ctx *generatorContext, trieId *trie.ID, prefix []byte, kind string, origin []byte, max int, valueConvertFn func([]byte) ([]byte, error)) (*proofResult, error) {
+func (g *generator) proveRange(ctx *generatorContext, trieId *trie.ID, prefix []byte, kind string, origin []byte, max int, valueConvertFn func([]byte) ([]byte, error)) (*proofResult, error) {
 	var (
 		keys     [][]byte
 		vals     [][]byte
@@ -245,9 +301,9 @@ func (dl *diskLayer) proveRange(ctx *generatorContext, trieId *trie.ID, prefix [
 		return &proofResult{keys: keys, vals: vals}, nil
 	}
 	// Snap state is chunked, generate edge proofs for verification.
-	tr, err := trie.New(trieId, dl.triedb)
+	tr, err := trie.New(trieId, g.triedb)
 	if err != nil {
-		ctx.stats.Log("Trie missing, state snapshotting paused", dl.root, dl.genMarker)
+		log.Info("Trie missing, state snapshotting paused", "state", ctx.root, "kind", kind, "root", trieId.Root)
 		return nil, errMissingTrie
 	}
 	// Generate the Merkle proofs for the first and last element
@@ -305,9 +361,9 @@ type onStateCallback func(key []byte, val []byte, write bool, delete bool) error
 // generateRange generates the state segment with particular prefix. Generation can
 // either verify the correctness of existing state through range-proof and skip
 // generation, or iterate trie to regenerate state on demand.
-func (dl *diskLayer) generateRange(ctx *generatorContext, trieId *trie.ID, prefix []byte, kind string, origin []byte, max int, onState onStateCallback, valueConvertFn func([]byte) ([]byte, error)) (bool, []byte, error) {
+func (g *generator) generateRange(ctx *generatorContext, trieId *trie.ID, prefix []byte, kind string, origin []byte, max int, onState onStateCallback, valueConvertFn func([]byte) ([]byte, error)) (bool, []byte, error) {
 	// Use range prover to check the validity of the flat state in the range
-	result, err := dl.proveRange(ctx, trieId, prefix, kind, origin, max, valueConvertFn)
+	result, err := g.proveRange(ctx, trieId, prefix, kind, origin, max, valueConvertFn)
 	if err != nil {
 		return false, nil, err
 	}
@@ -373,9 +429,9 @@ func (dl *diskLayer) generateRange(ctx *generatorContext, trieId *trie.ID, prefi
 	// if it's already opened with some nodes resolved.
 	tr := result.tr
 	if tr == nil {
-		tr, err = trie.New(trieId, dl.triedb)
+		tr, err = trie.New(trieId, g.triedb)
 		if err != nil {
-			ctx.stats.Log("Trie missing, state snapshotting paused", dl.root, dl.genMarker)
+			log.Info("Trie missing, state snapshotting paused", "state", ctx.root, "kind", kind, "root", trieId.Root)
 			return false, nil, errMissingTrie
 		}
 	}
@@ -473,32 +529,34 @@ func (dl *diskLayer) generateRange(ctx *generatorContext, trieId *trie.ID, prefi
 
 // checkAndFlush checks if an interruption signal is received or the
 // batch size has exceeded the allowance.
-func (dl *diskLayer) checkAndFlush(ctx *generatorContext, current []byte) error {
+func (g *generator) checkAndFlush(ctx *generatorContext, current []byte) error {
-	var abort chan *generatorStats
+	var abort chan struct{}
 	select {
-	case abort = <-dl.genAbort:
+	case abort = <-g.abort:
 	default:
 	}
 	if ctx.batch.ValueSize() > ethdb.IdealBatchSize || abort != nil {
-		if bytes.Compare(current, dl.genMarker) < 0 {
+		if bytes.Compare(current, ctx.marker) < 0 {
-			log.Error("Snapshot generator went backwards", "current", fmt.Sprintf("%x", current), "genMarker", fmt.Sprintf("%x", dl.genMarker))
+			log.Error("Snapshot generator went backwards", "current", fmt.Sprintf("%x", current), "genMarker", fmt.Sprintf("%x", ctx.marker))
 		}
-		// Flush out the batch anyway no matter it's empty or not.
+		// Persist the progress marker regardless of whether the batch is empty or not.
-		// It's possible that all the states are recovered and the
+		// It may happen that all the flat states in the database are correct, so the
-		// generation indeed makes progress.
+		// generator indeed makes progress even if there is nothing to commit.
-		journalProgress(ctx.batch, current, ctx.stats)
+		journalProgress(ctx.batch, current, g.stats)
 
+		// Flush out the database writes atomically
 		if err := ctx.batch.Write(); err != nil {
 			return err
 		}
 		ctx.batch.Reset()
 
-		dl.lock.Lock()
+		// Update the generation progress marker. This will also cascade the progress
-		dl.genMarker = current
+		// to the associated disk layer, allowing access to the newly generated parts.
-		dl.lock.Unlock()
+		ctx.setGenMarker(current)
 
+		// Abort the generation if it's required
 		if abort != nil {
-			ctx.stats.Log("Aborting state snapshot generation", dl.root, current)
+			g.stats.log("Aborting state snapshot generation", ctx.root, ctx.marker)
 			return newAbortErr(abort) // bubble up an error for interruption
 		}
 		// Don't hold the iterators too long, release them to let compactor works
@@ -506,7 +564,7 @@ func (dl *diskLayer) checkAndFlush(ctx *generatorContext, current []byte) error
 		ctx.reopenIterator(snapStorage)
 	}
 	if time.Since(ctx.logged) > 8*time.Second {
-		ctx.stats.Log("Generating state snapshot", dl.root, current)
+		g.stats.log("Generating state snapshot", ctx.root, ctx.marker)
 		ctx.logged = time.Now()
 	}
 	return nil
@@ -514,7 +572,7 @@ func (dl *diskLayer) checkAndFlush(ctx *generatorContext, current []byte) error
 
 // generateStorages generates the missing storage slots of the specific contract.
 // It's supposed to restart the generation from the given origin position.
-func generateStorages(ctx *generatorContext, dl *diskLayer, stateRoot common.Hash, account common.Hash, storageRoot common.Hash, storeMarker []byte) error {
+func (g *generator) generateStorages(ctx *generatorContext, account common.Hash, storageRoot common.Hash, storeMarker []byte) error {
 	onStorage := func(key []byte, val []byte, write bool, delete bool) error {
 		defer func(start time.Time) {
 			snapStorageWriteCounter.Inc(time.Since(start).Nanoseconds())
@@ -531,11 +589,11 @@ func generateStorages(ctx *generatorContext, dl *diskLayer, stateRoot common.Has
 		} else {
 			snapRecoveredStorageMeter.Mark(1)
 		}
-		ctx.stats.storage += common.StorageSize(1 + 2*common.HashLength + len(val))
+		g.stats.storage += common.StorageSize(1 + 2*common.HashLength + len(val))
-		ctx.stats.slots++
+		g.stats.slots++
 
 		// If we've exceeded our batch allowance or termination was requested, flush to disk
-		if err := dl.checkAndFlush(ctx, append(account[:], key...)); err != nil {
+		if err := g.checkAndFlush(ctx, append(account[:], key...)); err != nil {
 			return err
 		}
 		return nil
@@ -543,8 +601,8 @@ func generateStorages(ctx *generatorContext, dl *diskLayer, stateRoot common.Has
 	// Loop for re-generating the missing storage slots.
 	var origin = common.CopyBytes(storeMarker)
 	for {
-		id := trie.StorageTrieID(stateRoot, account, storageRoot)
+		id := trie.StorageTrieID(ctx.root, account, storageRoot)
-		exhausted, last, err := dl.generateRange(ctx, id, append(rawdb.SnapshotStoragePrefix, account.Bytes()...), snapStorage, origin, storageCheckRange, onStorage, nil)
+		exhausted, last, err := g.generateRange(ctx, id, append(rawdb.SnapshotStoragePrefix, account.Bytes()...), snapStorage, origin, storageCheckRange, onStorage, nil)
 		if err != nil {
 			return err // The procedure it aborted, either by external signal or internal error.
 		}
@@ -562,11 +620,11 @@ func generateStorages(ctx *generatorContext, dl *diskLayer, stateRoot common.Has
 // generateAccounts generates the missing snapshot accounts as well as their
 // storage slots in the main trie. It's supposed to restart the generation
 // from the given origin position.
-func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) error {
+func (g *generator) generateAccounts(ctx *generatorContext, accMarker []byte) error {
 	onAccount := func(key []byte, val []byte, write bool, delete bool) error {
 		// Make sure to clear all dangling storages before this account
 		account := common.BytesToHash(key)
-		ctx.removeStorageBefore(account)
+		g.stats.dangling += ctx.removeStorageBefore(account)
 
 		start := time.Now()
 		if delete {
@@ -599,17 +657,17 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 				rawdb.WriteAccountSnapshot(ctx.batch, account, data)
 				snapGeneratedAccountMeter.Mark(1)
 			}
-			ctx.stats.storage += common.StorageSize(1 + common.HashLength + dataLen)
+			g.stats.storage += common.StorageSize(1 + common.HashLength + dataLen)
-			ctx.stats.accounts++
+			g.stats.accounts++
 		}
 		// If the snap generation goes here after interrupted, genMarker may go backward
 		// when last genMarker is consisted of accountHash and storageHash
 		marker := account[:]
-		if accMarker != nil && bytes.Equal(marker, accMarker) && len(dl.genMarker) > common.HashLength {
+		if accMarker != nil && bytes.Equal(marker, accMarker) && len(ctx.marker) > common.HashLength {
-			marker = dl.genMarker[:]
+			marker = ctx.marker
 		}
 		// If we've exceeded our batch allowance or termination was requested, flush to disk
-		if err := dl.checkAndFlush(ctx, marker); err != nil {
+		if err := g.checkAndFlush(ctx, marker); err != nil {
 			return err
 		}
 		snapAccountWriteCounter.Inc(time.Since(start).Nanoseconds()) // let's count flush time as well
@@ -620,10 +678,10 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 			ctx.removeStorageAt(account)
 		} else {
 			var storeMarker []byte
-			if accMarker != nil && bytes.Equal(account[:], accMarker) && len(dl.genMarker) > common.HashLength {
+			if accMarker != nil && bytes.Equal(account[:], accMarker) && len(ctx.marker) > common.HashLength {
-				storeMarker = dl.genMarker[common.HashLength:]
+				storeMarker = ctx.marker[common.HashLength:]
 			}
-			if err := generateStorages(ctx, dl, dl.root, account, acc.Root, storeMarker); err != nil {
+			if err := g.generateStorages(ctx, account, acc.Root, storeMarker); err != nil {
 				return err
 			}
 		}
@@ -633,8 +691,8 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 	}
 	origin := common.CopyBytes(accMarker)
 	for {
-		id := trie.StateTrieID(dl.root)
+		id := trie.StateTrieID(ctx.root)
-		exhausted, last, err := dl.generateRange(ctx, id, rawdb.SnapshotAccountPrefix, snapAccount, origin, accountCheckRange, onAccount, types.FullAccountRLP)
+		exhausted, last, err := g.generateRange(ctx, id, rawdb.SnapshotAccountPrefix, snapAccount, origin, accountCheckRange, onAccount, types.FullAccountRLP)
 		if err != nil {
 			return err // The procedure it aborted, either by external signal or internal error.
 		}
@@ -643,7 +701,7 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 		// Last step, cleanup the storages after the last account.
 		// All the left storages should be treated as dangling.
 		if origin == nil || exhausted {
-			ctx.removeStorageLeft()
+			g.stats.dangling += ctx.removeStorageLeft()
 			break
 		}
 	}
@@ -654,15 +712,9 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 // constructing the state snapshot. All the arguments are purely for statistics
 // gathering and logging, since the method surfs the blocks as they arrive, often
 // being restarted.
-func (dl *diskLayer) generate(stats *generatorStats) {
+func (g *generator) generate(ctx *generatorContext) {
-	var (
+	g.stats.log("Resuming state snapshot generation", ctx.root, ctx.marker)
-		accMarker []byte
+	defer ctx.close()
-		abort     chan *generatorStats
-	)
-	if len(dl.genMarker) > 0 { // []byte{} is the start, use nil for that
-		accMarker = dl.genMarker[:common.HashLength]
-	}
-	stats.Log("Resuming state snapshot generation", dl.root, dl.genMarker)
 
 	// Initialize the global generator context. The snapshot iterators are
 	// opened at the interrupted position because the assumption is held
@@ -672,45 +724,44 @@ func (dl *diskLayer) generate(stats *generatorStats) {
 	// For the account or storage slot at the interruption, they will be
 	// processed twice by the generator(they are already processed in the
 	// last run) but it's fine.
-	ctx := newGeneratorContext(stats, dl.diskdb, accMarker, dl.genMarker)
+	var (
-	defer ctx.close()
+		accMarker, _ = splitMarker(ctx.marker)
-
+		abort        chan struct{}
-	if err := generateAccounts(ctx, dl, accMarker); err != nil {
+	)
+	if err := g.generateAccounts(ctx, accMarker); err != nil {
 		// Extract the received interruption signal if exists
 		if aerr, ok := err.(*abortErr); ok {
 			abort = aerr.abort
 		}
 		// Aborted by internal error, wait the signal
 		if abort == nil {
-			abort = <-dl.genAbort
+			abort = <-g.abort
 		}
-		abort <- stats
+		close(abort)
 		return
 	}
 	// Snapshot fully generated, set the marker to nil.
 	// Note even there is nothing to commit, persist the
 	// generator anyway to mark the snapshot is complete.
-	journalProgress(ctx.batch, nil, stats)
+	journalProgress(ctx.batch, nil, g.stats)
 	if err := ctx.batch.Write(); err != nil {
 		log.Error("Failed to flush batch", "err", err)
-
+		abort = <-g.abort
-		abort = <-dl.genAbort
+		close(abort)
-		abort <- stats
 		return
 	}
 	ctx.batch.Reset()
 
-	log.Info("Generated state snapshot", "accounts", stats.accounts, "slots", stats.slots,
+	log.Info("Generated state snapshot", "accounts", g.stats.accounts, "slots", g.stats.slots,
-		"storage", stats.storage, "dangling", stats.dangling, "elapsed", common.PrettyDuration(time.Since(stats.start)))
+		"storage", g.stats.storage, "dangling", g.stats.dangling, "elapsed", common.PrettyDuration(time.Since(g.stats.start)))
 
-	dl.lock.Lock()
+	// Update the generation progress marker
-	dl.genMarker = nil
+	ctx.setGenMarker(nil)
-	close(dl.genPending)
+	close(g.done)
-	dl.lock.Unlock()
 
 	// Someone will be looking for us, wait it out
-	abort = <-dl.genAbort
+	abort = <-g.abort
-	abort <- nil
+	close(abort)
 }
 
 // increaseKey increase the input key by one bit. Return nil if the entire
@@ -728,10 +779,10 @@ func increaseKey(key []byte) []byte {
 // abortErr wraps an interruption signal received to represent the
 // generation is aborted by external processes.
 type abortErr struct {
-	abort chan *generatorStats
+	abort chan struct{}
 }
 
-func newAbortErr(abort chan *generatorStats) error {
+func newAbortErr(abort chan struct{}) error {
 	return &abortErr{abort: abort}
 }
 

core/state/snapshot/generate_test.go

@@ -71,7 +71,7 @@ func testGeneration(t *testing.T, scheme string) {
 		t.Fatalf("have %#x want %#x", have, want)
 	}
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -80,9 +80,7 @@ func testGeneration(t *testing.T, scheme string) {
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation with existent flat state.
@@ -112,7 +110,7 @@ func testGenerateExistentState(t *testing.T, scheme string) {
 
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -121,9 +119,7 @@ func testGenerateExistentState(t *testing.T, scheme string) {
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 func checkSnapRoot(t *testing.T, snap *diskLayer, trieRoot common.Hash) {
@@ -330,17 +326,16 @@ func testGenerateExistentStateWithWrongStorage(t *testing.T, scheme string) {
 	t.Logf("Root: %#x\n", root) // Root = 0x8746cce9fd9c658b2cfd639878ed6584b7a2b3e73bb40f607fcfa156002429a0
 
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation with existent flat state, where the flat state
@@ -392,7 +387,7 @@ func testGenerateExistentStateWithWrongAccounts(t *testing.T, scheme string) {
 	t.Logf("Root: %#x\n", root) // Root = 0x825891472281463511e7ebcc7f109e4f9200c20fa384754e11fd605cd98464e8
 
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -401,9 +396,7 @@ func testGenerateExistentStateWithWrongAccounts(t *testing.T, scheme string) {
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation errors out correctly in case of a missing trie
@@ -433,7 +426,7 @@ func testGenerateCorruptAccountTrie(t *testing.T, scheme string) {
 
 	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 		t.Errorf("Snapshot generated against corrupt account trie")
 
@@ -441,9 +434,7 @@ func testGenerateCorruptAccountTrie(t *testing.T, scheme string) {
 		// Not generated fast enough, hopefully blocked inside on missing trie node fail
 	}
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation errors out correctly in case of a missing root
@@ -477,7 +468,7 @@ func testGenerateMissingStorageTrie(t *testing.T, scheme string) {
 
 	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 		t.Errorf("Snapshot generated against corrupt storage trie")
 
@@ -485,9 +476,7 @@ func testGenerateMissingStorageTrie(t *testing.T, scheme string) {
 		// Not generated fast enough, hopefully blocked inside on missing trie node fail
 	}
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation errors out correctly in case of a missing trie
@@ -519,7 +508,7 @@ func testGenerateCorruptStorageTrie(t *testing.T, scheme string) {
 
 	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 		t.Errorf("Snapshot generated against corrupt storage trie")
 
@@ -527,9 +516,7 @@ func testGenerateCorruptStorageTrie(t *testing.T, scheme string) {
 		// Not generated fast enough, hopefully blocked inside on missing trie node fail
 	}
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation when an extra account with storage exists in the snap state.
@@ -583,7 +570,7 @@ func testGenerateWithExtraAccounts(t *testing.T, scheme string) {
 	}
 	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -592,9 +579,8 @@ func testGenerateWithExtraAccounts(t *testing.T, scheme string) {
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
+
-	<-stop
 	// If we now inspect the snap db, there should exist no extraneous storage items
 	if data := rawdb.ReadStorageSnapshot(helper.diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
 		t.Fatalf("expected slot to be removed, got %v", string(data))
@@ -645,17 +631,16 @@ func testGenerateWithManyExtraAccounts(t *testing.T, scheme string) {
 	}
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests this case
@@ -694,17 +679,16 @@ func testGenerateWithExtraBeforeAndAfter(t *testing.T, scheme string) {
 	}
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // TestGenerateWithMalformedSnapdata tests what happes if we have some junk
@@ -734,17 +718,17 @@ func testGenerateWithMalformedSnapdata(t *testing.T, scheme string) {
 	}
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
+
-	<-stop
 	// If we now inspect the snap db, there should exist no extraneous storage items
 	if data := rawdb.ReadStorageSnapshot(helper.diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
 		t.Fatalf("expected slot to be removed, got %v", string(data))
@@ -771,17 +755,16 @@ func testGenerateFromEmptySnap(t *testing.T, scheme string) {
 	t.Logf("Root: %#x\n", root) // Root: 0x6f7af6d2e1a1bf2b84a3beb3f8b64388465fbc1e274ca5d5d3fc787ca78f59e4
 
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation with existent flat state, where the flat state
@@ -822,17 +805,16 @@ func testGenerateWithIncompleteStorage(t *testing.T, scheme string) {
 	t.Logf("Root: %#x\n", root) // Root: 0xca73f6f05ba4ca3024ef340ef3dfca8fdabc1b677ff13f5a9571fd49c16e67ff
 
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 func incKey(key []byte) []byte {
@@ -917,7 +899,7 @@ func testGenerateCompleteSnapshotWithDanglingStorage(t *testing.T, scheme string
 
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -926,9 +908,7 @@ func testGenerateCompleteSnapshotWithDanglingStorage(t *testing.T, scheme string
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }
 
 // Tests that snapshot generation with dangling storages. Dangling storage means
@@ -954,7 +934,7 @@ func testGenerateBrokenSnapshotWithDanglingStorage(t *testing.T, scheme string)
 
 	root, snap := helper.CommitAndGenerate()
 	select {
-	case <-snap.genPending:
+	case <-snap.generator.done:
 		// Snapshot generation succeeded
 
 	case <-time.After(3 * time.Second):
@@ -963,7 +943,5 @@ func testGenerateBrokenSnapshotWithDanglingStorage(t *testing.T, scheme string)
 	checkSnapRoot(t, snap, root)
 
 	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
+	snap.generator.stop()
-	snap.genAbort <- stop
-	<-stop
 }

core/state/snapshot/holdable_iterator_test.go

@@ -118,7 +118,7 @@ func TestReopenIterator(t *testing.T) {
 		}
 	}
 	// Iterate over the database with the given configs and verify the results
-	ctx, idx := newGeneratorContext(&generatorStats{}, db, nil, nil), -1
+	ctx, idx := newGeneratorContext(common.Hash{}, nil, nil, db), -1
 
 	idx++
 	ctx.account.Next()

core/state/snapshot/journal.go

@@ -102,6 +102,10 @@ func loadAndParseJournal(db ethdb.KeyValueStore, base *diskLayer) (snapshot, jou
 	if err := rlp.DecodeBytes(generatorBlob, &generator); err != nil {
 		return nil, journalGenerator{}, fmt.Errorf("failed to decode snapshot generator: %v", err)
 	}
+	// Nilness is lost after rlp decoding, explicitly set it back to empty.
+	if !generator.Done && generator.Marker == nil {
+		generator.Marker = []byte{}
+	}
 	// Retrieve the diff layer journal. It's possible that the journal is
 	// not existent, e.g. the disk layer is generating while that the Geth
 	// crashes without persisting the diff journal.
@@ -134,7 +138,6 @@ func loadSnapshot(diskdb ethdb.KeyValueStore, triedb *triedb.Database, root comm
 	}
 	base := &diskLayer{
 		diskdb: diskdb,
-		triedb: triedb,
 		cache:  fastcache.New(cache * 1024 * 1024),
 		root:   baseRoot,
 	}
@@ -167,27 +170,22 @@ func loadSnapshot(diskdb ethdb.KeyValueStore, triedb *triedb.Database, root comm
 	// Load the disk layer status from the generator if it's not complete
 	if !generator.Done {
 		base.genMarker = generator.Marker
-		if base.genMarker == nil {
-			base.genMarker = []byte{}
-		}
 	}
 	// Everything loaded correctly, resume any suspended operations
 	// if the background generation is allowed
 	if !generator.Done && !noBuild {
-		base.genPending = make(chan struct{})
-		base.genAbort = make(chan chan *generatorStats)
-
 		var origin uint64
 		if len(generator.Marker) >= 8 {
 			origin = binary.BigEndian.Uint64(generator.Marker)
 		}
-		go base.generate(&generatorStats{
+		base.generator = newGenerator(diskdb, triedb, &generatorStats{
 			origin:   origin,
 			start:    time.Now(),
 			accounts: generator.Accounts,
 			slots:    generator.Slots,
 			storage:  common.StorageSize(generator.Storage),
 		})
+		base.generator.run(baseRoot, generator.Marker, base.setGenMarker)
 	}
 	return snapshot, false, nil
 }
@@ -196,14 +194,12 @@ func loadSnapshot(diskdb ethdb.KeyValueStore, triedb *triedb.Database, root comm
 // the progress into the database.
 func (dl *diskLayer) Journal(buffer *bytes.Buffer) (common.Hash, error) {
 	// If the snapshot is currently being generated, abort it
-	var stats *generatorStats
+	if dl.generator != nil {
-	if dl.genAbort != nil {
+		dl.generator.stop()
-		abort := make(chan *generatorStats)
-		dl.genAbort <- abort
 
-		if stats = <-abort; stats != nil {
+		// safe to access dl.genMarker as the only mutator dl.generator
-			stats.Log("Journalling in-progress snapshot", dl.root, dl.genMarker)
+		// has been terminated.
-		}
+		dl.generator.stats.log("Journalling in-progress snapshot", dl.root, dl.genMarker)
 	}
 	// Ensure the layer didn't get stale
 	dl.lock.RLock()
@@ -212,9 +208,6 @@ func (dl *diskLayer) Journal(buffer *bytes.Buffer) (common.Hash, error) {
 	if dl.stale {
 		return common.Hash{}, ErrSnapshotStale
 	}
-	// Ensure the generator stats is written even if none was ran this cycle
-	journalProgress(dl.diskdb, dl.genMarker, stats)
-
 	log.Debug("Journalled disk layer", "root", dl.root)
 	return dl.root, nil
 }

core/state/snapshot/snapshot.go

@@ -235,7 +235,9 @@ func (t *Tree) waitBuild() {
 	t.lock.RLock()
 	for _, layer := range t.layers {
 		if layer, ok := layer.(*diskLayer); ok {
-			done = layer.genPending
+			if layer.generator != nil {
+				done = layer.generator.done
+			}
 			break
 		}
 	}
@@ -259,15 +261,11 @@ func (t *Tree) Disable() {
 		switch layer := layer.(type) {
 		case *diskLayer:
 			// If the base layer is generating, abort it
-			if layer.genAbort != nil {
+			if layer.generator != nil {
-				abort := make(chan *generatorStats)
+				layer.generator.stop()
-				layer.genAbort <- abort
-				<-abort
 			}
 			// Layer should be inactive now, mark it as stale
-			layer.lock.Lock()
+			layer.markStale()
-			layer.stale = true
-			layer.lock.Unlock()
 			layer.Release()
 
 		case *diffLayer:
@@ -385,7 +383,8 @@ func (t *Tree) Cap(root common.Hash, layers int) error {
 	if !ok {
 		return fmt.Errorf("snapshot [%#x] is disk layer", root)
 	}
-	// If the generator is still running, use a more aggressive cap
+	// If the generator is still running, use a more aggressive cap.
+	// Hold the read lock as the genMarker is not thread safe.
 	diff.origin.lock.RLock()
 	if diff.origin.genMarker != nil && layers > 8 {
 		layers = 8
@@ -498,7 +497,7 @@ func (t *Tree) cap(diff *diffLayer, layers int) *diskLayer {
 			// there's a snapshot being generated currently. In that case, the trie
 			// will move from underneath the generator so we **must** merge all the
 			// partial data down into the snapshot and restart the generation.
-			if flattened.parent.(*diskLayer).genAbort == nil {
+			if flattened.parent.(*diskLayer).generator == nil {
 				return nil
 			}
 		}
@@ -526,26 +525,21 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 	var (
 		base  = bottom.parent.(*diskLayer)
 		batch = base.diskdb.NewBatch()
-		stats *generatorStats
 	)
 	// If the disk layer is running a snapshot generator, abort it
-	if base.genAbort != nil {
+	if base.generator != nil {
-		abort := make(chan *generatorStats)
+		base.generator.stop()
-		base.genAbort <- abort
-		stats = <-abort
 	}
 	// Put the deletion in the batch writer, flush all updates in the final step.
 	rawdb.DeleteSnapshotRoot(batch)
 
 	// Mark the original base as stale as we're going to create a new wrapper
-	base.lock.Lock()
+	base.markStale()
-	if base.stale {
-		panic("parent disk layer is stale") // we've committed into the same base from two children, boo
-	}
-	base.stale = true
-	base.lock.Unlock()
 
 	// Destroy all the destructed accounts from the database
+	//
+	// Note it's safe to access base.genMarker here as the only mutator
+	// base.generator has been terminated.
 	for hash := range bottom.destructSet {
 		// Skip any account not covered yet by the snapshot
 		if base.genMarker != nil && bytes.Compare(hash[:], base.genMarker) > 0 {
@@ -627,9 +621,6 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 	// Update the snapshot block marker and write any remainder data
 	rawdb.WriteSnapshotRoot(batch, bottom.root)
 
-	// Write out the generator progress marker and report
-	journalProgress(batch, base.genMarker, stats)
-
 	// Flush all the updates in the single db operation. Ensure the
 	// disk layer transition is atomic.
 	if err := batch.Write(); err != nil {
@@ -640,19 +631,16 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 		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.
 	//
-	// Note, the `base.genAbort` comparison is not used normally, it's checked
+	// Note, the base.generator comparison is not used normally, it's checked
 	// to allow the tests to play with the marker without triggering this path.
-	if base.genMarker != nil && base.genAbort != nil {
+	if res.genMarker != nil && base.generator != nil {
-		res.genMarker = base.genMarker
+		res.generator = base.generator
-		res.genAbort = make(chan chan *generatorStats)
+		res.generator.run(res.root, res.genMarker, res.setGenMarker)
-		go res.generate(stats)
 	}
 	return res
 }
@@ -724,15 +712,11 @@ func (t *Tree) Rebuild(root common.Hash) {
 		switch layer := layer.(type) {
 		case *diskLayer:
 			// If the base layer is generating, abort it and save
-			if layer.genAbort != nil {
+			if layer.generator != nil {
-				abort := make(chan *generatorStats)
+				layer.generator.stop()
-				layer.genAbort <- abort
-				<-abort
 			}
 			// Layer should be inactive now, mark it as stale
-			layer.lock.Lock()
+			layer.markStale()
-			layer.stale = true
-			layer.lock.Unlock()
 			layer.Release()
 
 		case *diffLayer: