use StateAtBlock and reference states when recreating missing state

arbitrum/apibackend.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
@@ -445,20 +446,57 @@ func (a *APIBackend) stateAndHeaderFromHeader(ctx context.Context, header *types
 	}
 	bc := a.BlockChain()
 	stateFor := func(header *types.Header) (*state.StateDB, error) {
-		return bc.StateAt(header.Root)
+		if header.Root != (common.Hash{}) {
+			// Try referencing the root, if it isn't in dirties cache then Reference will have no effect
+			bc.StateCache().TrieDB().Reference(header.Root, common.Hash{})
+		}
+		statedb, err := state.New(header.Root, bc.StateCache(), bc.Snapshots())
+		if err != nil {
+			return nil, err
+		}
+		if header.Root != (common.Hash{}) {
+			// we are setting finalizer instead of returning a StateReleaseFunc to avoid changing ethapi.Backend interface to minimize diff to upstream
+			headerRoot := header.Root
+			runtime.SetFinalizer(statedb, func(_ *state.StateDB) {
+				bc.StateCache().TrieDB().Dereference(headerRoot)
+			})
+		}
+		return statedb, err
 	}
-	state, lastHeader, err := FindLastAvailableState(ctx, bc, stateFor, header, nil, a.b.config.MaxRecreateStateDepth)
+	lastState, lastHeader, err := FindLastAvailableState(ctx, bc, stateFor, header, nil, a.b.config.MaxRecreateStateDepth)
 	if err != nil {
 		return nil, nil, err
 	}
 	if lastHeader == header {
-		return state, header, nil
+		return lastState, header, nil
 	}
-	state, err = AdvanceStateUpToBlock(ctx, bc, state, header, lastHeader, nil)
+	if lastHeader.Root != (common.Hash{}) {
+		defer bc.StateCache().TrieDB().Dereference(lastHeader.Root)
+	}
+	targetBlock := bc.GetBlockByNumber(header.Number.Uint64())
+	if targetBlock == nil {
+		return nil, nil, errors.New("target block not found")
+	}
+	lastBlock := bc.GetBlockByNumber(lastHeader.Number.Uint64())
+	if lastBlock == nil {
+		return nil, nil, errors.New("last block not found")
+	}
+	reexec := uint64(0)
+	checkLive := false
+	preferDisk := true
+	statedb, release, err := eth.NewArbEthereum(a.b.arb.BlockChain(), a.ChainDb()).StateAtBlock(ctx, targetBlock, reexec, lastState, lastBlock, checkLive, preferDisk)
 	if err != nil {
 		return nil, nil, err
 	}
-	return state, header, err
+	// we are setting finalizer instead of returning a StateReleaseFunc to avoid changing ethapi.Backend interface to minimize diff to upstream
+	// to set a finalizer we need to allocated the obj in current block
+	statedb, err = state.New(header.Root, statedb.Database(), nil)
+	if header.Root != (common.Hash{}) {
+		runtime.SetFinalizer(statedb, func(_ *state.StateDB) {
+			release()
+		})
+	}
+	return statedb, header, err
 }
 
 func (a *APIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.BlockNumber) (*state.StateDB, *types.Header, error) {
@@ -476,7 +514,7 @@ func (a *APIBackend) StateAtBlock(ctx context.Context, block *types.Block, reexe
 		return nil, nil, types.ErrUseFallback
 	}
 	// DEV: This assumes that `StateAtBlock` only accesses the blockchain and chainDb fields
-	return eth.NewArbEthereum(a.b.arb.BlockChain(), a.ChainDb()).StateAtBlock(ctx, block, reexec, base, checkLive, preferDisk)
+	return eth.NewArbEthereum(a.b.arb.BlockChain(), a.ChainDb()).StateAtBlock(ctx, block, reexec, base, nil, checkLive, preferDisk)
 }
 
 func (a *APIBackend) StateAtTransaction(ctx context.Context, block *types.Block, txIndex int, reexec uint64) (*core.Message, vm.BlockContext, *state.StateDB, tracers.StateReleaseFunc, error) {

arbitrum/recreatestate.go

@@ -80,24 +80,3 @@ func AdvanceStateByBlock(ctx context.Context, bc *core.BlockChain, state *state.
 	}
 	return state, block, nil
 }
