initial import

core/discov/internal/balancer.go

@@ -0,0 +1,103 @@
+package internal
+
+import "sync"
+
+type (
+	DialFn  func(server string) (interface{}, error)
+	CloseFn func(server string, conn interface{}) error
+
+	Balancer interface {
+		AddConn(kv KV) error
+		IsEmpty() bool
+		Next(key ...string) (interface{}, bool)
+		RemoveKey(key string)
+		initialize()
+		setListener(listener Listener)
+	}
+
+	serverConn struct {
+		key  string
+		conn interface{}
+	}
+
+	baseBalancer struct {
+		exclusive bool
+		servers   map[string][]string
+		mapping   map[string]string
+		lock      sync.Mutex
+		dialFn    DialFn
+		closeFn   CloseFn
+		listener  Listener
+	}
+)
+
+func newBaseBalancer(dialFn DialFn, closeFn CloseFn, exclusive bool) *baseBalancer {
+	return &baseBalancer{
+		exclusive: exclusive,
+		servers:   make(map[string][]string),
+		mapping:   make(map[string]string),
+		dialFn:    dialFn,
+		closeFn:   closeFn,
+	}
+}
+
+// addKv adds the kv, returns if there are already other keys associate with the server
+func (b *baseBalancer) addKv(key, value string) ([]string, bool) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	keys := b.servers[value]
+	previous := append([]string(nil), keys...)
+	early := len(keys) > 0
+	if b.exclusive && early {
+		for _, each := range keys {
+			b.doRemoveKv(each)
+		}
+	}
+	b.servers[value] = append(b.servers[value], key)
+	b.mapping[key] = value
+
+	if early {
+		return previous, true
+	} else {
+		return nil, false
+	}
+}
+
+func (b *baseBalancer) doRemoveKv(key string) (server string, keepConn bool) {
+	server, ok := b.mapping[key]
+	if !ok {
+		return "", true
+	}
+
+	delete(b.mapping, key)
+	keys := b.servers[server]
+	remain := keys[:0]
+
+	for _, k := range keys {
+		if k != key {
+			remain = append(remain, k)
+		}
+	}
+
+	if len(remain) > 0 {
+		b.servers[server] = remain
+		return server, true
+	} else {
+		delete(b.servers, server)
+		return server, false
+	}
+}
+
+func (b *baseBalancer) removeKv(key string) (server string, keepConn bool) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	return b.doRemoveKv(key)
+}
+
+func (b *baseBalancer) setListener(listener Listener) {
+	b.lock.Lock()
+	b.listener = listener
+	b.lock.Unlock()
+}

core/discov/internal/balancer_test.go

@@ -0,0 +1,5 @@
+package internal
+
+type mockConn struct {
+	server string
+}

core/discov/internal/consistentbalancer.go

@@ -0,0 +1,152 @@
+package internal
+
+import (
+	"zero/core/hash"
+	"zero/core/logx"
+)
+
+type consistentBalancer struct {
+	*baseBalancer
+	conns     map[string]interface{}
+	buckets   *hash.ConsistentHash
+	bucketKey func(KV) string
+}
+
+func NewConsistentBalancer(dialFn DialFn, closeFn CloseFn, keyer func(kv KV) string) *consistentBalancer {
+	// we don't support exclusive mode for consistent Balancer, to avoid complexity,
+	// because there are few scenarios, use it on your own risks.
+	balancer := &consistentBalancer{
+		conns:     make(map[string]interface{}),
+		buckets:   hash.NewConsistentHash(),
+		bucketKey: keyer,
+	}
+	balancer.baseBalancer = newBaseBalancer(dialFn, closeFn, false)
+	return balancer
+}
+
+func (b *consistentBalancer) AddConn(kv KV) error {
+	// not adding kv and conn within a transaction, but it doesn't matter
+	// we just rollback the kv addition if dial failed
+	var conn interface{}
+	prev, found := b.addKv(kv.Key, kv.Val)
+	if found {
+		conn = b.handlePrevious(prev)
+	}
+
+	if conn == nil {
+		var err error
+		conn, err = b.dialFn(kv.Val)
+		if err != nil {
+			b.removeKv(kv.Key)
+			return err
+		}
+	}
+
+	bucketKey := b.bucketKey(kv)
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	b.conns[bucketKey] = conn
+	b.buckets.Add(bucketKey)
+	b.notify(bucketKey)
+
+	logx.Infof("added server, key: %s, server: %s", bucketKey, kv.Val)
+
+	return nil
+}
+
+func (b *consistentBalancer) getConn(key string) (interface{}, bool) {
+	b.lock.Lock()
+	conn, ok := b.conns[key]
+	b.lock.Unlock()
+
+	return conn, ok
+}
+
+func (b *consistentBalancer) handlePrevious(prev []string) interface{} {
+	if len(prev) == 0 {
+		return nil
+	}
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	// if not exclusive, only need to randomly find one connection
+	for key, conn := range b.conns {
+		if key == prev[0] {
+			return conn
+		}
+	}
+
+	return nil
+}
+
+func (b *consistentBalancer) initialize() {
+}
+
+func (b *consistentBalancer) notify(key string) {
+	if b.listener == nil {
+		return
+	}
+
+	var keys []string
+	var values []string
+	for k := range b.conns {
+		keys = append(keys, k)
+	}
+	for _, v := range b.mapping {
+		values = append(values, v)
+	}
+
+	b.listener.OnUpdate(keys, values, key)
+}
+
+func (b *consistentBalancer) RemoveKey(key string) {
+	kv := KV{Key: key}
+	server, keep := b.removeKv(key)
+	kv.Val = server
+	bucketKey := b.bucketKey(kv)
+	b.buckets.Remove(b.bucketKey(kv))
+
+	// wrap the query & removal in a function to make sure the quick lock/unlock
+	conn, ok := func() (interface{}, bool) {
+		b.lock.Lock()
+		defer b.lock.Unlock()
+
+		conn, ok := b.conns[bucketKey]
+		if ok {
+			delete(b.conns, bucketKey)
+		}
+
+		return conn, ok
+	}()
+	if ok && !keep {
+		logx.Infof("removing server, key: %s", kv.Key)
+		if err := b.closeFn(server, conn); err != nil {
+			logx.Error(err)
+		}
+	}
+
+	// notify without new key
+	b.notify("")
+}
+
+func (b *consistentBalancer) IsEmpty() bool {
+	b.lock.Lock()
+	empty := len(b.conns) == 0
+	b.lock.Unlock()
+
+	return empty
+}
+
+func (b *consistentBalancer) Next(keys ...string) (interface{}, bool) {
+	if len(keys) != 1 {
+		return nil, false
+	}
+
+	key := keys[0]
+	if node, ok := b.buckets.Get(key); !ok {
+		return nil, false
+	} else {
+		return b.getConn(node.(string))
+	}
+}

core/discov/internal/consistentbalancer_test.go

@@ -0,0 +1,178 @@
+package internal
+
+import (
+	"errors"
+	"sort"
+	"strconv"
+	"testing"
+
+	"zero/core/mathx"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConsistent_addConn(t *testing.T) {
+	b := NewConsistentBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return errors.New("error")
+	}, func(kv KV) string {
+		return kv.Key
+	})
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, map[string]interface{}{
+		"thekey1": mockConn{server: "thevalue"},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+	}, b.mapping)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey2",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, map[string]interface{}{
+		"thekey1": mockConn{server: "thevalue"},
+		"thekey2": mockConn{server: "thevalue"},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1", "thekey2"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+		"thekey2": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey1")
+	assert.EqualValues(t, map[string]interface{}{
+		"thekey2": mockConn{server: "thevalue"},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey2"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey2": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey2")
+	assert.Equal(t, 0, len(b.conns))
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+	assert.True(t, b.IsEmpty())
+}
+
+func TestConsistent_addConnError(t *testing.T) {
+	b := NewConsistentBalancer(func(server string) (interface{}, error) {
+		return nil, errors.New("error")
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, func(kv KV) string {
+		return kv.Key
+	})
+	assert.NotNil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.Equal(t, 0, len(b.conns))
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+}
+
+func TestConsistent_next(t *testing.T) {
+	b := NewConsistentBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return errors.New("error")
+	}, func(kv KV) string {
+		return kv.Key
+	})
+	b.initialize()
+
+	_, ok := b.Next("any")
+	assert.False(t, ok)
+
+	const size = 100
+	for i := 0; i < size; i++ {
+		assert.Nil(t, b.AddConn(KV{
+			Key: "thekey/" + strconv.Itoa(i),
+			Val: "thevalue/" + strconv.Itoa(i),
+		}))
+	}
+
+	m := make(map[interface{}]int)
+	const total = 10000
+	for i := 0; i < total; i++ {
+		val, ok := b.Next(strconv.Itoa(i))
+		assert.True(t, ok)
+		m[val]++
+	}
+
+	entropy := mathx.CalcEntropy(m, total)
+	assert.Equal(t, size, len(m))
+	assert.True(t, entropy > .95)
+
+	for i := 0; i < size; i++ {
+		b.RemoveKey("thekey/" + strconv.Itoa(i))
+	}
+	_, ok = b.Next()
+	assert.False(t, ok)
+}
+
+func TestConsistentBalancer_Listener(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	b := NewConsistentBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, func(kv KV) string {
+		return kv.Key
+	})
+	assert.Nil(t, b.AddConn(KV{
+		Key: "key1",
+		Val: "val1",
+	}))
+	assert.Nil(t, b.AddConn(KV{
+		Key: "key2",
+		Val: "val2",
+	}))
+
+	listener := NewMockListener(ctrl)
+	listener.EXPECT().OnUpdate(gomock.Any(), gomock.Any(), "key2").Do(func(keys, vals, _ interface{}) {
+		sort.Strings(keys.([]string))
+		sort.Strings(vals.([]string))
+		assert.EqualValues(t, []string{"key1", "key2"}, keys)
+		assert.EqualValues(t, []string{"val1", "val2"}, vals)
+	})
+	b.setListener(listener)
+	b.notify("key2")
+}
+
+func TestConsistentBalancer_remove(t *testing.T) {
+	b := NewConsistentBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, func(kv KV) string {
+		return kv.Key
+	})
+
+	assert.Nil(t, b.handlePrevious(nil))
+	assert.Nil(t, b.handlePrevious([]string{"any"}))
+}

core/discov/internal/etcdclient.go

@@ -0,0 +1,21 @@
+//go:generate mockgen -package internal -destination etcdclient_mock.go -source etcdclient.go EtcdClient
+package internal
+
+import (
+	"context"
+
+	"go.etcd.io/etcd/clientv3"
+	"google.golang.org/grpc"
+)
+
+type EtcdClient interface {
+	ActiveConnection() *grpc.ClientConn
+	Close() error
+	Ctx() context.Context
+	Get(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error)
+	Grant(ctx context.Context, ttl int64) (*clientv3.LeaseGrantResponse, error)
+	KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error)
+	Put(ctx context.Context, key, val string, opts ...clientv3.OpOption) (*clientv3.PutResponse, error)
+	Revoke(ctx context.Context, id clientv3.LeaseID) (*clientv3.LeaseRevokeResponse, error)
+	Watch(ctx context.Context, key string, opts ...clientv3.OpOption) clientv3.WatchChan
+}

core/discov/internal/etcdclient_mock.go

@@ -0,0 +1,182 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: etcdclient.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	context "context"
+	gomock "github.com/golang/mock/gomock"
+	clientv3 "go.etcd.io/etcd/clientv3"
+	grpc "google.golang.org/grpc"
+	reflect "reflect"
+)
+
+// MockEtcdClient is a mock of EtcdClient interface
+type MockEtcdClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockEtcdClientMockRecorder
+}
+
+// MockEtcdClientMockRecorder is the mock recorder for MockEtcdClient
+type MockEtcdClientMockRecorder struct {
+	mock *MockEtcdClient
+}
+
+// NewMockEtcdClient creates a new mock instance
+func NewMockEtcdClient(ctrl *gomock.Controller) *MockEtcdClient {
+	mock := &MockEtcdClient{ctrl: ctrl}
+	mock.recorder = &MockEtcdClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockEtcdClient) EXPECT() *MockEtcdClientMockRecorder {
+	return m.recorder
+}
+
+// ActiveConnection mocks base method
+func (m *MockEtcdClient) ActiveConnection() *grpc.ClientConn {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ActiveConnection")
+	ret0, _ := ret[0].(*grpc.ClientConn)
+	return ret0
+}
+
+// ActiveConnection indicates an expected call of ActiveConnection
+func (mr *MockEtcdClientMockRecorder) ActiveConnection() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ActiveConnection", reflect.TypeOf((*MockEtcdClient)(nil).ActiveConnection))
+}
+
+// Close mocks base method
+func (m *MockEtcdClient) Close() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Close")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Close indicates an expected call of Close
+func (mr *MockEtcdClientMockRecorder) Close() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockEtcdClient)(nil).Close))
+}
+
+// Ctx mocks base method
+func (m *MockEtcdClient) Ctx() context.Context {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Ctx")
+	ret0, _ := ret[0].(context.Context)
+	return ret0
+}
+
+// Ctx indicates an expected call of Ctx
+func (mr *MockEtcdClientMockRecorder) Ctx() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ctx", reflect.TypeOf((*MockEtcdClient)(nil).Ctx))
+}
+
+// Get mocks base method
+func (m *MockEtcdClient) Get(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{ctx, key}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Get", varargs...)
+	ret0, _ := ret[0].(*clientv3.GetResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Get indicates an expected call of Get
+func (mr *MockEtcdClientMockRecorder) Get(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{ctx, key}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockEtcdClient)(nil).Get), varargs...)
+}
+
+// Grant mocks base method
+func (m *MockEtcdClient) Grant(ctx context.Context, ttl int64) (*clientv3.LeaseGrantResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Grant", ctx, ttl)
+	ret0, _ := ret[0].(*clientv3.LeaseGrantResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Grant indicates an expected call of Grant
+func (mr *MockEtcdClientMockRecorder) Grant(ctx, ttl interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Grant", reflect.TypeOf((*MockEtcdClient)(nil).Grant), ctx, ttl)
+}
+
+// KeepAlive mocks base method
+func (m *MockEtcdClient) KeepAlive(ctx context.Context, id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "KeepAlive", ctx, id)
+	ret0, _ := ret[0].(<-chan *clientv3.LeaseKeepAliveResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// KeepAlive indicates an expected call of KeepAlive
+func (mr *MockEtcdClientMockRecorder) KeepAlive(ctx, id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "KeepAlive", reflect.TypeOf((*MockEtcdClient)(nil).KeepAlive), ctx, id)
+}
+
+// Put mocks base method
+func (m *MockEtcdClient) Put(ctx context.Context, key, val string, opts ...clientv3.OpOption) (*clientv3.PutResponse, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{ctx, key, val}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Put", varargs...)
+	ret0, _ := ret[0].(*clientv3.PutResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Put indicates an expected call of Put
+func (mr *MockEtcdClientMockRecorder) Put(ctx, key, val interface{}, opts ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{ctx, key, val}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockEtcdClient)(nil).Put), varargs...)
+}
+
+// Revoke mocks base method
+func (m *MockEtcdClient) Revoke(ctx context.Context, id clientv3.LeaseID) (*clientv3.LeaseRevokeResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Revoke", ctx, id)
+	ret0, _ := ret[0].(*clientv3.LeaseRevokeResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Revoke indicates an expected call of Revoke
+func (mr *MockEtcdClientMockRecorder) Revoke(ctx, id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Revoke", reflect.TypeOf((*MockEtcdClient)(nil).Revoke), ctx, id)
+}
+
+// Watch mocks base method
+func (m *MockEtcdClient) Watch(ctx context.Context, key string, opts ...clientv3.OpOption) clientv3.WatchChan {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{ctx, key}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Watch", varargs...)
+	ret0, _ := ret[0].(clientv3.WatchChan)
+	return ret0
+}
+
+// Watch indicates an expected call of Watch
+func (mr *MockEtcdClientMockRecorder) Watch(ctx, key interface{}, opts ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{ctx, key}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockEtcdClient)(nil).Watch), varargs...)
+}

core/discov/internal/listener.go

@@ -0,0 +1,6 @@
+//go:generate mockgen -package internal -destination listener_mock.go -source listener.go Listener
+package internal
+
+type Listener interface {
+	OnUpdate(keys []string, values []string, newKey string)
+}

core/discov/internal/listener_mock.go

@@ -0,0 +1,45 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: listener.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockListener is a mock of Listener interface
+type MockListener struct {
+	ctrl     *gomock.Controller
+	recorder *MockListenerMockRecorder
+}
+
+// MockListenerMockRecorder is the mock recorder for MockListener
+type MockListenerMockRecorder struct {
+	mock *MockListener
+}
+
+// NewMockListener creates a new mock instance
+func NewMockListener(ctrl *gomock.Controller) *MockListener {
+	mock := &MockListener{ctrl: ctrl}
+	mock.recorder = &MockListenerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockListener) EXPECT() *MockListenerMockRecorder {
+	return m.recorder
+}
+
+// OnUpdate mocks base method
+func (m *MockListener) OnUpdate(keys, values []string, newKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "OnUpdate", keys, values, newKey)
+}
+
+// OnUpdate indicates an expected call of OnUpdate
+func (mr *MockListenerMockRecorder) OnUpdate(keys, values, newKey interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnUpdate", reflect.TypeOf((*MockListener)(nil).OnUpdate), keys, values, newKey)
+}

core/discov/internal/registry.go

@@ -0,0 +1,310 @@
+package internal
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"zero/core/contextx"
+	"zero/core/lang"
+	"zero/core/logx"
+	"zero/core/syncx"
+	"zero/core/threading"
+
+	"go.etcd.io/etcd/clientv3"
+)
+
+var (
+	registryInstance = Registry{
+		clusters: make(map[string]*cluster),
+	}
+	connManager = syncx.NewResourceManager()
+)
+
+type Registry struct {
+	clusters map[string]*cluster
+	lock     sync.Mutex
+}
+
+func GetRegistry() *Registry {
+	return &registryInstance
+}
+
+func (r *Registry) getCluster(endpoints []string) *cluster {
+	clusterKey := getClusterKey(endpoints)
+	r.lock.Lock()
+	defer r.lock.Unlock()
+	c, ok := r.clusters[clusterKey]
+	if !ok {
+		c = newCluster(endpoints)
+		r.clusters[clusterKey] = c
+	}
+
+	return c
+}
+
+func (r *Registry) GetConn(endpoints []string) (EtcdClient, error) {
+	return r.getCluster(endpoints).getClient()
+}
+
+func (r *Registry) Monitor(endpoints []string, key string, l UpdateListener) error {
+	return r.getCluster(endpoints).monitor(key, l)
+}
+
+type cluster struct {
+	endpoints  []string
+	key        string
+	values     map[string]map[string]string
+	listeners  map[string][]UpdateListener
+	watchGroup *threading.RoutineGroup
+	done       chan lang.PlaceholderType
+	lock       sync.Mutex
+}
+
+func newCluster(endpoints []string) *cluster {
+	return &cluster{
+		endpoints:  endpoints,
+		key:        getClusterKey(endpoints),
+		values:     make(map[string]map[string]string),
+		listeners:  make(map[string][]UpdateListener),
+		watchGroup: threading.NewRoutineGroup(),
+		done:       make(chan lang.PlaceholderType),
+	}
+}
+
+func (c *cluster) context(cli EtcdClient) context.Context {
+	return contextx.ValueOnlyFrom(cli.Ctx())
+}
+
+func (c *cluster) getClient() (EtcdClient, error) {
+	val, err := connManager.GetResource(c.key, func() (io.Closer, error) {
+		return c.newClient()
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return val.(EtcdClient), nil
+}
+
+func (c *cluster) handleChanges(key string, kvs []KV) {
+	var add []KV
+	var remove []KV
+	c.lock.Lock()
+	listeners := append([]UpdateListener(nil), c.listeners[key]...)
+	vals, ok := c.values[key]
+	if !ok {
+		add = kvs
+		vals = make(map[string]string)
+		for _, kv := range kvs {
+			vals[kv.Key] = kv.Val
+		}
+		c.values[key] = vals
+	} else {
+		m := make(map[string]string)
+		for _, kv := range kvs {
+			m[kv.Key] = kv.Val
+		}
+		for k, v := range vals {
+			if val, ok := m[k]; !ok || v != val {
+				remove = append(remove, KV{
+					Key: k,
+					Val: v,
+				})
+			}
+		}
+		for k, v := range m {
+			if val, ok := vals[k]; !ok || v != val {
+				add = append(add, KV{
+					Key: k,
+					Val: v,
+				})
+			}
+		}
+		c.values[key] = m
+	}
+	c.lock.Unlock()
+
+	for _, kv := range add {
+		for _, l := range listeners {
+			l.OnAdd(kv)
+		}
+	}
+	for _, kv := range remove {
+		for _, l := range listeners {
+			l.OnDelete(kv)
+		}
+	}
+}
+
+func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
+	c.lock.Lock()
+	listeners := append([]UpdateListener(nil), c.listeners[key]...)
+	c.lock.Unlock()
+
+	for _, ev := range events {
+		switch ev.Type {
+		case clientv3.EventTypePut:
+			c.lock.Lock()
+			if vals, ok := c.values[key]; ok {
+				vals[string(ev.Kv.Key)] = string(ev.Kv.Value)
+			} else {
+				c.values[key] = map[string]string{string(ev.Kv.Key): string(ev.Kv.Value)}
+			}
+			c.lock.Unlock()
+			for _, l := range listeners {
+				l.OnAdd(KV{
+					Key: string(ev.Kv.Key),
+					Val: string(ev.Kv.Value),
+				})
+			}
+		case clientv3.EventTypeDelete:
+			if vals, ok := c.values[key]; ok {
+				delete(vals, string(ev.Kv.Key))
+			}
+			for _, l := range listeners {
+				l.OnDelete(KV{
+					Key: string(ev.Kv.Key),
+					Val: string(ev.Kv.Value),
+				})
+			}
+		default:
+			logx.Errorf("Unknown event type: %v", ev.Type)
+		}
+	}
+}
+
+func (c *cluster) load(cli EtcdClient, key string) {
+	var resp *clientv3.GetResponse
+	for {
+		var err error
+		ctx, cancel := context.WithTimeout(c.context(cli), RequestTimeout)
+		resp, err = cli.Get(ctx, makeKeyPrefix(key), clientv3.WithPrefix())
+		cancel()
+		if err == nil {
+			break
+		}
+
+		logx.Error(err)
+		time.Sleep(coolDownInterval)
+	}
+
+	var kvs []KV
+	c.lock.Lock()
+	for _, ev := range resp.Kvs {
+		kvs = append(kvs, KV{
+			Key: string(ev.Key),
+			Val: string(ev.Value),
+		})
+	}
+	c.lock.Unlock()
+
+	c.handleChanges(key, kvs)
+}
+
+func (c *cluster) monitor(key string, l UpdateListener) error {
+	c.lock.Lock()
+	c.listeners[key] = append(c.listeners[key], l)
+	c.lock.Unlock()
+
+	cli, err := c.getClient()
+	if err != nil {
+		return err
+	}
+
+	c.load(cli, key)
+	c.watchGroup.Run(func() {
+		c.watch(cli, key)
+	})
+
+	return nil
+}
+
+func (c *cluster) newClient() (EtcdClient, error) {
+	cli, err := NewClient(c.endpoints)
+	if err != nil {
+		return nil, err
+	}
+
+	go c.watchConnState(cli)
+
+	return cli, nil
+}
+
+func (c *cluster) reload(cli EtcdClient) {
+	c.lock.Lock()
+	close(c.done)
+	c.watchGroup.Wait()
+	c.done = make(chan lang.PlaceholderType)
+	c.watchGroup = threading.NewRoutineGroup()
+	var keys []string
+	for k := range c.listeners {
+		keys = append(keys, k)
+	}
+	c.lock.Unlock()
+
+	for _, key := range keys {
+		k := key
+		c.watchGroup.Run(func() {
+			c.load(cli, k)
+			c.watch(cli, k)
+		})
+	}
+}
+
+func (c *cluster) watch(cli EtcdClient, key string) {
+	rch := cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key), clientv3.WithPrefix())
+	for {
+		select {
+		case wresp, ok := <-rch:
+			if !ok {
+				logx.Error("etcd monitor chan has been closed")
+				return
+			}
+			if wresp.Canceled {
+				logx.Error("etcd monitor chan has been canceled")
+				return
+			}
+			if wresp.Err() != nil {
+				logx.Error(fmt.Sprintf("etcd monitor chan error: %v", wresp.Err()))
+				return
+			}
+
+			c.handleWatchEvents(key, wresp.Events)
+		case <-c.done:
+			return
+		}
+	}
+}
+
+func (c *cluster) watchConnState(cli EtcdClient) {
+	watcher := newStateWatcher()
+	watcher.addListener(func() {
+		go c.reload(cli)
+	})
+	watcher.watch(cli.ActiveConnection())
+}
+
+func DialClient(endpoints []string) (EtcdClient, error) {
+	return clientv3.New(clientv3.Config{
+		Endpoints:            endpoints,
+		AutoSyncInterval:     autoSyncInterval,
+		DialTimeout:          DialTimeout,
+		DialKeepAliveTime:    dialKeepAliveTime,
+		DialKeepAliveTimeout: DialTimeout,
+		RejectOldCluster:     true,
+	})
+}
+
+func getClusterKey(endpoints []string) string {
+	sort.Strings(endpoints)
+	return strings.Join(endpoints, endpointsSeparator)
+}
+
+func makeKeyPrefix(key string) string {
+	return fmt.Sprintf("%s%c", key, Delimiter)
+}

