From b05b82c531b198ab2d14def4325c5e42ce9fc094 Mon Sep 17 00:00:00 2001
From: Vladimir Dementyev <dementiev.vm@gmail.com>
Date: Tue, 16 Nov 2021 14:47:47 +0300
Subject: [PATCH] feat: add grpc client max payload size config options

Closes #143
---
 CHANGELOG.md   |  2 ++
 cli/options.go |  4 ++++
 rpc/config.go  |  4 ++++
 rpc/rpc.go     | 21 ++++++++++++++++++---
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b68d16fe..e0d8cbc4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## master
 
+- Add `rpc_max_call_recv_size` and `rpc_max_call_send_size` options to allow modifying the corresponding limits for gRPC client connection. ([@palkan][])
+
 ## 1.1.3 (2021-09-16)
 
 - Fixed potential deadlocks in Hub. ([@palkan][])
diff --git a/cli/options.go b/cli/options.go
index 01a0115b..3f429fe0 100644
--- a/cli/options.go
+++ b/cli/options.go
@@ -68,6 +68,8 @@ func init() {
 	fs.StringVar(&defaults.RPC.Host, "rpc_host", "localhost:50051", "")
 	fs.IntVar(&defaults.RPC.Concurrency, "rpc_concurrency", 28, "")
 	fs.BoolVar(&defaults.RPC.EnableTLS, "rpc_enable_tls", false, "")
+	fs.IntVar(&defaults.RPC.MaxRecvSize, "rpc_max_call_recv_size", 0, "")
+	fs.IntVar(&defaults.RPC.MaxSendSize, "rpc_max_call_send_size", 0, "")
 	fs.StringVar(&headers, "headers", "cookie", "")
 
 	fs.IntVar(&defaults.WS.ReadBufferSize, "read_buffer_size", 1024, "")
@@ -157,6 +159,8 @@ OPTIONS
   --rpc_host                             RPC service address, default: localhost:50051, env: ANYCABLE_RPC_HOST
   --rpc_concurrency                      Max number of concurrent RPC request; should be slightly less than the RPC server concurrency, default: 28, env: ANYCABLE_RPC_CONCURRENCY
   --rpc_enable_tls                       Enable client-side TLS with the RPC server, default: false, env: ANYCABLE_RPC_ENABLE_TLS
+  --rpc_max_call_recv_size               Override default MaxCallRecvMsgSize for RPC client (bytes), default: none, env: ANYCABLE_RPC_MAX_CALL_RECV_SIZE
+  --rpc_max_call_send_size               Override default MaxCallSendMsgSize for RPC client (bytes), default: none, env: ANYCABLE_RPC_MAX_CALL_SEND_SIZE
   --headers                              List of headers to proxy to RPC, default: cookie, env: ANYCABLE_HEADERS
 
   --disconnect_rate                      Max number of Disconnect calls per second, default: 100, env: ANYCABLE_DISCONNECT_RATE
diff --git a/rpc/config.go b/rpc/config.go
index 0ff436a6..860d9c59 100644
--- a/rpc/config.go
+++ b/rpc/config.go
@@ -10,6 +10,10 @@ type Config struct {
 	Concurrency int
 	// Enable client-side TLS on RPC connections?
 	EnableTLS bool
+	// Max recieve msg size (bytes)
+	MaxRecvSize int
+	// Max send msg size (bytes)
+	MaxSendSize int
 }
 
 // NewConfig builds a new config
diff --git a/rpc/rpc.go b/rpc/rpc.go
index 1e43835a..2445549d 100644
--- a/rpc/rpc.go
+++ b/rpc/rpc.go
@@ -17,10 +17,10 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
-	"google.golang.org/grpc/credentials"
 
 	"crypto/tls"
 )
@@ -140,7 +140,7 @@ func (c *Controller) Start() error {
 	if enableTLS {
 		tlsConfig := &tls.Config{
 			InsecureSkipVerify: false,
-			MinVersion: tls.VersionTLS12,
+			MinVersion:         tls.VersionTLS12,
 		}
 
 		dialOptions = append(dialOptions, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
@@ -148,9 +148,24 @@ func (c *Controller) Start() error {
 		dialOptions = append(dialOptions, grpc.WithInsecure())
 	}
 
+	var callOptions = []grpc.CallOption{}
+
+	// Zero is the default
+	if c.config.MaxRecvSize != 0 {
+		callOptions = append(callOptions, grpc.MaxCallRecvMsgSize(c.config.MaxRecvSize))
+	}
+
+	if c.config.MaxSendSize != 0 {
+		callOptions = append(callOptions, grpc.MaxCallSendMsgSize(c.config.MaxSendSize))
+	}
+
+	if len(callOptions) > 0 {
+		dialOptions = append(dialOptions, grpc.WithDefaultCallOptions(callOptions...))
+	}
+
 	conn, err := grpc.Dial(
 		host,
-		dialOptions...
+		dialOptions...,
 	)
 
 	c.initSemaphore(capacity)