Add context cancellation to SS DB layer

evmrpc/initialize_block_test.go

@@ -0,0 +1,107 @@
+package evmrpc
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/big"
+	"testing"
+
+	ethtypes "github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/trie"
+	"github.com/stretchr/testify/require"
+
+	"github.com/sei-protocol/sei-chain/app/legacyabci"
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	abci "github.com/sei-protocol/sei-chain/sei-tendermint/abci/types"
+	"github.com/sei-protocol/sei-chain/sei-tendermint/rpc/coretypes"
+	tmtypes "github.com/sei-protocol/sei-chain/sei-tendermint/types"
+)
+
+func TestReleaseOnContextPanic(t *testing.T) {
+	t.Parallel()
+
+	var released int
+	release := func() { released++ }
+
+	err := releaseOnContextPanic(release, context.Canceled)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Equal(t, 1, released)
+
+	err = releaseOnContextPanic(release, context.DeadlineExceeded)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Equal(t, 2, released)
+
+	err = releaseOnContextPanic(release, fmt.Errorf("skip: %w", context.DeadlineExceeded))
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Equal(t, 3, released)
+
+	require.Panics(t, func() {
+		_ = releaseOnContextPanic(release, errors.New("pebble seek failed"))
+	})
+	require.Equal(t, 4, released)
+
+	require.Panics(t, func() {
+		_ = releaseOnContextPanic(release, "not an error")
+	})
+	require.Equal(t, 5, released)
+}
+
+func TestInitializeBlockReleasesLeaseOnBeginBlockDeadline(t *testing.T) {
+	orig := runTraceBeginBlock
+	t.Cleanup(func() { runTraceBeginBlock = orig })
+	runTraceBeginBlock = func(sdk.Context, int64, []abci.VoteInfo, []abci.Misbehavior, legacyabci.BeginBlockKeepers) {
+		panic(context.DeadlineExceeded)
+	}
+
+	var released int
+	backend, block := newInitializeBlockTestBackend(t)
+	_, _, release, err := backend.initializeBlock(t.Context(), block, func(int64) (sdk.Context, func()) {
+		return sdk.Context{}, func() { released++ }
+	})
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Equal(t, 1, released, "base snapshot lease must be released on BeginBlock abort")
+	release()
+	require.Equal(t, 1, released, "returned release must be a no-op after recover")
+}
+
+func TestInitializeBlockReleasesLeaseOnUnrelatedBeginBlockPanic(t *testing.T) {
+	orig := runTraceBeginBlock
+	t.Cleanup(func() { runTraceBeginBlock = orig })
+	runTraceBeginBlock = func(sdk.Context, int64, []abci.VoteInfo, []abci.Misbehavior, legacyabci.BeginBlockKeepers) {
+		panic("boom")
+	}
+
+	var released int
+	backend, block := newInitializeBlockTestBackend(t)
+	require.Panics(t, func() {
+		_, _, _, _ = backend.initializeBlock(t.Context(), block, func(int64) (sdk.Context, func()) {
+			return sdk.Context{}, func() { released++ }
+		})
+	})
+	require.Equal(t, 1, released)
+}
+
+func newInitializeBlockTestBackend(t *testing.T) (*Backend, *ethtypes.Block) {
+	t.Helper()
+	tm := &fakeTMClient{
+		status: &coretypes.ResultStatus{SyncInfo: coretypes.SyncInfo{LatestBlockHeight: 10, EarliestBlockHeight: 1}},
+		blocksByHeight: map[int64]*coretypes.ResultBlock{
+			8: {
+				Block: &tmtypes.Block{
+					Header:     tmtypes.Header{Height: 8},
+					LastCommit: &tmtypes.Commit{},
+				},
+			},
+		},
+	}
+	return &Backend{
+			tmClient:   tm,
+			watermarks: newTestWatermarkManager(tm, 10, nil, 10),
+		}, ethtypes.NewBlock(
+			&ethtypes.Header{Number: big.NewInt(8), Time: 1, Difficulty: big.NewInt(0)},
+			&ethtypes.Body{},
+			nil,
+			trie.NewStackTrie(nil),
+		)
+}

evmrpc/simulate.go

@@ -669,6 +669,9 @@ func (b *Backend) replayTransactionTillIndex(ctx context.Context, block *ethtype
 		}
 		_ = b.app.DeliverTx(sdkCtx, abci.RequestDeliverTxV2{Tx: tx}, sdkTx, sha256.Sum256(tx))
 	}
+	if err := ctx.Err(); err != nil {
+		return nil, nil, emptyRelease, err
+	}
 	success = true
 	return state.NewDBImpl(sdkCtx.WithIsEVM(true), b.keeper, true), tmBlock.Block.Txs, release, nil
 }
