GRPC Streaming Batching

protocol/app/app.go

@@ -458,6 +458,9 @@ func New(
 			if app.SlinkyClient != nil {
 				app.SlinkyClient.Stop()
 			}
+			if app.GrpcStreamingManager != nil {
+				app.GrpcStreamingManager.Stop()
+			}
 			return nil
 		},
 	)
@@ -1931,7 +1934,11 @@ func getGrpcStreamingManagerFromOptions(
 ) (manager streamingtypes.GrpcStreamingManager) {
 	if appFlags.GrpcStreamingEnabled {
 		logger.Info("GRPC streaming is enabled")
-		return streaming.NewGrpcStreamingManager(logger)
+		return streaming.NewGrpcStreamingManager(
+			logger,
+			appFlags.GrpcStreamingFlushIntervalMs,
+			appFlags.GrpcStreamingMaxBufferSize,
+		)
 	}
 	return streaming.NewNoopGrpcStreamingManager()
 }

protocol/app/flags/flags.go

@@ -21,8 +21,11 @@ type Flags struct {
 	GrpcEnable  bool
 
 	// Grpc Streaming
-	GrpcStreamingEnabled bool
-	VEOracleEnabled      bool // Slinky Vote Extensions
+	GrpcStreamingEnabled         bool
+	GrpcStreamingFlushIntervalMs uint32
+	GrpcStreamingMaxBufferSize   uint32
+
+	VEOracleEnabled bool // Slinky Vote Extensions
 }
 
 // List of CLI flags.
@@ -37,7 +40,9 @@ const (
 	GrpcEnable  = "grpc.enable"
 
 	// Grpc Streaming
-	GrpcStreamingEnabled = "grpc-streaming-enabled"
+	GrpcStreamingEnabled         = "grpc-streaming-enabled"
+	GrpcStreamingFlushIntervalMs = "grpc-streaming-flush-interval-ms"
+	GrpcStreamingMaxBufferSize   = "grpc-streaming-max-buffer-size"
 
 	// Slinky VEs enabled
 	VEOracleEnabled = "slinky-vote-extension-oracle-enabled"
@@ -50,8 +55,11 @@ const (
 	DefaultNonValidatingFullNode = false
 	DefaultDdErrorTrackingFormat = false
 
-	DefaultGrpcStreamingEnabled = false
-	DefaultVEOracleEnabled      = true
+	DefaultGrpcStreamingEnabled         = false
+	DefaultGrpcStreamingFlushIntervalMs = 50
+	DefaultGrpcStreamingMaxBufferSize   = 10000
+
+	DefaultVEOracleEnabled = true
 )
 
 // AddFlagsToCmd adds flags to app initialization.
@@ -85,6 +93,16 @@ func AddFlagsToCmd(cmd *cobra.Command) {
 		DefaultGrpcStreamingEnabled,
 		"Whether to enable grpc streaming for full nodes",
 	)
+	cmd.Flags().Uint32(
+		GrpcStreamingFlushIntervalMs,
+		DefaultGrpcStreamingFlushIntervalMs,
+		"Flush interval (in ms) for grpc streaming",
+	)
+	cmd.Flags().Uint32(
+		GrpcStreamingMaxBufferSize,
+		DefaultGrpcStreamingMaxBufferSize,
+		"Maximum buffer size before grpc streaming cancels all connections",
+	)
 	cmd.Flags().Bool(
 		VEOracleEnabled,
 		DefaultVEOracleEnabled,
@@ -104,6 +122,12 @@ func (f *Flags) Validate() error {
 		if !f.GrpcEnable {
 			return fmt.Errorf("grpc.enable must be set to true - grpc streaming requires gRPC server")
 		}
+		if f.GrpcStreamingMaxBufferSize == 0 {
+			return fmt.Errorf("grpc streaming buffer size must be positive number")
+		}
+		if f.GrpcStreamingFlushIntervalMs == 0 {
+			return fmt.Errorf("grpc streaming flush interval must be positive number")
+		}
 	}
 	return nil
 }
@@ -124,8 +148,11 @@ func GetFlagValuesFromOptions(
 		GrpcAddress: config.DefaultGRPCAddress,
 		GrpcEnable:  true,
 
-		GrpcStreamingEnabled: DefaultGrpcStreamingEnabled,
-		VEOracleEnabled:      true,
+		GrpcStreamingEnabled:         DefaultGrpcStreamingEnabled,
+		GrpcStreamingFlushIntervalMs: DefaultGrpcStreamingFlushIntervalMs,
+		GrpcStreamingMaxBufferSize:   DefaultGrpcStreamingMaxBufferSize,
+
+		VEOracleEnabled: true,
 	}
 
 	// Populate the flags if they exist.
@@ -171,6 +198,18 @@ func GetFlagValuesFromOptions(
 		}
 	}
 
+	if option := appOpts.Get(GrpcStreamingFlushIntervalMs); option != nil {
+		if v, err := cast.ToUint32E(option); err == nil {
+			result.GrpcStreamingFlushIntervalMs = v
+		}
+	}
+
+	if option := appOpts.Get(GrpcStreamingMaxBufferSize); option != nil {
+		if v, err := cast.ToUint32E(option); err == nil {
+			result.GrpcStreamingMaxBufferSize = v
+		}
+	}
+
 	if option := appOpts.Get(VEOracleEnabled); option != nil {
 		if v, err := cast.ToBoolE(option); err == nil {
 			result.VEOracleEnabled = v

protocol/app/flags/flags_test.go

@@ -32,6 +32,12 @@ func TestAddFlagsToCommand(t *testing.T) {
 		fmt.Sprintf("Has %s flag", flags.GrpcStreamingEnabled): {
 			flagName: flags.GrpcStreamingEnabled,
 		},
+		fmt.Sprintf("Has %s flag", flags.GrpcStreamingFlushIntervalMs): {
+			flagName: flags.GrpcStreamingFlushIntervalMs,
+		},
+		fmt.Sprintf("Has %s flag", flags.GrpcStreamingMaxBufferSize): {
+			flagName: flags.GrpcStreamingMaxBufferSize,
+		},
 	}
 
 	for name, tc := range tests {
@@ -63,9 +69,11 @@ func TestValidate(t *testing.T) {
 		},
 		"success - gRPC streaming enabled for validating nodes": {
 			flags: flags.Flags{
-				NonValidatingFullNode: false,
-				GrpcEnable:            true,
-				GrpcStreamingEnabled:  true,
+				NonValidatingFullNode:        false,
+				GrpcEnable:                   true,
+				GrpcStreamingEnabled:         true,
+				GrpcStreamingFlushIntervalMs: 100,
+				GrpcStreamingMaxBufferSize:   10000,
 			},
 		},
 		"failure - gRPC disabled": {
@@ -82,6 +90,26 @@ func TestValidate(t *testing.T) {
 			},
 			expectedErr: fmt.Errorf("grpc.enable must be set to true - grpc streaming requires gRPC server"),
 		},
+		"failure - gRPC streaming enabled with zero buffer size": {
+			flags: flags.Flags{
+				NonValidatingFullNode:        true,
+				GrpcEnable:                   true,
+				GrpcStreamingEnabled:         true,
+				GrpcStreamingFlushIntervalMs: 100,
+				GrpcStreamingMaxBufferSize:   0,
+			},
+			expectedErr: fmt.Errorf("grpc streaming buffer size must be positive number"),
+		},
+		"failure - gRPC streaming enabled with zero flush interval ms": {
+			flags: flags.Flags{
+				NonValidatingFullNode:        true,
+				GrpcEnable:                   true,
+				GrpcStreamingEnabled:         true,
+				GrpcStreamingFlushIntervalMs: 0,
+				GrpcStreamingMaxBufferSize:   10000,
+			},
+			expectedErr: fmt.Errorf("grpc streaming flush interval must be positive number"),
+		},
 	}
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
@@ -107,6 +135,8 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 		expectedGrpcAddress               string
 		expectedGrpcEnable                bool
 		expectedGrpcStreamingEnable       bool
+		expectedGrpcStreamingFlushMs      uint32
+		expectedGrpcStreamingBufferSize   uint32
 	}{
 		"Sets to default if unset": {
 			expectedNonValidatingFullNodeFlag: false,
@@ -115,22 +145,28 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 			expectedGrpcAddress:               "localhost:9090",
 			expectedGrpcEnable:                true,
 			expectedGrpcStreamingEnable:       false,
+			expectedGrpcStreamingFlushMs:      50,
+			expectedGrpcStreamingBufferSize:   10000,
 		},
 		"Sets values from options": {
 			optsMap: map[string]any{
-				flags.NonValidatingFullNodeFlag: true,
-				flags.DdAgentHost:               "agentHostTest",
-				flags.DdTraceAgentPort:          uint16(777),
-				flags.GrpcEnable:                false,
-				flags.GrpcAddress:               "localhost:9091",
-				flags.GrpcStreamingEnabled:      "true",
+				flags.NonValidatingFullNodeFlag:    true,
+				flags.DdAgentHost:                  "agentHostTest",
+				flags.DdTraceAgentPort:             uint16(777),
+				flags.GrpcEnable:                   false,
+				flags.GrpcAddress:                  "localhost:9091",
+				flags.GrpcStreamingEnabled:         "true",
+				flags.GrpcStreamingFlushIntervalMs: uint32(408),
+				flags.GrpcStreamingMaxBufferSize:   uint32(650),
 			},
 			expectedNonValidatingFullNodeFlag: true,
 			expectedDdAgentHost:               "agentHostTest",
 			expectedDdTraceAgentPort:          777,
 			expectedGrpcEnable:                false,
 			expectedGrpcAddress:               "localhost:9091",
 			expectedGrpcStreamingEnable:       true,
+			expectedGrpcStreamingFlushMs:      408,
+			expectedGrpcStreamingBufferSize:   650,
 		},
 	}
 
@@ -168,6 +204,21 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 				tc.expectedGrpcAddress,
 				flags.GrpcAddress,
 			)
