Fix passing default grpc call options in Kubernetes client

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>
serathius · Jul 24, 2024 · e3eed53 · e3eed53
1 parent 9a6c9ae
commit e3eed53
Showing 1 changed file with 25 additions and 4 deletions.
diff --git a/client/v3/kubernetes/client.go b/client/v3/kubernetes/client.go
@@ -16,6 +16,10 @@ package kubernetes
 
 import (
 	"context"
+	"fmt"
+	"math"
+
+	"google.golang.org/grpc"
 
 	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -33,6 +37,22 @@ func New(cfg clientv3.Config) (*Client, error) {
 		Client: c,
 		kv:     clientv3.RetryKVClient(c),
 	}
+	kc.callOpts = []grpc.CallOption{
+		grpc.WaitForReady(true),
+		grpc.MaxCallSendMsgSize(2 * 1024 * 1024),
+		grpc.MaxCallRecvMsgSize(math.MaxInt32),
+	}
+	if cfg.MaxCallSendMsgSize > 0 || cfg.MaxCallRecvMsgSize > 0 {
+		if cfg.MaxCallRecvMsgSize > 0 && cfg.MaxCallSendMsgSize > cfg.MaxCallRecvMsgSize {
+			return nil, fmt.Errorf("gRPC message recv limit (%d bytes) must be greater than send limit (%d bytes)", cfg.MaxCallRecvMsgSize, cfg.MaxCallSendMsgSize)
+		}
+		if cfg.MaxCallSendMsgSize > 0 {
+			kc.callOpts[1] = grpc.MaxCallSendMsgSize(cfg.MaxCallSendMsgSize)
+		}
+		if cfg.MaxCallRecvMsgSize > 0 {
+			kc.callOpts[2] = grpc.MaxCallRecvMsgSize(cfg.MaxCallRecvMsgSize)
+		}
+	}
 	kc.Kubernetes = kc
 	return kc, nil
 }
@@ -41,12 +61,13 @@ type Client struct {
 	*clientv3.Client
 	Kubernetes Interface
 	kv         pb.KVClient
+	callOpts   []grpc.CallOption
 }
 
 var _ Interface = (*Client)(nil)
 
 func (k Client) Get(ctx context.Context, key string, opts GetOptions) (resp GetResponse, err error) {
-	rangeResp, err := k.kv.Range(ctx, getRequest(key, opts.Revision))
+	rangeResp, err := k.kv.Range(ctx, getRequest(key, opts.Revision), k.callOpts...)
 	if err != nil {
 		return resp, clientv3.ContextError(ctx, err)
 	}
@@ -66,7 +87,7 @@ func (k Client) List(ctx context.Context, prefix string, opts ListOptions) (resp
 		RangeEnd: []byte(rangeEnd),
 		Limit:    opts.Limit,
 		Revision: opts.Revision,
-	})
+	}, k.callOpts...)
 	if err != nil {
 		return resp, clientv3.ContextError(ctx, err)
 	}
@@ -81,7 +102,7 @@ func (k Client) Count(ctx context.Context, prefix string, _ CountOptions) (int64
 		Key:       []byte(prefix),
 		RangeEnd:  []byte(clientv3.GetPrefixRangeEnd(prefix)),
 		CountOnly: true,
-	})
+	}, k.callOpts...)
 	if err != nil {
 		return 0, clientv3.ContextError(ctx, err)
 	}
@@ -145,7 +166,7 @@ func (k Client) optimisticTxn(ctx context.Context, key string, expectedRevision
 	if onFailure != nil {
 		txn.Failure = []*pb.RequestOp{onFailure}
 	}
-	return k.kv.Txn(ctx, txn)
+	return k.kv.Txn(ctx, txn, k.callOpts...)
 }
 
 func getRequest(key string, revision int64) *pb.RangeRequest {