feat: Optimistic Execution

baseapp/abci.go

@@ -494,6 +494,8 @@ func (app *BaseApp) ProcessProposal(req *abci.RequestProcessProposal) (resp *abc
 	// processed the first block, as we want to avoid overwriting the finalizeState
 	// after state changes during InitChain.
 	if req.Height > app.initialHeight {
+		// abort any running OE
+		app.optimisticExec.Abort()
 		app.setState(execModeFinalize, header)
 	}
 
@@ -534,6 +536,16 @@ func (app *BaseApp) ProcessProposal(req *abci.RequestProcessProposal) (resp *abc
 		return &abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_REJECT}, nil
 	}
 
+	// Only execute optimistic execution if OE is enabled and the block height
+	// is greater than the initial height. During the first block we'll be carrying
+	// state from InitChain, so it would be impossible for us to easily revert.
+	// After the first block has been processed, the next blocks will get executed
+	// optimistically, so that when the ABCI client calls `FinalizeBlock` the app
+	// can have a response ready.
+	if app.optimisticExec.Enabled() && req.Height > app.initialHeight {
+		app.optimisticExec.Execute(req)
+	}
+
 	return resp, nil
 }
 
@@ -647,17 +659,11 @@ func (app *BaseApp) VerifyVoteExtension(req *abci.RequestVerifyVoteExtension) (r
 	return resp, err
 }
 
-// FinalizeBlock will execute the block proposal provided by RequestFinalizeBlock.
-// Specifically, it will execute an application's BeginBlock (if defined), followed
-// by the transactions in the proposal, finally followed by the application's
-// EndBlock (if defined).
-//
-// For each raw transaction, i.e. a byte slice, BaseApp will only execute it if
-// it adheres to the sdk.Tx interface. Otherwise, the raw transaction will be
-// skipped. This is to support compatibility with proposers injecting vote
-// extensions into the proposal, which should not themselves be executed in cases
-// where they adhere to the sdk.Tx interface.
-func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.ResponseFinalizeBlock, error) {
+// internalFinalizeBlock it's what actually executes the block, called by the Optimistic
+// Execution flow or by the FinalizeBlock ABCI method. The context received is
+// only used to handle early cancellation, for anything related to state app.finalizeBlockState.ctx
+// must be used.
+func (app *BaseApp) internalFinalizeBlock(ctx context.Context, req *abci.RequestFinalizeBlock) (*abci.ResponseFinalizeBlock, error) {
 	var events []abci.Event
 
 	if err := app.checkHalt(req.Height, req.Time); err != nil {
@@ -732,6 +738,15 @@ func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.Respons
 		return nil, err
 	}
 
+	// First check for an abort signal after beginBlock, as it's the first place
+	// we spend any significant amount of time.
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		// continue
+	}
+
 	events = append(events, beginBlock.Events...)
 
 	// Iterate over all raw transactions in the proposal and attempt to execute
@@ -758,6 +773,14 @@ func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.Respons
 			)
 		}
 
+		// check after every tx if we should abort
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			// continue
+		}
+
 		txResults = append(txResults, response)
 	}
 
@@ -770,6 +793,14 @@ func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.Respons
 		return nil, err
 	}
 
+	// check after endBlock if we should abort, to avoid propagating the result
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+		// continue
+	}
+
 	events = append(events, endBlock.Events...)
 	cp := app.GetConsensusParams(app.finalizeBlockState.ctx)
 
@@ -778,10 +809,47 @@ func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.Respons
 		TxResults:             txResults,
 		ValidatorUpdates:      endBlock.ValidatorUpdates,
 		ConsensusParamUpdates: &cp,
-		AppHash:               app.workingHash(),
 	}, nil
 }
 