core/discov/internal/registry_test.go

@@ -0,0 +1,245 @@
+package internal
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"zero/core/contextx"
+	"zero/core/stringx"
+
+	"github.com/coreos/etcd/mvcc/mvccpb"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"go.etcd.io/etcd/clientv3"
+)
+
+var mockLock sync.Mutex
+
+func setMockClient(cli EtcdClient) func() {
+	mockLock.Lock()
+	NewClient = func([]string) (EtcdClient, error) {
+		return cli, nil
+	}
+	return func() {
+		NewClient = DialClient
+		mockLock.Unlock()
+	}
+}
+
+func TestGetCluster(t *testing.T) {
+	c1 := GetRegistry().getCluster([]string{"first"})
+	c2 := GetRegistry().getCluster([]string{"second"})
+	c3 := GetRegistry().getCluster([]string{"first"})
+	assert.Equal(t, c1, c3)
+	assert.NotEqual(t, c1, c2)
+}
+
+func TestGetClusterKey(t *testing.T) {
+	assert.Equal(t, getClusterKey([]string{"localhost:1234", "remotehost:5678"}),
+		getClusterKey([]string{"remotehost:5678", "localhost:1234"}))
+}
+
+func TestCluster_HandleChanges(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	l := NewMockUpdateListener(ctrl)
+	l.EXPECT().OnAdd(KV{
+		Key: "first",
+		Val: "1",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "second",
+		Val: "2",
+	})
+	l.EXPECT().OnDelete(KV{
+		Key: "first",
+		Val: "1",
+	})
+	l.EXPECT().OnDelete(KV{
+		Key: "second",
+		Val: "2",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "third",
+		Val: "3",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "fourth",
+		Val: "4",
+	})
+	c := newCluster([]string{"any"})
+	c.listeners["any"] = []UpdateListener{l}
+	c.handleChanges("any", []KV{
+		{
+			Key: "first",
+			Val: "1",
+		},
+		{
+			Key: "second",
+			Val: "2",
+		},
+	})
+	assert.EqualValues(t, map[string]string{
+		"first":  "1",
+		"second": "2",
+	}, c.values["any"])
+	c.handleChanges("any", []KV{
+		{
+			Key: "third",
+			Val: "3",
+		},
+		{
+			Key: "fourth",
+			Val: "4",
+		},
+	})
+	assert.EqualValues(t, map[string]string{
+		"third":  "3",
+		"fourth": "4",
+	}, c.values["any"])
+}
+
+func TestCluster_Load(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	cli.EXPECT().Get(gomock.Any(), "any/", gomock.Any()).Return(&clientv3.GetResponse{
+		Kvs: []*mvccpb.KeyValue{
+			{
+				Key:   []byte("hello"),
+				Value: []byte("world"),
+			},
+		},
+	}, nil)
+	cli.EXPECT().Ctx().Return(context.Background())
+	c := &cluster{
+		values: make(map[string]map[string]string),
+	}
+	c.load(cli, "any")
+}
+
+func TestCluster_Watch(t *testing.T) {
+	tests := []struct {
+		name      string
+		method    int
+		eventType mvccpb.Event_EventType
+	}{
+		{
+			name:      "add",
+			eventType: clientv3.EventTypePut,
+		},
+		{
+			name:      "delete",
+			eventType: clientv3.EventTypeDelete,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+			cli := NewMockEtcdClient(ctrl)
+			restore := setMockClient(cli)
+			defer restore()
+			ch := make(chan clientv3.WatchResponse)
+			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+			cli.EXPECT().Ctx().Return(context.Background())
+			var wg sync.WaitGroup
+			wg.Add(1)
+			c := &cluster{
+				listeners: make(map[string][]UpdateListener),
+				values:    make(map[string]map[string]string),
+			}
+			listener := NewMockUpdateListener(ctrl)
+			c.listeners["any"] = []UpdateListener{listener}
+			listener.EXPECT().OnAdd(gomock.Any()).Do(func(kv KV) {
+				assert.Equal(t, "hello", kv.Key)
+				assert.Equal(t, "world", kv.Val)
+				wg.Done()
+			}).MaxTimes(1)
+			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ interface{}) {
+				wg.Done()
+			}).MaxTimes(1)
+			go c.watch(cli, "any")
+			ch <- clientv3.WatchResponse{
+				Events: []*clientv3.Event{
+					{
+						Type: test.eventType,
+						Kv: &mvccpb.KeyValue{
+							Key:   []byte("hello"),
+							Value: []byte("world"),
+						},
+					},
+				},
+			}
+			wg.Wait()
+		})
+	}
+}
+
+func TestClusterWatch_RespFailures(t *testing.T) {
+	resps := []clientv3.WatchResponse{
+		{
+			Canceled: true,
+		},
+		{
+			// cause resp.Err() != nil
+			CompactRevision: 1,
+		},
+	}
+
+	for _, resp := range resps {
+		t.Run(stringx.Rand(), func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+			cli := NewMockEtcdClient(ctrl)
+			restore := setMockClient(cli)
+			defer restore()
+			ch := make(chan clientv3.WatchResponse)
+			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+			cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
+			c := new(cluster)
+			go func() {
+				ch <- resp
+			}()
+			c.watch(cli, "any")
+		})
+	}
+}
+
+func TestClusterWatch_CloseChan(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	ch := make(chan clientv3.WatchResponse)
+	cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+	cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
+	c := new(cluster)
+	go func() {
+		close(ch)
+	}()
+	c.watch(cli, "any")
+}
+
+func TestValueOnlyContext(t *testing.T) {
+	ctx := contextx.ValueOnlyFrom(context.Background())
+	ctx.Done()
+	assert.Nil(t, ctx.Err())
+}
+
+type mockedSharedCalls struct {
+	fn func() (interface{}, error)
+}
+
+func (c mockedSharedCalls) Do(_ string, fn func() (interface{}, error)) (interface{}, error) {
+	return c.fn()
+}
+
+func (c mockedSharedCalls) DoEx(_ string, fn func() (interface{}, error)) (interface{}, bool, error) {
+	val, err := c.fn()
+	return val, true, err
+}

