batch updates + flush out changes at EndBlocker and PrepareCheckState…

… and DeliverTx
dydxprotocol · May 13, 2024 · f50ea03 · f50ea03
1 parent ba9e7f3
commit f50ea03
Show file tree

Hide file tree

Showing 8 changed files with 153 additions and 63 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,10 @@ func (app *App) EndBlocker(ctx sdk.Context) (sdk.EndBlock, error) {
 	}
 	block := app.IndexerEventManager.ProduceBlock(ctx)
 	app.IndexerEventManager.SendOnchainData(block)
+
+	if app.GrpcStreamingManager.Enabled() {
+		app.GrpcStreamingManager.FlushStreamUpdates(uint32(ctx.BlockHeight()), ctx.ExecMode())
+	}
 	return response, err
 }
 
@@ -1527,6 +1535,10 @@ func (app *App) PrepareCheckStater(ctx sdk.Context) {
 	if err := app.ModuleManager.PrepareCheckState(ctx); err != nil {
 		panic(err)
 	}
+
+	if app.GrpcStreamingManager.Enabled() {
+		app.GrpcStreamingManager.FlushStreamUpdates(uint32(ctx.BlockHeight()), ctx.ExecMode())
+	}
 }
 
 // InitChainer application update at chain initialization.
@@ -1765,8 +1777,14 @@ func getGrpcStreamingManagerFromOptions(
 	logger log.Logger,
 ) (manager streamingtypes.GrpcStreamingManager) {
 	if appFlags.GrpcStreamingEnabled {
-		logger.Info("GRPC streaming is enabled")
-		return streaming.NewGrpcStreamingManager()
+		flushIntervalMs := uint32(appFlags.GrpcStreamingFlushIntervalMs)
+		logger.Info(
+			fmt.Sprintf(
+				"GRPC streaming is enabled with flush interval %+v",
+				flushIntervalMs,
+			),
+		)
+		return streaming.NewGrpcStreamingManager(uint32(appFlags.GrpcStreamingFlushIntervalMs))
 	}
 	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
+	GrpcStreamingFlushIntervalMs uint16
 }
 
 // List of CLI flags.
@@ -36,7 +37,8 @@ const (
 	GrpcEnable  = "grpc.enable"
 
 	// Grpc Streaming
-	GrpcStreamingEnabled = "grpc-streaming-enabled"
+	GrpcStreamingEnabled         = "grpc-streaming-enabled"
+	GrpcStreamingFlushIntervalMs = "grpc-streaming-flush-interval-ms"
 )
 
 // Default values.
@@ -46,7 +48,8 @@ const (
 	DefaultNonValidatingFullNode = false
 	DefaultDdErrorTrackingFormat = false
 
-	DefaultGrpcStreamingEnabled = false
+	DefaultGrpcStreamingEnabled         = false
+	DefaultGrpcStreamingFlushIntervalMs = 10
 )
 
 // 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(
+		GrpcStreamingFlushIntervalMs,
+		DefaultGrpcStreamingFlushIntervalMs,
+		"Interval on which to flush grpc stream updates",
+	)
 }
 
 // Validate checks that the flags are valid.
@@ -114,7 +122,8 @@ func GetFlagValuesFromOptions(
 		GrpcAddress: config.DefaultGRPCAddress,
 		GrpcEnable:  true,
 
-		GrpcStreamingEnabled: DefaultGrpcStreamingEnabled,
+		GrpcStreamingEnabled:         DefaultGrpcStreamingEnabled,
+		GrpcStreamingFlushIntervalMs: DefaultGrpcStreamingFlushIntervalMs,
 	}
 
 	// Populate the flags if they exist.
@@ -160,5 +169,11 @@ func GetFlagValuesFromOptions(
 		}
 	}
 
+	if option := appOpts.Get(GrpcStreamingFlushIntervalMs); option != nil {
+		if v, err := cast.ToUint16E(option); err == nil {
+			result.GrpcStreamingFlushIntervalMs = 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.GrpcStreamingFlushIntervalMs): {
+			flagName: flags.GrpcStreamingFlushIntervalMs,
+		},
 	}
 
 	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,
+				GrpcStreamingFlushIntervalMs: 15,
 			},
 		},
 		"failure - gRPC disabled": {
@@ -107,6 +111,7 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 		expectedGrpcAddress               string
 		expectedGrpcEnable                bool
 		expectedGrpcStreamingEnable       bool
+		expectedGrpcStreamingFlushMs      uint16
 	}{
 		"Sets to default if unset": {
 			expectedNonValidatingFullNodeFlag: false,
@@ -115,22 +120,25 @@ func TestGetFlagValuesFromOptions(t *testing.T) {
 			expectedGrpcAddress:               "localhost:9090",
 			expectedGrpcEnable:                true,
 			expectedGrpcStreamingEnable:       false,
+			expectedGrpcStreamingFlushMs:      10,
 		},
 		"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: "15",
 			},
 			expectedNonValidatingFullNodeFlag: true,
 			expectedDdAgentHost:               "agentHostTest",
 			expectedDdTraceAgentPort:          777,
 			expectedGrpcEnable:                false,
 			expectedGrpcAddress:               "localhost:9091",
 			expectedGrpcStreamingEnable:       true,
+			expectedGrpcStreamingFlushMs:      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.expectedGrpcStreamingFlushMs,
+				flags.GrpcStreamingFlushIntervalMs,
+			)
 		})
 	}
 }
diff --git a/protocol/lib/metrics/metric_keys.go b/protocol/lib/metrics/metric_keys.go
@@ -60,6 +60,7 @@ const (
 	FullNodeGrpc                    = "full_node_grpc"
 	GrpcSendOrderbookUpdatesLatency = "grpc_send_orderbook_updates_latency"
 	GrpcSendOrderbookFillsLatency   = "grpc_send_orderbook_fills_latency"
+	GrpcFlushStreamUpdatesLatency   = "grpc_flush_stream_updates_latency"
 	EndBlocker                      = "end_blocker"
 	EndBlockerLag                   = "end_blocker_lag"
 )