go-zero/core/discov/subscriber.go

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
}