execution: abort cache transitions after interrupted read-ahead drain

execution/cache/cache_test.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/erigontech/erigon/common"
 	"github.com/erigontech/erigon/common/crypto"
+	"github.com/erigontech/erigon/common/dbg"
 	"github.com/erigontech/erigon/db/kv"
 )
 
@@ -890,6 +891,28 @@ func TestDomainCache_PutIfAbsentAtomicWithPut(t *testing.T) {
 	}
 }
 
+// The drain-before-unwind convention (drainReadAhead ordered before every
+// epoch bump) is enforced here: a cache-populating warmup still in flight at
+// Unwind time can stamp a dead-fork value with the post-unwind epoch.
+func TestStateCache_UnwindAssertsWarmupInFlight(t *testing.T) {
+	old := dbg.AssertStateCache
+	dbg.AssertStateCache = true
+	t.Cleanup(func() { dbg.AssertStateCache = old })
+
+	b := 1 * datasize.MB
+	sc := NewStateCache(b, b, b, b)
+	sc.WarmupStarted()
+	require.Panics(t, func() { sc.Unwind(10) }, "epoch bump with a warmup in flight must fail loud")
+	sc.WarmupDone()
+	require.NotPanics(t, func() { sc.Unwind(10) })
+
+	// Without the assert flag the gauge is inert.
+	dbg.AssertStateCache = false
+	sc.WarmupStarted()
+	defer sc.WarmupDone()
+	require.NotPanics(t, func() { sc.Unwind(10) })
+}
+
 // A Delete racing an update-in-place put must not double-subtract the
 // displaced entry's size: freelru's OnEvict subtracts it for the Remove, and
 // put's update delta subtracts it again unless the two writers share the

execution/cache/state_cache.go

