address comments

dydxprotocol · Oct 2, 2024 · ff9eb5f · ff9eb5f
1 parent 01d9cc2
commit ff9eb5f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/protocol/finalizeblock/event_stager.go b/protocol/finalizeblock/event_stager.go
@@ -35,11 +35,14 @@ func NewEventStager[T proto.Message](
 	}
 }
 
-// GetStagedFinalizeBlockEventsFromStore retrieves all staged events from the store.
-func (s EventStager[T]) GetStagedFinalizeBlockEventsFromStore(
-	store storetypes.KVStore,
+// GetStagedFinalizeBlockEvents retrieves all staged events from the store.
+func (s EventStager[T]) GetStagedFinalizeBlockEvents(
+	ctx sdk.Context,
 	newStagedEvent func() T,
 ) []T {
+	noGasCtx := ctx.WithGasMeter(ante_types.NewFreeInfiniteGasMeter())
+	store := noGasCtx.TransientStore(s.transientStoreKey)
+
 	count := s.getStagedEventsCount(store)
 	events := make([]T, count)
 	store = prefix.NewStore(store, []byte(s.stagedEventKeyPrefix))
@@ -65,10 +68,9 @@ func (s EventStager[T]) getStagedEventsCount(
 // StageFinalizeBlockEvent stages an event in the transient store.
 func (s EventStager[T]) StageFinalizeBlockEvent(
 	ctx sdk.Context,
-	eventBytes []byte,
+	stagedEvent T,
 ) {
 	noGasCtx := ctx.WithGasMeter(ante_types.NewFreeInfiniteGasMeter())
-
 	store := noGasCtx.TransientStore(s.transientStoreKey)
 
 	// Increment events count.
@@ -77,5 +79,5 @@ func (s EventStager[T]) StageFinalizeBlockEvent(
 
 	// Store events keyed by index.
 	store = prefix.NewStore(store, []byte(s.stagedEventKeyPrefix))
-	store.Set(lib.Uint32ToKey(count), eventBytes)
+	store.Set(lib.Uint32ToKey(count), s.cdc.MustMarshal(stagedEvent))
 }
diff --git a/protocol/streaming/full_node_streaming_manager.go b/protocol/streaming/full_node_streaming_manager.go
@@ -407,18 +407,16 @@ func (sm *FullNodeStreamingManagerImpl) SendSubaccountUpdate(
 	}
 	sm.finalizeBlockStager.StageFinalizeBlockEvent(
 		ctx,
-		sm.cdc.MustMarshal(&stagedEvent),
+		&stagedEvent,
 	)
 }
 
 // Retrieve all events staged during `FinalizeBlock`.
 func (sm *FullNodeStreamingManagerImpl) GetStagedFinalizeBlockEvents(
 	ctx sdk.Context,
 ) []clobtypes.StagedFinalizeBlockEvent {
-	noGasCtx := ctx.WithGasMeter(ante_types.NewFreeInfiniteGasMeter())
-	store := noGasCtx.TransientStore(sm.streamingManagerTransientStoreKey)
-	events := sm.finalizeBlockStager.GetStagedFinalizeBlockEventsFromStore(
-		store,
+	events := sm.finalizeBlockStager.GetStagedFinalizeBlockEvents(
+		ctx,
 		func() *clobtypes.StagedFinalizeBlockEvent {
 			return &clobtypes.StagedFinalizeBlockEvent{}
 		},
@@ -557,7 +555,7 @@ func (sm *FullNodeStreamingManagerImpl) SendOrderbookUpdates(
 	}
 	sm.finalizeBlockStager.StageFinalizeBlockEvent(
 		ctx,
-		sm.cdc.MustMarshal(&stagedEvent),
+		&stagedEvent,
 	)
 }
 
@@ -632,7 +630,7 @@ func (sm *FullNodeStreamingManagerImpl) SendOrderbookFillUpdate(
 
 	sm.finalizeBlockStager.StageFinalizeBlockEvent(
 		ctx,
-		sm.cdc.MustMarshal(&stagedEvent),
+		&stagedEvent,
 	)
 }