use finalizer instead of returning state release function

Author: magicxyyz

arbitrum/apibackend.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
@@ -229,11 +230,10 @@ func (a *APIBackend) FeeHistory(
 
 	// use the most recent average compute rate for all blocks
 	// note: while we could query this value for each block, it'd be prohibitively expensive
-	state, _, release, err := a.StateAndHeaderByNumber(ctx, newestBlock)
+	state, _, err := a.StateAndHeaderByNumber(ctx, newestBlock)
 	if err != nil {
 		return common.Big0, nil, nil, nil, err
 	}
-	defer release()
 	speedLimit, err := core.GetArbOSSpeedLimitPerSecond(state)
 	if err != nil {
 		return common.Big0, nil, nil, nil, err
@@ -434,59 +434,77 @@ func (a *APIBackend) BlockByNumberOrHash(ctx context.Context, blockNrOrHash rpc.
 	return nil, errors.New("invalid arguments; neither block nor hash specified")
 }
 
-func (a *APIBackend) stateAndHeaderFromHeader(ctx context.Context, header *types.Header, err error) (*state.StateDB, *types.Header, ethapi.StateReleaseFunc, error) {
+func (a *APIBackend) stateAndHeaderFromHeader(ctx context.Context, header *types.Header, err error) (*state.StateDB, *types.Header, error) {
 	if err != nil {
-		return nil, header, nil, err
+		return nil, header, err
 	}
 	if header == nil {
-		return nil, nil, nil, errors.New("header not found")
+		return nil, nil, errors.New("header not found")
 	}
 	if !a.BlockChain().Config().IsArbitrumNitro(header.Number) {
-		return nil, header, nil, types.ErrUseFallback
+		return nil, header, types.ErrUseFallback
 	}
 	bc := a.BlockChain()
 	stateFor := func(header *types.Header) (*state.StateDB, error) {
 		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{})
 		}
-		state, err := bc.StateAt(header.Root)
-		return state, err
+		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
 	}
 	lastState, lastHeader, err := FindLastAvailableState(ctx, bc, stateFor, header, nil, a.b.config.MaxRecreateStateDepth)
 	if err != nil {
-		return nil, nil, nil, err
-	}
-	release := func() {
-		if lastHeader.Root != (common.Hash{}) {
-			bc.StateCache().TrieDB().Dereference(lastHeader.Root)
-		}
+		return nil, nil, err
 	}
 	if lastHeader == header {
-		return lastState, header, release, nil
+		return lastState, header, nil
+	}
+	if lastHeader.Root != (common.Hash{}) {
+		defer bc.StateCache().TrieDB().Dereference(lastHeader.Root)
 	}
-	defer release()
 	targetBlock := bc.GetBlockByNumber(header.Number.Uint64())
 	if targetBlock == nil {
-		return nil, nil, nil, errors.New("target block not found")
+		return nil, nil, errors.New("target block not found")
 	}
 	lastBlock := bc.GetBlockByNumber(lastHeader.Number.Uint64())
 	if lastBlock == nil {
-		return nil, nil, nil, errors.New("last block not found")
+		return nil, nil, errors.New("last block not found")
 	}
 	reexec := uint64(0)
 	checkLive := false
 	preferDisk := true
-	state, release, err := eth.NewArbEthereum(a.b.arb.BlockChain(), a.ChainDb()).StateAtBlock(ctx, targetBlock, reexec, lastState, lastBlock, checkLive, preferDisk)
-	return state, header, release, err
+	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
+	}
+	// 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, ethapi.StateReleaseFunc, error) {
+func (a *APIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.BlockNumber) (*state.StateDB, *types.Header, error) {
 	header, err := a.HeaderByNumber(ctx, number)
 	return a.stateAndHeaderFromHeader(ctx, header, err)
 }
 
-func (a *APIBackend) StateAndHeaderByNumberOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) (*state.StateDB, *types.Header, ethapi.StateReleaseFunc, error) {
+func (a *APIBackend) StateAndHeaderByNumberOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) (*state.StateDB, *types.Header, error) {
 	header, err := a.HeaderByNumberOrHash(ctx, blockNrOrHash)
 	return a.stateAndHeaderFromHeader(ctx, header, err)
 }

eth/api_backend.go

@@ -37,7 +37,6 @@ import (
 	"github.com/ethereum/go-ethereum/eth/tracers"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/event"
-	"github.com/ethereum/go-ethereum/internal/ethapi"
 	"github.com/ethereum/go-ethereum/miner"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/rpc"
@@ -199,46 +198,46 @@ func (b *EthAPIBackend) PendingBlockAndReceipts() (*types.Block, types.Receipts)
 	return b.eth.miner.PendingBlockAndReceipts()
 }
 
