Refactor max message sizes in comm client config

Signed-off-by: andrew-coleman <andrew_coleman@uk.ibm.com>

internal/pkg/comm/client.go

@@ -24,10 +24,6 @@ type GRPCClient struct {
 	dialOpts []grpc.DialOption
 	// Duration for which to block while established a new connection
 	timeout time.Duration
-	// Maximum message size the client can receive
-	maxRecvMsgSize int
-	// Maximum message size the client can send
-	maxSendMsgSize int
 }
 
 // NewGRPCClient creates a new implementation of GRPCClient given an address
@@ -57,8 +53,18 @@ func NewGRPCClient(config ClientConfig) (*GRPCClient, error) {
 	}
 	client.timeout = config.Timeout
 	// set send/recv message size to package defaults
-	client.maxRecvMsgSize = MaxRecvMsgSize
-	client.maxSendMsgSize = MaxSendMsgSize
+	maxRecvMsgSize := MaxRecvMsgSize
+	if config.MaxRecvMsgSize != 0 {
+		maxRecvMsgSize = config.MaxRecvMsgSize
+	}
+	maxSendMsgSize := MaxSendMsgSize
+	if config.MaxSendMsgSize != 0 {
+		maxSendMsgSize = config.MaxSendMsgSize
+	}
+	client.dialOpts = append(client.dialOpts, grpc.WithDefaultCallOptions(
+		grpc.MaxCallRecvMsgSize(maxRecvMsgSize),
+		grpc.MaxCallSendMsgSize(maxSendMsgSize),
+	))
 
 	return client, nil
 }
@@ -131,16 +137,6 @@ func (client *GRPCClient) MutualTLSRequired() bool {
 		len(client.tlsConfig.Certificates) > 0
 }
 
-// SetMaxRecvMsgSize sets the maximum message size the client can receive
-func (client *GRPCClient) SetMaxRecvMsgSize(size int) {
-	client.maxRecvMsgSize = size
-}
-
-// SetMaxSendMsgSize sets the maximum message size the client can send
-func (client *GRPCClient) SetMaxSendMsgSize(size int) {
-	client.maxSendMsgSize = size
-}
-
 // SetServerRootCAs sets the list of authorities used to verify server
 // certificates based on a list of PEM-encoded X509 certificate authorities
 func (client *GRPCClient) SetServerRootCAs(serverRoots [][]byte) error {
@@ -195,11 +191,6 @@ func (client *GRPCClient) NewConnection(address string, tlsOptions ...TLSOption)
 		dialOpts = append(dialOpts, grpc.WithInsecure())
 	}
 
-	dialOpts = append(dialOpts, grpc.WithDefaultCallOptions(
-		grpc.MaxCallRecvMsgSize(client.maxRecvMsgSize),
-		grpc.MaxCallSendMsgSize(client.maxSendMsgSize),
-	))
-
 	ctx, cancel := context.WithTimeout(context.Background(), client.timeout)
 	defer cancel()
 	conn, err := grpc.DialContext(ctx, address, dialOpts...)

internal/pkg/comm/client_test.go

@@ -478,25 +478,19 @@ func TestSetMessageSize(t *testing.T) {
 		},
 	}
 
-	// set up test client
-	client, err := comm.NewGRPCClient(comm.ClientConfig{
-		Timeout: testTimeout,
-	})
-	if err != nil {
-		t.Fatalf("error creating test client: %v", err)
-	}
 	// run tests
 	for _, test := range tests {
 		test := test
 		address := lis.Addr().String()
 		t.Run(test.name, func(t *testing.T) {
 			t.Log(test.name)
-			if test.maxRecvSize > 0 {
-				client.SetMaxRecvMsgSize(test.maxRecvSize)
-			}
-			if test.maxSendSize > 0 {
-				client.SetMaxSendMsgSize(test.maxSendSize)
-			}
+			// set up test client
+			client, err := comm.NewGRPCClient(comm.ClientConfig{
+				Timeout:        testTimeout,
+				MaxRecvMsgSize: test.maxRecvSize,
+				MaxSendMsgSize: test.maxSendSize,
+			})
+			require.NoError(t, err, "error creating test client")
 			conn, err := client.NewConnection(address)
 			assert.NoError(t, err)
 			defer conn.Close()

internal/pkg/comm/config.go

@@ -77,6 +77,10 @@ type ClientConfig struct {
 	Timeout time.Duration
 	// AsyncConnect makes connection creation non blocking
 	AsyncConnect bool
+	// Maximum message size the client can receive
+	MaxRecvMsgSize int
+	// Maximum message size the client can send
+	MaxSendMsgSize int
 }
 
 // Clone clones this ClientConfig