cmd/evm/internal/t8ntool: stream t8n alloc to ease heavy memory cases (#34785)

cmd/evm/internal/t8ntool/execution.go

@@ -17,9 +17,11 @@
 package t8ntool
 
 import (
+	"encoding/json"
 	"fmt"
 	stdmath "math"
 	"math/big"
+	"os"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -47,6 +49,9 @@ type Prestate struct {
 	Env        stEnv                         `json:"env"`
 	Pre        types.GenesisAlloc            `json:"pre"`
 	TreeLeaves map[common.Hash]hexutil.Bytes `json:"vkt,omitempty"`
+	// AllocPath, when non-empty, causes Apply to stream the alloc from disk
+	// instead of reading Pre, so the full map never materializes in memory.
+	AllocPath string `json:"-"`
 }
 
 //go:generate go run github.com/fjl/gencodec -type ExecutionResult -field-override executionResultMarshaling -out gen_execresult.go
@@ -146,8 +151,19 @@ func (pre *Prestate) Apply(vmConfig vm.Config, chainConfig *params.ChainConfig,
 		return h
 	}
 	var (
-		isEIP4762   = chainConfig.IsUBT(big.NewInt(int64(pre.Env.Number)), pre.Env.Timestamp)
-		statedb     = MakePreState(rawdb.NewMemoryDatabase(), pre.Pre, isEIP4762)
+		isEIP4762 = chainConfig.IsUBT(big.NewInt(int64(pre.Env.Number)), pre.Env.Timestamp)
+		statedb   *state.StateDB
+	)
+	if pre.AllocPath != "" {
+		var err error
+		statedb, err = MakePreStateStreaming(rawdb.NewMemoryDatabase(), pre.AllocPath, isEIP4762)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+	} else {
+		statedb = MakePreState(rawdb.NewMemoryDatabase(), pre.Pre, isEIP4762)
+	}
+	var (
 		signer      = types.MakeSigner(chainConfig, new(big.Int).SetUint64(pre.Env.Number), pre.Env.Timestamp)
 		gaspool     = core.NewGasPool(pre.Env.GasLimit)
 		blockHash   = common.Hash{0x13, 0x37}
@@ -414,6 +430,76 @@ func MakePreState(db ethdb.Database, accounts types.GenesisAlloc, isBintrie bool
 	return statedb
 }
 
+// MakePreStateStreaming is like MakePreState, but decodes the alloc from disk
+// one account at a time so the full map is never held in memory.
+func MakePreStateStreaming(db ethdb.Database, allocPath string, isBintrie bool) (*state.StateDB, error) {
+	tdb := triedb.NewDatabase(db, &triedb.Config{Preimages: true, IsUBT: isBintrie})
+	sdb := state.NewDatabase(tdb, nil)
+
+	root := types.EmptyRootHash
+	if isBintrie {
+		root = types.EmptyBinaryHash
+	}
+	statedb, err := state.New(root, sdb)
+	if err != nil {
+		return nil, NewError(ErrorEVM, fmt.Errorf("failed to create initial statedb: %v", err))
+	}
+
+	f, err := os.Open(allocPath)
+	if err != nil {
+		return nil, NewError(ErrorIO, fmt.Errorf("failed reading alloc file: %v", err))
+	}
+	defer f.Close()
+
+	dec := json.NewDecoder(f)
+	tok, err := dec.Token()
+	if err != nil {
+		return nil, NewError(ErrorJson, fmt.Errorf("failed reading alloc opening token: %v", err))
+	}
+	if d, ok := tok.(json.Delim); !ok || d != '{' {
+		return nil, NewError(ErrorJson, fmt.Errorf("expected alloc object, got %v", tok))
+	}
+	for dec.More() {
+		keyTok, err := dec.Token()
+		if err != nil {
+			return nil, NewError(ErrorJson, fmt.Errorf("failed reading alloc key: %v", err))
+		}
+		keyStr, ok := keyTok.(string)
+		if !ok {
+			return nil, NewError(ErrorJson, fmt.Errorf("alloc key not a string: %v", keyTok))
+		}
+		addr := common.HexToAddress(keyStr)
+		var acct types.Account
+		if err := dec.Decode(&acct); err != nil {
+			return nil, NewError(ErrorJson, fmt.Errorf("failed decoding account %s: %v", keyStr, err))
+		}
+		statedb.SetCode(addr, acct.Code, tracing.CodeChangeUnspecified)
+		statedb.SetNonce(addr, acct.Nonce, tracing.NonceChangeGenesis)
+		if acct.Balance != nil {
+			statedb.SetBalance(addr, uint256.MustFromBig(acct.Balance), tracing.BalanceIncreaseGenesisBalance)
+		}
+		for k, v := range acct.Storage {
+			statedb.SetState(addr, k, v)
+		}
+	}
+	if _, err := dec.Token(); err != nil {
+		return nil, NewError(ErrorJson, fmt.Errorf("failed reading alloc closing token: %v", err))
+	}
+
+	root, err = statedb.Commit(0, false, false)
+	if err != nil {
+		return nil, NewError(ErrorEVM, fmt.Errorf("failed to commit initial state: %v", err))
+	}
+	if isBintrie {
+		return statedb, nil
+	}
+	statedb, err = state.New(root, sdb)
+	if err != nil {
+		return nil, NewError(ErrorEVM, fmt.Errorf("failed to reopen state after commit: %v", err))
+	}
+	return statedb, nil
+}
+
 func rlpHash(x any) (h common.Hash) {
 	hw := keccak.NewLegacyKeccak256()
 	rlp.Encode(hw, x)

cmd/evm/internal/t8ntool/transition.go

@@ -17,6 +17,7 @@
 package t8ntool
 
 import (
+	"bufio"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -115,11 +116,10 @@ func Transition(ctx *cli.Context) error {
 		}
 	}
 	if allocStr != stdinSelector {
-		if err := readFile(allocStr, "alloc", &inputData.Alloc); err != nil {
-			return err
-		}
+		prestate.AllocPath = allocStr
+	} else {
+		prestate.Pre = inputData.Alloc
 	}
-	prestate.Pre = inputData.Alloc
 
 	if btStr != stdinSelector && btStr != "" {
 		if err := readFile(btStr, "BT", &inputData.BT); err != nil {
@@ -223,22 +223,57 @@ func Transition(ctx *cli.Context) error {
 			return err
 		}
 	}
-	// Dump the execution result
+	// Dump the execution result.
 	var (
-		collector = make(Alloc)
+		collector Alloc
 		btleaves  map[common.Hash]hexutil.Bytes
 	)
 	isBinary := chainConfig.IsUBT(big.NewInt(int64(prestate.Env.Number)), prestate.Env.Timestamp)
-	if !isBinary {
+	allocOutput := ctx.String(OutputAllocFlag.Name)
+	switch {
+	case !isBinary && allocOutput != "" && allocOutput != "stdout" && allocOutput != "stderr":
+		// Stream directly to the output file to avoid materializing the
+		// whole post-state in memory. dispatchOutput is told to skip alloc
+		// by clearing the output name.
+		if err := writeStreamedAlloc(filepath.Join(baseDir, allocOutput), s); err != nil {
+			return err
+		}
+		allocOutput = ""
+	case !isBinary:
+		collector = make(Alloc)
 		s.DumpToCollector(collector, nil)
-	} else {
+	default:
 		btleaves = make(map[common.Hash]hexutil.Bytes)
 		if err := s.DumpBinTrieLeaves(btleaves); err != nil {
 			return err
 		}
 	}
+	return dispatchOutput(ctx, baseDir, result, collector, allocOutput, body, btleaves)
+}
 
-	return dispatchOutput(ctx, baseDir, result, collector, body, btleaves)
+// writeStreamedAlloc writes the post-state alloc to path one account at a
+// time, producing the same JSON shape as saveFile on an Alloc map.
+func writeStreamedAlloc(path string, s *state.StateDB) error {
+	f, err := os.Create(path)
+	if err != nil {
+		return NewError(ErrorIO, fmt.Errorf("failed creating alloc output file: %v", err))
+	}
+	bw := bufio.NewWriter(f)
+	sa := newStreamingAlloc(bw)
+	s.DumpToCollector(sa, nil)
+	if err := sa.Close(); err != nil {
+		f.Close()
+		return NewError(ErrorIO, fmt.Errorf("failed writing alloc output: %v", err))
+	}
+	if err := bw.Flush(); err != nil {
+		f.Close()
+		return NewError(ErrorIO, fmt.Errorf("failed flushing alloc output: %v", err))
+	}
+	if err := f.Close(); err != nil {
+		return NewError(ErrorIO, fmt.Errorf("failed closing alloc output file: %v", err))
+	}
+	log.Info("Wrote file", "file", path)
+	return nil
 }
 
 func applyLondonChecks(env *stEnv, chainConfig *params.ChainConfig) error {
@@ -327,6 +362,10 @@ func (g Alloc) OnAccount(addr *common.Address, dumpAccount state.DumpAccount) {
 	if addr == nil {
 		return
 	}
+	g[*addr] = dumpAccountToTypesAccount(dumpAccount)
+}
+
+func dumpAccountToTypesAccount(dumpAccount state.DumpAccount) types.Account {
 	balance, _ := new(big.Int).SetString(dumpAccount.Balance, 0)
 	var storage map[common.Hash]common.Hash
 	if dumpAccount.Storage != nil {
@@ -335,13 +374,64 @@ func (g Alloc) OnAccount(addr *common.Address, dumpAccount state.DumpAccount) {
 			storage[k] = common.HexToHash(v)
 		}
 	}
-	genesisAccount := types.Account{
+	return types.Account{
 		Code:    dumpAccount.Code,
 		Storage: storage,
 		Balance: balance,
 		Nonce:   dumpAccount.Nonce,
 	}
-	g[*addr] = genesisAccount
+}
+
+// streamingAlloc is a DumpCollector that writes each account to w as it is
+// visited, emitting a single JSON object keyed by address. Close must be
+// called to emit the closing brace.
+type streamingAlloc struct {
+	w        io.Writer
+	wroteOne bool
+	err      error
+}
+
+func newStreamingAlloc(w io.Writer) *streamingAlloc {
+	return &streamingAlloc{w: w}
+}
+
+func (s *streamingAlloc) write(b []byte) {
+	if s.err != nil {
+		return
+	}
+	_, s.err = s.w.Write(b)
+}
+
+func (s *streamingAlloc) OnRoot(common.Hash) {
+	s.write([]byte{'{'})
+}
+
+func (s *streamingAlloc) OnAccount(addr *common.Address, dumpAccount state.DumpAccount) {
+	if s.err != nil || addr == nil {
+		return
+	}
+	keyJSON, err := json.Marshal(*addr)
+	if err != nil {
+		s.err = err
+		return
+	}
+	valueJSON, err := json.Marshal(dumpAccountToTypesAccount(dumpAccount))
+	if err != nil {
+		s.err = err
+		return
+	}
+	if s.wroteOne {
+		s.write([]byte{','})
+	}
+	s.write(keyJSON)
+	s.write([]byte{':'})
+	s.write(valueJSON)
+	s.wroteOne = true
+}
+
+func (s *streamingAlloc) Close() error {
+	s.write([]byte{'}'})
+	return s.err
 }
 
 // saveFile marshals the object to the given file
@@ -359,8 +449,9 @@ func saveFile(baseDir, filename string, data interface{}) error {
 }
 
 // dispatchOutput writes the output data to either stderr or stdout, or to the specified
-// files
-func dispatchOutput(ctx *cli.Context, baseDir string, result *ExecutionResult, alloc Alloc, body hexutil.Bytes, bt map[common.Hash]hexutil.Bytes) error {
+// files. An empty allocOutput skips the alloc dispatch, which is used when the
+// alloc has already been streamed to disk by the caller.
+func dispatchOutput(ctx *cli.Context, baseDir string, result *ExecutionResult, alloc Alloc, allocOutput string, body hexutil.Bytes, bt map[common.Hash]hexutil.Bytes) error {
 	stdOutObject := make(map[string]interface{})
 	stdErrObject := make(map[string]interface{})
 	dispatch := func(baseDir, fName, name string, obj interface{}) error {
@@ -378,7 +469,7 @@ func dispatchOutput(ctx *cli.Context, baseDir string, result *ExecutionResult, a
 		}
 		return nil
 	}
-	if err := dispatch(baseDir, ctx.String(OutputAllocFlag.Name), "alloc", alloc); err != nil {
+	if err := dispatch(baseDir, allocOutput, "alloc", alloc); err != nil {
 		return err
 	}
 	if err := dispatch(baseDir, ctx.String(OutputResultFlag.Name), "result", result); err != nil {