[WIP] execution: fix and optimise BAL warmup

common/dbg/experiments.go

@@ -136,6 +136,8 @@ var (
 	DisableAdaptivePin   = EnvBool("DISABLE_ADAPTIVE_PIN", false)
 	AssertStateCache     = EnvBool("ASSERT_STATE_CACHE", false)
 	ReadAhead            = EnvBool("READ_AHEAD", true)
+	ReadAheadWorkers     = EnvInt("READ_AHEAD_WORKERS", runtime.NumCPU())
+	ReadAheadWait        = EnvBool("READ_AHEAD_WAIT", false)
 	// FilesAsyncIO warms cold state .kv pages via io_uring before the mmap read, so
 	// a would-be blocking page fault becomes a non-blocking read that releases the
 	// goroutine's P. Linux + io_uring only; self-disables (reads use ordinary faults)

execution/exec/blocks_read_ahead.go

@@ -83,10 +83,13 @@ func (bra *BlockReadAheader) SetStateCache(sc *cache.StateCache) {
 // (codeHash→bytes) + size-cache layers via PutCodeWithHashIfAbsent, keyed by
 // the code's own keccak hash so every cached pair is self-consistent.
 type cachePopulatingGetter struct {
-	g        kv.TemporalGetter
-	sc       *cache.StateCache
-	progress func(kv.Domain) uint64 // domain progress source for stamping negative fills
-	stepSize uint64                 // for the read txNum upper bound (last txNum of the read's step)
+	g                 kv.TemporalGetter
+	sc                *cache.StateCache
+	progress          func(kv.Domain) uint64 // domain progress source for stamping negative fills
+	stepSize          uint64                 // for the read txNum upper bound (last txNum of the read's step)
+	lastCodeAddr      common.Address
+	lastCodeHash      common.Hash
+	lastCodeHashKnown bool
 }
 
 func newCachePopulatingGetter(tx kv.TemporalTx, sc *cache.StateCache) *cachePopulatingGetter {
@@ -96,6 +99,16 @@ func newCachePopulatingGetter(tx kv.TemporalTx, sc *cache.StateCache) *cachePopu
 
 func (cpg *cachePopulatingGetter) GetLatest(name kv.Domain, k []byte) ([]byte, kv.Step, error) {
 	v, step, err := cpg.g.GetLatest(name, k)
+	if name == kv.AccountsDomain {
+		cpg.lastCodeHashKnown = false
+		if err == nil && len(k) == len(cpg.lastCodeAddr) {
+			if codeHash := accounts.DeserialiseV3CodeHash(v); len(codeHash) == len(cpg.lastCodeHash) {
+				cpg.lastCodeAddr = common.BytesToAddress(k)
+				copy(cpg.lastCodeHash[:], codeHash)
+				cpg.lastCodeHashKnown = true
+			}
+		}
+	}
 	if err == nil && cpg.sc != nil {
 		// If-absent writes only: this runs in a fire-and-forget goroutine over a
 		// committed snapshot, so an unconditional Put racing an FCU flush's
@@ -129,6 +142,23 @@ func (cpg *cachePopulatingGetter) GetLatest(name kv.Domain, k []byte) ([]byte, k
 	return v, step, err
 }
 
+func (cpg *cachePopulatingGetter) GetCode(addr []byte, _ uint64) ([]byte, bool, error) {
+	// A warmup worker calls ReadAccountData immediately before ReadAccountCode.
+	// The account read provides the code hash, which lets the code read probe the
+	// code cache before falling back to the database. This avoids repeated database
+	// reads for accounts sharing identical code.
+	if cpg.sc != nil && cpg.lastCodeHashKnown && bytes.Equal(addr, cpg.lastCodeAddr[:]) {
+		if code, ok := cpg.sc.GetCodeByHash(cpg.lastCodeHash[:]); ok {
+			return code, true, nil
+		}
+	}
+	code, _, err := cpg.GetLatest(kv.CodeDomain, addr)
+	if err != nil {
+		return nil, false, err
+	}
+	return code, len(code) > 0, nil
+}
+
 func (cpg *cachePopulatingGetter) HasPrefix(name kv.Domain, prefix []byte) ([]byte, []byte, bool, error) {
 	return cpg.g.HasPrefix(name, prefix)
 }
@@ -137,7 +167,7 @@ func (cpg *cachePopulatingGetter) StepsInFiles(entitySet ...kv.Domain) kv.Step {
 	return cpg.g.StepsInFiles(entitySet...)
 }
 
-func (bra *BlockReadAheader) AddHeaderAndBody(ctx context.Context, db kv.RoDB, header *types.Header, body *types.Body) {
+func (bra *BlockReadAheader) AddHeaderAndBody(ctx context.Context, db kv.RoDB, tx kv.Getter, header *types.Header, body *types.Body) {
 	blockHash := header.Hash()
 	bra.headers.Add(blockHash, header)
 	bra.bodies.Add(blockHash, body)
@@ -146,9 +176,20 @@ func (bra *BlockReadAheader) AddHeaderAndBody(ctx context.Context, db kv.RoDB, h
 		if !bra.warming.CompareAndSwap(false, true) {
 			return
 		}
+		var bal types.BlockAccessList
+		balBytes, err := tx.GetOne(kv.BlockAccessList, dbutils.BlockBodyKey(header.Number.Uint64(), blockHash))
+		if err != nil {
+			log.Warn("[warmBody] failed to read BAL", "blockNum", header.Number.Uint64(), "blockHash", blockHash, "err", err)
+		} else if len(balBytes) > 0 {
+			bal, err = types.DecodeBlockAccessListBytes(balBytes)
+			if err != nil {
+				log.Warn("[warmBody] failed to decode BAL", "blockNum", header.Number.Uint64(), "blockHash", blockHash, "err", err)
+			}
+		}
 		bra.warmWg.Go(func() {
-			bra.warmBody(ctx, db, header, body, 8) // use 8 workers for warming
+			bra.warmBody(ctx, db, body, bal, dbg.ReadAheadWorkers)
 		})
+		bra.waitForWarmupIfConfigured(ctx)
 	}
 }
 
@@ -167,134 +208,140 @@ func (bra *BlockReadAheader) WaitForWarmup(ctx context.Context) {
 	}
 }
 
+func (bra *BlockReadAheader) waitForWarmupIfConfigured(ctx context.Context) {
+	if dbg.ReadAheadWait {
+		bra.WaitForWarmup(ctx)
+	}
+}
+
 func (bra *BlockReadAheader) AddSenders(senders []byte, blockHash common.Hash) {
 	if _, ok := bra.bodies.Get(blockHash); !ok {
 		return
 	}
 	bra.senders.Add(blockHash, bytes.Clone(senders))
 }
 
-// warmBody warms state for all transactions in a body using multiple workers.
-// 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) {
-	defer bra.warming.Store(false)
+type balWarmupTaskKind uint8
+
+const (
+	balWarmAccount balWarmupTaskKind = iota
+	balWarmStorageChanges
+	balWarmStorageReads
+)
+
+type balWarmupTask struct {
+	accountIndex int
+	kind         balWarmupTaskKind
+	slotIndex    int
+}
+
+func makeBALWarmupPlan(bal types.BlockAccessList, workers int) ([]balWarmupTask, int) {
+	taskCount := len(bal)
+	for _, account := range bal {
+		taskCount += len(account.StorageChanges) + len(account.StorageReads)
+	}
+	tasks := make([]balWarmupTask, 0, taskCount)
+	for accountIndex, account := range bal {
+		tasks = append(tasks, balWarmupTask{accountIndex: accountIndex, kind: balWarmAccount})
+		for slotIndex := range account.StorageChanges {
+			tasks = append(tasks, balWarmupTask{accountIndex: accountIndex, kind: balWarmStorageChanges, slotIndex: slotIndex})
+		}
+		for slotIndex := range account.StorageReads {
+			tasks = append(tasks, balWarmupTask{accountIndex: accountIndex, kind: balWarmStorageReads, slotIndex: slotIndex})
+		}
+	}
+	return tasks, min(workers, len(tasks))
+}
 
+func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, body *types.Body, bal types.BlockAccessList, workers int) {
+	defer bra.warming.Store(false)
 	if !dbg.ReadAhead {
 		return
 	}
-
 	if workers <= 0 {
 		workers = 1
 	}
+	if len(bal) > 0 {
+		bra.warmBAL(ctx, db, bal, workers)
+		return
+	}
+	bra.warmTxns(ctx, db, body.Transactions, workers)
+}
 
+func (bra *BlockReadAheader) warmBAL(ctx context.Context, db kv.RoDB, bal types.BlockAccessList, workers int) {
+	tasks, balWorkers := makeBALWarmupPlan(bal, workers)
+	var nextTask atomic.Uint64
 	var wg errgroup.Group
-
-	// If BAL exists in DB, use BAL warming (more complete)
-	var bal types.BlockAccessList
-	if header != nil && db != nil {
-		tx, err := db.BeginRo(ctx)
-		if err != nil {
-			log.Warn("[warmBody] failed to open tx for BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
-		} else {
-			data, err := tx.GetOne(kv.BlockAccessList, dbutils.BlockBodyKey(header.Number.Uint64(), header.Hash()))
+	for w := range balWorkers {
+		workerID := w
+		wg.Go(func() error {
+			startTime := time.Now()
+			tx, err := db.BeginRo(ctx)
 			if err != nil {
-				log.Warn("[warmBody] failed to read BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
-			} else if len(data) > 0 {
-				bal, err = types.DecodeBlockAccessListBytes(data)
-				if err != nil {
-					log.Warn("[warmBody] failed to decode BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
-				}
+				return err
 			}
-			tx.Rollback()
-		}
-	}
-
-	balLen := len(bal)
-	if balLen > 0 {
-		balWorkers := min(workers, balLen)
-
-		// Pre-divide work: each worker gets a dedicated range of BAL entries
-		entriesPerWorker := (balLen + balWorkers - 1) / balWorkers
-
-		for w := range balWorkers {
-			start := w * entriesPerWorker
-			end := min(start+entriesPerWorker, balLen)
-			if start >= balLen {
-				break
+			defer tx.Rollback()
+			ttx, ok := tx.(kv.TemporalTx)
+			if !ok {
+				return nil
 			}
-
-			// Capture loop variables for closure
-			workerStart, workerEnd, workerID := start, end, w
-			wg.Go(func() error {
-				startTime := time.Now()
-				tx, err := db.BeginRo(ctx)
-				if err != nil {
-					return err
-				}
-				defer tx.Rollback()
-
-				ttx, ok := tx.(kv.TemporalTx)
-				if !ok {
-					return nil
+			var getter kv.TemporalGetter = ttx
+			if bra.stateCache != nil {
+				getter = newCachePopulatingGetter(ttx, bra.stateCache)
+			}
+			stateReader := state.NewReaderV3(getter)
+			tasksProcessed := 0
+			for {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				default:
 				}
-				var getter kv.TemporalGetter = ttx
-				if bra.stateCache != nil {
-					getter = newCachePopulatingGetter(ttx, bra.stateCache)
+				taskIndex := int(nextTask.Add(1) - 1)
+				if taskIndex >= len(tasks) {
+					break
 				}
-				stateReader := state.NewReaderV3(getter)
-
-				for idx := workerStart; idx < workerEnd; idx++ {
-					select {
-					case <-ctx.Done():
-						return ctx.Err()
-					default:
-					}
-
-					acctChanges := bal[idx]
+				task := tasks[taskIndex]
+				acctChanges := bal[task.accountIndex]
+				switch task.kind {
+				case balWarmAccount:
 					acct, _ := stateReader.ReadAccountData(acctChanges.Address)
-					// Warm code if account has code or if there are code changes.
-					if acct != nil && !acct.CodeHash.IsEmpty() {
+					if (acct != nil && !acct.CodeHash.IsEmpty()) || len(acctChanges.CodeChanges) > 0 {
 						stateReader.ReadAccountCode(acctChanges.Address)
-					} else if len(acctChanges.CodeChanges) > 0 {
-						stateReader.ReadAccountCode(acctChanges.Address)
-					}
-					for _, slotChanges := range acctChanges.StorageChanges {
-						stateReader.ReadAccountStorage(acctChanges.Address, slotChanges.Slot)
-					}
-					for _, slot := range acctChanges.StorageReads {
-						stateReader.ReadAccountStorage(acctChanges.Address, slot)
 					}
+				case balWarmStorageChanges:
+					slot := acctChanges.StorageChanges[task.slotIndex].Slot
+					stateReader.ReadAccountStorage(acctChanges.Address, slot)
+				case balWarmStorageReads:
+					slot := acctChanges.StorageReads[task.slotIndex]
+					stateReader.ReadAccountStorage(acctChanges.Address, slot)
 				}
-				log.Debug("[warmBody] BAL worker finished", "worker", workerID, "entries", workerEnd-workerStart, "elapsed", time.Since(startTime))
-				return nil
-			})
-		}
-		wg.Wait()
-		return
+				tasksProcessed++
+			}
+			log.Debug("[warmBAL] worker finished", "worker", workerID, "tasks", tasksProcessed, "elapsed", time.Since(startTime))
+			return nil
+		})
 	}
-	// Fallback: per-transaction warming when no BAL
-	txns := body.Transactions
+	wg.Wait()
+}
+
+func (bra *BlockReadAheader) warmTxns(ctx context.Context, db kv.RoDB, txns types.Transactions, workers int) {
 	if len(txns) == 0 {
 		return
 	}
-
 	txnLen := len(txns)
 	if workers > txnLen {
 		workers = txnLen
 	}
-
 	// Pre-divide work: each worker gets a dedicated range of transactions
 	txnsPerWorker := (txnLen + workers - 1) / workers
-
+	var wg errgroup.Group
 	for w := 0; w < workers; w++ {
 		start := w * txnsPerWorker
 		end := min(start+txnsPerWorker, txnLen)
 		if start >= txnLen {
 			break
 		}
-
 		// Capture loop variables for closure
 		workerStart, workerEnd, workerID := start, end, w
 		wg.Go(func() error {
@@ -304,36 +351,29 @@ func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, header *t
 				return err
 			}
 			defer tx.Rollback()
-
 			ttx, ok := tx.(kv.TemporalTx)
 			if !ok {
 				return nil
 			}
 			var getter kv.TemporalGetter = ttx
-			var cpg *cachePopulatingGetter
 			if bra.stateCache != nil {
-				cpg = newCachePopulatingGetter(ttx, bra.stateCache)
-				getter = cpg
+				getter = newCachePopulatingGetter(ttx, bra.stateCache)
 			}
 			stateReader := state.NewReaderV3(getter)
-
 			for txIdx := workerStart; txIdx < workerEnd; txIdx++ {
 				select {
 				case <-ctx.Done():
 					return ctx.Err()
 				default:
 				}
-
 				txn := txns[txIdx]
-
 				// Warm To account and its code if it has one
 				if toAddr := txn.GetTo(); toAddr != nil {
 					to := accounts.InternAddress(*toAddr)
 					if acct, _ := stateReader.ReadAccountData(to); acct != nil && !acct.CodeHash.IsEmpty() {
 						stateReader.ReadAccountCode(to)
 					}
 				}
-
 				// Warm transaction access list accounts and their code
 				for _, entry := range txn.GetAccessList() {
 					addr := accounts.InternAddress(entry.Address)
@@ -345,11 +385,10 @@ func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, header *t
 					}
 				}
 			}
-			log.Debug("[warmBody] TX worker finished", "worker", workerID, "txns", workerEnd-workerStart, "elapsed", time.Since(startTime))
+			log.Debug("[warmTxns] worker finished", "worker", workerID, "txns", workerEnd-workerStart, "elapsed", time.Since(startTime))
 			return nil
 		})
 	}
-
 	wg.Wait()
 }