initial import

This commit is contained in:
kevin
2020-07-26 17:09:05 +08:00
commit 7e3a369a8f
647 changed files with 54754 additions and 0 deletions

40
core/discov/clients.go Normal file
View File

@@ -0,0 +1,40 @@
package discov
import (
"fmt"
"strings"
"zero/core/discov/internal"
)
const (
indexOfKey = iota
indexOfId
)
const timeToLive int64 = 10
var TimeToLive = timeToLive
func extract(etcdKey string, index int) (string, bool) {
if index < 0 {
return "", false
}
fields := strings.FieldsFunc(etcdKey, func(ch rune) bool {
return ch == internal.Delimiter
})
if index >= len(fields) {
return "", false
}
return fields[index], true
}
func extractId(etcdKey string) (string, bool) {
return extract(etcdKey, indexOfId)
}
func makeEtcdKey(key string, id int64) string {
return fmt.Sprintf("%s%c%d", key, internal.Delimiter, id)
}

View File

@@ -0,0 +1,36 @@
package discov
import (
"sync"
"testing"
"zero/core/discov/internal"
"github.com/stretchr/testify/assert"
)
var mockLock sync.Mutex
func setMockClient(cli internal.EtcdClient) func() {
mockLock.Lock()
internal.NewClient = func([]string) (internal.EtcdClient, error) {
return cli, nil
}
return func() {
internal.NewClient = internal.DialClient
mockLock.Unlock()
}
}
func TestExtract(t *testing.T) {
id, ok := extractId("key/123/val")
assert.True(t, ok)
assert.Equal(t, "123", id)
_, ok = extract("any", -1)
assert.False(t, ok)
}
func TestMakeKey(t *testing.T) {
assert.Equal(t, "key/123", makeEtcdKey("key", 123))
}

18
core/discov/config.go Normal file
View File

@@ -0,0 +1,18 @@
package discov
import "errors"
type EtcdConf struct {
Hosts []string
Key string
}
func (c EtcdConf) Validate() error {
if len(c.Hosts) == 0 {
return errors.New("empty etcd hosts")
} else if len(c.Key) == 0 {
return errors.New("empty etcd key")
} else {
return nil
}
}

View File

@@ -0,0 +1,46 @@
package discov
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestConfig(t *testing.T) {
tests := []struct {
EtcdConf
pass bool
}{
{
EtcdConf: EtcdConf{},
pass: false,
},
{
EtcdConf: EtcdConf{
Key: "any",
},
pass: false,
},
{
EtcdConf: EtcdConf{
Hosts: []string{"any"},
},
pass: false,
},
{
EtcdConf: EtcdConf{
Hosts: []string{"any"},
Key: "key",
},
pass: true,
},
}
for _, test := range tests {
if test.pass {
assert.Nil(t, test.EtcdConf.Validate())
} else {
assert.NotNil(t, test.EtcdConf.Validate())
}
}
}

47
core/discov/facade.go Normal file
View File

@@ -0,0 +1,47 @@
package discov
import (
"zero/core/discov/internal"
"zero/core/lang"
)
type (
Facade struct {
endpoints []string
registry *internal.Registry
}
FacadeListener interface {
OnAdd(key, val string)
OnDelete(key string)
}
)
func NewFacade(endpoints []string) Facade {
return Facade{
endpoints: endpoints,
registry: internal.GetRegistry(),
}
}
func (f Facade) Client() internal.EtcdClient {
conn, err := f.registry.GetConn(f.endpoints)
lang.Must(err)
return conn
}
func (f Facade) Monitor(key string, l FacadeListener) {
f.registry.Monitor(f.endpoints, key, listenerAdapter{l})
}
type listenerAdapter struct {
l FacadeListener
}
func (la listenerAdapter) OnAdd(kv internal.KV) {
la.l.OnAdd(kv.Key, kv.Val)
}
func (la listenerAdapter) OnDelete(kv internal.KV) {
la.l.OnDelete(kv.Key)
}

View File

@@ -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()
}

View File

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

View File

@@ -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))
}
}

View File

@@ -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"}))
}

