constraining # of arb'able txs per block

x/protorev/genesis.go

@@ -18,11 +18,16 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState)
 	k.SetParams(ctx, genState.Params)
 
 	// Init module state
-	k.SetProtoRevEnabled(ctx, genState.Params.Enabled)
-	k.SetDaysSinceModuleGenesis(ctx, 0)
+	k.SetProtoRevEnabled(ctx, genState.Params.Enabled)                              // configure the module to be enabled or disabled
+	k.SetDaysSinceModuleGenesis(ctx, 0)                                             // configure the number of days the module has been active to be 0
+	k.SetLatestBlockHeight(ctx, uint64(ctx.BlockHeight()))                          // configure the block height the module was initialized at
+	k.SetCurrentTxCount(ctx, 0)                                                     // configure the current tx count to be 0
+	if err := k.SetMaxTxsPerBlock(ctx, types.MaxTransactionsPerBlock); err != nil { // configure the max txs per block
+		panic(err)
+	}
 
 	// Default we only allow 3 pools to be arbitraged against per tx
-	if err := k.SetMaxPools(ctx, 3); err != nil {
+	if err := k.SetMaxPoolsPerTx(ctx, 3); err != nil {
 		panic(err)
 	}
 

x/protorev/keeper/posthandler.go

@@ -1,6 +1,8 @@
 package keeper
 
 import (
+	"fmt"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	gammtypes "github.com/osmosis-labs/osmosis/v13/x/gamm/types"
@@ -25,39 +27,39 @@ func NewProtoRevDecorator(protoRevDecorator Keeper) ProtoRevDecorator {
 // This posthandler will first check if there were any swaps in the tx. If so, collect all of the pools, build three
 // pool routes for cyclic arbitrage, and then execute the optimal route if it exists.
 func (protoRevDec ProtoRevDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
+	if ctx.GasMeter().Limit() == 0 {
+		return next(ctx, tx, simulate)
+	}
+
 	// Create a cache context to execute the posthandler such that
 	// 1. If there is an error, then the cache context is discarded
 	// 2. If there is no error, then the cache context is written to the main context
 	cacheCtx, write := ctx.CacheContext()
+	cacheCtx = cacheCtx.WithGasMeter(sdk.NewInfiniteGasMeter())
 
-	txGasWanted := cacheCtx.GasMeter().Limit()
-	if txGasWanted == 0 {
+	// Check if the protorev posthandler can be executed
+	if err := protoRevDec.ProtoRevKeeper.AnteHandlerCheck(cacheCtx); err != nil {
 		return next(ctx, tx, simulate)
 	}
 
-	// Change the gas meter to be an infinite gas meter which allows the entire posthandler to work without consuming any gas
-	cacheCtx = cacheCtx.WithGasMeter(sdk.NewInfiniteGasMeter())
-
-	// Only execute the protoRev module if it is enabled
-	if enabled, err := protoRevDec.ProtoRevKeeper.GetProtoRevEnabled(cacheCtx); err != nil || !enabled {
+	swappedPools := ExtractSwappedPools(tx)
+	// If there were no swaps, then we don't need to do anything
+	if len(swappedPools) == 0 {
 		return next(ctx, tx, simulate)
 	}
 
-	// Get the max number of pools to iterate through
-	maxPoolsToIterate, err := protoRevDec.ProtoRevKeeper.GetMaxPools(cacheCtx)
+	maxPoolsToIterate, err := protoRevDec.ProtoRevKeeper.GetMaxPoolsPerTx(cacheCtx)
 	if err != nil {
 		return next(ctx, tx, simulate)
 	}
 
-	// Find routes for every single pool that was swapped on (up to maxPoolsToIterate pools per tx)
-	swappedPools := ExtractSwappedPools(tx)
+	// Find every pool that was swapped on up to maxPoolsToIterate pools
 	if len(swappedPools) > int(maxPoolsToIterate) {
 		swappedPools = swappedPools[:int(maxPoolsToIterate)]
 	}
 
 	tradeErr := error(nil)
 	for _, swap := range swappedPools {
-		// If there was is an error executing the trade, break and set tradeErr
 		if err := protoRevDec.ProtoRevKeeper.ProtoRevTrade(cacheCtx, swap); err != nil {
 			tradeErr = err
 			break
@@ -66,6 +68,13 @@ func (protoRevDec ProtoRevDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simu
 
 	// If there was no error, write the cache context to the main context
 	if tradeErr == nil {
+		// Increment the current tx count
+		currentTxCount, err := protoRevDec.ProtoRevKeeper.GetCurrentTxCount(cacheCtx)
+		if err != nil {
+			return next(ctx, tx, simulate)
+		}
+		protoRevDec.ProtoRevKeeper.SetCurrentTxCount(cacheCtx, currentTxCount+1)
+
 		write()
 		ctx.EventManager().EmitEvents(cacheCtx.EventManager().Events())
 	} else {
@@ -75,6 +84,42 @@ func (protoRevDec ProtoRevDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simu
 	return next(ctx, tx, simulate)
 }
 
+// AnteHandlerCheck checks if the module is enabled and if the number of transactions to be processed per block has been reached.
+func (k Keeper) AnteHandlerCheck(ctx sdk.Context) error {
+	// Only execute the posthandler if the module is enabled
+	if enabled, err := k.GetProtoRevEnabled(ctx); err != nil || !enabled {
+		return fmt.Errorf("protoRev is not enabled")
+	}
+
+	latestBlockHeight, err := k.GetLatestBlockHeight(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get last seen block height")
+	}
+
+	currentTxCount, err := k.GetCurrentTxCount(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get current tx count")
+	}
+
+	maxTxCount, err := k.GetMaxTxsPerBlock(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get max txs per block")
+	}
+
+	// Only execute the posthandler if the number of transactions to be processed per block has not been reached
+	blockHeight := uint64(ctx.BlockHeight())
+	if blockHeight == latestBlockHeight {
+		if currentTxCount >= maxTxCount {
+			return fmt.Errorf("max tx count reached")
+		}
+	} else {
+		// Reset the current tx count
+		k.SetCurrentTxCount(ctx, 0)
+		k.SetLatestBlockHeight(ctx, blockHeight)
+	}
+	return nil
+}
+
 // ProtoRevTrade wraps around the build routes, iterate routes, and execute trade functionality to execute an cyclic arbitrage trade
 // if it exists. It returns an error if there was an error executing the trade and a boolean if the trade was executed.
 func (k Keeper) ProtoRevTrade(ctx sdk.Context, swap SwapToBackrun) error {

x/protorev/keeper/posthandler_test.go

@@ -16,14 +16,15 @@ import (
 func (suite *KeeperTestSuite) TestAnteHandle() {
 
 	type param struct {
-		msgs                []sdk.Msg
-		txFee               sdk.Coins
-		minGasPrices        sdk.DecCoins
-		gasLimit            uint64
-		isCheckTx           bool
-		baseDenomGas        bool
-		expectedNumOfTrades sdk.Int
-		expectedProfits     []*sdk.Coin
+		msgs                     []sdk.Msg
+		txFee                    sdk.Coins
+		minGasPrices             sdk.DecCoins
+		gasLimit                 uint64
+		isCheckTx                bool
+		baseDenomGas             bool
+		expectedNumOfTrades      sdk.Int
+		expectedProfits          []*sdk.Coin
+		expectedTransactionCount uint64
 	}
 
 	txBuilder := suite.clientCtx.TxConfig.NewTxBuilder()
@@ -42,14 +43,15 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 		{
 			name: "Random Msg - Expect Nothing to Happen",
 			params: param{
-				msgs:                []sdk.Msg{testdata.NewTestMsg(addr0)},
-				txFee:               sdk.NewCoins(sdk.NewCoin("uosmo", sdk.NewInt(10000))),
-				minGasPrices:        sdk.NewDecCoins(),
-				gasLimit:            500000,
-				isCheckTx:           false,
-				baseDenomGas:        true,
-				expectedNumOfTrades: sdk.ZeroInt(),
-				expectedProfits:     []*sdk.Coin{},
+				msgs:                     []sdk.Msg{testdata.NewTestMsg(addr0)},
+				txFee:                    sdk.NewCoins(sdk.NewCoin("uosmo", sdk.NewInt(10000))),
+				minGasPrices:             sdk.NewDecCoins(),
+				gasLimit:                 500000,
+				isCheckTx:                false,
+				baseDenomGas:             true,
+				expectedNumOfTrades:      sdk.ZeroInt(),
+				expectedProfits:          []*sdk.Coin{},
+				expectedTransactionCount: 0,
 			},
 			expectPass: true,
 		},
@@ -69,13 +71,14 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 						TokenOutMinAmount: sdk.NewInt(10000),
 					},
 				},
-				txFee:               sdk.NewCoins(sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000))),
-				minGasPrices:        sdk.NewDecCoins(),
-				gasLimit:            0,
-				isCheckTx:           false,
-				baseDenomGas:        true,
-				expectedNumOfTrades: sdk.ZeroInt(),
-				expectedProfits:     []*sdk.Coin{},
+				txFee:                    sdk.NewCoins(sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000))),
+				minGasPrices:             sdk.NewDecCoins(),
+				gasLimit:                 0,
+				isCheckTx:                false,
+				baseDenomGas:             true,
+				expectedNumOfTrades:      sdk.ZeroInt(),
+				expectedProfits:          []*sdk.Coin{},
+				expectedTransactionCount: 0,
 			},
 			expectPass: true,
 		},
@@ -95,13 +98,14 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 						TokenOutMinAmount: sdk.NewInt(1),
 					},
 				},
-				txFee:               sdk.NewCoins(sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000))),
-				minGasPrices:        sdk.NewDecCoins(),
-				gasLimit:            500000,
-				isCheckTx:           false,
-				baseDenomGas:        true,
-				expectedNumOfTrades: sdk.ZeroInt(),
-				expectedProfits:     []*sdk.Coin{},
+				txFee:                    sdk.NewCoins(sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000))),
+				minGasPrices:             sdk.NewDecCoins(),
+				gasLimit:                 500000,
+				isCheckTx:                false,
+				baseDenomGas:             true,
+				expectedNumOfTrades:      sdk.ZeroInt(),
+				expectedProfits:          []*sdk.Coin{},
+				expectedTransactionCount: 1,
 			},
 			expectPass: true,
 		},
