go-zero/core/discov/subclient.go

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