Revert "merge finalized subaccount updates #2230" (#2348)

Author: jonfung-dydx

protocol/lib/metrics/metric_keys.go

@@ -66,19 +66,19 @@ const (
 	GateWithdrawalsIfNegativeTncSubaccountSeenLatency = "gate_withdrawals_if_negative_tnc_subaccount_seen_latency"
 
 	// Full node grpc
-	FullNodeGrpc                              = "full_node_grpc"
-	GrpcSendOrderbookUpdatesLatency           = "grpc_send_orderbook_updates_latency"
-	GrpcSendOrderbookSnapshotLatency          = "grpc_send_orderbook_snapshot_latency"
-	GrpcSendSubaccountSnapshotLatency         = "grpc_send_subaccount_snapshot_latency"
-	GrpcSendOrderbookFillsLatency             = "grpc_send_orderbook_fills_latency"
-	GrpcSendFinalizedSubaccountUpdatesLatency = "grpc_send_finalized_subaccount_updates_latency"
-	GrpcAddUpdateToBufferCount                = "grpc_add_update_to_buffer_count"
-	GrpcAddToSubscriptionChannelCount         = "grpc_add_to_subscription_channel_count"
-	GrpcSendResponseToSubscriberCount         = "grpc_send_response_to_subscriber_count"
-	GrpcStreamSubscriberCount                 = "grpc_stream_subscriber_count"
-	GrpcStreamNumUpdatesBuffered              = "grpc_stream_num_updates_buffered"
-	GrpcFlushUpdatesLatency                   = "grpc_flush_updates_latency"
-	GrpcSubscriptionChannelLength             = "grpc_subscription_channel_length"
+	FullNodeGrpc                      = "full_node_grpc"
+	GrpcSendOrderbookUpdatesLatency   = "grpc_send_orderbook_updates_latency"
+	GrpcSendOrderbookSnapshotLatency  = "grpc_send_orderbook_snapshot_latency"
+	GrpcSendSubaccountSnapshotLatency = "grpc_send_subaccount_snapshot_latency"
+	GrpcSendOrderbookFillsLatency     = "grpc_send_orderbook_fills_latency"
+	GrpcSendSubaccountUpdatesLatency  = "grpc_send_subaccount_updates_latency"
+	GrpcAddUpdateToBufferCount        = "grpc_add_update_to_buffer_count"
+	GrpcAddToSubscriptionChannelCount = "grpc_add_to_subscription_channel_count"
+	GrpcSendResponseToSubscriberCount = "grpc_send_response_to_subscriber_count"
+	GrpcStreamSubscriberCount         = "grpc_stream_subscriber_count"
+	GrpcStreamNumUpdatesBuffered      = "grpc_stream_num_updates_buffered"
+	GrpcFlushUpdatesLatency           = "grpc_flush_updates_latency"
+	GrpcSubscriptionChannelLength     = "grpc_subscription_channel_length"
 
 	EndBlocker    = "end_blocker"
 	EndBlockerLag = "end_blocker_lag"

protocol/mocks/ClobKeeper.go

@@ -3,7 +3,6 @@ package streaming
 import (
 	"fmt"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
@@ -56,8 +55,8 @@ type FullNodeStreamingManagerImpl struct {
 type OrderbookSubscription struct {
 	subscriptionId uint32
 
-	// Whether the subscription is initialized with snapshot.
-	initialized *atomic.Bool
+	// Initialize the subscription with orderbook snapshots.
+	initialize *sync.Once
 
 	// Clob pair ids to subscribe to.
 	clobPairIds []uint32
@@ -76,10 +75,6 @@ type OrderbookSubscription struct {
 	nextSnapshotBlock uint32
 }
 
-func (sub *OrderbookSubscription) IsInitialized() bool {
-	return sub.initialized.Load()
-}
-
 func NewFullNodeStreamingManager(
 	logger log.Logger,
 	flushIntervalMs uint32,
@@ -164,7 +159,7 @@ func (sm *FullNodeStreamingManagerImpl) Subscribe(
 	}
 	subscription := &OrderbookSubscription{
 		subscriptionId: sm.nextSubscriptionId,
-		initialized:    &atomic.Bool{}, // False by default.
+		initialize:     &sync.Once{},
 		clobPairIds:    clobPairIds,
 		subaccountIds:  sIds,
 		messageSender:  messageSender,
@@ -516,23 +511,19 @@ func (sm *FullNodeStreamingManagerImpl) SendTakerOrderStatus(
 	)
 }
 
-// SendFinalizedSubaccountUpdates groups subaccount updates by their subaccount ids and
+// SendSubaccountUpdates groups subaccount updates by their subaccount ids and
 // sends messages to the subscribers.
-func (sm *FullNodeStreamingManagerImpl) SendFinalizedSubaccountUpdates(
+func (sm *FullNodeStreamingManagerImpl) SendSubaccountUpdates(
 	subaccountUpdates []satypes.StreamSubaccountUpdate,
 	blockHeight uint32,
 	execMode sdk.ExecMode,
 ) {
 	defer metrics.ModuleMeasureSince(
 		metrics.FullNodeGrpc,
-		metrics.GrpcSendFinalizedSubaccountUpdatesLatency,
+		metrics.GrpcSendSubaccountUpdatesLatency,
 		time.Now(),
 	)
 
-	if execMode != sdk.ExecModeFinalize {
-		panic("SendFinalizedSubaccountUpdates should only be called in ExecModeFinalize")
-	}
-
 	// Group subaccount updates by subaccount id.
 	streamUpdates := make([]clobtypes.StreamUpdate, 0)
 	subaccountIds := make([]*satypes.SubaccountId, 0)
@@ -679,33 +670,9 @@ func (sm *FullNodeStreamingManagerImpl) FlushStreamUpdatesWithLock() {
 	sm.EmitMetrics()
 }
 
-func (sm *FullNodeStreamingManagerImpl) GetSubaccountSnapshotsForInitStreams(
-	getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
-) map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate {
-	sm.Lock()
-	defer sm.Unlock()
-
-	ret := make(map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate)
-	for _, subscription := range sm.orderbookSubscriptions {
-		// If the subscription has been initialized, no need to grab the subaccount snapshot.
-		if alreadyInitialized := subscription.initialized.Load(); alreadyInitialized {
-			continue
-		}
-
-		for _, subaccountId := range subscription.subaccountIds {
-			if _, exists := ret[subaccountId]; exists {
-				continue
-			}
-
-			ret[subaccountId] = getSubaccountSnapshot(subaccountId)
-		}
-	}
-	return ret
-}
-
 func (sm *FullNodeStreamingManagerImpl) InitializeNewStreams(
 	getOrderbookSnapshot func(clobPairId clobtypes.ClobPairId) *clobtypes.OffchainUpdates,
-	subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate,
+	getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
 	blockHeight uint32,
 	execMode sdk.ExecMode,
 ) {
@@ -719,40 +686,31 @@ func (sm *FullNodeStreamingManagerImpl) InitializeNewStreams(
 	updatesByClobPairId := make(map[uint32]*clobtypes.OffchainUpdates)
 
 	for subscriptionId, subscription := range sm.orderbookSubscriptions {
-		if alreadyInitialized := subscription.initialized.Swap(true); !alreadyInitialized {
-			allUpdates := clobtypes.NewOffchainUpdates()
-			for _, clobPairId := range subscription.clobPairIds {
-				if _, ok := updatesByClobPairId[clobPairId]; !ok {
-					updatesByClobPairId[clobPairId] = getOrderbookSnapshot(clobtypes.ClobPairId(clobPairId))
-				}
-				allUpdates.Append(updatesByClobPairId[clobPairId])
-			}
-
-			saUpdates := []*satypes.StreamSubaccountUpdate{}
-			for _, subaccountId := range subscription.subaccountIds {
-				// The subaccount snapshot may not exist due to the following race condition
-				// 1. At beginning of PrepareCheckState we get snapshot for all subscribed subaccounts.
-				// 2. A new subaccount is subscribed to by a new subscription.
-				// 3. InitializeNewStreams is called.
-				// Then the new subaccount would not be included in the snapshot.
-				// We are okay with this behavior.
-				if saUpdate, ok := subaccountSnapshots[subaccountId]; ok {
-					saUpdates = append(saUpdates, saUpdate)
-				}
-			}
-
-			sm.SendCombinedSnapshot(allUpdates, saUpdates, subscriptionId, blockHeight, execMode)
-
-			if sm.snapshotBlockInterval != 0 {
-				subscription.nextSnapshotBlock = blockHeight + sm.snapshotBlockInterval
-			}
-		}
-
-		// If the snapshot block interval is enabled and the next block is a snapshot block,
-		// reset the `atomic.Bool` so snapshots are sent for the next block.
+		// If the snapshot block interval is enabled, reset the sync.Once in order to
+		// re-send snapshots out.
 		if sm.snapshotBlockInterval > 0 &&
-			blockHeight+1 == subscription.nextSnapshotBlock {
-			subscription.initialized = &atomic.Bool{} // False by default.
+			blockHeight == subscription.nextSnapshotBlock {
+			subscription.initialize = &sync.Once{}
 		}
+
+		subscription.initialize.Do(
+			func() {
+				allUpdates := clobtypes.NewOffchainUpdates()
+				for _, clobPairId := range subscription.clobPairIds {
+					if _, ok := updatesByClobPairId[clobPairId]; !ok {
+						updatesByClobPairId[clobPairId] = getOrderbookSnapshot(clobtypes.ClobPairId(clobPairId))
+					}
+					allUpdates.Append(updatesByClobPairId[clobPairId])
+				}
+				saUpdates := []*satypes.StreamSubaccountUpdate{}
+				for _, subaccountId := range subscription.subaccountIds {
+					saUpdates = append(saUpdates, getSubaccountSnapshot(subaccountId))
+				}
+				sm.SendCombinedSnapshot(allUpdates, saUpdates, subscriptionId, blockHeight, execMode)
+				if sm.snapshotBlockInterval != 0 {
+					subscription.nextSnapshotBlock = blockHeight + sm.snapshotBlockInterval
+				}
+			},
+		)
 	}
 }

protocol/streaming/full_node_streaming_manager.go

@@ -51,7 +51,7 @@ func (sm *NoopGrpcStreamingManager) SendTakerOrderStatus(
 ) {
 }
 
-func (sm *NoopGrpcStreamingManager) SendFinalizedSubaccountUpdates(
+func (sm *NoopGrpcStreamingManager) SendSubaccountUpdates(
 	subaccountUpdates []satypes.StreamSubaccountUpdate,
 	blockHeight uint32,
 	execMode sdk.ExecMode,
@@ -62,15 +62,9 @@ func (sm *NoopGrpcStreamingManager) TracksSubaccountId(id satypes.SubaccountId)
 	return false
 }
 
-func (sm *NoopGrpcStreamingManager) GetSubaccountSnapshotsForInitStreams(
-	getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
-) map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate {
-	return nil
-}
-
 func (sm *NoopGrpcStreamingManager) InitializeNewStreams(
 	getOrderbookSnapshot func(clobPairId clobtypes.ClobPairId) *clobtypes.OffchainUpdates,
-	subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate,
+	getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
 	blockHeight uint32,
 	execMode sdk.ExecMode,
 ) {

protocol/streaming/noop_streaming_manager.go

@@ -22,13 +22,10 @@ type FullNodeStreamingManager interface {
 	// L3+ Orderbook updates.
 	InitializeNewStreams(
 		getOrderbookSnapshot func(clobPairId clobtypes.ClobPairId) *clobtypes.OffchainUpdates,
-		subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate,
+		getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
 		blockHeight uint32,
 		execMode sdk.ExecMode,
 	)
-	GetSubaccountSnapshotsForInitStreams(
-		getSubaccountSnapshot func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate,
-	) map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate
 	SendOrderbookUpdates(
 		offchainUpdates *clobtypes.OffchainUpdates,
 		blockHeight uint32,
@@ -45,7 +42,7 @@ type FullNodeStreamingManager interface {
 		blockHeight uint32,
 		execMode sdk.ExecMode,
 	)
-	SendFinalizedSubaccountUpdates(
+	SendSubaccountUpdates(
 		subaccountUpdates []satypes.StreamSubaccountUpdate,
 		blockHeight uint32,
 		execMode sdk.ExecMode,

protocol/streaming/types/interface.go

@@ -14,7 +14,6 @@ import (
 	"github.com/dydxprotocol/v4-chain/protocol/lib/metrics"
 	"github.com/dydxprotocol/v4-chain/protocol/x/clob/keeper"
 	"github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
-	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
 )
 
 // PreBlocker executes all ABCI PreBlock logic respective to the clob module.
@@ -147,15 +146,6 @@ func PrepareCheckState(
 		log.BlockHeight, ctx.BlockHeight()+1,
 	)
 
-	// We just committed block `h`, preparing `CheckState` of `h+1`
-	// Before we modify the `CheckState`, we first take the snapshot of
-	// the subscribed subaccounts at the end of block `h`. This we send finalized state of
-	// the subaccounts below in `InitializeNewStreams`.
-	var subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate
-	if keeper.GetFullNodeStreamingManager().Enabled() {
-		subaccountSnapshots = keeper.GetSubaccountSnapshotsForInitStreams(ctx)
-	}
-
 	// Get the events generated from processing the matches in the latest block.
 	processProposerMatchesEvents := keeper.GetProcessProposerMatchesEvents(ctx)
 	if ctx.BlockHeight() != int64(processProposerMatchesEvents.BlockHeight) {
@@ -269,11 +259,7 @@ func PrepareCheckState(
 	)
 
 	// Initialize new streams with orderbook snapshots, if any.
-	keeper.InitializeNewStreams(
-		ctx,
-		// Use the subaccount snapshot at the top of function to initialize the streams.
-		subaccountSnapshots,
-	)
+	keeper.InitializeNewStreams(ctx)
 
 	// Set per-orderbook gauges.
 	keeper.MemClob.SetMemclobGauges(ctx)

protocol/x/clob/abci.go

@@ -3,9 +3,8 @@ package keeper
 import (
 	"errors"
 	"fmt"
-	"sync/atomic"
-
 	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+	"sync/atomic"
 
 	"cosmossdk.io/log"
 	"cosmossdk.io/store/prefix"
@@ -257,31 +256,9 @@ func (k *Keeper) SetAnteHandler(anteHandler sdk.AnteHandler) {
 	k.antehandler = anteHandler
 }
 
-func (k Keeper) GetSubaccountSnapshotsForInitStreams(
-	ctx sdk.Context,
-) (
-	subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate,
-) {
-	lib.AssertCheckTxMode(ctx)
-
-	return k.GetFullNodeStreamingManager().GetSubaccountSnapshotsForInitStreams(
-		func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate {
-			subaccountUpdate := k.subaccountsKeeper.GetStreamSubaccountUpdate(
-				ctx,
-				subaccountId,
-				true,
-			)
-			return &subaccountUpdate
-		},
-	)
-}
-
 // InitializeNewStreams initializes new streams for all uninitialized clob pairs
 // by sending the corresponding orderbook snapshots.
-func (k Keeper) InitializeNewStreams(
-	ctx sdk.Context,
-	subaccountSnapshots map[satypes.SubaccountId]*satypes.StreamSubaccountUpdate,
-) {
+func (k Keeper) InitializeNewStreams(ctx sdk.Context) {
 	streamingManager := k.GetFullNodeStreamingManager()
 
 	streamingManager.InitializeNewStreams(
@@ -291,7 +268,14 @@ func (k Keeper) InitializeNewStreams(
 				clobPairId,
 			)
 		},
-		subaccountSnapshots,
+		func(subaccountId satypes.SubaccountId) *satypes.StreamSubaccountUpdate {
+			subaccountUpdate := k.subaccountsKeeper.GetStreamSubaccountUpdate(
+				ctx,
+				subaccountId,
+				true,
+			)
+			return &subaccountUpdate
+		},
 		lib.MustConvertIntegerToUint32(ctx.BlockHeight()),
 		ctx.ExecMode(),
 	)