Skip to content

Commit 18146ee

Browse files
authored
Merge pull request #574 from giantwu/fix-grpc-client-banlance
Fix grpc client banlance
2 parents 4a663a1 + 59a6a44 commit 18146ee

3 files changed

Lines changed: 40 additions & 1 deletion

File tree

client/dtmgrpc/dtmgimp/grpc_clients.go

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ package dtmgimp
88

99
import (
1010
"fmt"
11+
"os"
1112
"sync"
1213

1314
"github.com/dtm-labs/dtm/client/dtmcli/dtmimp"
@@ -39,6 +40,9 @@ var normalClients, rawClients sync.Map
3940
// ClientInterceptors declares grpc.UnaryClientInterceptors slice
4041
var ClientInterceptors = []grpc.UnaryClientInterceptor{}
4142

43+
// GrpcServiceConfigGetter is a function to get gRPC service config, can be set by server config
44+
var GrpcServiceConfigGetter func() string
45+
4246
// MustGetDtmClient 1
4347
func MustGetDtmClient(grpcServer string) dtmgpb.DtmClient {
4448
return dtmgpb.NewDtmClient(MustGetGrpcConn(grpcServer, false))
@@ -61,7 +65,12 @@ func GetGrpcConn(grpcServer string, isRaw bool) (conn *grpc.ClientConn, rerr err
6165
interceptors := append(ClientInterceptors, GrpcClientLog)
6266
interceptors = append(interceptors, dtmdriver.Middlewares.Grpc...)
6367
inOpt := grpc.WithChainUnaryInterceptor(interceptors...)
64-
conn, rerr := grpc.Dial(grpcServer, inOpt, grpc.WithTransportCredentials(insecure.NewCredentials()), opts)
68+
grpcServiceConfig, hasConfig := getGrpcServiceConfig()
69+
dialOpts := []grpc.DialOption{inOpt, grpc.WithTransportCredentials(insecure.NewCredentials()), opts}
70+
if hasConfig {
71+
dialOpts = append(dialOpts, grpc.WithDefaultServiceConfig(grpcServiceConfig))
72+
}
73+
conn, rerr := grpc.Dial(grpcServer, dialOpts...)
6574
if rerr == nil {
6675
clients.Store(grpcServer, conn)
6776
v = conn
@@ -77,3 +86,20 @@ func MustGetGrpcConn(grpcServer string, isRaw bool) *grpc.ClientConn {
7786
dtmimp.E2P(err)
7887
return conn
7988
}
89+
90+
// getGrpcServiceConfig returns the gRPC service config from config getter or environment variable, and a bool indicating if config is set
91+
func getGrpcServiceConfig() (string, bool) {
92+
// First try to get from config getter (set by server)
93+
if GrpcServiceConfigGetter != nil {
94+
if config := GrpcServiceConfigGetter(); config != "" {
95+
return config, true
96+
}
97+
}
98+
// Fallback to environment variable
99+
config := os.Getenv("GRPC_SERVICE_CONFIG")
100+
if config != "" {
101+
return config, true
102+
}
103+
// No config set
104+
return "", false
105+
}

conf.sample.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,3 +70,6 @@
7070
# AlertRetryLimit: 3 # default 3; if a transaction branch has been retried 3 times, the AlertHook will be called
7171
# AlertWebHook: '' # default ''; sample: 'http://localhost:8080/dtm-hook'. this hook will be called like this:
7272
## curl -H "Content-Type: application/json" -d '{"gid":"xxxx","status":"submitted","retry_count":3}' http://localhost:8080/dtm-hook
73+
74+
# GrpcBalancer: 'round_robin' # default round_robin,For more information about service configs, see:https://github.com/grpc/grpc/blob/master/doc/service_config.md
75+
# GrpcServiceConfig: '{"loadBalancingConfig": [{"round_robin":{}}]}' # default gRPC service config for client connections

dtmsvr/config/config.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"io/ioutil"
66

77
"github.com/dtm-labs/dtm/client/dtmcli"
8+
"github.com/dtm-labs/dtm/client/dtmgrpc/dtmgimp"
89
"github.com/dtm-labs/logger"
910
"gopkg.in/yaml.v3"
1011
)
@@ -104,6 +105,7 @@ type Type struct {
104105
AlertRetryLimit int64 `yaml:"AlertRetryLimit" default:"3"`
105106
AlertWebHook string `yaml:"AlertWebHook"`
106107
AdminBasePath string `yaml:"AdminBasePath"`
108+
GrpcServiceConfig string `yaml:"GrpcServiceConfig" default:"{\"loadBalancingConfig\": [{\"round_robin\":{}}]}"`
107109
}
108110

109111
// Config config
@@ -124,4 +126,12 @@ func MustLoadConfig(confFile string) {
124126
err = checkConfig(&Config)
125127
logger.FatalfIf(err != nil, `config error: '%v'.
126128
please visit http://d.dtm.pub to see the config document.`, err)
129+
// Set gRPC service config for client library
130+
if Config.GrpcServiceConfig == "" {
131+
Config.GrpcServiceConfig = `{"loadBalancingConfig": [{"round_robin":{}}]}`
132+
}
133+
// Set config getter function for client library to read config directly
134+
dtmgimp.GrpcServiceConfigGetter = func() string {
135+
return Config.GrpcServiceConfig
136+
}
127137
}

0 commit comments

Comments
 (0)