diff --git a/client/dtmgrpc/dtmgimp/grpc_clients.go b/client/dtmgrpc/dtmgimp/grpc_clients.go index 2291c52..94effbe 100644 --- a/client/dtmgrpc/dtmgimp/grpc_clients.go +++ b/client/dtmgrpc/dtmgimp/grpc_clients.go @@ -8,6 +8,7 @@ package dtmgimp import ( "fmt" + "os" "sync" "github.com/dtm-labs/dtm/client/dtmcli/dtmimp" @@ -39,6 +40,9 @@ var normalClients, rawClients sync.Map // ClientInterceptors declares grpc.UnaryClientInterceptors slice var ClientInterceptors = []grpc.UnaryClientInterceptor{} +// GrpcServiceConfigGetter is a function to get gRPC service config, can be set by server config +var GrpcServiceConfigGetter func() string + // MustGetDtmClient 1 func MustGetDtmClient(grpcServer string) dtmgpb.DtmClient { return dtmgpb.NewDtmClient(MustGetGrpcConn(grpcServer, false)) @@ -61,7 +65,12 @@ func GetGrpcConn(grpcServer string, isRaw bool) (conn *grpc.ClientConn, rerr err interceptors := append(ClientInterceptors, GrpcClientLog) interceptors = append(interceptors, dtmdriver.Middlewares.Grpc...) inOpt := grpc.WithChainUnaryInterceptor(interceptors...) - conn, rerr := grpc.Dial(grpcServer, inOpt, grpc.WithTransportCredentials(insecure.NewCredentials()), opts) + grpcServiceConfig, hasConfig := getGrpcServiceConfig() + dialOpts := []grpc.DialOption{inOpt, grpc.WithTransportCredentials(insecure.NewCredentials()), opts} + if hasConfig { + dialOpts = append(dialOpts, grpc.WithDefaultServiceConfig(grpcServiceConfig)) + } + conn, rerr := grpc.Dial(grpcServer, dialOpts...) if rerr == nil { clients.Store(grpcServer, conn) v = conn @@ -77,3 +86,20 @@ func MustGetGrpcConn(grpcServer string, isRaw bool) *grpc.ClientConn { dtmimp.E2P(err) return conn } + +// getGrpcServiceConfig returns the gRPC service config from config getter or environment variable, and a bool indicating if config is set +func getGrpcServiceConfig() (string, bool) { + // First try to get from config getter (set by server) + if GrpcServiceConfigGetter != nil { + if config := GrpcServiceConfigGetter(); config != "" { + return config, true + } + } + // Fallback to environment variable + config := os.Getenv("GRPC_SERVICE_CONFIG") + if config != "" { + return config, true + } + // No config set + return "", false +} diff --git a/conf.sample.yml b/conf.sample.yml index 6da326f..c23079b 100644 --- a/conf.sample.yml +++ b/conf.sample.yml @@ -70,3 +70,6 @@ # AlertRetryLimit: 3 # default 3; if a transaction branch has been retried 3 times, the AlertHook will be called # AlertWebHook: '' # default ''; sample: 'http://localhost:8080/dtm-hook'. this hook will be called like this: ## curl -H "Content-Type: application/json" -d '{"gid":"xxxx","status":"submitted","retry_count":3}' http://localhost:8080/dtm-hook + +# GrpcBalancer: 'round_robin' # default round_robin,For more information about service configs, see:https://github.com/grpc/grpc/blob/master/doc/service_config.md +# GrpcServiceConfig: '{"loadBalancingConfig": [{"round_robin":{}}]}' # default gRPC service config for client connections diff --git a/dtmsvr/config/config.go b/dtmsvr/config/config.go index 35815f2..e27bed1 100644 --- a/dtmsvr/config/config.go +++ b/dtmsvr/config/config.go @@ -5,6 +5,7 @@ import ( "io/ioutil" "github.com/dtm-labs/dtm/client/dtmcli" + "github.com/dtm-labs/dtm/client/dtmgrpc/dtmgimp" "github.com/dtm-labs/logger" "gopkg.in/yaml.v3" ) @@ -104,6 +105,7 @@ type Type struct { AlertRetryLimit int64 `yaml:"AlertRetryLimit" default:"3"` AlertWebHook string `yaml:"AlertWebHook"` AdminBasePath string `yaml:"AdminBasePath"` + GrpcServiceConfig string `yaml:"GrpcServiceConfig" default:"{\"loadBalancingConfig\": [{\"round_robin\":{}}]}"` } // Config config @@ -124,4 +126,12 @@ func MustLoadConfig(confFile string) { err = checkConfig(&Config) logger.FatalfIf(err != nil, `config error: '%v'. please visit http://d.dtm.pub to see the config document.`, err) + // Set gRPC service config for client library + if Config.GrpcServiceConfig == "" { + Config.GrpcServiceConfig = `{"loadBalancingConfig": [{"round_robin":{}}]}` + } + // Set config getter function for client library to read config directly + dtmgimp.GrpcServiceConfigGetter = func() string { + return Config.GrpcServiceConfig + } }