From 54460fb7712455930caf3410a60131384fcc7097 Mon Sep 17 00:00:00 2001
From: LaurenceLiZhixin <382673304@qq.com>
Date: Mon, 6 Dec 2021 22:29:26 +0800
Subject: [PATCH] Ftr: add grpc max size support

---
 pkg/config/config.go        | 30 ++++++++++++++++++++++++++++++
 pkg/triple/dubbo3_client.go |  9 +++++++++
 pkg/triple/dubbo3_server.go |  8 ++++++++
 3 files changed, 47 insertions(+)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index fae8cd6..9609df5 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -40,6 +40,12 @@ type Option struct {
 	HeaderGroup      string
 	HeaderAppVersion string
 
+	// grpc opts
+	GRPCMaxCallSendMsgSize   int
+	GRPCMaxServerSendMsgSize int
+	GRPCMaxCallRecvMsgSize   int
+	GRPCMaxServerRecvMsgSize int
+
 	// tracing
 	JaegerAddress     string
 	JaegerServiceName string
@@ -166,6 +172,30 @@ func WithNumWorker(numWorkers uint32) OptionFunction {
 	}
 }
 
+func WithGRPCMaxCallSendMessageSize(maxCallSendMsgSize int) OptionFunction {
+	return func(o *Option) {
+		o.GRPCMaxCallSendMsgSize = maxCallSendMsgSize
+	}
+}
+
+func WithGRPCMaxCallRecvMessageSize(maxCallRecvMsgSize int) OptionFunction {
+	return func(o *Option) {
+		o.GRPCMaxCallRecvMsgSize = maxCallRecvMsgSize
+	}
+}
+
+func WithGRPCMaxServerSendMessageSize(maxServerSendMsgSize int) OptionFunction {
+	return func(o *Option) {
+		o.GRPCMaxServerSendMsgSize = maxServerSendMsgSize
+	}
+}
+
+func WithGRPCMaxServerRecvMessageSize(maxServerRecvMsgSize int) OptionFunction {
+	return func(o *Option) {
+		o.GRPCMaxServerRecvMsgSize = maxServerRecvMsgSize
+	}
+}
+
 func WithJaegerConfig(address, serviceName string, useAgent bool) OptionFunction {
 	return func(o *Option) {
 		o.JaegerAddress = address
diff --git a/pkg/triple/dubbo3_client.go b/pkg/triple/dubbo3_client.go
index d610ed8..5091d9d 100644
--- a/pkg/triple/dubbo3_client.go
+++ b/pkg/triple/dubbo3_client.go
@@ -85,6 +85,15 @@ func NewTripleClient(impl interface{}, opt *config.Option) (*TripleClient, error
 		)
 	}
 
+	defaultCallOpts := make([]grpc.CallOption, 0)
+	if opt.GRPCMaxCallSendMsgSize != 0 {
+		defaultCallOpts = append(defaultCallOpts, grpc.MaxCallSendMsgSize(opt.GRPCMaxCallSendMsgSize))
+	}
+	if opt.GRPCMaxCallRecvMsgSize != 0 {
+		defaultCallOpts = append(defaultCallOpts, grpc.MaxCallRecvMsgSize(opt.GRPCMaxCallRecvMsgSize))
+	}
+	dialOpts = append(dialOpts, grpc.WithDefaultCallOptions(defaultCallOpts...))
+
 	if opt.CodecType == constant.PBCodecName {
 		// put dubbo3 network logic to tripleConn, creat pb stub invoker
 		tripleClient.stubInvoker = reflect.ValueOf(getInvoker(impl, newTripleConn(opt.Location, dialOpts...)))
diff --git a/pkg/triple/dubbo3_server.go b/pkg/triple/dubbo3_server.go
index 76eafc0..1fc6d5a 100644
--- a/pkg/triple/dubbo3_server.go
+++ b/pkg/triple/dubbo3_server.go
@@ -34,6 +34,7 @@ import (
 	"github.com/dubbogo/grpc-go/encoding/msgpack"
 	"github.com/dubbogo/grpc-go/encoding/proto_wrapper_api"
 	"github.com/dubbogo/grpc-go/encoding/raw_proto"
+
 	perrors "github.com/pkg/errors"
 )
 
@@ -204,6 +205,13 @@ func newGrpcServerWithCodec(opt *config.Option) *grpc.Server {
 		)
 	}
 
+	if opt.GRPCMaxServerRecvMsgSize != 0 {
+		serverOpts = append(serverOpts, grpc.MaxRecvMsgSize(opt.GRPCMaxServerSendMsgSize))
+	}
+	if opt.GRPCMaxCallSendMsgSize != 0 {
+		serverOpts = append(serverOpts, grpc.MaxSendMsgSize(opt.GRPCMaxServerRecvMsgSize))
+	}
+
 	var err error
 	switch opt.CodecType {
 	case constant.PBCodecName: