Add a channel buffer to decouple abci and grpc streaming (backport #1530

) (#1595) Co-authored-by: Jonathan Fung <[email protected]> Co-authored-by: Jonathan Fung <[email protected]>
dydxprotocol · May 28, 2024 · f15e98f · f15e98f
1 parent efa59b4
commit f15e98f
Show file tree

Hide file tree

Showing 8 changed files with 222 additions and 113 deletions.
diff --git a/protocol/app/app.go b/protocol/app/app.go
@@ -3,6 +3,7 @@ package app
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"math/big"
 	"net/http"
@@ -425,6 +426,9 @@ func New(
 			if app.Server != nil {
 				app.Server.Stop()
 			}
+			if app.GrpcStreamingManager != nil {
+				app.GrpcStreamingManager.Stop()
+			}
 			return nil
 		},
 	)
@@ -1510,6 +1514,7 @@ func (app *App) EndBlocker(ctx sdk.Context) (sdk.EndBlock, error) {
 	}
 	block := app.IndexerEventManager.ProduceBlock(ctx)
 	app.IndexerEventManager.SendOnchainData(block)
+	app.GrpcStreamingManager.EmitMetrics()
 	return response, err
 }
 
@@ -1527,6 +1532,7 @@ func (app *App) PrepareCheckStater(ctx sdk.Context) {
 	if err := app.ModuleManager.PrepareCheckState(ctx); err != nil {
 		panic(err)
 	}
+	app.GrpcStreamingManager.EmitMetrics()
 }
 
 // InitChainer application update at chain initialization.
@@ -1765,8 +1771,9 @@ func getGrpcStreamingManagerFromOptions(
 	logger log.Logger,
 ) (manager streamingtypes.GrpcStreamingManager) {
 	if appFlags.GrpcStreamingEnabled {
-		logger.Info("GRPC streaming is enabled")
-		return streaming.NewGrpcStreamingManager()
+		grpcStreamingBufferSize := uint32(appFlags.GrpcStreamingBufferSize)
+		logger.Info(fmt.Sprintf("GRPC streaming is enabled with buffer size %d", grpcStreamingBufferSize))
+		return streaming.NewGrpcStreamingManager(logger, grpcStreamingBufferSize)
 	}
 	return streaming.NewNoopGrpcStreamingManager()
 }
diff --git a/protocol/app/flags/flags.go b/protocol/app/flags/flags.go
@@ -21,7 +21,8 @@ type Flags struct {
 	GrpcEnable  bool
 
 	// Grpc Streaming
-	GrpcStreamingEnabled bool
+	GrpcStreamingEnabled    bool
+	GrpcStreamingBufferSize uint16
 }
 
 // List of CLI flags.
@@ -36,7 +37,8 @@ const (
 	GrpcEnable  = "grpc.enable"
 
 	// Grpc Streaming
-	GrpcStreamingEnabled = "grpc-streaming-enabled"
+	GrpcStreamingEnabled    = "grpc-streaming-enabled"
+	GrpcStreamingBufferSize = "grpc-streaming-buffer-size"
 )
 
 // Default values.
@@ -46,7 +48,8 @@ const (
 	DefaultNonValidatingFullNode = false
 	DefaultDdErrorTrackingFormat = false
 
-	DefaultGrpcStreamingEnabled = false
+	DefaultGrpcStreamingEnabled    = false
+	DefaultGrpcStreamingBufferSize = 1000
 )
 
 // AddFlagsToCmd adds flags to app initialization.
@@ -80,6 +83,11 @@ func AddFlagsToCmd(cmd *cobra.Command) {
 		DefaultGrpcStreamingEnabled,
 		"Whether to enable grpc streaming for full nodes",
 	)
+	cmd.Flags().Uint16(
+		GrpcStreamingBufferSize,
+		DefaultGrpcStreamingBufferSize,
+		"Protocol-side buffer channel size to store grpc stream updates before dropping messages",
+	)
 }
 
 // Validate checks that the flags are valid.
@@ -114,7 +122,8 @@ func GetFlagValuesFromOptions(
 		GrpcAddress: config.DefaultGRPCAddress,
 		GrpcEnable:  true,
 
-		GrpcStreamingEnabled: DefaultGrpcStreamingEnabled,
+		GrpcStreamingEnabled:    DefaultGrpcStreamingEnabled,
+		GrpcStreamingBufferSize: DefaultGrpcStreamingBufferSize,
 	}
 
 	// Populate the flags if they exist.
@@ -160,5 +169,11 @@ func GetFlagValuesFromOptions(
 		}
 	}
 
+	if option := appOpts.Get(GrpcStreamingBufferSize); option != nil {
+		if v, err := cast.ToUint16E(option); err == nil {
+			result.GrpcStreamingBufferSize = v
+		}
+	}
+
 	return result
 }
diff --git a/protocol/app/flags/flags_test.go b/protocol/app/flags/flags_test.go
@@ -32,6 +32,9 @@ func TestAddFlagsToCommand(t *testing.T) {
 		fmt.Sprintf("Has %s flag", flags.GrpcStreamingEnabled): {
 			flagName: flags.GrpcStreamingEnabled,
 		},
+		fmt.Sprintf("Has %s flag", flags.GrpcStreamingBufferSize): {
+			flagName: flags.GrpcStreamingBufferSize,
+		},
 	}
 
 	for name, tc := range tests {
@@ -63,9 +66,10 @@ 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,
+				GrpcStreamingBufferSize: 15,
 			},
 		},
 		"failure - gRPC disabled": {
@@ -107,6 +111,7 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 		expectedGrpcAddress               string
 		expectedGrpcEnable                bool
 		expectedGrpcStreamingEnable       bool
+		expectedGrpcStreamingBufferSize   uint16
 	}{
 		"Sets to default if unset": {
 			expectedNonValidatingFullNodeFlag: false,
@@ -115,6 +120,7 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 			expectedGrpcAddress:               "localhost:9090",
 			expectedGrpcEnable:                true,
 			expectedGrpcStreamingEnable:       false,
+			expectedGrpcStreamingBufferSize:   1000,
 		},
 		"Sets values from options": {
 			optsMap: map[string]any{
@@ -124,13 +130,15 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 				flags.GrpcEnable:                false,
 				flags.GrpcAddress:               "localhost:9091",
 				flags.GrpcStreamingEnabled:      "true",
+				flags.GrpcStreamingBufferSize:   "15",
 			},
 			expectedNonValidatingFullNodeFlag: true,
 			expectedDdAgentHost:               "agentHostTest",
 			expectedDdTraceAgentPort:          777,
 			expectedGrpcEnable:                false,
 			expectedGrpcAddress:               "localhost:9091",
 			expectedGrpcStreamingEnable:       true,
+			expectedGrpcStreamingBufferSize:   15,
 		},
 	}
 
@@ -168,6 +176,21 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 				tc.expectedGrpcAddress,
 				flags.GrpcAddress,
 			)
+			require.Equal(
+				t,
+				tc.expectedGrpcAddress,
+				flags.GrpcAddress,
+			)
+			require.Equal(
+				t,
+				tc.expectedGrpcStreamingEnable,
+				flags.GrpcStreamingEnabled,
+			)
+			require.Equal(
+				t,
+				tc.expectedGrpcStreamingBufferSize,
+				flags.GrpcStreamingBufferSize,
+			)
 		})
 	}
 }
diff --git a/protocol/lib/metrics/metric_keys.go b/protocol/lib/metrics/metric_keys.go
@@ -60,6 +60,8 @@ const (
 	FullNodeGrpc                    = "full_node_grpc"
 	GrpcSendOrderbookUpdatesLatency = "grpc_send_orderbook_updates_latency"
 	GrpcSendOrderbookFillsLatency   = "grpc_send_orderbook_fills_latency"
+	GrpcStreamingBufferSize         = "grpc_streaming_buffer_size"
+	GrpcStreamingNumConnections     = "grpc_streaming_num_connections"
 	EndBlocker                      = "end_blocker"
 	EndBlockerLag                   = "end_blocker_lag"
 )