core/discov/internal/roundrobinbalancer.go

@@ -0,0 +1,148 @@
+package internal
+
+import (
+	"math/rand"
+	"time"
+
+	"zero/core/logx"
+)
+
+type roundRobinBalancer struct {
+	*baseBalancer
+	conns []serverConn
+	index int
+}
+
+func NewRoundRobinBalancer(dialFn DialFn, closeFn CloseFn, exclusive bool) *roundRobinBalancer {
+	balancer := new(roundRobinBalancer)
+	balancer.baseBalancer = newBaseBalancer(dialFn, closeFn, exclusive)
+	return balancer
+}
+
+func (b *roundRobinBalancer) AddConn(kv KV) error {
+	var conn interface{}
+	prev, found := b.addKv(kv.Key, kv.Val)
+	if found {
+		conn = b.handlePrevious(prev, kv.Val)
+	}
+
+	if conn == nil {
+		var err error
+		conn, err = b.dialFn(kv.Val)
+		if err != nil {
+			b.removeKv(kv.Key)
+			return err
+		}
+	}
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	b.conns = append(b.conns, serverConn{
+		key:  kv.Key,
+		conn: conn,
+	})
+	b.notify(kv.Key)
+
+	return nil
+}
+
+func (b *roundRobinBalancer) handlePrevious(prev []string, server string) interface{} {
+	if len(prev) == 0 {
+		return nil
+	}
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if b.exclusive {
+		for _, item := range prev {
+			conns := b.conns[:0]
+			for _, each := range b.conns {
+				if each.key == item {
+					if err := b.closeFn(server, each.conn); err != nil {
+						logx.Error(err)
+					}
+				} else {
+					conns = append(conns, each)
+				}
+			}
+			b.conns = conns
+		}
+	} else {
+		for _, each := range b.conns {
+			if each.key == prev[0] {
+				return each.conn
+			}
+		}
+	}
+
+	return nil
+}
+
+func (b *roundRobinBalancer) initialize() {
+	rand.Seed(time.Now().UnixNano())
+	if len(b.conns) > 0 {
+		b.index = rand.Intn(len(b.conns))
+	}
+}
+
+func (b *roundRobinBalancer) IsEmpty() bool {
+	b.lock.Lock()
+	empty := len(b.conns) == 0
+	b.lock.Unlock()
+
+	return empty
+}
+
+func (b *roundRobinBalancer) Next(...string) (interface{}, bool) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if len(b.conns) == 0 {
+		return nil, false
+	}
+
+	b.index = (b.index + 1) % len(b.conns)
+	return b.conns[b.index].conn, true
+}
+
+func (b *roundRobinBalancer) notify(key string) {
+	if b.listener == nil {
+		return
+	}
+
+	// b.servers has the format of map[conn][]key
+	var keys []string
+	var values []string
+	for k, v := range b.servers {
+		values = append(values, k)
+		keys = append(keys, v...)
+	}
+
+	b.listener.OnUpdate(keys, values, key)
+}
+
+func (b *roundRobinBalancer) RemoveKey(key string) {
+	server, keep := b.removeKv(key)
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	conns := b.conns[:0]
+	for _, conn := range b.conns {
+		if conn.key == key {
+			// there are other keys assocated with the conn, don't close the conn.
+			if keep {
+				continue
+			}
+			if err := b.closeFn(server, conn.conn); err != nil {
+				logx.Error(err)
+			}
+		} else {
+			conns = append(conns, conn)
+		}
+	}
+	b.conns = conns
+	// notify without new key
+	b.notify("")
+}

core/discov/internal/roundrobinbalancer_test.go