-func (b *EthAPIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.BlockNumber) (*state.StateDB, *types.Header, ethapi.StateReleaseFunc, error) {
+func (b *EthAPIBackend) StateAndHeaderByNumber(ctx context.Context, number rpc.BlockNumber) (*state.StateDB, *types.Header, error) {
 	// Pending state is only known by the miner
 	if number == rpc.PendingBlockNumber {
 		block, state := b.eth.miner.Pending()
 		if block == nil || state == nil {
-			return nil, nil, nil, errors.New("pending state is not available")
+			return nil, nil, errors.New("pending state is not available")
 		}
-		return state, block.Header(), ethapi.NoOpStateRelease, nil
+		return state, block.Header(), nil
 	}
 	// Otherwise resolve the block number and return its state
 	header, err := b.HeaderByNumber(ctx, number)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	if header == nil {
-		return nil, nil, nil, errors.New("header not found")
+		return nil, nil, errors.New("header not found")
 	}
 	stateDb, err := b.eth.BlockChain().StateAt(header.Root)
-	return stateDb, header, ethapi.NoOpStateRelease, err
+	return stateDb, header, err
 }
 
-func (b *EthAPIBackend) StateAndHeaderByNumberOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) (*state.StateDB, *types.Header, ethapi.StateReleaseFunc, error) {
+func (b *EthAPIBackend) StateAndHeaderByNumberOrHash(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) (*state.StateDB, *types.Header, error) {
 	if blockNr, ok := blockNrOrHash.Number(); ok {
 		return b.StateAndHeaderByNumber(ctx, blockNr)
 	}
 	if hash, ok := blockNrOrHash.Hash(); ok {
 		header, err := b.HeaderByHash(ctx, hash)
 		if err != nil {
-			return nil, nil, nil, err
+			return nil, nil, err
 		}
 		if header == nil {
-			return nil, nil, nil, errors.New("header for hash not found")
+			return nil, nil, errors.New("header for hash not found")
 		}
 		if blockNrOrHash.RequireCanonical && b.eth.blockchain.GetCanonicalHash(header.Number.Uint64()) != hash {
-			return nil, nil, nil, errors.New("hash is not currently canonical")
+			return nil, nil, errors.New("hash is not currently canonical")
 		}
 		stateDb, err := b.eth.BlockChain().StateAt(header.Root)
-		return stateDb, header, ethapi.NoOpStateRelease, err
+		return stateDb, header, err
 	}
-	return nil, nil, nil, errors.New("invalid arguments; neither block nor hash specified")
+	return nil, nil, errors.New("invalid arguments; neither block nor hash specified")
 }
 
 func (b *EthAPIBackend) GetReceipts(ctx context.Context, hash common.Hash) (types.Receipts, error) {

eth/catalyst/api_test.go

@@ -1143,11 +1143,10 @@ func TestWithdrawals(t *testing.T) {
 	}
 
 	// 11: verify withdrawals were processed.
-	db, _, release, err := ethservice.APIBackend.StateAndHeaderByNumber(context.Background(), rpc.BlockNumber(execData.ExecutionPayload.Number))
+	db, _, err := ethservice.APIBackend.StateAndHeaderByNumber(context.Background(), rpc.BlockNumber(execData.ExecutionPayload.Number))
 	if err != nil {
 		t.Fatalf("unable to load db: %v", err)
 	}
-	defer release()
 	for i, w := range blockParams.Withdrawals {
 		// w.Amount is in gwei, balance in wei
 		if db.GetBalance(w.Address).Uint64() != w.Amount*params.GWei {

graphql/graphql.go

@@ -86,21 +86,20 @@ type Account struct {
 }
 
 // getState fetches the StateDB object for an account.
-func (a *Account) getState(ctx context.Context) (*state.StateDB, ethapi.StateReleaseFunc, error) {
-	state, _, release, err := a.r.backend.StateAndHeaderByNumberOrHash(ctx, a.blockNrOrHash)
-	return state, release, err
+func (a *Account) getState(ctx context.Context) (*state.StateDB, error) {
+	state, _, err := a.r.backend.StateAndHeaderByNumberOrHash(ctx, a.blockNrOrHash)
+	return state, err
 }
 
 func (a *Account) Address(ctx context.Context) (common.Address, error) {
 	return a.address, nil
 }
 
 func (a *Account) Balance(ctx context.Context) (hexutil.Big, error) {
-	state, release, err := a.getState(ctx)
+	state, err := a.getState(ctx)
 	if err != nil {
 		return hexutil.Big{}, err
 	}
-	defer release()
 	balance := state.GetBalance(a.address)
 	if balance == nil {
 		return hexutil.Big{}, fmt.Errorf("failed to load balance %x", a.address)
@@ -117,29 +116,26 @@ func (a *Account) TransactionCount(ctx context.Context) (hexutil.Uint64, error)
 		}
 		return hexutil.Uint64(nonce), nil
 	}
-	state, release, err := a.getState(ctx)
+	state, err := a.getState(ctx)
 	if err != nil {
 		return 0, err
 	}
-	defer release()
 	return hexutil.Uint64(state.GetNonce(a.address)), nil
 }
 
 func (a *Account) Code(ctx context.Context) (hexutil.Bytes, error) {
-	state, release, err := a.getState(ctx)
+	state, err := a.getState(ctx)
 	if err != nil {
 		return hexutil.Bytes{}, err
 	}
-	defer release()
 	return state.GetCode(a.address), nil
 }
 
 func (a *Account) Storage(ctx context.Context, args struct{ Slot common.Hash }) (common.Hash, error) {
-	state, release, err := a.getState(ctx)
+	state, err := a.getState(ctx)
 	if err != nil {
 		return common.Hash{}, err
 	}
-	defer release()
 	return state.GetState(a.address, args.Slot), nil
 }