use StateAtBlock and reference states when recreating missing state

arbitrum/apibackend.go

@@ -14,6 +14,8 @@ import (
 	"github.com/ethereum/go-ethereum/eth"
 	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/metrics"
+	"github.com/ethereum/go-ethereum/trie"
 
 	"github.com/ethereum/go-ethereum/accounts"
 	"github.com/ethereum/go-ethereum/common"
@@ -22,6 +24,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/bloombits"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/eth/filters"
@@ -32,6 +35,11 @@ import (
 	"github.com/ethereum/go-ethereum/rpc"
 )
 
+var (
+	liveStatesCounter      = metrics.NewRegisteredCounter("arb/apibackend/states/live", nil)
+	recreatedStatesCounter = metrics.NewRegisteredCounter("arb/apibackend/states/recreated", nil)
+)
+
 type APIBackend struct {
 	b *Backend
 
@@ -234,6 +242,7 @@ func (a *APIBackend) FeeHistory(
 		return common.Big0, nil, nil, nil, err
 	}
 	speedLimit, err := core.GetArbOSSpeedLimitPerSecond(state)
+	state.Release()
 	if err != nil {
 		return common.Big0, nil, nil, nil, err
 	}
@@ -444,21 +453,72 @@ func (a *APIBackend) stateAndHeaderFromHeader(ctx context.Context, header *types
 		return nil, header, types.ErrUseFallback
 	}
 	bc := a.BlockChain()
-	stateFor := func(header *types.Header) (*state.StateDB, error) {
-		return bc.StateAt(header.Root)
+	stateFor := func(db state.Database, snapshots *snapshot.Tree) func(header *types.Header) (*state.StateDB, StateReleaseFunc, error) {
+		return func(header *types.Header) (*state.StateDB, StateReleaseFunc, error) {
+			if header.Root != (common.Hash{}) {
+				// Try referencing the root, if it isn't in dirties cache then Reference will have no effect
+				db.TrieDB().Reference(header.Root, common.Hash{})
+			}
+			statedb, err := state.New(header.Root, db, snapshots)
+			if err != nil {
+				return nil, nil, err
+			}
+			if header.Root != (common.Hash{}) {
+				headerRoot := header.Root
+				return statedb, func() { db.TrieDB().Dereference(headerRoot) }, nil
+			}
+			return statedb, NoopStateRelease, nil
+		}
 	}
-	state, lastHeader, err := FindLastAvailableState(ctx, bc, stateFor, header, nil, a.b.config.MaxRecreateStateDepth)
+	liveState, liveStateRelease, err := stateFor(bc.StateCache(), bc.Snapshots())(header)
+	if err == nil {
+		liveStatesCounter.Inc(1)
+		liveState.SetRelease(func() {
+			liveStateRelease()
+		})
+		return liveState, header, nil
+	}
+	// else err != nil => we don't need to call liveStateRelease
+
+	// Create an ephemeral trie.Database for isolating the live one
+	// note: triedb cleans cache is disabled in trie.HashDefaults
+	// note: only states committed to diskdb can be found as we're creating new triedb
+	// note: snapshots are not used here
+	ephemeral := state.NewDatabaseWithConfig(a.ChainDb(), trie.HashDefaults)
+	lastState, lastHeader, lastStateRelease, err := FindLastAvailableState(ctx, bc, stateFor(ephemeral, nil), header, nil, a.b.config.MaxRecreateStateDepth)
 	if err != nil {
 		return nil, nil, err
 	}
+	// make sure that we haven't found the state in diskdb
 	if lastHeader == header {
-		return state, header, nil
-	}
-	state, err = AdvanceStateUpToBlock(ctx, bc, state, header, lastHeader, nil)
+		liveStatesCounter.Inc(1)
+		lastState.SetRelease(func() {
+			lastStateRelease()
+		})
+		return lastState, header, nil
+	}
+	defer lastStateRelease()
+	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 := false // preferDisk is ignored in this case
+	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 nil, nil, fmt.Errorf("failed to recreate state: %w", 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
+	recreatedStatesCounter.Inc(1)
+	statedb.SetRelease(func() {
+		release()
+	})
+	return statedb, header, err
 }
 
 func (a *APIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.BlockNumber) (*state.StateDB, *types.Header, error) {
@@ -476,7 +536,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/recordingdb.go

@@ -309,7 +309,12 @@ func (r *RecordingDatabase) PreimagesFromRecording(chainContextIf core.ChainCont
 }
 
 func (r *RecordingDatabase) GetOrRecreateState(ctx context.Context, header *types.Header, logFunc StateBuildingLogFunction) (*state.StateDB, error) {
-	state, currentHeader, err := FindLastAvailableState(ctx, r.bc, r.StateFor, header, logFunc, -1)
+	stateFor := func(header *types.Header) (*state.StateDB, StateReleaseFunc, error) {
+		state, err := r.StateFor(header)
+		// we don't use the release functor pattern here yet
+		return state, NoopStateRelease, err
+	}
+	state, currentHeader, _, err := FindLastAvailableState(ctx, r.bc, stateFor, header, logFunc, -1)
 	if err != nil {
 		return nil, err
 	}

arbitrum/recreatestate.go

@@ -9,58 +9,64 @@ import (
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
+	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/pkg/errors"
 )
 
 var (
 	ErrDepthLimitExceeded = errors.New("state recreation l2 gas depth limit exceeded")
 )
 
+type StateReleaseFunc tracers.StateReleaseFunc
+
+var NoopStateRelease StateReleaseFunc = func() {}
+
 type StateBuildingLogFunction func(targetHeader, header *types.Header, hasState bool)
-type StateForHeaderFunction func(header *types.Header) (*state.StateDB, error)
+type StateForHeaderFunction func(header *types.Header) (*state.StateDB, StateReleaseFunc, error)
 
 // finds last available state and header checking it first for targetHeader then looking backwards
 // if maxDepthInL2Gas is positive, it constitutes a limit for cumulative l2 gas used of the traversed blocks
 // else if maxDepthInL2Gas is -1, the traversal depth is not limited
 // otherwise only targetHeader state is checked and no search is performed
-func FindLastAvailableState(ctx context.Context, bc *core.BlockChain, stateFor StateForHeaderFunction, targetHeader *types.Header, logFunc StateBuildingLogFunction, maxDepthInL2Gas int64) (*state.StateDB, *types.Header, error) {
+func FindLastAvailableState(ctx context.Context, bc *core.BlockChain, stateFor StateForHeaderFunction, targetHeader *types.Header, logFunc StateBuildingLogFunction, maxDepthInL2Gas int64) (*state.StateDB, *types.Header, StateReleaseFunc, error) {
 	genesis := bc.Config().ArbitrumChainParams.GenesisBlockNum
 	currentHeader := targetHeader
 	var state *state.StateDB
 	var err error
 	var l2GasUsed uint64
+	release := NoopStateRelease
 	for ctx.Err() == nil {
 		lastHeader := currentHeader
-		state, err = stateFor(currentHeader)
+		state, release, err = stateFor(currentHeader)
 		if err == nil {
 			break
 		}
 		if maxDepthInL2Gas > 0 {
 			receipts := bc.GetReceiptsByHash(currentHeader.Hash())
 			if receipts == nil {
-				return nil, lastHeader, fmt.Errorf("failed to get receipts for hash %v", currentHeader.Hash())
+				return nil, lastHeader, nil, fmt.Errorf("failed to get receipts for hash %v", currentHeader.Hash())
 			}
 			for _, receipt := range receipts {
 				l2GasUsed += receipt.GasUsed - receipt.GasUsedForL1
 			}
 			if l2GasUsed > uint64(maxDepthInL2Gas) {
-				return nil, lastHeader, ErrDepthLimitExceeded
+				return nil, lastHeader, nil, ErrDepthLimitExceeded
 			}
 		} else if maxDepthInL2Gas != InfiniteMaxRecreateStateDepth {
-			return nil, lastHeader, err
+			return nil, lastHeader, nil, err
 		}
 		if logFunc != nil {
 			logFunc(targetHeader, currentHeader, false)
 		}
 		if currentHeader.Number.Uint64() <= genesis {
-			return nil, lastHeader, errors.Wrap(err, fmt.Sprintf("moved beyond genesis looking for state %d, genesis %d", targetHeader.Number.Uint64(), genesis))
+			return nil, lastHeader, nil, errors.Wrap(err, fmt.Sprintf("moved beyond genesis looking for state %d, genesis %d", targetHeader.Number.Uint64(), genesis))
 		}
 		currentHeader = bc.GetHeader(currentHeader.ParentHash, currentHeader.Number.Uint64()-1)
 		if currentHeader == nil {
-			return nil, lastHeader, fmt.Errorf("chain doesn't contain parent of block %d hash %v", lastHeader.Number, lastHeader.Hash())
+			return nil, lastHeader, nil, fmt.Errorf("chain doesn't contain parent of block %d hash %v", lastHeader.Number, lastHeader.Hash())
 		}
 	}
-	return state, currentHeader, ctx.Err()
+	return state, currentHeader, release, ctx.Err()
 }
 
 func AdvanceStateByBlock(ctx context.Context, bc *core.BlockChain, state *state.StateDB, targetHeader *types.Header, blockToRecreate uint64, prevBlockHash common.Hash, logFunc StateBuildingLogFunction) (*state.StateDB, *types.Block, error) {

core/blockchain.go

@@ -1071,7 +1071,8 @@ func (bc *BlockChain) Stop() {
 		//  - HEAD:     So we don't need to reprocess any blocks in the general case
 		//  - HEAD-1:   So we don't do large reorgs if our HEAD becomes an uncle
 		//  - HEAD-127: So we have a hard limit on the number of blocks reexecuted
-		if !bc.cacheConfig.TrieDirtyDisabled {
+		// It applies for both full node and sparse archive node
+		if !bc.cacheConfig.TrieDirtyDisabled || bc.cacheConfig.MaxNumberOfBlocksToSkipStateSaving > 0 || bc.cacheConfig.MaxAmountOfGasToSkipStateSaving > 0 {
 			triedb := bc.triedb
 
 			for _, offset := range []uint64{0, 1, bc.cacheConfig.TriesInMemory - 1, math.MaxUint64} {
@@ -1496,7 +1497,7 @@ func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.
 		return nil
 	}
 	// If we're running an archive node, flush
-	// If MaxNumberOfBlocksToSkipStateSaving or MaxAmountOfGasToSkipStateSaving is not zero, then flushing of some blocks will be skipped:
+	// Sparse archive: if MaxNumberOfBlocksToSkipStateSaving or MaxAmountOfGasToSkipStateSaving is not zero, then flushing of some blocks will be skipped:
 	// * at most MaxNumberOfBlocksToSkipStateSaving block state commits will be skipped
 	// * sum of gas used in skipped blocks will be at most MaxAmountOfGasToSkipStateSaving
 	archiveNode := bc.cacheConfig.TrieDirtyDisabled
@@ -1526,7 +1527,7 @@ func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.
 		// we are skipping saving the trie to diskdb, so we need to keep the trie in memory and garbage collect it later
 	}
 
-	// Full node or archive node that's not keeping all states, do proper garbage collection
+	// Full node or sparse archive node that's not keeping all states, do proper garbage collection
 	bc.triedb.Reference(root, common.Hash{}) // metadata reference to keep trie alive
 	bc.triegc.Push(trieGcEntry{root, block.Header().Time}, -int64(block.NumberU64()))
 

core/state/statedb.go

@@ -145,9 +145,28 @@ type StateDB struct {
 	// Testing hooks
 	onCommit func(states *triestate.Set) // Hook invoked when commit is performed
 
+	// arbitrum: cleanup hook
+	release  func()
+	released bool
+
 	deterministic bool
 }
 
+func (s *StateDB) SetRelease(release func()) {
+	if s.release == nil {
+		s.release = release
+	} else {
+		log.Error("StateDB.SetRelease called more then once, may cause memory leak")
+	}
+}
+
+func (s *StateDB) Release() {
+	if s.release != nil && !s.released {
+		s.release()
+		s.released = true
+	}
+}
+
 // New creates a new state from a given trie.
 func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error) {
 	tr, err := db.OpenTrie(root)
@@ -175,6 +194,9 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 		accessList:           newAccessList(),
 		transientStorage:     newTransientStorage(),
 		hasher:               crypto.NewKeccakState(),
+
+		release:  nil,
+		released: false,
 	}
 	if sdb.snaps != nil {
 		sdb.snap = sdb.snaps.Snapshot(root)

eth/api_backend.go

@@ -414,7 +414,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) {