+			require.Equal(
+				t,
+				tc.expectedGrpcStreamingEnable,
+				flags.GrpcStreamingEnabled,
+			)
+			require.Equal(
+				t,
+				tc.expectedGrpcStreamingFlushMs,
+				flags.GrpcStreamingFlushIntervalMs,
+			)
+			require.Equal(
+				t,
+				tc.expectedGrpcStreamingBufferSize,
+				flags.GrpcStreamingMaxBufferSize,
+			)
 		})
 	}
 }

protocol/streaming/grpc/grpc_streaming_manager.go

@@ -25,6 +25,15 @@ type GrpcStreamingManagerImpl struct {
 	// orderbookSubscriptions maps subscription IDs to their respective orderbook subscriptions.
 	orderbookSubscriptions map[uint32]*OrderbookSubscription
 	nextSubscriptionId     uint32
+
+	// grpc stream will batch and flush out messages every 10 ms.
+	ticker *time.Ticker
+	done   chan bool
+	// map of clob pair id to stream updates.
+	streamUpdateCache map[uint32][]clobtypes.StreamUpdate
+	numUpdatesInCache uint32
+
+	maxUpdatesInCache uint32
 }
 
 // OrderbookSubscription represents a active subscription to the orderbook updates stream.
@@ -41,12 +50,36 @@ type OrderbookSubscription struct {
 
 func NewGrpcStreamingManager(
 	logger log.Logger,
+	flushIntervalMs uint32,
+	maxUpdatesInCache uint32,
 ) *GrpcStreamingManagerImpl {
 	logger = logger.With(log.ModuleKey, "grpc-streaming")
-	return &GrpcStreamingManagerImpl{
+	grpcStreamingManager := &GrpcStreamingManagerImpl{
 		logger:                 logger,
 		orderbookSubscriptions: make(map[uint32]*OrderbookSubscription),
+		nextSubscriptionId:     0,
+
+		ticker:            time.NewTicker(time.Duration(flushIntervalMs) * time.Millisecond),
+		done:              make(chan bool),
+		streamUpdateCache: make(map[uint32][]clobtypes.StreamUpdate),
+		numUpdatesInCache: 0,
+
+		maxUpdatesInCache: maxUpdatesInCache,
 	}
+
+	// Start the goroutine for pushing order updates through
+	go func() {
+		for {
+			select {
+			case <-grpcStreamingManager.ticker.C:
+				grpcStreamingManager.FlushStreamUpdates()
+			case <-grpcStreamingManager.done:
+				return
+			}
+		}
+	}()
+
+	return grpcStreamingManager
 }
 
 func (sm *GrpcStreamingManagerImpl) Enabled() bool {
@@ -88,6 +121,10 @@ func (sm *GrpcStreamingManagerImpl) Subscribe(
 	return nil
 }
 
+func (sm *GrpcStreamingManagerImpl) Stop() {
+	sm.done <- true
+}
+
 // SendOrderbookUpdates groups updates by their clob pair ids and
 // sends messages to the subscribers.
 func (sm *GrpcStreamingManagerImpl) SendOrderbookUpdates(
@@ -133,9 +170,7 @@ func (sm *GrpcStreamingManagerImpl) SendOrderbookUpdates(
 		}
 	}
 
-	sm.sendStreamUpdate(
-		updatesByClobPairId,
-	)
+	sm.AddUpdatesToCache(updatesByClobPairId, uint32(len(updates)))
 }
 
 // SendOrderbookFillUpdates groups fills by their clob pair ids and
@@ -172,29 +207,49 @@ func (sm *GrpcStreamingManagerImpl) SendOrderbookFillUpdates(
 		updatesByClobPairId[clobPairId] = append(updatesByClobPairId[clobPairId], streamUpdate)
 	}
 
-	sm.sendStreamUpdate(
-		updatesByClobPairId,
-	)
+	sm.AddUpdatesToCache(updatesByClobPairId, uint32(len(orderbookFills)))
 }
 
-// sendStreamUpdate takes in a map of clob pair id to stream updates and emits them to subscribers.
-func (sm *GrpcStreamingManagerImpl) sendStreamUpdate(
+func (sm *GrpcStreamingManagerImpl) AddUpdatesToCache(
 	updatesByClobPairId map[uint32][]clobtypes.StreamUpdate,
+	numUpdatesToAdd uint32,
 ) {
+	sm.Lock()
+	defer sm.Unlock()
+
+	for clobPairId, streamUpdates := range updatesByClobPairId {
+		sm.streamUpdateCache[clobPairId] = append(sm.streamUpdateCache[clobPairId], streamUpdates...)
+	}
+	sm.numUpdatesInCache += numUpdatesToAdd
+
+	// Remove all subscriptions and wipe the buffer if buffer overflows.
+	if sm.numUpdatesInCache > sm.maxUpdatesInCache {
+		sm.logger.Error("GRPC Streaming buffer full capacity. Dropping messages and all subscriptions. " +
+			"Disconnect all clients and increase buffer size via the grpc-stream-buffer-size flag.")
+		for id := range sm.orderbookSubscriptions {
+			delete(sm.orderbookSubscriptions, id)
+		}
+		clear(sm.streamUpdateCache)
+		sm.numUpdatesInCache = 0
+	}
+}
+
+// FlushStreamUpdates takes in a map of clob pair id to stream updates and emits them to subscribers.
+func (sm *GrpcStreamingManagerImpl) FlushStreamUpdates() {
+	sm.Lock()
+	defer sm.Unlock()
+
 	metrics.IncrCounter(
 		metrics.GrpcEmitProtocolUpdateCount,
 		1,
 	)
 
-	sm.Lock()
-	defer sm.Unlock()
-
 	// Send updates to subscribers.
 	idsToRemove := make([]uint32, 0)
 	for id, subscription := range sm.orderbookSubscriptions {
 		streamUpdatesForSubscription := make([]clobtypes.StreamUpdate, 0)
 		for _, clobPairId := range subscription.clobPairIds {
-			if update, ok := updatesByClobPairId[clobPairId]; ok {
+			if update, ok := sm.streamUpdateCache[clobPairId]; ok {
 				streamUpdatesForSubscription = append(streamUpdatesForSubscription, update...)
 			}
 		}
@@ -209,6 +264,7 @@ func (sm *GrpcStreamingManagerImpl) sendStreamUpdate(
 					Updates: streamUpdatesForSubscription,
 				},
 			); err != nil {
+				sm.logger.Error("Error sending out update", "err", err)
 				idsToRemove = append(idsToRemove, id)
 			}
 		}
@@ -219,6 +275,10 @@ func (sm *GrpcStreamingManagerImpl) sendStreamUpdate(
 	for _, id := range idsToRemove {
 		delete(sm.orderbookSubscriptions, id)
 	}
+
+	clear(sm.streamUpdateCache)
+	sm.numUpdatesInCache = 0
+
 	sm.EmitMetrics()
 }