View File

@@ -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
}

View File

@@ -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...)
}

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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
}

View File

@@ -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("")
}

View File

@@ -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)
}

View File

@@ -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()
}
}
}

View File

@@ -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)
}

View File

@@ -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()
}

View File

@@ -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)
}
)

View File

@@ -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)
}

View File

@@ -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
)

View File

@@ -0,0 +1,4 @@
apiVersion: v1
kind: Namespace
metadata:
name: discov

View File

@@ -0,0 +1,378 @@
apiVersion: v1
kind: Service
metadata:
name: etcd
namespace: discov
spec:
ports:
- name: etcd-port
port: 2379
protocol: TCP
targetPort: 2379
selector:
app: etcd
---
apiVersion: v1
kind: Pod
metadata:
labels:
app: etcd
etcd_node: etcd0
name: etcd0
namespace: discov
spec:
containers:
- command:
- /usr/local/bin/etcd
- --name
- etcd0
- --initial-advertise-peer-urls
- http://etcd0:2380
- --listen-peer-urls
- http://0.0.0.0:2380
- --listen-client-urls
- http://0.0.0.0:2379
- --advertise-client-urls
- http://etcd0:2379
- --initial-cluster
- etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
- --initial-cluster-state
- new
image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
name: etcd0
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
imagePullSecrets:
- name: aliyun
affinity:
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values:
- etcd
topologyKey: "kubernetes.io/hostname"
restartPolicy: Always
---
apiVersion: v1
kind: Service
metadata:
labels:
etcd_node: etcd0
name: etcd0
namespace: discov
spec:
ports:
- name: client
port: 2379
protocol: TCP
targetPort: 2379
- name: server
port: 2380
protocol: TCP
targetPort: 2380
selector:
etcd_node: etcd0
---
apiVersion: v1
kind: Pod
metadata:
labels:
app: etcd
etcd_node: etcd1
name: etcd1
namespace: discov
spec:
containers:
- command:
- /usr/local/bin/etcd
- --name
- etcd1
- --initial-advertise-peer-urls
- http://etcd1:2380
- --listen-peer-urls
- http://0.0.0.0:2380
- --listen-client-urls
- http://0.0.0.0:2379
- --advertise-client-urls
- http://etcd1:2379
- --initial-cluster
- etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
- --initial-cluster-state
- new
image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
name: etcd1
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
imagePullSecrets:
- name: aliyun
affinity:
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values:
- etcd
topologyKey: "kubernetes.io/hostname"
restartPolicy: Always
---
apiVersion: v1
kind: Service
metadata:
labels:
etcd_node: etcd1
name: etcd1
namespace: discov
spec:
ports:
- name: client
port: 2379
protocol: TCP
targetPort: 2379
- name: server
port: 2380
protocol: TCP
targetPort: 2380
selector:
etcd_node: etcd1
---
apiVersion: v1
kind: Pod
metadata:
labels:
app: etcd
etcd_node: etcd2
name: etcd2
namespace: discov
spec:
containers:
- command:
- /usr/local/bin/etcd
- --name
- etcd2
- --initial-advertise-peer-urls
- http://etcd2:2380
- --listen-peer-urls
- http://0.0.0.0:2380
- --listen-client-urls
- http://0.0.0.0:2379
- --advertise-client-urls
- http://etcd2:2379
- --initial-cluster
- etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
- --initial-cluster-state
- new
image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
name: etcd2
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
imagePullSecrets:
- name: aliyun
affinity:
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values:
- etcd
topologyKey: "kubernetes.io/hostname"
restartPolicy: Always
---
apiVersion: v1
kind: Service
metadata:
labels:
etcd_node: etcd2
name: etcd2
namespace: discov
spec:
ports:
- name: client
port: 2379
protocol: TCP
targetPort: 2379
- name: server
port: 2380
protocol: TCP
targetPort: 2380
selector:
etcd_node: etcd2
---
apiVersion: v1
kind: Pod
metadata:
labels:
app: etcd
etcd_node: etcd3
name: etcd3
namespace: discov
spec:
containers:
- command:
- /usr/local/bin/etcd
- --name
- etcd3
- --initial-advertise-peer-urls
- http://etcd3:2380
- --listen-peer-urls
- http://0.0.0.0:2380
- --listen-client-urls
- http://0.0.0.0:2379
- --advertise-client-urls
- http://etcd3:2379
- --initial-cluster
- etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
- --initial-cluster-state
- new
image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
name: etcd3
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
imagePullSecrets:
- name: aliyun
affinity:
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values:
- etcd
topologyKey: "kubernetes.io/hostname"
restartPolicy: Always
---
apiVersion: v1
kind: Service
metadata:
labels:
etcd_node: etcd3
name: etcd3
namespace: discov
spec:
ports:
- name: client
port: 2379
protocol: TCP
targetPort: 2379
- name: server
port: 2380
protocol: TCP
targetPort: 2380
selector:
etcd_node: etcd3
---
apiVersion: v1
kind: Pod
metadata:
labels:
app: etcd
etcd_node: etcd4
name: etcd4
namespace: discov
spec:
containers:
- command:
- /usr/local/bin/etcd
- --name
- etcd4
- --initial-advertise-peer-urls
- http://etcd4:2380
- --listen-peer-urls
- http://0.0.0.0:2380
- --listen-client-urls
- http://0.0.0.0:2379
- --advertise-client-urls
- http://etcd4:2379
- --initial-cluster
- etcd0=http://etcd0:2380,etcd1=http://etcd1:2380,etcd2=http://etcd2:2380,etcd3=http://etcd3:2380,etcd4=http://etcd4:2380
- --initial-cluster-state
- new
image: registry-vpc.cn-hangzhou.aliyuncs.com/xapp/etcd:latest
name: etcd4
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
imagePullSecrets:
- name: aliyun
affinity:
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values:
- etcd
topologyKey: "kubernetes.io/hostname"
restartPolicy: Always
---
apiVersion: v1
kind: Service
metadata:
labels:
etcd_node: etcd4
name: etcd4
namespace: discov
spec:
ports:
- name: client
port: 2379
protocol: TCP
targetPort: 2379
- name: server
port: 2380
protocol: TCP
targetPort: 2380
selector:
etcd_node: etcd4