@@ -0,0 +1,321 @@
+package internal
+
+import (
+	"errors"
+	"sort"
+	"strconv"
+	"testing"
+
+	"zero/core/logx"
+	"zero/core/mathx"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+)
+
+func init() {
+	logx.Disable()
+}
+
+func TestRoundRobin_addConn(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return errors.New("error")
+	}, false)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey1",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+	}, b.mapping)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey2",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey1",
+			conn: mockConn{server: "thevalue"},
+		},
+		{
+			key:  "thekey2",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1", "thekey2"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+		"thekey2": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey1")
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey2",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey2"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey2": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey2")
+	assert.EqualValues(t, []serverConn{}, b.conns)
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+	assert.True(t, b.IsEmpty())
+}
+
+func TestRoundRobin_addConnExclusive(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, true)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey1",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+	}, b.mapping)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey2",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey2",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey2"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey2": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey1")
+	b.RemoveKey("thekey2")
+	assert.EqualValues(t, []serverConn{}, b.conns)
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+	assert.True(t, b.IsEmpty())
+}
+
+func TestRoundRobin_addConnDupExclusive(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return errors.New("error")
+	}, true)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey1",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"thevalue": {"thekey1"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey1": "thevalue",
+	}, b.mapping)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey",
+		Val: "anothervalue",
+	}))
+	assert.Nil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.EqualValues(t, []serverConn{
+		{
+			key:  "thekey",
+			conn: mockConn{server: "anothervalue"},
+		},
+		{
+			key:  "thekey1",
+			conn: mockConn{server: "thevalue"},
+		},
+	}, b.conns)
+	assert.EqualValues(t, map[string][]string{
+		"anothervalue": {"thekey"},
+		"thevalue":     {"thekey1"},
+	}, b.servers)
+	assert.EqualValues(t, map[string]string{
+		"thekey":  "anothervalue",
+		"thekey1": "thevalue",
+	}, b.mapping)
+	assert.False(t, b.IsEmpty())
+
+	b.RemoveKey("thekey")
+	b.RemoveKey("thekey1")
+	assert.EqualValues(t, []serverConn{}, b.conns)
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+	assert.True(t, b.IsEmpty())
+}
+
+func TestRoundRobin_addConnError(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return nil, errors.New("error")
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, true)
+	assert.NotNil(t, b.AddConn(KV{
+		Key: "thekey1",
+		Val: "thevalue",
+	}))
+	assert.Nil(t, b.conns)
+	assert.EqualValues(t, map[string][]string{}, b.servers)
+	assert.EqualValues(t, map[string]string{}, b.mapping)
+}
+
+func TestRoundRobin_initialize(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, true)
+	for i := 0; i < 100; i++ {
+		assert.Nil(t, b.AddConn(KV{
+			Key: "thekey/" + strconv.Itoa(i),
+			Val: "thevalue/" + strconv.Itoa(i),
+		}))
+	}
+
+	m := make(map[int]int)
+	const total = 1000
+	for i := 0; i < total; i++ {
+		b.initialize()
+		m[b.index]++
+	}
+
+	mi := make(map[interface{}]int, len(m))
+	for k, v := range m {
+		mi[k] = v
+	}
+	entropy := mathx.CalcEntropy(mi, total)
+	assert.True(t, entropy > .95)
+}
+
+func TestRoundRobin_next(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return errors.New("error")
+	}, true)
+	const size = 100
+	for i := 0; i < size; i++ {
+		assert.Nil(t, b.AddConn(KV{
+			Key: "thekey/" + strconv.Itoa(i),
+			Val: "thevalue/" + strconv.Itoa(i),
+		}))
+	}
+
+	m := make(map[interface{}]int)
+	const total = 10000
+	for i := 0; i < total; i++ {
+		val, ok := b.Next()
+		assert.True(t, ok)
+		m[val]++
+	}
+
+	entropy := mathx.CalcEntropy(m, total)
+	assert.Equal(t, size, len(m))
+	assert.True(t, entropy > .95)
+
+	for i := 0; i < size; i++ {
+		b.RemoveKey("thekey/" + strconv.Itoa(i))
+	}
+	_, ok := b.Next()
+	assert.False(t, ok)
+}
+
+func TestRoundRobinBalancer_Listener(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, true)
+	assert.Nil(t, b.AddConn(KV{
+		Key: "key1",
+		Val: "val1",
+	}))
+	assert.Nil(t, b.AddConn(KV{
+		Key: "key2",
+		Val: "val2",
+	}))
+
+	listener := NewMockListener(ctrl)
+	listener.EXPECT().OnUpdate(gomock.Any(), gomock.Any(), "key2").Do(func(keys, vals, _ interface{}) {
+		sort.Strings(vals.([]string))
+		sort.Strings(keys.([]string))
+		assert.EqualValues(t, []string{"key1", "key2"}, keys)
+		assert.EqualValues(t, []string{"val1", "val2"}, vals)
+	})
+	b.setListener(listener)
+	b.notify("key2")
+}
+
+func TestRoundRobinBalancer_remove(t *testing.T) {
+	b := NewRoundRobinBalancer(func(server string) (interface{}, error) {
+		return mockConn{
+			server: server,
+		}, nil
+	}, func(server string, conn interface{}) error {
+		return nil
+	}, true)
+
+	assert.Nil(t, b.handlePrevious(nil, "any"))
+	_, ok := b.doRemoveKv("any")
+	assert.True(t, ok)
+}

core/discov/internal/statewatcher.go