@@ -683,13 +686,14 @@ func (b *Backend) StateAtBlock(ctx context.Context, block *ethtypes.Block, reexe
 	return statedb, release, nil
 }
 
-func (b *Backend) initializeBlock(ctx context.Context, block *ethtypes.Block, ctxProvider TraceContextProvider) (sdk.Context, *coretypes.ResultBlock, tracers.StateReleaseFunc, error) {
+func (b *Backend) initializeBlock(ctx context.Context, block *ethtypes.Block, ctxProvider TraceContextProvider) (sdkCtx sdk.Context, tmBlock *coretypes.ResultBlock, release tracers.StateReleaseFunc, err error) {
 	emptyRelease := func() {}
+	release = emptyRelease
 	// get the parent block using block.parentHash
 	prevBlockHeight := max(block.Number().Int64()-1, 0)
 
 	blockNumber := block.Number().Int64()
-	tmBlock, err := blockByNumberRespectingWatermarks(ctx, b.tmClient, b.watermarks, &blockNumber, 1)
+	tmBlock, err = blockByNumberRespectingWatermarks(ctx, b.tmClient, b.watermarks, &blockNumber, 1)
 	if err != nil {
 		return sdk.Context{}, nil, emptyRelease, fmt.Errorf("cannot find block %d from tendermint", blockNumber)
 	}
@@ -701,17 +705,57 @@ func (b *Backend) initializeBlock(ctx context.Context, block *ethtypes.Block, ct
 	reqBeginBlock := tmBlock.Block.ToReqBeginBlock(res.Validators)
 	reqBeginBlock.Simulate = true
 	baseCtx, baseRelease := ctxProvider(prevBlockHeight)
-	sdkCtx := baseCtx.WithBlockHeight(blockNumber).WithBlockTime(tmBlock.Block.Time)
-	legacyabci.BeginBlock(sdkCtx, blockNumber, reqBeginBlock.LastCommitInfo.Votes, tmBlock.Block.Evidence.ToABCI(), b.beginBlockKeepers)
-	nextCtx, nextRelease := ctxProvider(sdkCtx.BlockHeight())
+	var nextRelease func()
+	var released bool
+	release = func() {
+		if released {
+			return
+		}
+		released = true
+		if nextRelease != nil {
+			nextRelease()
+		}
+		baseRelease()
+	}
+	defer func() {
+		if r := recover(); r != nil {
+			err = releaseOnContextPanic(release, r)
+			release = emptyRelease
+			sdkCtx = sdk.Context{}
+			tmBlock = nil
+		}
+	}()
+	sdkCtx = baseCtx.WithBlockHeight(blockNumber).WithBlockTime(tmBlock.Block.Time)
+	if ctx != nil {
+		// The RPC/trace deadline must be on the SDK context so KVStore
+		// iteration can pass it into the SS MVCC skip loops.
+		sdkCtx = sdkCtx.WithContext(ctx)
+	}
+	runTraceBeginBlock(sdkCtx, blockNumber, reqBeginBlock.LastCommitInfo.Votes, tmBlock.Block.Evidence.ToABCI(), b.beginBlockKeepers)
+	var nextCtx sdk.Context
+	nextCtx, nextRelease = ctxProvider(sdkCtx.BlockHeight())
 	sdkCtx = sdkCtx.WithNextMs(
 		nextCtx.MultiStore(),
 		[]string{"oracle", "oracle_mem"},
 	)
-	return sdkCtx, tmBlock, func() {
-		nextRelease()
-		baseRelease()
-	}, nil
+	return sdkCtx, tmBlock, release, nil
+}
+
+// runTraceBeginBlock is the BeginBlock used when reconstructing historical
+// state for traces.
+var runTraceBeginBlock = legacyabci.BeginBlock
+
+// releaseOnContextPanic releases leased snapshots and returns a cancel or
+// deadline panic as an error. Other panics are re-raised after release.
+func releaseOnContextPanic(release func(), recovered any) error {
+	if release != nil {
+		release()
+	}
+	err, ok := recovered.(error)
+	if ok && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) {
+		return err
+	}
+	panic(recovered)
 }
 
 func (b *Backend) GetEVM(_ context.Context, msg *core.Message, stateDB vm.StateDB, h *ethtypes.Header, vmConfig *vm.Config, blockCtx *vm.BlockContext) *vm.EVM {

evmrpc/trace_baker.go

@@ -191,7 +191,7 @@ func (b *TraceBaker) bakeBlockOneTracer(height int64, tracer string) bool {
 
 	tracerName := tracer
 	results, err := b.tracersAPI.TraceBlockByNumber(ctx, rpc.BlockNumber(height), &gethtracers.TraceConfig{Tracer: &tracerName})
-	if err != nil {
+	if _, err = resultUnlessExpired(ctx, results, err); err != nil {
 		atomic.AddUint64(&b.failed, 1)
 		bakerLogger.Debug("trace baker block trace failed", "height", height, "tracer", tracer, "err", err)
 		return false

evmrpc/trace_baker_test.go

@@ -20,16 +20,20 @@ import (
 
 // fakeTracerAPI drives the baker with controllable per-call results.
 type fakeTracerAPI struct {
-	mu      sync.Mutex
-	calls   int32
-	results map[int64][]*gethtracers.TxTraceResult // keyed by height
-	errs    map[int64]error
-	gate    chan struct{}           // when set, blocks each call until released
-	gates   map[int64]chan struct{} // optional per-height gates
+	mu                 sync.Mutex
+	calls              int32
+	results            map[int64][]*gethtracers.TxTraceResult // keyed by height
+	errs               map[int64]error
+	gate               chan struct{}           // when set, blocks each call until released
+	gates              map[int64]chan struct{} // optional per-height gates
+	synthesizeOnCancel bool                    // wait for ctx.Done, then return results with a nil error
 }
 
-func (f *fakeTracerAPI) TraceBlockByNumber(_ context.Context, number rpc.BlockNumber, _ *gethtracers.TraceConfig) ([]*gethtracers.TxTraceResult, error) {
+func (f *fakeTracerAPI) TraceBlockByNumber(ctx context.Context, number rpc.BlockNumber, _ *gethtracers.TraceConfig) ([]*gethtracers.TxTraceResult, error) {
 	atomic.AddInt32(&f.calls, 1)
+	if f.synthesizeOnCancel {
+		<-ctx.Done()
+	}
 	f.mu.Lock()
 	gate := f.gate
 	if f.gates != nil {
@@ -151,6 +155,38 @@ func TestTraceBakerErrorBecomesFailedCount(t *testing.T) {
 	require.Equal(t, uint64(0), b.BakedCount(), "errors should not count as baked")
 }
 
+func TestTraceBakerDoesNotCacheExpiredSynthesizedTrace(t *testing.T) {
+	cache, err := keeper.NewTraceDB(t.TempDir())
+	require.NoError(t, err)
+	defer cache.Close()
+
+	tx := common.HexToHash("0xee")
+	api := &fakeTracerAPI{
+		synthesizeOnCancel: true,
+		results: map[int64][]*gethtracers.TxTraceResult{
+			11: {{TxHash: tx, Result: map[string]string{"error": "store access cancelled"}}},
+		},
+	}
+	b := NewTraceBaker(nil, cache, TraceBakerConfig{
+		Workers:     1,
+		QueueSize:   8,
+		BakeTimeout: time.Millisecond,
+	})
+	b.tracersAPI = api
+
+	require.False(t, b.bakeBlockOneTracer(11, "callTracer"))
+	require.Equal(t, uint64(1), b.FailedCount())
+	require.Equal(t, uint64(0), b.BakedCount())
+
+	_, ok, err := cache.Get(11, "callTracer", tx)
+	require.NoError(t, err)
+	require.False(t, ok, "expired bake must not persist a per-tx row")
+
+	_, ok, err = cache.GetBlock(11, "callTracer")
+	require.NoError(t, err)
+	require.False(t, ok, "expired bake must not persist a block row")
+}
+
 func TestTraceBakerSkipsNilOrErroredTxResults(t *testing.T) {
 	// Per-tx errors come back as {Result: nil}; baker must skip them.
 	cache, err := keeper.NewTraceDB(t.TempDir())

evmrpc/trace_cancellation_test.go

@@ -0,0 +1,29 @@
+package evmrpc
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestResultUnlessExpiredReportsTheDeadline(t *testing.T) {
+	live := t.Context()
+
+	traced := map[string]string{"gas": "0x1"}
+	result, err := resultUnlessExpired(live, traced, nil)
+	require.NoError(t, err)
+	require.Equal(t, traced, result)
+
+	underlying := errors.New("tracer failed")
+	result, err = resultUnlessExpired(live, nil, underlying)
+	require.ErrorIs(t, err, underlying)
+	require.Nil(t, result)
+
+	expired, cancelExpired := context.WithCancel(t.Context())
+	cancelExpired()
+	result, err = resultUnlessExpired(expired, map[string]string{"error": "store access cancelled"}, nil)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, result)
+}

evmrpc/trace_profile.go

@@ -116,7 +116,7 @@ func (api *DebugAPI) TraceTransactionProfile(ctx context.Context, hash common.Ha
 	}
 	api.clampDefaultStructLogLimit(config)
 	traceResult, err := api.profiledTraceTx(ctx, tx, msg, txctx, blockCtx, statedb, config, nil, false, &phases.traceExecutionPhaseDurations)
-	if err != nil {
+	if _, err = resultUnlessExpired(ctx, traceResult, err); err != nil {
 		return nil, err
 	}