@@ -18,7 +18,9 @@ package cache
 
 import (
 	"bytes"
+	"fmt"
 	"strings"
+	"sync/atomic"
 
 	"github.com/c2h5oh/datasize"
 
@@ -52,6 +54,12 @@ const (
 // Code uses CodeCache (two-level for deduplication).
 type StateCache struct {
 	caches [kv.DomainLen]Cache
+
+	// warmupsInFlight counts fire-and-forget cache-populating prefetches
+	// (WarmupStarted/WarmupDone). Unwind asserts it is zero: a prefetch put
+	// racing the epoch bump could stamp a dead-fork value with the post-unwind
+	// epoch and have it served as canonical.
+	warmupsInFlight atomic.Int64
 }
 
 // NewStateCache creates a new StateCache with the specified byte capacities.
@@ -280,14 +288,29 @@ func (c *StateCache) Close() {
 // GenericCaches and the CodeCache, all layers) bumps an epoch + lowers a floor
 // and drops stale entries lazily on read. This is the sole cache-invalidation
 // path on unwind — the executor never touches the cache during forward execution.
+//
+// Callers must drain any in-flight cache-populating warmup first (see
+// WarmupStarted); the assert converts that convention into a loud failure.
 func (c *StateCache) Unwind(unwindToTxNum uint64) {
+	if dbg.AssertStateCache {
+		if n := c.warmupsInFlight.Load(); n != 0 {
+			panic(fmt.Sprintf("StateCache.Unwind with %d cache-populating warmup(s) in flight — missing drain before the epoch bump", n))
+		}
+	}
 	for _, cache := range c.caches {
 		if cache != nil {
 			cache.Unwind(unwindToTxNum)
 		}
 	}
 }
 
+// WarmupStarted and WarmupDone bracket a fire-and-forget cache-populating
+// prefetch; see warmupsInFlight.
+func (c *StateCache) WarmupStarted() { c.warmupsInFlight.Add(1) }
+
+// WarmupDone is the counterpart of WarmupStarted.
+func (c *StateCache) WarmupDone() { c.warmupsInFlight.Add(-1) }
+
 // GetCache returns the cache for the given domain.
 // Returns nil if the domain is not supported.
 func (c *StateCache) GetCache(domain kv.Domain) Cache {

execution/exec/blocks_read_ahead.go

@@ -146,24 +146,34 @@
 		if !bra.warming.CompareAndSwap(false, true) {
 			return
 		}
+		sc := bra.stateCache
 		bra.warmWg.Go(func() {
-			bra.warmBody(ctx, db, header, body, 8) // use 8 workers for warming
+			if sc != nil {
+				sc.WarmupStarted()
+				defer sc.WarmupDone()
+			}
+			bra.warmBody(ctx, db, sc, header, body, 8) // use 8 workers for warming
 		})
 	}
 }
 
-// WaitForWarmup blocks until any in-flight warmBody goroutine finishes or
-// the context is cancelled. Call before closing the database to avoid
-// waitTxsAllDoneOnClose hangs.
-func (bra *BlockReadAheader) WaitForWarmup(ctx context.Context) {
+// WaitForWarmup blocks until any in-flight warmBody goroutine finishes or the
+// context is cancelled, reporting whether the warmup fully drained. False
+// means a warmup may still be running — callers about to bump the cache epoch
+// must treat it as a failed precondition. Call before closing the database to
+// avoid waitTxsAllDoneOnClose hangs (that caller may ignore the result: it only
+// needs a bounded wait).
+func (bra *BlockReadAheader) WaitForWarmup(ctx context.Context) bool {
 	done := make(chan struct{})
 	go func() {
 		bra.warmWg.Wait()
 		close(done)
 	}()
 	select {
 	case <-done:
+		return true
 	case <-ctx.Done():
+		return false
 	}
 }
 
@@ -178,7 +188,9 @@
 // It reads: To accounts, To account code, To account storage from access lists,
 // and block-level access lists. Each worker creates its own transaction.
 // Only one warmBody can run at a time - concurrent calls are no-ops.
-func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, header *types.Header, body *types.Body, workers int) {
+// sc is the launch-time cache snapshot (see AddHeaderAndBody), nil to warm the
+// OS page cache only.
+func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, sc *cache.StateCache, header *types.Header, body *types.Body, workers int) {
 	defer bra.warming.Store(false)
 
 	if !dbg.ReadAhead {
@@ -240,8 +252,8 @@
 					return nil
 				}
 				var getter kv.TemporalGetter = ttx
-				if bra.stateCache != nil {
-					getter = newCachePopulatingGetter(ttx, bra.stateCache)
+				if sc != nil {
+					getter = newCachePopulatingGetter(ttx, sc)
 				}
 				stateReader := state.NewReaderV3(getter)
 
@@ -311,8 +323,8 @@
 			}
 			var getter kv.TemporalGetter = ttx
 			var cpg *cachePopulatingGetter
-			if bra.stateCache != nil {
-				cpg = newCachePopulatingGetter(ttx, bra.stateCache)
+			if sc != nil {
+				cpg = newCachePopulatingGetter(ttx, sc)
 				getter = cpg
 			}
 			stateReader := state.NewReaderV3(getter)

execution/exec/blocks_read_ahead_test.go

@@ -17,6 +17,7 @@
 package exec
 
 import (
+	"context"
 	"testing"
 
 	"github.com/c2h5oh/datasize"
@@ -27,6 +28,22 @@ import (
 	"github.com/erigontech/erigon/execution/cache"
 )
 
+// A cancelled wait can return while a warmup is still in flight — the gauge
+// convention only holds for a drained return, so callers about to bump the
+// cache epoch must be able to tell the two apart.
+func TestWaitForWarmupReportsDrained(t *testing.T) {
+	bra := &BlockReadAheader{}
+	require.True(t, bra.WaitForWarmup(context.Background()), "nothing in flight — drained")
+
+	bra.warmWg.Add(1)
+	cancelled, cancel := context.WithCancel(context.Background())
+	cancel()
+	require.False(t, bra.WaitForWarmup(cancelled), "cancelled wait with a live warmup must report undrained")
+
+	bra.warmWg.Done()
+	require.True(t, bra.WaitForWarmup(context.Background()), "drained after the warmup finished")
+}
+
 // stubTemporalGetter stands in for the committed-state snapshot a warmup
 // goroutine reads: every GetLatest returns the same fixed value.
 type stubTemporalGetter struct {

execution/execmodule/exec_module.go

@@ -384,22 +384,24 @@ func (e *ExecModule) canonicalHash(ctx context.Context, tx kv.Tx, blockNumber ui
 	return canonical, nil
 }
 
-// drainReadAhead blocks until any in-flight block-assembly warmup finishes.
-// warmBody is fire-and-forget and populates the shared state/branch caches; if
-// it is still running when an unwind bumps the cache epoch, it can Put a
-// pre-unwind (dead-fork) value stamped with the post-unwind epoch — IsStale then
-// returns false and the stale value is served as canonical (wrong root). A
-// laggard Put can likewise land after a flush's cache-apply and pin the
-// pre-flush snapshot. Call before any unwind epoch-bump or flush cache-apply.
-func (e *ExecModule) drainReadAhead() {
+// drainReadAhead blocks until any in-flight block-assembly warmup finishes,
+// reporting whether it fully drained — false only when the module context is
+// cancelled (shutdown). warmBody is fire-and-forget and populates the shared
+// state/branch caches; if it is still running when an unwind bumps the cache
+// epoch, it can Put a pre-unwind (dead-fork) value stamped with the post-unwind
+// epoch — IsStale then returns false and the stale value is served as canonical
+// (wrong root). A laggard Put can likewise land after a flush's cache-apply and
+// pin the pre-flush snapshot. Call before any unwind epoch-bump or flush
+// cache-apply, and do not proceed to them on false.
+func (e *ExecModule) drainReadAhead() bool {
 	if e.readAheader == nil {
-		return
+		return true
 	}
 	ctx := e.bacgroundCtx
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	e.readAheader.WaitForWarmup(ctx)
+	return e.readAheader.WaitForWarmup(ctx)
 }
 
 func (e *ExecModule) unwindToCommonCanonical(sd *execctx.SharedDomains, tx kv.TemporalRwTx, header *types.Header) error {
@@ -430,7 +432,9 @@ func (e *ExecModule) unwindToCommonCanonical(sd *execctx.SharedDomains, tx kv.Te
 		return err
 	}
 
-	e.drainReadAhead()
+	if !e.drainReadAhead() {
+		return fmt.Errorf("read-ahead drain interrupted before unwind: %w", e.bacgroundCtx.Err())
+	}
 	if err := e.pipelineExecutor.UnwindTo(unwindPoint, stagedsync.ExecUnwind, tx); err != nil {
 		return err
 	}

execution/execmodule/forkchoice.go

@@ -363,8 +363,10 @@ func (e *ExecModule) updateForkChoice(ctx context.Context, originalBlockHash, sa
 	// Drain any warmup a preceding newPayload spawned: its Puts reflect a
 	// pre-FCU snapshot and must land before this FCU's unwind epoch-bump and
 	// flush cache-apply, not after them (no new warmup starts while we hold
-	// the semaphore).
-	e.drainReadAhead()
+	// the semaphore). An interrupted drain means shutdown — bail.
+	if !e.drainReadAhead() {
+		return sendForkchoiceErrorWithoutWaiting(e.logger, outcomeCh, e.bacgroundCtx.Err(), false)
+	}
 
 	var validationError string
 

execution/execmodule/set_head.go

@@ -113,8 +113,10 @@ func (e *ExecModule) SetHead(ctx context.Context, targetBlock uint64) error {
 
 	// Drain in-flight warmup before the unwind bumps the cache epoch, so a
 	// fire-and-forget warmup can't Put a dead-fork value stamped with the new
-	// epoch (cross-fork contamination).
-	e.drainReadAhead()
+	// epoch (cross-fork contamination). An interrupted drain means shutdown.
+	if !e.drainReadAhead() {
+		return fmt.Errorf("read-ahead drain interrupted before unwind: %w", e.bacgroundCtx.Err())
+	}
 
 	// Set the unwind point and run the unwind
 	if err := e.pipelineExecutor.UnwindTo(targetBlock, stagedsync.StagedUnwind, tx); err != nil {