go-zero/core/discov/subscriber.go

package discov

import (
	"sync"
	"sync/atomic"

	"github.com/tal-tech/go-zero/core/discov/internal"
	"github.com/tal-tech/go-zero/core/syncx"
)

type (
	// SubOption defines the method to customize a Subscriber.
	SubOption func(sub *Subscriber)

	// A Subscriber is used to subscribe the given key on a etcd cluster.
	Subscriber struct {
		endpoints []string
		exclusive bool
		items     *container
	}
)

// NewSubscriber returns a Subscriber.
// endpoints is the hosts of the etcd cluster.
// key is the key to subscribe.
// opts are used to customize the Subscriber.
func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscriber, error) {
	sub := &Subscriber{
		endpoints: endpoints,
	}
	for _, opt := range opts {
		opt(sub)
	}
	sub.items = newContainer(sub.exclusive)

	if err := internal.GetRegistry().Monitor(endpoints, key, sub.items); err != nil {
		return nil, err
	}

	return sub, nil
}

// AddListener adds listener to s.
func (s *Subscriber) AddListener(listener func()) {
	s.items.addListener(listener)
}

// Values returns all the subscription values.
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(sub *Subscriber) {
		sub.exclusive = true
	}
}

// WithSubEtcdAccount customizes the Subscriber with given etcd username/password.
func WithSubEtcdAccount(user, pass string) SubOption {
	return func(sub *Subscriber) {
		internal.AddAccount(sub.endpoints, user, pass)
	}
}

type container struct {
	exclusive bool
	values    map[string][]string
	mapping   map[string]string
	snapshot  atomic.Value
	dirty     *syncx.AtomicBool
	listeners []func()
	lock      sync.Mutex
}

func newContainer(exclusive bool) *container {
	return &container{
		exclusive: exclusive,
		values:    make(map[string][]string),
		mapping:   make(map[string]string),
		dirty:     syncx.ForAtomicBool(true),
	}
}

func (c *container) OnAdd(kv internal.KV) {
	c.addKv(kv.Key, kv.Val)
	c.notifyChange()
}

func (c *container) OnDelete(kv internal.KV) {
	c.removeKey(kv.Key)
	c.notifyChange()
}

// 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()

	c.dirty.Set(true)
	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
	}

	return nil, false
}

func (c *container) addListener(listener func()) {
	c.lock.Lock()
	c.listeners = append(c.listeners, listener)
	c.lock.Unlock()
}

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 {
	if !c.dirty.True() {
		return c.snapshot.Load().([]string)
	}

	c.lock.Lock()
	defer c.lock.Unlock()

	var vals []string
	for each := range c.values {
		vals = append(vals, each)
	}
	c.snapshot.Store(vals)
	c.dirty.Set(false)

	return vals
}

func (c *container) notifyChange() {
	c.lock.Lock()
	listeners := append(([]func())(nil), c.listeners...)
	c.lock.Unlock()

	for _, listener := range listeners {
		listener()
	}
}

// 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.dirty.Set(true)
	c.doRemoveKey(key)
}