143
core/discov/publisher.go Normal file
View File

@@ -0,0 +1,143 @@
package discov
import (
"zero/core/discov/internal"
"zero/core/lang"
"zero/core/logx"
"zero/core/proc"
"zero/core/syncx"
"zero/core/threading"
"go.etcd.io/etcd/clientv3"
)
type (
PublisherOption func(client *Publisher)
Publisher struct {
endpoints []string
key string
fullKey string
id int64
value string
lease clientv3.LeaseID
quit *syncx.DoneChan
pauseChan chan lang.PlaceholderType
resumeChan chan lang.PlaceholderType
}
)
func NewPublisher(endpoints []string, key, value string, opts ...PublisherOption) *Publisher {
publisher := &Publisher{
endpoints: endpoints,
key: key,
value: value,
quit: syncx.NewDoneChan(),
pauseChan: make(chan lang.PlaceholderType),
resumeChan: make(chan lang.PlaceholderType),
}
for _, opt := range opts {
opt(publisher)
}
return publisher
}
func (p *Publisher) KeepAlive() error {
cli, err := internal.GetRegistry().GetConn(p.endpoints)
if err != nil {
return err
}
p.lease, err = p.register(cli)
if err != nil {
return err
}
proc.AddWrapUpListener(func() {
p.Stop()
})
return p.keepAliveAsync(cli)
}
func (p *Publisher) Pause() {
p.pauseChan <- lang.Placeholder
}
func (p *Publisher) Resume() {
p.resumeChan <- lang.Placeholder
}
func (p *Publisher) Stop() {
p.quit.Close()
}
func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
ch, err := cli.KeepAlive(cli.Ctx(), p.lease)
if err != nil {
return err
}
threading.GoSafe(func() {
for {
select {
case _, ok := <-ch:
if !ok {
p.revoke(cli)
if err := p.KeepAlive(); err != nil {
logx.Errorf("KeepAlive: %s", err.Error())
}
return
}
case <-p.pauseChan:
logx.Infof("paused etcd renew, key: %s, value: %s", p.key, p.value)
p.revoke(cli)
select {
case <-p.resumeChan:
if err := p.KeepAlive(); err != nil {
logx.Errorf("KeepAlive: %s", err.Error())
}
return
case <-p.quit.Done():
return
}
case <-p.quit.Done():
p.revoke(cli)
return
}
}
})
return nil
}
func (p *Publisher) register(client internal.EtcdClient) (clientv3.LeaseID, error) {
resp, err := client.Grant(client.Ctx(), TimeToLive)
if err != nil {
return clientv3.NoLease, err
}
lease := resp.ID
if p.id > 0 {
p.fullKey = makeEtcdKey(p.key, p.id)
} else {
p.fullKey = makeEtcdKey(p.key, int64(lease))
}
_, err = client.Put(client.Ctx(), p.fullKey, p.value, clientv3.WithLease(lease))
return lease, err
}
func (p *Publisher) revoke(cli internal.EtcdClient) {
if _, err := cli.Revoke(cli.Ctx(), p.lease); err != nil {
logx.Error(err)
}
}
func WithId(id int64) PublisherOption {
return func(publisher *Publisher) {
publisher.id = id
}
}