@@ -0,0 +1,58 @@
+//go:generate mockgen -package internal -destination statewatcher_mock.go -source statewatcher.go etcdConn
+package internal
+
+import (
+	"context"
+	"sync"
+
+	"google.golang.org/grpc/connectivity"
+)
+
+type (
+	etcdConn interface {
+		GetState() connectivity.State
+		WaitForStateChange(ctx context.Context, sourceState connectivity.State) bool
+	}
+
+	stateWatcher struct {
+		disconnected bool
+		currentState connectivity.State
+		listeners    []func()
+		lock         sync.Mutex
+	}
+)
+
+func newStateWatcher() *stateWatcher {
+	return new(stateWatcher)
+}
+
+func (sw *stateWatcher) addListener(l func()) {
+	sw.lock.Lock()
+	sw.listeners = append(sw.listeners, l)
+	sw.lock.Unlock()
+}
+
+func (sw *stateWatcher) watch(conn etcdConn) {
+	sw.currentState = conn.GetState()
+	for {
+		if conn.WaitForStateChange(context.Background(), sw.currentState) {
+			newState := conn.GetState()
+			sw.lock.Lock()
+			sw.currentState = newState
+
+			switch newState {
+			case connectivity.TransientFailure, connectivity.Shutdown:
+				sw.disconnected = true
+			case connectivity.Ready:
+				if sw.disconnected {
+					sw.disconnected = false
+					for _, l := range sw.listeners {
+						l()
+					}
+				}
+			}
+
+			sw.lock.Unlock()
+		}
+	}
+}

core/discov/internal/statewatcher_mock.go

@@ -0,0 +1,63 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: statewatcher.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	context "context"
+	gomock "github.com/golang/mock/gomock"
+	connectivity "google.golang.org/grpc/connectivity"
+	reflect "reflect"
+)
+
+// MocketcdConn is a mock of etcdConn interface
+type MocketcdConn struct {
+	ctrl     *gomock.Controller
+	recorder *MocketcdConnMockRecorder
+}
+
+// MocketcdConnMockRecorder is the mock recorder for MocketcdConn
+type MocketcdConnMockRecorder struct {
+	mock *MocketcdConn
+}
+
+// NewMocketcdConn creates a new mock instance
+func NewMocketcdConn(ctrl *gomock.Controller) *MocketcdConn {
+	mock := &MocketcdConn{ctrl: ctrl}
+	mock.recorder = &MocketcdConnMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MocketcdConn) EXPECT() *MocketcdConnMockRecorder {
+	return m.recorder
+}
+
+// GetState mocks base method
+func (m *MocketcdConn) GetState() connectivity.State {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetState")
+	ret0, _ := ret[0].(connectivity.State)
+	return ret0
+}
+
+// GetState indicates an expected call of GetState
+func (mr *MocketcdConnMockRecorder) GetState() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetState", reflect.TypeOf((*MocketcdConn)(nil).GetState))
+}
+
+// WaitForStateChange mocks base method
+func (m *MocketcdConn) WaitForStateChange(ctx context.Context, sourceState connectivity.State) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForStateChange", ctx, sourceState)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// WaitForStateChange indicates an expected call of WaitForStateChange
+func (mr *MocketcdConnMockRecorder) WaitForStateChange(ctx, sourceState interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForStateChange", reflect.TypeOf((*MocketcdConn)(nil).WaitForStateChange), ctx, sourceState)
+}

core/discov/internal/statewatcher_test.go

@@ -0,0 +1,27 @@
+package internal
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"google.golang.org/grpc/connectivity"
+)
+
+func TestStateWatcher_watch(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	watcher := newStateWatcher()
+	var wg sync.WaitGroup
+	wg.Add(1)
+	watcher.addListener(func() {
+		wg.Done()
+	})
+	conn := NewMocketcdConn(ctrl)
+	conn.EXPECT().GetState().Return(connectivity.Ready)
+	conn.EXPECT().GetState().Return(connectivity.TransientFailure)
+	conn.EXPECT().GetState().Return(connectivity.Ready).AnyTimes()
+	conn.EXPECT().WaitForStateChange(gomock.Any(), gomock.Any()).Return(true).AnyTimes()
+	go watcher.watch(conn)
+	wg.Wait()
+}

core/discov/internal/updatelistener.go

@@ -0,0 +1,14 @@
+//go:generate mockgen -package internal -destination updatelistener_mock.go -source updatelistener.go UpdateListener
+package internal
+
+type (
+	KV struct {
+		Key string
+		Val string
+	}
+
+	UpdateListener interface {
+		OnAdd(kv KV)
+		OnDelete(kv KV)
+	}
+)

core/discov/internal/updatelistener_mock.go

@@ -0,0 +1,57 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: updatelistener.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockUpdateListener is a mock of UpdateListener interface
+type MockUpdateListener struct {
+	ctrl     *gomock.Controller
+	recorder *MockUpdateListenerMockRecorder
+}
+
+// MockUpdateListenerMockRecorder is the mock recorder for MockUpdateListener
+type MockUpdateListenerMockRecorder struct {
+	mock *MockUpdateListener
+}
+
+// NewMockUpdateListener creates a new mock instance
+func NewMockUpdateListener(ctrl *gomock.Controller) *MockUpdateListener {
+	mock := &MockUpdateListener{ctrl: ctrl}
+	mock.recorder = &MockUpdateListenerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockUpdateListener) EXPECT() *MockUpdateListenerMockRecorder {
+	return m.recorder
+}
+
+// OnAdd mocks base method
+func (m *MockUpdateListener) OnAdd(kv KV) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "OnAdd", kv)
+}
+
+// OnAdd indicates an expected call of OnAdd
+func (mr *MockUpdateListenerMockRecorder) OnAdd(kv interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnAdd", reflect.TypeOf((*MockUpdateListener)(nil).OnAdd), kv)
+}
+
+// OnDelete mocks base method
+func (m *MockUpdateListener) OnDelete(kv KV) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "OnDelete", kv)
+}
+
+// OnDelete indicates an expected call of OnDelete
+func (mr *MockUpdateListenerMockRecorder) OnDelete(kv interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnDelete", reflect.TypeOf((*MockUpdateListener)(nil).OnDelete), kv)
+}

core/discov/internal/vars.go

@@ -0,0 +1,19 @@
+package internal
+
+import "time"
+
+const (
+	autoSyncInterval   = time.Minute
+	coolDownInterval   = time.Second
+	dialTimeout        = 5 * time.Second
+	dialKeepAliveTime  = 5 * time.Second
+	requestTimeout     = 3 * time.Second
+	Delimiter          = '/'
+	endpointsSeparator = ","
+)
+
+var (
+	DialTimeout    = dialTimeout
+	RequestTimeout = requestTimeout
+	NewClient      = DialClient
+)