@@ -133,6 +137,7 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 						Amount: sdk.NewInt(24848),
 					},
 				},
+				expectedTransactionCount: 2,
 			},
 			expectPass: true,
 		},
@@ -168,6 +173,7 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 						Amount: sdk.NewInt(24848),
 					},
 				},
+				expectedTransactionCount: 3,
 			},
 			expectPass: true,
 		},
@@ -203,14 +209,14 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 						Amount: sdk.NewInt(56609900),
 					},
 				},
+				expectedTransactionCount: 4,
 			},
 			expectPass: true,
 		},
 	}
 
 	for _, tc := range tests {
 		suite.Run(tc.name, func() {
-
 			suite.Ctx = suite.Ctx.WithIsCheckTx(tc.params.isCheckTx)
 			suite.Ctx = suite.Ctx.WithGasMeter(sdk.NewInfiniteGasMeter())
 			suite.Ctx = suite.Ctx.WithMinGasPrices(tc.params.minGasPrices)
@@ -265,6 +271,11 @@ func (suite *KeeperTestSuite) TestAnteHandle() {
 				profits := suite.App.AppKeepers.ProtoRevKeeper.GetAllProfits(suite.Ctx)
 				suite.Require().Equal(tc.params.expectedProfits, profits)
 
+				// Check that the current count of transactions that could have been arb'd is correct
+				currentTxCount, err := suite.App.AppKeepers.ProtoRevKeeper.GetCurrentTxCount(suite.Ctx)
+				suite.Require().NoError(err)
+				suite.Require().Equal(tc.params.expectedTransactionCount, currentTxCount)
+
 			} else {
 				suite.Require().Error(err)
 			}

x/protorev/keeper/protorev.go

@@ -241,6 +241,46 @@ func (k Keeper) SetProtoRevEnabled(ctx sdk.Context, enabled bool) {
 	store.Set(types.KeyPrefixProtoRevEnabled, bz)
 }
 
+// GetCurrentTxCount returns the number of transactions that had at least one swap msg for the current block height
+func (k Keeper) GetCurrentTxCount(ctx sdk.Context) (uint64, error) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixCurrentTxCount)
+	bz := store.Get(types.KeyPrefixCurrentTxCount)
+	if bz == nil {
+		// This should never happen as the module is initialized on genesis and reset in the post handler
+		return 0, fmt.Errorf("current tx count has not been set in state")
+	}
+
+	res := sdk.BigEndianToUint64(bz)
+
+	return res, nil
+}
+
+// SetCurrentTxCount sets the number of transactions that had at least one swap msg for the current block height
+func (k Keeper) SetCurrentTxCount(ctx sdk.Context, txCount uint64) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixCurrentTxCount)
+	store.Set(types.KeyPrefixCurrentTxCount, sdk.Uint64ToBigEndian(txCount))
+}
+
+// GetLatestBlockHeight returns the latest block height that protorev was run on
+func (k Keeper) GetLatestBlockHeight(ctx sdk.Context) (uint64, error) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixLatestBlockHeight)
+	bz := store.Get(types.KeyPrefixLatestBlockHeight)
+	if bz == nil {
+		// This should never happen as the module is initialized on genesis and reset in the post handler
+		return 0, fmt.Errorf("block height has not been set in state")
+	}
+
+	res := sdk.BigEndianToUint64(bz)
+
+	return res, nil
+}
+
+// SetLatestBlockHeight sets the latest block height that protorev was run on
+func (k Keeper) SetLatestBlockHeight(ctx sdk.Context, blockHeight uint64) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixLatestBlockHeight)
+	store.Set(types.KeyPrefixLatestBlockHeight, sdk.Uint64ToBigEndian(blockHeight))
+}
+
 // ---------------------- Admin Stores  ---------------------- //
 
 // GetAdminAccount returns the admin account for protorev