View File

@@ -0,0 +1,151 @@
package discov
import (
"errors"
"sync"
"testing"
"zero/core/discov/internal"
"zero/core/logx"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
)
func init() {
logx.Disable()
}
func TestPublisher_register(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
ID: id,
}, nil)
cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
pub := NewPublisher(nil, "thekey", "thevalue")
_, err := pub.register(cli)
assert.Nil(t, err)
}
func TestPublisher_registerWithId(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id = 2
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
ID: 1,
}, nil)
cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
pub := NewPublisher(nil, "thekey", "thevalue", WithId(id))
_, err := pub.register(cli)
assert.Nil(t, err)
}
func TestPublisher_registerError(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(nil, errors.New("error"))
pub := NewPublisher(nil, "thekey", "thevalue")
val, err := pub.register(cli)
assert.NotNil(t, err)
assert.Equal(t, clientv3.NoLease, val)
}
func TestPublisher_revoke(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id clientv3.LeaseID = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().Revoke(gomock.Any(), id)
pub := NewPublisher(nil, "thekey", "thevalue")
pub.lease = id
pub.revoke(cli)
}
func TestPublisher_revokeError(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id clientv3.LeaseID = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().Revoke(gomock.Any(), id).Return(nil, errors.New("error"))
pub := NewPublisher(nil, "thekey", "thevalue")
pub.lease = id
pub.revoke(cli)
}
func TestPublisher_keepAliveAsyncError(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id clientv3.LeaseID = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().KeepAlive(gomock.Any(), id).Return(nil, errors.New("error"))
pub := NewPublisher(nil, "thekey", "thevalue")
pub.lease = id
assert.NotNil(t, pub.keepAliveAsync(cli))
}
func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id clientv3.LeaseID = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().KeepAlive(gomock.Any(), id)
var wg sync.WaitGroup
wg.Add(1)
cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
wg.Done()
})
pub := NewPublisher(nil, "thekey", "thevalue")
pub.lease = id
pub.Stop()
assert.Nil(t, pub.keepAliveAsync(cli))
wg.Wait()
}
func TestPublisher_keepAliveAsyncPause(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
const id clientv3.LeaseID = 1
cli := internal.NewMockEtcdClient(ctrl)
restore := setMockClient(cli)
defer restore()
cli.EXPECT().Ctx().AnyTimes()
cli.EXPECT().KeepAlive(gomock.Any(), id)
pub := NewPublisher(nil, "thekey", "thevalue")
var wg sync.WaitGroup
wg.Add(1)
cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
pub.Stop()
wg.Done()
})
pub.lease = id
assert.Nil(t, pub.keepAliveAsync(cli))
pub.Pause()
wg.Wait()
}

