initial import

core/discov/subclient.go

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