@@ -277,10 +317,10 @@ func (k Keeper) SetDeveloperAccount(ctx sdk.Context, developerAccount sdk.AccAdd
 	store.Set(types.KeyPrefixDeveloperAccount, developerAccount.Bytes())
 }
 
-// GetMaxPools returns the max number of pools that can be iterated after a swap
-func (k Keeper) GetMaxPools(ctx sdk.Context) (uint64, error) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxPools)
-	bz := store.Get(types.KeyPrefixMaxPools)
+// GetMaxPoolsPerTx returns the max number of pools that can be iterated per transaction
+func (k Keeper) GetMaxPoolsPerTx(ctx sdk.Context) (uint64, error) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxPoolsPerTx)
+	bz := store.Get(types.KeyPrefixMaxPoolsPerTx)
 	if bz == nil {
 		// This should never happen as the module is initialized on genesis
 		return 0, fmt.Errorf("max pools configuration has not been set in state")
@@ -290,15 +330,41 @@ func (k Keeper) GetMaxPools(ctx sdk.Context) (uint64, error) {
 	return res, nil
 }
 
-// SetMaxPools sets the max number of pools that can be iterated after a swap.
-func (k Keeper) SetMaxPools(ctx sdk.Context, maxPools uint64) error {
-	if maxPools == 0 || maxPools > types.MaxIterablePools {
-		return fmt.Errorf("max pools must be between 1 and 5")
+// SetMaxPoolsPerTxPerTx sets the max number of pools that can be iterated per transaction
+func (k Keeper) SetMaxPoolsPerTx(ctx sdk.Context, maxPools uint64) error {
+	if maxPools == 0 || maxPools > types.MaxIterablePoolsPerTx {
+		return fmt.Errorf("max pools must be between 1 and %d", types.MaxIterablePoolsPerTx)
 	}
 
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxPools)
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxPoolsPerTx)
 	bz := sdk.Uint64ToBigEndian(maxPools)
-	store.Set(types.KeyPrefixMaxPools, bz)
+	store.Set(types.KeyPrefixMaxPoolsPerTx, bz)
+
+	return nil
+}
+
+// GetMaxTxsPerBlock returns the max number of transactions that can be back-run in a block
+func (k Keeper) GetMaxTxsPerBlock(ctx sdk.Context) (uint64, error) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxTxsPerBlock)
+	bz := store.Get(types.KeyPrefixMaxTxsPerBlock)
+	if bz == nil {
+		// This should never happen as the module is initialized on genesis
+		return 0, fmt.Errorf("max txs configuration has not been set in state")
+	}
+
+	res := sdk.BigEndianToUint64(bz)
+	return res, nil
+}
+
+// SetMaxTxsPerBlock sets the max number of transactions that can be back-run in a block
+func (k Keeper) SetMaxTxsPerBlock(ctx sdk.Context, maxTxs uint64) error {
+	if maxTxs == 0 || maxTxs > types.MaxTransactionsPerBlock {
+		return fmt.Errorf("max txs must be between 1 and %d", types.MaxTransactionsPerBlock)
+	}
+
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixMaxTxsPerBlock)
+	bz := sdk.Uint64ToBigEndian(maxTxs)
+	store.Set(types.KeyPrefixMaxTxsPerBlock, bz)
 
 	return nil
 }