35
core/discov/renewer.go Normal file
View File

@@ -0,0 +1,35 @@
package discov
import "zero/core/logx"
type (
Renewer interface {
Start()
Stop()
Pause()
Resume()
}
etcdRenewer struct {
*Publisher
}
)
func NewRenewer(endpoints []string, key, value string, renewId int64) Renewer {
var publisher *Publisher
if renewId > 0 {
publisher = NewPublisher(endpoints, key, value, WithId(renewId))
} else {
publisher = NewPublisher(endpoints, key, value)
}
return &etcdRenewer{
Publisher: publisher,
}
}
func (sr *etcdRenewer) Start() {
if err := sr.KeepAlive(); err != nil {
logx.Error(err)
}
}

186
core/discov/subclient.go Normal file
View File

@@ -0,0 +1,186 @@
package discov
import (
"sync"
"zero/core/discov/internal"
"zero/core/logx"
)
const (
_ = iota // keyBasedBalance, default
idBasedBalance
)
type (
Listener internal.Listener
subClient struct {
balancer internal.Balancer
lock sync.Mutex
cond *sync.Cond
listeners []internal.Listener
}
balanceOptions struct {
balanceType int
}
BalanceOption func(*balanceOptions)
RoundRobinSubClient struct {
*subClient
}
ConsistentSubClient struct {
*subClient
}
BatchConsistentSubClient struct {
*ConsistentSubClient
}
)
func NewRoundRobinSubClient(endpoints []string, key string, dialFn internal.DialFn, closeFn internal.CloseFn,
opts ...SubOption) (*RoundRobinSubClient, error) {
var subOpts subOptions
for _, opt := range opts {
opt(&subOpts)
}
cli, err := newSubClient(endpoints, key, internal.NewRoundRobinBalancer(dialFn, closeFn, subOpts.exclusive))
if err != nil {
return nil, err
}
return &RoundRobinSubClient{
subClient: cli,
}, nil
}
func NewConsistentSubClient(endpoints []string, key string, dialFn internal.DialFn,
closeFn internal.CloseFn, opts ...BalanceOption) (*ConsistentSubClient, error) {
var balanceOpts balanceOptions
for _, opt := range opts {
opt(&balanceOpts)
}
var keyer func(internal.KV) string
switch balanceOpts.balanceType {
case idBasedBalance:
keyer = func(kv internal.KV) string {
if id, ok := extractId(kv.Key); ok {
return id
} else {
return kv.Key
}
}
default:
keyer = func(kv internal.KV) string {
return kv.Val
}
}
cli, err := newSubClient(endpoints, key, internal.NewConsistentBalancer(dialFn, closeFn, keyer))
if err != nil {
return nil, err
}
return &ConsistentSubClient{
subClient: cli,
}, nil
}
func NewBatchConsistentSubClient(endpoints []string, key string, dialFn internal.DialFn, closeFn internal.CloseFn,
opts ...BalanceOption) (*BatchConsistentSubClient, error) {
cli, err := NewConsistentSubClient(endpoints, key, dialFn, closeFn, opts...)
if err != nil {
return nil, err
}
return &BatchConsistentSubClient{
ConsistentSubClient: cli,
}, nil
}
func newSubClient(endpoints []string, key string, balancer internal.Balancer) (*subClient, error) {
client := &subClient{
balancer: balancer,
}
client.cond = sync.NewCond(&client.lock)
if err := internal.GetRegistry().Monitor(endpoints, key, client); err != nil {
return nil, err
}
return client, nil
}
func (c *subClient) AddListener(listener internal.Listener) {
c.lock.Lock()
c.listeners = append(c.listeners, listener)
c.lock.Unlock()
}
func (c *subClient) OnAdd(kv internal.KV) {
c.lock.Lock()
defer c.lock.Unlock()
if err := c.balancer.AddConn(kv); err != nil {
logx.Error(err)
} else {
c.cond.Broadcast()
}
}
func (c *subClient) OnDelete(kv internal.KV) {
c.balancer.RemoveKey(kv.Key)
}
func (c *subClient) WaitForServers() {
logx.Error("Waiting for alive servers")
c.lock.Lock()
defer c.lock.Unlock()
if c.balancer.IsEmpty() {
c.cond.Wait()
}
}
func (c *subClient) onAdd(keys []string, servers []string, newKey string) {
// guarded by locked outside
for _, listener := range c.listeners {
listener.OnUpdate(keys, servers, newKey)
}
}
func (c *RoundRobinSubClient) Next() (interface{}, bool) {
return c.balancer.Next()
}
func (c *ConsistentSubClient) Next(key string) (interface{}, bool) {
return c.balancer.Next(key)
}
func (bc *BatchConsistentSubClient) Next(keys []string) (map[interface{}][]string, bool) {
if len(keys) == 0 {
return nil, false
}
result := make(map[interface{}][]string)
for _, key := range keys {
dest, ok := bc.ConsistentSubClient.Next(key)
if !ok {
return nil, false
}
result[dest] = append(result[dest], key)
}
return result, true
}
func BalanceWithId() BalanceOption {
return func(opts *balanceOptions) {
opts.balanceType = idBasedBalance
}
}