-
-func AdvanceStateUpToBlock(ctx context.Context, bc *core.BlockChain, state *state.StateDB, targetHeader *types.Header, lastAvailableHeader *types.Header, logFunc StateBuildingLogFunction) (*state.StateDB, error) {
-	returnedBlockNumber := targetHeader.Number.Uint64()
-	blockToRecreate := lastAvailableHeader.Number.Uint64() + 1
-	prevHash := lastAvailableHeader.Hash()
-	for ctx.Err() == nil {
-		state, block, err := AdvanceStateByBlock(ctx, bc, state, targetHeader, blockToRecreate, prevHash, logFunc)
-		if err != nil {
-			return nil, err
-		}
-		prevHash = block.Hash()
-		if blockToRecreate >= returnedBlockNumber {
-			if block.Hash() != targetHeader.Hash() {
-				return nil, fmt.Errorf("blockHash doesn't match when recreating number: %d expected: %v got: %v", blockToRecreate, targetHeader.Hash(), block.Hash())
-			}
-			return state, nil
-		}
-		blockToRecreate++
-	}
-	return nil, ctx.Err()
-}

eth/api_backend.go

@@ -423,7 +423,7 @@ func (b *EthAPIBackend) StartMining() error {
 }
 
 func (b *EthAPIBackend) StateAtBlock(ctx context.Context, block *types.Block, reexec uint64, base *state.StateDB, readOnly bool, preferDisk bool) (*state.StateDB, tracers.StateReleaseFunc, error) {
-	return b.eth.StateAtBlock(ctx, block, reexec, base, readOnly, preferDisk)
+	return b.eth.StateAtBlock(ctx, block, reexec, base, nil, readOnly, preferDisk)
 }
 
 func (b *EthAPIBackend) StateAtTransaction(ctx context.Context, block *types.Block, txIndex int, reexec uint64) (*core.Message, vm.BlockContext, *state.StateDB, tracers.StateReleaseFunc, error) {

eth/state_accessor.go

@@ -47,13 +47,14 @@ import (
 //   - reexec:     The maximum number of blocks to reprocess trying to obtain the desired state
 //   - base:       If the caller is tracing multiple blocks, the caller can provide the parent
 //     state continuously from the callsite.
+//   - baseBlock:  Arbitrum specific: caller can provide the block from which reprocessing should start. Previous argument (base) is assumed to be the state at the block. If base is not provided, baseBlock is ignored.
 //   - readOnly:   If true, then the live 'blockchain' state database is used. No mutation should
 //     be made from caller, e.g. perform Commit or other 'save-to-disk' changes.
 //     Otherwise, the trash generated by caller may be persisted permanently.
 //   - preferDisk: this arg can be used by the caller to signal that even though the 'base' is
 //     provided, it would be preferable to start from a fresh state, if we have it
 //     on disk.
-func (eth *Ethereum) StateAtBlock(ctx context.Context, block *types.Block, reexec uint64, base *state.StateDB, readOnly bool, preferDisk bool) (statedb *state.StateDB, release tracers.StateReleaseFunc, err error) {
+func (eth *Ethereum) StateAtBlock(ctx context.Context, block *types.Block, reexec uint64, base *state.StateDB, baseBlock *types.Block, readOnly bool, preferDisk bool) (statedb *state.StateDB, release tracers.StateReleaseFunc, err error) {
 	var (
 		current  *types.Block
 		database state.Database
@@ -66,8 +67,10 @@ func (eth *Ethereum) StateAtBlock(ctx context.Context, block *types.Block, reexe
 		// The state is available in live database, create a reference
 		// on top to prevent garbage collection and return a release
 		// function to deref it.
+
+		// Try referencing the root, if it isn't in dirties cache then Reference will have no effect
+		eth.blockchain.StateCache().TrieDB().Reference(block.Root(), common.Hash{})
 		if statedb, err = eth.blockchain.StateAt(block.Root()); err == nil {
-			statedb.Database().TrieDB().Reference(block.Root(), common.Hash{})
 			return statedb, func() {
 				statedb.Database().TrieDB().Dereference(block.Root())
 			}, nil
@@ -95,7 +98,11 @@ func (eth *Ethereum) StateAtBlock(ctx context.Context, block *types.Block, reexe
 		}
 		// The optional base statedb is given, mark the start point as parent block
 		statedb, database, report = base, base.Database(), false
-		current = eth.blockchain.GetBlock(block.ParentHash(), block.NumberU64()-1)
+		if baseBlock == nil {
+			current = eth.blockchain.GetBlock(block.ParentHash(), block.NumberU64()-1)
+		} else {
+			current = baseBlock
+		}
 	} else {
 		// Otherwise, try to reexec blocks until we find a state or reach our limit
 		current = block
@@ -214,7 +221,7 @@ func (eth *Ethereum) stateAtTransaction(ctx context.Context, block *types.Block,
 	}
 	// Lookup the statedb of parent block from the live database,
 	// otherwise regenerate it on the flight.
-	statedb, release, err := eth.StateAtBlock(ctx, parent, reexec, nil, true, false)
+	statedb, release, err := eth.StateAtBlock(ctx, parent, reexec, nil, nil, true, false)
 	if err != nil {
 		return nil, vm.BlockContext{}, nil, nil, err
 	}