diff --git a/pkg/discoveryregistry/zookeeper/conf.go b/pkg/discoveryregistry/zookeeper/conf.go index 04ee6b179..55532a370 100644 --- a/pkg/discoveryregistry/zookeeper/conf.go +++ b/pkg/discoveryregistry/zookeeper/conf.go @@ -1,4 +1,4 @@ -package openKeeper +package zookeeper import ( "github.com/go-zookeeper/zk" diff --git a/pkg/discoveryregistry/zookeeper/discover.go b/pkg/discoveryregistry/zookeeper/discover.go index c50ec2654..01900e08e 100644 --- a/pkg/discoveryregistry/zookeeper/discover.go +++ b/pkg/discoveryregistry/zookeeper/discover.go @@ -1,4 +1,4 @@ -package openKeeper +package zookeeper import ( "context" @@ -107,7 +107,7 @@ func (s *ZkClient) GetConn(ctx context.Context, serviceName string, opts ...grpc if len(conns) == 0 { return nil, ErrConnIsNil } - return conns[0], nil + return s.getConnBalance(conns), nil } func (s *ZkClient) GetFirstConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (*grpc.ClientConn, error) { diff --git a/pkg/discoveryregistry/zookeeper/load_balancing.go b/pkg/discoveryregistry/zookeeper/load_balancing.go new file mode 100644 index 000000000..c20639d31 --- /dev/null +++ b/pkg/discoveryregistry/zookeeper/load_balancing.go @@ -0,0 +1,23 @@ +package zookeeper + +import ( + "sync" + + "google.golang.org/grpc" +) + +type RoundRobin struct { + index int + lock sync.Mutex +} + +func (r *RoundRobin) getConnBalance(conns []*grpc.ClientConn) (conn *grpc.ClientConn) { + r.lock.Lock() + defer r.lock.Unlock() + if r.index < len(conns)-1 { + r.index++ + } else { + r.index = 0 + } + return conns[r.index] +} diff --git a/pkg/discoveryregistry/zookeeper/zk.go b/pkg/discoveryregistry/zookeeper/zk.go index d89567c8a..2fd9e14c7 100644 --- a/pkg/discoveryregistry/zookeeper/zk.go +++ b/pkg/discoveryregistry/zookeeper/zk.go @@ -1,4 +1,4 @@ -package openKeeper +package zookeeper import ( "context" @@ -38,6 +38,7 @@ type ZkClient struct { resolvers map[string]*Resolver localConns map[string][]resolver.Address balancerName string + RoundRobin } type ZkOption func(*ZkClient)