151
core/discov/subscriber.go Normal file
View File

@@ -0,0 +1,151 @@
package discov
import (
"sync"
"zero/core/discov/internal"
)
type (
subOptions struct {
exclusive bool
}
SubOption func(opts *subOptions)
Subscriber struct {
items *container
}
)
func NewSubscriber(endpoints []string, key string, opts ...SubOption) *Subscriber {
var subOpts subOptions
for _, opt := range opts {
opt(&subOpts)
}
subscriber := &Subscriber{
items: newContainer(subOpts.exclusive),
}
internal.GetRegistry().Monitor(endpoints, key, subscriber.items)
return subscriber
}
func (s *Subscriber) Values() []string {
return s.items.getValues()
}
// exclusive means that key value can only be 1:1,
// which means later added value will remove the keys associated with the same value previously.
func Exclusive() SubOption {
return func(opts *subOptions) {
opts.exclusive = true
}
}
type container struct {
exclusive bool
values map[string][]string
mapping map[string]string
lock sync.Mutex
}
func newContainer(exclusive bool) *container {
return &container{
exclusive: exclusive,
values: make(map[string][]string),
mapping: make(map[string]string),
}
}
func (c *container) OnAdd(kv internal.KV) {
c.addKv(kv.Key, kv.Val)
}
func (c *container) OnDelete(kv internal.KV) {
c.removeKey(kv.Key)
}
// addKv adds the kv, returns if there are already other keys associate with the value
func (c *container) addKv(key, value string) ([]string, bool) {
c.lock.Lock()
defer c.lock.Unlock()
keys := c.values[value]
previous := append([]string(nil), keys...)
early := len(keys) > 0
if c.exclusive && early {
for _, each := range keys {
c.doRemoveKey(each)
}
}
c.values[value] = append(c.values[value], key)
c.mapping[key] = value
if early {
return previous, true
} else {
return nil, false
}
}
func (c *container) doRemoveKey(key string) {
server, ok := c.mapping[key]
if !ok {
return
}
delete(c.mapping, key)
keys := c.values[server]
remain := keys[:0]
for _, k := range keys {
if k != key {
remain = append(remain, k)
}
}
if len(remain) > 0 {
c.values[server] = remain
} else {
delete(c.values, server)
}
}
func (c *container) getValues() []string {
c.lock.Lock()
defer c.lock.Unlock()
var vs []string
for each := range c.values {
vs = append(vs, each)
}
return vs
}
// removeKey removes the kv, returns true if there are still other keys associate with the value
func (c *container) removeKey(key string) {
c.lock.Lock()
defer c.lock.Unlock()
c.doRemoveKey(key)
}
func (c *container) removeVal(val string) (empty bool) {
c.lock.Lock()
defer c.lock.Unlock()
for k := range c.values {
if k == val {
delete(c.values, k)
}
}
for k, v := range c.mapping {
if v == val {
delete(c.mapping, k)
}
}
return len(c.values) == 0
}