Store pruneable orders in a key-per-order format rather than key-per-…

…height

protocol/app/upgrades/v5.0.0/upgrade.go

@@ -108,6 +108,9 @@ func CreateUpgradeHandler(
 		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
 		sdkCtx.Logger().Info(fmt.Sprintf("Running %s Upgrade...", UpgradeName))
 
+		// Migrate pruneable orders to new format
+		clobKeeper.MigratePruneableOrders(sdkCtx)
+
 		// Set all perpetuals to cross market type
 		perpetualsUpgrade(sdkCtx, perpetualsKeeper)
 

protocol/x/clob/abci_test.go

@@ -22,8 +22,6 @@ import (
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
 
-	"cosmossdk.io/store/prefix"
-	storetypes "cosmossdk.io/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/dydxprotocol/v4-chain/protocol/mocks"
 	keepertest "github.com/dydxprotocol/v4-chain/protocol/testutil/keeper"
@@ -51,9 +49,7 @@ func assertFillAmountAndPruneState(
 	t *testing.T,
 	k *keeper.Keeper,
 	ctx sdk.Context,
-	storeKey storetypes.StoreKey,
 	expectedFillAmounts map[types.OrderId]satypes.BaseQuantums,
-	expectedPruneableBlockHeights map[uint32][]types.OrderId,
 	expectedPrunedOrders map[types.OrderId]bool,
 ) {
 	for orderId, expectedFillAmount := range expectedFillAmounts {
@@ -66,28 +62,6 @@ func assertFillAmountAndPruneState(
 		exists, _, _ := k.GetOrderFillAmount(ctx, orderId)
 		require.False(t, exists)
 	}
-
-	for blockHeight, orderIds := range expectedPruneableBlockHeights {
-		// Verify that expected `blockHeightToPotentiallyPrunableOrders` were deleted.
-		blockHeightToPotentiallyPrunableOrdersStore := prefix.NewStore(
-			ctx.KVStore(storeKey),
-			[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
-		)
-
-		potentiallyPrunableOrdersBytes := blockHeightToPotentiallyPrunableOrdersStore.Get(
-			lib.Uint32ToKey(blockHeight),
-		)
-
-		var potentiallyPrunableOrders = &types.PotentiallyPrunableOrders{}
-		err := potentiallyPrunableOrders.Unmarshal(potentiallyPrunableOrdersBytes)
-		require.NoError(t, err)
-
-		require.ElementsMatch(
-			t,
-			potentiallyPrunableOrders.OrderIds,
-			orderIds,
-		)
-	}
 }
 
 func TestEndBlocker_Failure(t *testing.T) {
@@ -184,7 +158,6 @@ func TestEndBlocker_Success(t *testing.T) {
 
 		// Expectations.
 		expectedFillAmounts                  map[types.OrderId]satypes.BaseQuantums
-		expectedPruneableBlockHeights        map[uint32][]types.OrderId
 		expectedPrunedOrders                 map[types.OrderId]bool
 		expectedStatefulPlacementInState     map[types.OrderId]bool
 		expectedStatefulOrderTimeSlice       map[time.Time][]types.OrderId
@@ -802,9 +775,7 @@ func TestEndBlocker_Success(t *testing.T) {
 				t,
 				ks.ClobKeeper,
 				ctx,
-				ks.StoreKey,
 				tc.expectedFillAmounts,
-				tc.expectedPruneableBlockHeights,
 				tc.expectedPrunedOrders,
 			)
 

protocol/x/clob/keeper/order_state.go

@@ -1,6 +1,9 @@
 package keeper
 
 import (
+	"bytes"
+	"encoding/binary"
+
 	"cosmossdk.io/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
@@ -125,14 +128,32 @@ func (k Keeper) GetOrderFillAmount(
 	return true, satypes.BaseQuantums(orderFillState.FillAmount), orderFillState.PrunableBlockHeight
 }
 
-// AddOrdersForPruning creates or updates a slice of `orderIds` to state for potential future pruning from state.
-// These orders will be checked for pruning from state at `prunableBlockHeight`. If the `orderIds` slice provided
-// contains duplicates, the duplicates will be ignored.
+func (k Keeper) GetPruneableOrdersStore(ctx sdk.Context, height uint32) prefix.Store {
+	var buf bytes.Buffer
+	buf.Write([]byte(types.PrunableOrdersKeyPrefix))
+	buf.Write(lib.Uint32ToKey(height))
+	buf.Write([]byte(":"))
+	return prefix.NewStore(ctx.KVStore(k.storeKey), buf.Bytes())
+}
+
+// AddOrdersForPruning creates or updates `orderIds` to state for potential future pruning from state.
 func (k Keeper) AddOrdersForPruning(ctx sdk.Context, orderIds []types.OrderId, prunableBlockHeight uint32) {
+	store := k.GetPruneableOrdersStore(ctx, prunableBlockHeight)
+	for _, orderId := range orderIds {
+		store.Set(
+			orderId.ToStateKey(),
+			k.cdc.MustMarshal(&orderId),
+		)
+	}
+}
+
+// LegacyAddOrdersForPruning is the old key-per-height way of storing orders to prune.
+// DO NOT USE. Retained for testing purposes.
+func (k Keeper) LegacyAddOrdersForPruning(ctx sdk.Context, orderIds []types.OrderId, prunableBlockHeight uint32) {
 	// Retrieve an instance of the store.
 	store := prefix.NewStore(
 		ctx.KVStore(k.storeKey),
-		[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
+		[]byte(types.LegacyBlockHeightToPotentiallyPrunableOrdersPrefix),
 	)
 
 	// Retrieve the `PotentiallyPrunableOrders` bytes from the store.
@@ -187,27 +208,13 @@ func (k Keeper) AddOrdersForPruning(ctx sdk.Context, orderIds []types.OrderId, p
 // Note: An order is only deemed prunable if the `prunableBlockHeight` on the `OrderFillState` is less than or equal
 // to the provided `blockHeight` passed this method. Returns a slice of unique `OrderIds` which were pruned from state.
 func (k Keeper) PruneOrdersForBlockHeight(ctx sdk.Context, blockHeight uint32) (prunedOrderIds []types.OrderId) {
-	// Retrieve an instance of the stores.
-	blockHeightToPotentiallyPrunableOrdersStore := prefix.NewStore(
-		ctx.KVStore(k.storeKey),
-		[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
-	)
-
-	// Retrieve the raw bytes of the `prunableOrders`.
-	potentiallyPrunableOrderBytes := blockHeightToPotentiallyPrunableOrdersStore.Get(
-		lib.Uint32ToKey(blockHeight),
-	)
-
-	// If there are no prunable orders for this block, then there is nothing to do. Early return.
-	if potentiallyPrunableOrderBytes == nil {
-		return
-	}
+	potentiallyPrunableOrdersStore := k.GetPruneableOrdersStore(ctx, blockHeight)
+	it := potentiallyPrunableOrdersStore.Iterator(nil, nil)
+	defer it.Close()
 
-	var potentiallyPrunableOrders types.PotentiallyPrunableOrders
-	k.cdc.MustUnmarshal(potentiallyPrunableOrderBytes, &potentiallyPrunableOrders)
-
-	for _, orderId := range potentiallyPrunableOrders.OrderIds {
-		// Check if the order can be pruned, and prune if so.
+	for ; it.Valid(); it.Next() {
+		var orderId types.OrderId
+		k.cdc.MustUnmarshal(it.Value(), &orderId)
 		exists, _, prunableBlockHeight := k.GetOrderFillAmount(ctx, orderId)
 		if exists && prunableBlockHeight <= blockHeight {
 			k.RemoveOrderFillAmount(ctx, orderId)
@@ -221,14 +228,31 @@ func (k Keeper) PruneOrdersForBlockHeight(ctx sdk.Context, blockHeight uint32) (
 				)
 			}
 		}
+		potentiallyPrunableOrdersStore.Delete(it.Key())
 	}
 
-	// Delete the key for prunable orders at this block height.
-	blockHeightToPotentiallyPrunableOrdersStore.Delete(
-		lib.Uint32ToKey(blockHeight),
+	return prunedOrderIds
+}
+
+// MigratePruneableOrders is used to migrate prunable orders from key-per-height to key-per-order format.
+func (k Keeper) MigratePruneableOrders(ctx sdk.Context) {
+	store := prefix.NewStore(
+		ctx.KVStore(k.storeKey),
+		[]byte(types.LegacyBlockHeightToPotentiallyPrunableOrdersPrefix),
 	)
+	it := store.Iterator(nil, nil)
+	defer it.Close()
 
-	return prunedOrderIds
+	for ; it.Valid(); it.Next() {
+		if it.Value() == nil {
+			continue
+		}
+
+		height := binary.BigEndian.Uint32(it.Value())
+		var potentiallyPrunableOrders types.PotentiallyPrunableOrders
+		k.cdc.MustUnmarshal(it.Value(), &potentiallyPrunableOrders)
+		k.AddOrdersForPruning(ctx, potentiallyPrunableOrders.OrderIds, height)
+	}
 }
 
 // RemoveOrderFillAmount removes the fill amount of an Order from state and the memstore.

protocol/x/clob/keeper/order_state_test.go

@@ -1,12 +1,9 @@
 package keeper_test
 
 import (
-	"sort"
 	"testing"
 
-	"cosmossdk.io/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	"github.com/dydxprotocol/v4-chain/protocol/mocks"
 	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
 	keepertest "github.com/dydxprotocol/v4-chain/protocol/testutil/keeper"
@@ -262,99 +259,6 @@ func TestOrderFillAmountInitMemStore_Success(t *testing.T) {
 	require.False(t, exists)
 }
 
-func TestAddOrdersForPruning_Determinism(t *testing.T) {
-	memClob := &mocks.MemClob{}
-	memClob.On("SetClobKeeper", mock.Anything).Return()
-
-	ks := keepertest.NewClobKeepersTestContext(
-		t,
-		memClob,
-		&mocks.BankKeeper{},
-		&mocks.IndexerEventManager{},
-	)
-
-	blockHeight := uint32(10)
-
-	store := prefix.NewStore(
-		ks.Ctx.KVStore(ks.StoreKey),
-		[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
-	)
-
-	orders := []types.OrderId{
-		constants.Order_Alice_Num0_Id0_Clob0_Buy5_Price10_GTB15.OrderId,
-		constants.Order_Bob_Num0_Id0_Clob1_Sell10_Price15_GTB20.OrderId,
-		constants.Order_Alice_Num1_Id0_Clob0_Sell10_Price15_GTB20.OrderId,
-		constants.Order_Alice_Num0_Id1_Clob0_Sell10_Price15_GTB15.OrderId,
-	}
-
-	expectedOrders := []types.OrderId{
-		constants.Order_Alice_Num0_Id0_Clob0_Buy5_Price10_GTB15.OrderId,
-		constants.Order_Alice_Num0_Id1_Clob0_Sell10_Price15_GTB15.OrderId,
-		constants.Order_Alice_Num1_Id0_Clob0_Sell10_Price15_GTB20.OrderId,
-		constants.Order_Bob_Num0_Id0_Clob1_Sell10_Price15_GTB20.OrderId,
-	}
-
-	for i := 0; i < 100; i++ {
-		ks.ClobKeeper.AddOrdersForPruning(
-			ks.Ctx,
-			orders,
-			blockHeight,
-		)
-
-		potentiallyPrunableOrdersBytes := store.Get(
-			lib.Uint32ToKey(blockHeight),
-		)
-
-		var potentiallyPrunableOrders = &types.PotentiallyPrunableOrders{}
-		err := potentiallyPrunableOrders.Unmarshal(potentiallyPrunableOrdersBytes)
-		require.NoError(t, err)
-
-		sort.Sort(types.SortedOrders(expectedOrders))
-		for i, o := range potentiallyPrunableOrders.OrderIds {
-			require.Equal(t, o, expectedOrders[i])
-		}
-	}
-}
-
-func TestAddOrdersForPruning_DuplicateOrderIds(t *testing.T) {
-	memClob := &mocks.MemClob{}
-	memClob.On("SetClobKeeper", mock.Anything).Return()
-	ks := keepertest.NewClobKeepersTestContext(
-		t,
-		memClob,
-		&mocks.BankKeeper{},
-		&mocks.IndexerEventManager{},
-	)
-
-	blockHeight := uint32(10)
-
-	ks.ClobKeeper.AddOrdersForPruning(
-		ks.Ctx,
-		[]types.OrderId{
-			constants.Order_Alice_Num0_Id1_Clob0_Sell10_Price15_GTB15.OrderId,
-			constants.Order_Alice_Num0_Id1_Clob0_Sell10_Price15_GTB15.OrderId,
-			constants.Order_Alice_Num0_Id2_Clob1_Sell5_Price10_GTB15.OrderId,
-			constants.Order_Alice_Num0_Id2_Clob1_Sell5_Price10_GTB15.OrderId,
-		},
-		blockHeight,
-	)
-
-	store := prefix.NewStore(
-		ks.Ctx.KVStore(ks.StoreKey),
-		[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
-	)
-
-	potentiallyPrunableOrdersBytes := store.Get(
-		lib.Uint32ToKey(blockHeight),
-	)
-
-	var potentiallyPrunableOrders = &types.PotentiallyPrunableOrders{}
-	err := potentiallyPrunableOrders.Unmarshal(potentiallyPrunableOrdersBytes)
-	require.NoError(t, err)
-
-	require.Len(t, potentiallyPrunableOrders.OrderIds, 2)
-}
-
 func TestPruning(t *testing.T) {
 	tests := map[string]struct {
 		// Setup.
@@ -571,22 +475,70 @@ func TestPruning(t *testing.T) {
 				require.Equal(t, prunableBlockHeight, tc.expectedPrunableBlockHeight)
 			}
 
-			// Verify that expected `blockHeightToPotentiallyPrunableOrdersStore` were deleted.
-			blockHeightToPotentiallyPrunableOrdersStore := prefix.NewStore(
-				ks.Ctx.KVStore(ks.StoreKey),
-				[]byte(types.BlockHeightToPotentiallyPrunableOrdersPrefix),
-			)
-
+			// Verify all prune order keys were deleted for specified heights
 			for _, blockHeight := range tc.expectedEmptyPotentiallyPrunableOrderBlockHeights {
-				has := blockHeightToPotentiallyPrunableOrdersStore.Has(
-					lib.Uint32ToKey(blockHeight),
-				)
-				require.False(t, has)
+				potentiallyPrunableOrdersStore := ks.ClobKeeper.GetPruneableOrdersStore(ks.Ctx, blockHeight)
+				it := potentiallyPrunableOrdersStore.Iterator(nil, nil)
+				defer it.Close()
+				require.False(t, it.Valid())
 			}
 		})
 	}
 }
 
+func TestMigratePruneableOrders(t *testing.T) {
+	memClob := &mocks.MemClob{}
+	memClob.On("SetClobKeeper", mock.Anything).Return()
+
+	ks := keepertest.NewClobKeepersTestContext(
+		t,
+		memClob,
+		&mocks.BankKeeper{},
+		&mocks.IndexerEventManager{},
+	)
+
+	ordersA := []types.OrderId{
+		constants.Order_Alice_Num0_Id0_Clob0_Buy5_Price10_GTB15.OrderId,
+		constants.Order_Alice_Num1_Id0_Clob0_Sell10_Price15_GTB20.OrderId,
+		constants.Order_Alice_Num0_Id1_Clob0_Sell10_Price15_GTB15.OrderId,
+	}
+	ordersB := []types.OrderId{
+		constants.Order_Alice_Num0_Id0_Clob0_Buy5_Price10_GTB15.OrderId,
+		constants.Order_Bob_Num0_Id0_Clob1_Sell10_Price15_GTB20.OrderId,
+		constants.Order_Alice_Num1_Id0_Clob0_Sell10_Price15_GTB20.OrderId,
+	}
+
+	ks.ClobKeeper.AddOrdersForPruning(
+		ks.Ctx,
+		ordersA,
+		10,
+	)
+	ks.ClobKeeper.AddOrdersForPruning(
+		ks.Ctx,
+		ordersB,
+		100,
+	)
+
+	ks.ClobKeeper.MigratePruneableOrders(ks.Ctx)
+
+	getPostMigrationOrdersAtHeight := func(height uint32) []types.OrderId {
+		postMigrationOrders := []types.OrderId{}
+		storeA := ks.ClobKeeper.GetPruneableOrdersStore(ks.Ctx, height)
+		it := storeA.Iterator(nil, nil)
+		defer it.Close()
+		for ; it.Valid(); it.Next() {
+			var orderId types.OrderId
+			err := orderId.Unmarshal(it.Value())
+			require.NoError(t, err)
+			postMigrationOrders = append(postMigrationOrders, orderId)
+		}
+		return postMigrationOrders
+	}
+
+	require.ElementsMatch(t, ordersA, getPostMigrationOrdersAtHeight(10))
+	require.ElementsMatch(t, ordersB, getPostMigrationOrdersAtHeight(100))
+}
+
 func TestRemoveOrderFillAmount(t *testing.T) {
 	tests := map[string]struct {
 		// Setup.