Allow setting grpc send and recv message sizes independently

Signed-off-by: Priya Bibra <[email protected]>
slackhq · Aug 12, 2024 · 79e613d · 79e613d
1 parent 404332b
commit 79e613d
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 5 deletions.
diff --git a/go/flags/endtoend/vtgate.txt b/go/flags/endtoend/vtgate.txt
@@ -54,6 +54,8 @@ Usage of vtgate:
       --grpc_key string                                                  server private key to use for gRPC connections, requires grpc_cert, enables TLS
       --grpc_max_connection_age duration                                 Maximum age of a client connection before GoAway is sent. (default 2562047h47m16.854775807s)
       --grpc_max_connection_age_grace duration                           Additional grace period after grpc_max_connection_age, after which connections are forcibly closed. (default 2562047h47m16.854775807s)
+      --grpc_max_message_recv_size int                                   Maximum allowed RPC message size when receiving. If 0, defaults to grpc_max_message_size.
+      --grpc_max_message_send_size int                                   Maximum allowed RPC message size when sending. If 0, defaults to grpc_max_message_size.
       --grpc_max_message_size int                                        Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'. (default 16777216)
       --grpc_port int                                                    Port to listen on for gRPC calls
       --grpc_prometheus                                                  Enable gRPC monitoring with Prometheus

diff --git a/go/flags/endtoend/vttablet.txt b/go/flags/endtoend/vttablet.txt
@@ -234,6 +234,8 @@ Usage of vttablet:
       --grpc_key string                                                  server private key to use for gRPC connections, requires grpc_cert, enables TLS
       --grpc_max_connection_age duration                                 Maximum age of a client connection before GoAway is sent. (default 2562047h47m16.854775807s)
       --grpc_max_connection_age_grace duration                           Additional grace period after grpc_max_connection_age, after which connections are forcibly closed. (default 2562047h47m16.854775807s)
+      --grpc_max_message_recv_size int                                   Maximum allowed RPC message size when receiving. If 0, defaults to grpc_max_message_size.
+      --grpc_max_message_send_size int                                   Maximum allowed RPC message size when sending. If 0, defaults to grpc_max_message_size.
       --grpc_max_message_size int                                        Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'. (default 16777216)
       --grpc_port int                                                    Port to listen on for gRPC calls
       --grpc_prometheus                                                  Enable gRPC monitoring with Prometheus

diff --git a/go/vt/grpcclient/client.go b/go/vt/grpcclient/client.go
@@ -77,8 +77,8 @@ func DialContext(ctx context.Context, target string, failFast FailFast, opts ...
 	grpccommon.EnableTracingOpt()
 	newopts := []grpc.DialOption{
 		grpc.WithDefaultCallOptions(
-			grpc.MaxCallRecvMsgSize(*grpccommon.MaxMessageSize),
-			grpc.MaxCallSendMsgSize(*grpccommon.MaxMessageSize),
+			grpc.MaxCallRecvMsgSize(grpccommon.MaxMessageRecvSize()),
+			grpc.MaxCallSendMsgSize(grpccommon.MaxMessageSendSize()),
 			grpc.WaitForReady(bool(!failFast)),
 		),
 	}

diff --git a/go/vt/grpccommon/options.go b/go/vt/grpccommon/options.go
@@ -31,6 +31,11 @@ var (
 	// accept. Larger messages will be rejected.
 	// Note: We're using 16 MiB as default value because that's the default in MySQL
 	MaxMessageSize = flag.Int("grpc_max_message_size", defaultMaxMessageSize, "Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'.")
+	// These options override MaxMessageSize if > 0, allowing us to control the max
+	// size sending independently from receiving.
+	MaxMsgRecvSize = flag.Int("grpc_max_message_recv_size", 0, "Maximum allowed RPC message size when receiving. If 0, defaults to grpc_max_message_size.")
+	MaxMsgSendSize = flag.Int("grpc_max_message_send_size", 0, "Maximum allowed RPC message size when sending. If 0, defaults to grpc_max_message_size.")
+
 	// EnableTracing sets a flag to enable grpc client/server tracing.
 	EnableTracing = flag.Bool("grpc_enable_tracing", false, "Enable GRPC tracing")
 
@@ -49,6 +54,20 @@ func EnableTracingOpt() {
 	})
 }
 
+func MaxMessageRecvSize() int {
+	if *MaxMsgRecvSize > 0 {
+		return *MaxMsgRecvSize
+	}
+	return *MaxMessageSize
+}
+
+func MaxMessageSendSize() int {
+	if *MaxMsgSendSize > 0 {
+		return *MaxMsgSendSize
+	}
+	return *MaxMessageSize
+}
+
 func init() {
 	stats.NewString("GrpcVersion").Set(grpc.Version)
 }
diff --git a/go/vt/servenv/grpc_server.go b/go/vt/servenv/grpc_server.go
@@ -159,9 +159,11 @@ func createGRPCServer() {
 	// grpc: received message length XXXXXXX exceeding the max size 4194304
 	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
 	// because it's not enforced on the client side.
-	log.Infof("Setting grpc max message size to %d", *grpccommon.MaxMessageSize)
-	opts = append(opts, grpc.MaxRecvMsgSize(*grpccommon.MaxMessageSize))
-	opts = append(opts, grpc.MaxSendMsgSize(*grpccommon.MaxMessageSize))
+	maxSendSize := grpccommon.MaxMessageSendSize()
+	maxRecvSize := grpccommon.MaxMessageRecvSize()
+	log.Infof("Setting grpc server max message sizes to %d (sending), %d (receiving)", maxSendSize, maxRecvSize)
+	opts = append(opts, grpc.MaxRecvMsgSize(maxRecvSize))
+	opts = append(opts, grpc.MaxSendMsgSize(maxSendSize))
 
 	if *GRPCInitialConnWindowSize != 0 {
 		log.Infof("Setting grpc server initial conn window size to %d", int32(*GRPCInitialConnWindowSize))