+// FinalizeBlock will execute the block proposal provided by RequestFinalizeBlock.
+// Specifically, it will execute an application's BeginBlock (if defined), followed
+// by the transactions in the proposal, finally followed by the application's
+// EndBlock (if defined).
+//
+// For each raw transaction, i.e. a byte slice, BaseApp will only execute it if
+// it adheres to the sdk.Tx interface. Otherwise, the raw transaction will be
+// skipped. This is to support compatibility with proposers injecting vote
+// extensions into the proposal, which should not themselves be executed in cases
+// where they adhere to the sdk.Tx interface.
+func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.ResponseFinalizeBlock, error) {
+	if app.optimisticExec.Initialized() {
+		// check if the hash we got is the same as the one we are executing
+		aborted := app.optimisticExec.AbortIfNeeded(req.Hash)
+		// Wait for the OE to finish, regardless of whether it was aborted or not
+		res, err := app.optimisticExec.WaitResult()
+
+		// only return if we are not aborting
+		if !aborted {
+			if res != nil {
+				res.AppHash = app.workingHash()
+			}
+			return res, err
+		}
+
+		// if it was aborted, we need to reset the state
+		app.finalizeBlockState = nil
+		app.optimisticExec.Reset()
+	}
+
+	// if no OE is running, just run the block (this is either a block replay or a OE that got aborted)
+	res, err := app.internalFinalizeBlock(context.Background(), req)
+	if res != nil {
+		res.AppHash = app.workingHash()
+	}
+	return res, err
+}
+
 // checkHalt checkes if height or time exceeds halt-height or halt-time respectively.
 func (app *BaseApp) checkHalt(height int64, time time.Time) error {
 	var halt bool

baseapp/abci_test.go

@@ -2197,3 +2197,44 @@ func TestBaseApp_VoteExtensions(t *testing.T) {
 	committedAvgPrice := suite.baseApp.NewContext(true).KVStore(capKey1).Get([]byte("avgPrice"))
 	require.Equal(t, avgPrice, committedAvgPrice)
 }
+
+func TestOptimisticExecution(t *testing.T) {
+	suite := NewBaseAppSuite(t, baseapp.SetOptimisticExecution())
+
+	_, err := suite.baseApp.InitChain(&abci.RequestInitChain{
+		ConsensusParams: &cmtproto.ConsensusParams{},
+	})
+	require.NoError(t, err)
+
+	// run 50 blocks
+	for i := 0; i < 50; i++ {
+		tx := newTxCounter(t, suite.txConfig, 0, 1)
+		txBytes, err := suite.txConfig.TxEncoder()(tx)
+		require.NoError(t, err)
+
+		reqProcProp := abci.RequestProcessProposal{
+			Txs:    [][]byte{txBytes},
+			Height: suite.baseApp.LastBlockHeight() + 1,
+			Hash:   []byte("some-hash" + strconv.FormatInt(suite.baseApp.LastBlockHeight()+1, 10)),
+		}
+
+		respProcProp, err := suite.baseApp.ProcessProposal(&reqProcProp)
+		require.Equal(t, abci.ResponseProcessProposal_ACCEPT, respProcProp.Status)
+		require.NoError(t, err)
+
+		reqFinalizeBlock := abci.RequestFinalizeBlock{
+			Height: reqProcProp.Height,
+			Txs:    reqProcProp.Txs,
+			Hash:   reqProcProp.Hash,
+		}
+
+		respFinalizeBlock, err := suite.baseApp.FinalizeBlock(&reqFinalizeBlock)
+		require.NoError(t, err)
+		require.Len(t, respFinalizeBlock.TxResults, 1)
+
+		_, err = suite.baseApp.Commit()
+		require.NoError(t, err)
+	}
+
+	require.Equal(t, int64(50), suite.baseApp.LastBlockHeight())
+}

baseapp/baseapp.go

@@ -23,6 +23,7 @@ import (
 	"cosmossdk.io/store/snapshots"
 	storetypes "cosmossdk.io/store/types"
 
+	"github.com/cosmos/cosmos-sdk/baseapp/oe"
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	servertypes "github.com/cosmos/cosmos-sdk/server/types"
@@ -187,6 +188,8 @@ type BaseApp struct {
 	chainID string
 
 	cdc codec.Codec
+
+	optimisticExec *oe.OptimisticExecution
 }
 
 // NewBaseApp returns a reference to an initialized BaseApp. It accepts a

baseapp/oe/optimistic_execution.go

@@ -0,0 +1,152 @@
+package oe
+
+import (
+	"bytes"
+	"context"
+	"encoding/hex"
+	"math/rand"
+	"sync"
+	"time"
+
+	abci "github.com/cometbft/cometbft/abci/types"
+
+	"cosmossdk.io/log"
+)
+
+// FinalizeBlockFunc is the function that is called by the OE to finalize the
+// block. It is the same as the one in the ABCI app.
+type FinalizeBlockFunc func(context.Context, *abci.RequestFinalizeBlock) (*abci.ResponseFinalizeBlock, error)
+
+// OptimisticExecution is a struct that contains the OE context. It is used to
+// run the FinalizeBlock function in a goroutine, and to abort it if needed.
+type OptimisticExecution struct {
+	finalizeBlockFunc FinalizeBlockFunc // ABCI FinalizeBlock function with a context
+	logger            log.Logger
+
+	mtx         sync.Mutex
+	stopCh      chan struct{}
+	request     *abci.RequestFinalizeBlock
+	response    *abci.ResponseFinalizeBlock
+	err         error
+	cancelFunc  func() // cancel function for the context
+	initialized bool   // A boolean value indicating whether the struct has been initialized
+
+	// debugging/testing options
+	abortRate int // number from 0 to 100 that determines the percentage of OE that should be aborted
+}
+
+func NewOptimisticExecution(logger log.Logger, fn FinalizeBlockFunc, opts ...func(*OptimisticExecution)) *OptimisticExecution {
+	oe := &OptimisticExecution{logger: logger, finalizeBlockFunc: fn}
+	for _, opt := range opts {
+		opt(oe)
+	}
+	return oe
+}
+
+func WithAbortRate(rate int) func(*OptimisticExecution) {
+	return func(oe *OptimisticExecution) {
+		oe.abortRate = rate
+	}
+}
+
+// Reset resets the OE context. Must be called whenever we want to invalidate
+// the current OE.
+func (oe *OptimisticExecution) Reset() {
+	oe.mtx.Lock()
+	defer oe.mtx.Unlock()
+	oe.request = nil
+	oe.response = nil
+	oe.err = nil
+	oe.initialized = false
+}
+
+func (oe *OptimisticExecution) Enabled() bool {
+	return oe != nil
+}
+
+// Initialized returns true if the OE was initialized, meaning that it contains
+// a request and it was run or it is running.
+func (oe *OptimisticExecution) Initialized() bool {
+	if oe == nil {
+		return false
+	}
+	oe.mtx.Lock()
+	defer oe.mtx.Unlock()
+
+	return oe.initialized
+}
+
+// Execute initializes the OE and starts it in a goroutine.
+func (oe *OptimisticExecution) Execute(req *abci.RequestProcessProposal) {
+	oe.mtx.Lock()
+	defer oe.mtx.Unlock()
+
+	oe.stopCh = make(chan struct{})
+	oe.request = &abci.RequestFinalizeBlock{
+		Txs:                req.Txs,
+		DecidedLastCommit:  req.ProposedLastCommit,
+		Misbehavior:        req.Misbehavior,
+		Hash:               req.Hash,
+		Height:             req.Height,
+		Time:               req.Time,
+		NextValidatorsHash: req.NextValidatorsHash,
+		ProposerAddress:    req.ProposerAddress,
+	}
+
+	oe.logger.Debug("OE started", "height", req.Height, "hash", hex.EncodeToString(req.Hash), "time", req.Time.String())
+	ctx, cancel := context.WithCancel(context.Background())
+	oe.cancelFunc = cancel
+	oe.initialized = true
+
+	go func() {
+		start := time.Now()
+		resp, err := oe.finalizeBlockFunc(ctx, oe.request)
+		oe.mtx.Lock()
+		executionTime := time.Since(start)
+		oe.logger.Debug("OE finished", "duration", executionTime.String(), "height", req.Height, "hash", hex.EncodeToString(req.Hash))
+		oe.response, oe.err = resp, err
+		close(oe.stopCh)
+		oe.mtx.Unlock()
+	}()
+}
+
+// AbortIfNeeded aborts the OE if the request hash is not the same as the one in
+// the running OE. Returns true if the OE was aborted.
+func (oe *OptimisticExecution) AbortIfNeeded(reqHash []byte) bool {
+	if oe == nil {
+		return false
+	}
+
+	oe.mtx.Lock()
+	defer oe.mtx.Unlock()
+
+	if !bytes.Equal(oe.request.Hash, reqHash) {
+		oe.logger.Error("OE aborted due to hash mismatch", "oe_hash", hex.EncodeToString(oe.request.Hash), "req_hash", hex.EncodeToString(reqHash), "oe_height", oe.request.Height, "req_height", oe.request.Height)
+		oe.cancelFunc()
+		return true
+	} else if oe.abortRate > 0 && rand.Intn(100) < oe.abortRate {
+		// this is for test purposes only, we can emulate a certain percentage of
+		// OE needed to be aborted.
+		oe.cancelFunc()
+		oe.logger.Error("OE aborted due to test abort rate")
+		return true
+	}
+
+	return false
+}
+
+// Abort aborts the OE unconditionally and waits for it to finish.
+func (oe *OptimisticExecution) Abort() {
+	if oe == nil || oe.cancelFunc == nil {
+		return
+	}
+
+	oe.cancelFunc()
+	<-oe.stopCh
+}
+
+// WaitResult waits for the OE to finish and returns the result.
+func (oe *OptimisticExecution) WaitResult() (*abci.ResponseFinalizeBlock, error) {
+	<-oe.stopCh
+	return oe.response, oe.err
+}

baseapp/options.go

@@ -13,6 +13,7 @@ import (
 	snapshottypes "cosmossdk.io/store/snapshots/types"
 	storetypes "cosmossdk.io/store/types"
 
+	"github.com/cosmos/cosmos-sdk/baseapp/oe"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -104,6 +105,13 @@ func SetChainID(chainID string) func(*BaseApp) {
 	return func(app *BaseApp) { app.chainID = chainID }
 }
 
+// SetOptimisticExecution enables optimistic execution.
+func SetOptimisticExecution(opts ...func(*oe.OptimisticExecution)) func(*BaseApp) {
+	return func(app *BaseApp) {
+		app.optimisticExec = oe.NewOptimisticExecution(app.logger, app.internalFinalizeBlock, opts...)
+	}
+}
+
 func (app *BaseApp) SetName(name string) {
 	if app.sealed {
 		panic("SetName() on sealed BaseApp")