Move memstore hydration initialized flag to memstore (#2681)

Author: jayy04

protocol/mocks/ClobKeeper.go

@@ -65,7 +65,7 @@ func (cd ClobDecorator) AnteHandle(
 	}
 
 	// Disable order placement and cancelation processing if the clob keeper is not initialized.
-	if !cd.clobKeeper.IsInitialized() {
+	if !cd.clobKeeper.IsInMemStructuresInitialized() {
 		return ctx, errorsmod.Wrap(
 			types.ErrClobNotInitialized,
 			"clob keeper is not initialized. Please wait for the next block.",

protocol/x/clob/ante/clob.go

@@ -53,7 +53,7 @@ func runTestCase(t *testing.T, tc TestCase) {
 	// Setup AnteHandler.
 	mockClobKeeper := &mocks.ClobKeeper{}
 	mockClobKeeper.On("Logger", mock.Anything).Return(log.NewNopLogger()).Maybe()
-	mockClobKeeper.On("IsInitialized").Return(true).Maybe()
+	mockClobKeeper.On("IsInMemStructuresInitialized").Return(true).Maybe()
 	cd := ante.NewClobDecorator(mockClobKeeper)
 	antehandler := sdk.ChainAnteDecorators(cd)
 	if tc.setupMocks != nil {

protocol/x/clob/ante/clob_test.go

@@ -1,7 +1,6 @@
 package keeper
 
 import (
-	"errors"
 	"fmt"
 	"math/big"
 	"sync/atomic"
@@ -54,8 +53,7 @@ type (
 		streamingManager         streamingtypes.FullNodeStreamingManager
 		finalizeBlockEventStager finalizeblock.EventStager[*types.ClobStagedFinalizeBlockEvent]
 
-		initialized         *atomic.Bool
-		memStoreInitialized *atomic.Bool
+		inMemStructuresInitialized *atomic.Bool
 
 		Flags flags.ClobFlags
 
@@ -104,29 +102,28 @@ func NewKeeper(
 	revshareKeeper types.RevShareKeeper,
 ) *Keeper {
 	keeper := &Keeper{
-		cdc:                     cdc,
-		storeKey:                storeKey,
-		memKey:                  memKey,
-		transientStoreKey:       transientStoreKey,
-		authorities:             lib.UniqueSliceToSet(authorities),
-		MemClob:                 memClob,
-		PerpetualIdToClobPairId: make(map[uint32][]types.ClobPairId),
-		subaccountsKeeper:       subaccountsKeeper,
-		assetsKeeper:            assetsKeeper,
-		blockTimeKeeper:         blockTimeKeeper,
-		bankKeeper:              bankKeeper,
-		feeTiersKeeper:          feeTiersKeeper,
-		perpetualsKeeper:        perpetualsKeeper,
-		pricesKeeper:            pricesKeeper,
-		statsKeeper:             statsKeeper,
-		rewardsKeeper:           rewardsKeeper,
-		affiliatesKeeper:        affiliatesKeeper,
-		accountPlusKeeper:       accountPlusKeeper,
-		indexerEventManager:     indexerEventManager,
-		streamingManager:        streamingManager,
-		memStoreInitialized:     &atomic.Bool{}, // False by default.
-		initialized:             &atomic.Bool{}, // False by default.
-		txDecoder:               txDecoder,
+		cdc:                        cdc,
+		storeKey:                   storeKey,
+		memKey:                     memKey,
+		transientStoreKey:          transientStoreKey,
+		authorities:                lib.UniqueSliceToSet(authorities),
+		MemClob:                    memClob,
+		PerpetualIdToClobPairId:    make(map[uint32][]types.ClobPairId),
+		subaccountsKeeper:          subaccountsKeeper,
+		assetsKeeper:               assetsKeeper,
+		blockTimeKeeper:            blockTimeKeeper,
+		bankKeeper:                 bankKeeper,
+		feeTiersKeeper:             feeTiersKeeper,
+		perpetualsKeeper:           perpetualsKeeper,
+		pricesKeeper:               pricesKeeper,
+		statsKeeper:                statsKeeper,
+		rewardsKeeper:              rewardsKeeper,
+		affiliatesKeeper:           affiliatesKeeper,
+		accountPlusKeeper:          accountPlusKeeper,
+		indexerEventManager:        indexerEventManager,
+		streamingManager:           streamingManager,
+		inMemStructuresInitialized: &atomic.Bool{}, // False by default.
+		txDecoder:                  txDecoder,
 		mevTelemetryConfig: MevTelemetryConfig{
 			Enabled:    clobFlags.MevTelemetryEnabled,
 			Hosts:      clobFlags.MevTelemetryHosts,
@@ -180,21 +177,23 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 func (k Keeper) InitializeForGenesis(ctx sdk.Context) {
 }
 
-// IsInitialized returns whether the clob keeper has been hydrated.
-func (k Keeper) IsInitialized() bool {
-	return k.initialized.Load()
+// IsInMemStructuresInitialized returns whether the clob keeper has been hydrated.
+func (k Keeper) IsInMemStructuresInitialized() bool {
+	return k.inMemStructuresInitialized.Load()
 }
 
 // Initialize hydrates the clob keeper with the necessary in memory data structures.
 func (k Keeper) Initialize(ctx sdk.Context) {
-	alreadyInitialized := k.initialized.Swap(true)
+	// Initialize memstore in clobKeeper with order fill amounts and stateful orders.
+	k.InitMemStore(ctx)
+
+	// Code below hydrates the in memory data structures and is not rolled back even if
+	// the block execution is discarded by OE. Therefore, they are only called once.
+	alreadyInitialized := k.inMemStructuresInitialized.Swap(true)
 	if alreadyInitialized {
 		return
 	}
 
-	// Initialize memstore in clobKeeper with order fill amounts and stateful orders.
-	k.InitMemStore(ctx)
-
 	// Branch the context for hydration.
 	// This means that new order matches from hydration will get added to the operations
 	// queue but the corresponding state changes will be discarded.
@@ -254,11 +253,14 @@ func (k Keeper) ProcessStagedFinalizeBlockEvents(ctx sdk.Context) {
 // InitMemStore initializes the memstore of the `clob` keeper.
 // This is called during app initialization in `app.go`, before any ABCI calls are received.
 func (k Keeper) InitMemStore(ctx sdk.Context) {
-	alreadyInitialized := k.memStoreInitialized.Swap(true)
+	alreadyInitialized := k.GetMemstoreInitialized(ctx)
 	if alreadyInitialized {
-		panic(errors.New("Memory store already initialized and is not intended to be invoked more then once."))
+		return
 	}
 
+	// Set memstore initialized flag.
+	k.SetMemstoreInitialized(ctx)
+
 	memStore := ctx.KVStore(k.memKey)
 	memStoreType := memStore.GetStoreType()
 	if memStoreType != storetypes.StoreTypeMemory {
@@ -307,6 +309,19 @@ func (k Keeper) InitMemStore(ctx sdk.Context) {
 	}
 }
 
+func (k Keeper) GetMemstoreInitialized(ctx sdk.Context) bool {
+	store := ctx.KVStore(k.memKey)
+	return store.Has([]byte(types.KeyMemstoreInitialized))
+}
+
+func (k Keeper) SetMemstoreInitialized(ctx sdk.Context) {
+	store := ctx.KVStore(k.memKey)
+	store.Set(
+		[]byte(types.KeyMemstoreInitialized),
+		[]byte{1},
+	)
+}
+
 // Sets the ante handler after it has been constructed. This breaks a cycle between
 // when the ante handler is constructed and when the clob keeper is constructed.
 func (k *Keeper) SetAnteHandler(anteHandler sdk.AnteHandler) {

protocol/x/clob/keeper/keeper.go

@@ -31,10 +31,7 @@ func TestInitMemStore_OnlyAllowedOnce(t *testing.T) {
 
 	ks.ClobKeeper.InitMemStore(ks.Ctx)
 
-	// Initializing a second time causes a panic
-	require.Panics(t, func() {
-		ks.ClobKeeper.InitMemStore(ks.Ctx)
-	})
+	require.True(t, ks.ClobKeeper.GetMemstoreInitialized(ks.Ctx))
 }
 
 func TestInitMemStore_StatefulOrderCount(t *testing.T) {

protocol/x/clob/keeper/keeper_test.go

@@ -13,7 +13,7 @@ type ClobKeeper interface {
 	LiquidationsKeeper
 	LiquidationsConfigKeeper
 
-	IsInitialized() bool
+	IsInMemStructuresInitialized() bool
 	Initialize(ctx sdk.Context)
 
 	AddOrderToOrderbookSubaccountUpdatesCheck(

protocol/x/clob/types/clob_keeper.go

@@ -79,6 +79,9 @@ const (
 
 // Memstore
 const (
+	// KeyMemstoreInitialized is the key to check if the memstore has been initialized.
+	KeyMemstoreInitialized = "MemstoreInit"
+
 	// ProcessProposerMatchesEventsKey is the key to retrieve information about how to update
 	// memclob state based on the latest block.
 	ProcessProposerMatchesEventsKey = "ProposerEvents"