chore: import upstream snapshot with attribution

registry/nats/nats.go

@@ -0,0 +1,417 @@
+// Package nats provides a NATS registry using broadcast queries
+package nats
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/nats-io/nats.go"
+	"go-micro.dev/v6/registry"
+)
+
+type natsRegistry struct {
+	addrs          []string
+	opts           registry.Options
+	nopts          nats.Options
+	queryTopic     string
+	watchTopic     string
+	registerAction string
+
+	sync.RWMutex
+	conn      *nats.Conn
+	services  map[string][]*registry.Service
+	listeners map[string]chan bool
+}
+
+var (
+	defaultQueryTopic     = "micro.nats.query"
+	defaultWatchTopic     = "micro.nats.watch"
+	defaultRegisterAction = "create"
+)
+
+func configure(n *natsRegistry, opts ...registry.Option) error {
+	for _, o := range opts {
+		o(&n.opts)
+	}
+
+	natsOptions := nats.GetDefaultOptions()
+	if n, ok := n.opts.Context.Value(optionsKey{}).(nats.Options); ok {
+		natsOptions = n
+	}
+
+	queryTopic := defaultQueryTopic
+	if qt, ok := n.opts.Context.Value(queryTopicKey{}).(string); ok {
+		queryTopic = qt
+	}
+
+	watchTopic := defaultWatchTopic
+	if wt, ok := n.opts.Context.Value(watchTopicKey{}).(string); ok {
+		watchTopic = wt
+	}
+
+	registerAction := defaultRegisterAction
+	if ra, ok := n.opts.Context.Value(registerActionKey{}).(string); ok {
+		registerAction = ra
+	}
+
+	// Options have higher priority than nats.Options
+	// only if Addrs, Secure or TLSConfig were not set through a Option
+	// we read them from nats.Option
+	if len(n.opts.Addrs) == 0 {
+		n.opts.Addrs = natsOptions.Servers
+	}
+
+	if !n.opts.Secure {
+		n.opts.Secure = natsOptions.Secure
+	}
+
+	if n.opts.TLSConfig == nil {
+		n.opts.TLSConfig = natsOptions.TLSConfig
+	}
+
+	// check & add nats:// prefix (this makes also sure that the addresses
+	// stored in natsaddrs and n.opts.Addrs are identical)
+	n.opts.Addrs = setAddrs(n.opts.Addrs)
+
+	n.addrs = n.opts.Addrs
+	n.nopts = natsOptions
+	n.queryTopic = queryTopic
+	n.watchTopic = watchTopic
+	n.registerAction = registerAction
+
+	return nil
+}
+
+func setAddrs(addrs []string) []string {
+	var cAddrs []string
+	for _, addr := range addrs {
+		if len(addr) == 0 {
+			continue
+		}
+		if !strings.HasPrefix(addr, "nats://") {
+			addr = "nats://" + addr
+		}
+		cAddrs = append(cAddrs, addr)
+	}
+	if len(cAddrs) == 0 {
+		cAddrs = []string{nats.DefaultURL}
+	}
+	return cAddrs
+}
+
+func (n *natsRegistry) newConn() (*nats.Conn, error) {
+	opts := n.nopts
+	opts.Servers = n.addrs
+	opts.Secure = n.opts.Secure
+	opts.TLSConfig = n.opts.TLSConfig
+
+	// secure might not be set
+	if opts.TLSConfig != nil {
+		opts.Secure = true
+	}
+
+	return opts.Connect()
+}
+
+func (n *natsRegistry) getConn() (*nats.Conn, error) {
+	n.Lock()
+	defer n.Unlock()
+
+	if n.conn != nil {
+		return n.conn, nil
+	}
+
+	c, err := n.newConn()
+	if err != nil {
+		return nil, err
+	}
+	n.conn = c
+
+	return n.conn, nil
+}
+
+func (n *natsRegistry) register(s *registry.Service) error {
+	conn, err := n.getConn()
+	if err != nil {
+		return err
+	}
+
+	n.Lock()
+	defer n.Unlock()
+
+	// cache service
+	n.services[s.Name] = addServices(n.services[s.Name], cp([]*registry.Service{s}))
+
+	// create query listener
+	if n.listeners[s.Name] == nil {
+		listener := make(chan bool)
+
+		// create a subscriber that responds to queries
+		sub, err := conn.Subscribe(n.queryTopic, func(m *nats.Msg) {
+			var result *registry.Result
+
+			if err := json.Unmarshal(m.Data, &result); err != nil {
+				return
+			}
+
+			var services []*registry.Service
+
+			switch result.Action {
+			// is this a get query and we own the service?
+			case "get":
+				if result.Service.Name != s.Name {
+					return
+				}
+				n.RLock()
+				services = cp(n.services[s.Name])
+				n.RUnlock()
+			// it's a list request, but we're still only a
+			// subscriber for this service... so just get this service
+			// totally suboptimal
+			case "list":
+				n.RLock()
+				services = cp(n.services[s.Name])
+				n.RUnlock()
+			default:
+				// does not match
+				return
+			}
+
+			// respond to query
+			for _, service := range services {
+				b, err := json.Marshal(service)
+				if err != nil {
+					continue
+				}
+				_ = conn.Publish(m.Reply, b)
+			}
+		})
+		if err != nil {
+			return err
+		}
+
+		// Unsubscribe if we're told to do so
+		go func() {
+			<-listener
+			_ = sub.Unsubscribe()
+		}()
+
+		n.listeners[s.Name] = listener
+	}
+
+	return nil
+}
+
+func (n *natsRegistry) deregister(s *registry.Service) error {
+	n.Lock()
+	defer n.Unlock()
+
+	services := delServices(n.services[s.Name], cp([]*registry.Service{s}))
+	if len(services) > 0 {
+		n.services[s.Name] = services
+		return nil
+	}
+
+	// delete cached service
+	delete(n.services, s.Name)
+
+	// delete query listener
+	if listener, lexists := n.listeners[s.Name]; lexists {
+		close(listener)
+		delete(n.listeners, s.Name)
+	}
+
+	return nil
+}
+
+func (n *natsRegistry) query(s string, quorum int) ([]*registry.Service, error) {
+	conn, err := n.getConn()
+	if err != nil {
+		return nil, err
+	}
+
+	var action string
+	var service *registry.Service
+
+	if len(s) > 0 {
+		action = "get"
+		service = &registry.Service{Name: s}
+	} else {
+		action = "list"
+	}
+
+	inbox := nats.NewInbox()
+
+	response := make(chan *registry.Service, 10)
+
+	sub, err := conn.Subscribe(inbox, func(m *nats.Msg) {
+		var service *registry.Service
+		if err := json.Unmarshal(m.Data, &service); err != nil {
+			return
+		}
+		select {
+		case response <- service:
+		case <-time.After(n.opts.Timeout):
+		}
+	})
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = sub.Unsubscribe() }()
+
+	b, err := json.Marshal(&registry.Result{Action: action, Service: service})
+	if err != nil {
+		return nil, err
+	}
+
+	if err := conn.PublishMsg(&nats.Msg{
+		Subject: n.queryTopic,
+		Reply:   inbox,
+		Data:    b,
+	}); err != nil {
+		return nil, err
+	}
+
+	timeoutChan := time.After(n.opts.Timeout)
+
+	serviceMap := make(map[string]*registry.Service)
+
+loop:
+	for {
+		select {
+		case service := <-response:
+			key := service.Name + "-" + service.Version
+			srv, ok := serviceMap[key]
+			if ok {
+				srv.Nodes = append(srv.Nodes, service.Nodes...)
+				serviceMap[key] = srv
+			} else {
+				serviceMap[key] = service
+			}
+
+			if quorum > 0 && len(serviceMap[key].Nodes) >= quorum {
+				break loop
+			}
+		case <-timeoutChan:
+			break loop
+		}
+	}
+
+	var services []*registry.Service
+	for _, service := range serviceMap {
+		services = append(services, service)
+	}
+	return services, nil
+}
+
+func (n *natsRegistry) Init(opts ...registry.Option) error {
+	return configure(n, opts...)
+}
+
+func (n *natsRegistry) Options() registry.Options {
+	return n.opts
+}
+
+func (n *natsRegistry) Register(s *registry.Service, opts ...registry.RegisterOption) error {
+	if err := n.register(s); err != nil {
+		return err
+	}
+
+	conn, err := n.getConn()
+	if err != nil {
+		return err
+	}
+
+	b, err := json.Marshal(&registry.Result{Action: n.registerAction, Service: s})
+	if err != nil {
+		return err
+	}
+
+	return conn.Publish(n.watchTopic, b)
+}
+
+func (n *natsRegistry) Deregister(s *registry.Service, opts ...registry.DeregisterOption) error {
+	if err := n.deregister(s); err != nil {
+		return err
+	}
+
+	conn, err := n.getConn()
+	if err != nil {
+		return err
+	}
+
+	b, err := json.Marshal(&registry.Result{Action: "delete", Service: s})
+	if err != nil {
+		return err
+	}
+	return conn.Publish(n.watchTopic, b)
+}
+
+func (n *natsRegistry) GetService(s string, opts ...registry.GetOption) ([]*registry.Service, error) {
+	services, err := n.query(s, getQuorum(n.opts))
+	if err != nil {
+		return nil, err
+	}
+	return services, nil
+}
+
+func (n *natsRegistry) ListServices(opts ...registry.ListOption) ([]*registry.Service, error) {
+	s, err := n.query("", 0)
+	if err != nil {
+		return nil, err
+	}
+
+	var services []*registry.Service
+	serviceMap := make(map[string]*registry.Service)
+
+	for _, v := range s {
+		serviceMap[v.Name] = &registry.Service{Name: v.Name, Version: v.Version}
+	}
+
+	for _, v := range serviceMap {
+		services = append(services, v)
+	}
+
+	return services, nil
+}
+
+func (n *natsRegistry) Watch(opts ...registry.WatchOption) (registry.Watcher, error) {
+	conn, err := n.getConn()
+	if err != nil {
+		return nil, err
+	}
+
+	sub, err := conn.SubscribeSync(n.watchTopic)
+	if err != nil {
+		return nil, err
+	}
+
+	var wo registry.WatchOptions
+	for _, o := range opts {
+		o(&wo)
+	}
+
+	return &natsWatcher{sub, wo}, nil
+}
+
+func (n *natsRegistry) String() string {
+	return "nats"
+}
+
+func NewNatsRegistry(opts ...registry.Option) registry.Registry {
+	options := registry.Options{
+		Timeout: time.Millisecond * 100,
+		Context: context.Background(),
+	}
+
+	n := &natsRegistry{
+		opts:      options,
+		services:  make(map[string][]*registry.Service),
+		listeners: make(map[string]chan bool),
+	}
+	_ = configure(n, opts...)
+	return n
+}

registry/nats/nats_assert_test.go

@@ -0,0 +1,18 @@
+package nats_test
+
+import (
+	"reflect"
+	"testing"
+)
+
+func assertNoError(tb testing.TB, actual error) {
+	if actual != nil {
+		tb.Errorf("expected no error, got %v", actual)
+	}
+}
+
+func assertEqual(tb testing.TB, expected, actual interface{}) {
+	if !reflect.DeepEqual(expected, actual) {
+		tb.Errorf("expected %v, got %v", expected, actual)
+	}
+}

registry/nats/nats_environment_test.go

@@ -0,0 +1,69 @@
+package nats_test
+
+import (
+	"os"
+	"testing"
+
+	log "go-micro.dev/v6/logger"
+	"go-micro.dev/v6/registry"
+	"go-micro.dev/v6/registry/nats"
+)
+
+type environment struct {
+	registryOne   registry.Registry
+	registryTwo   registry.Registry
+	registryThree registry.Registry
+
+	serviceOne registry.Service
+	serviceTwo registry.Service
+
+	nodeOne   registry.Node
+	nodeTwo   registry.Node
+	nodeThree registry.Node
+}
+
+var e environment
+
+func TestMain(m *testing.M) {
+	natsURL := os.Getenv("NATS_URL")
+	if natsURL == "" {
+		log.Infof("NATS_URL is undefined - skipping tests")
+		return
+	}
+
+	e.registryOne = nats.NewNatsRegistry(registry.Addrs(natsURL), nats.Quorum(1))
+	e.registryTwo = nats.NewNatsRegistry(registry.Addrs(natsURL), nats.Quorum(1))
+	e.registryThree = nats.NewNatsRegistry(registry.Addrs(natsURL), nats.Quorum(1))
+
+	e.serviceOne.Name = "one"
+	e.serviceOne.Version = "default"
+	e.serviceOne.Nodes = []*registry.Node{&e.nodeOne}
+
+	e.serviceTwo.Name = "two"
+	e.serviceTwo.Version = "default"
+	e.serviceTwo.Nodes = []*registry.Node{&e.nodeOne, &e.nodeTwo}
+
+	e.nodeOne.Id = "one"
+	e.nodeTwo.Id = "two"
+	e.nodeThree.Id = "three"
+
+	if err := e.registryOne.Register(&e.serviceOne); err != nil {
+		log.Fatal(err)
+	}
+
+	if err := e.registryOne.Register(&e.serviceTwo); err != nil {
+		log.Fatal(err)
+	}
+
+	result := m.Run()
+
+	if err := e.registryOne.Deregister(&e.serviceOne); err != nil {
+		log.Fatal(err)
+	}
+
+	if err := e.registryOne.Deregister(&e.serviceTwo); err != nil {
+		log.Fatal(err)
+	}
+
+	os.Exit(result)
+}

registry/nats/nats_options.go

@@ -0,0 +1,87 @@
+package nats
+
+import (
+	"context"
+
+	"github.com/nats-io/nats.go"
+	"go-micro.dev/v6/registry"
+)
+
+type contextQuorumKey struct{}
+type optionsKey struct{}
+type watchTopicKey struct{}
+type queryTopicKey struct{}
+type registerActionKey struct{}
+
+var (
+	DefaultQuorum = 0
+)
+
+func getQuorum(o registry.Options) int {
+	if o.Context == nil {
+		return DefaultQuorum
+	}
+
+	value := o.Context.Value(contextQuorumKey{})
+	if v, ok := value.(int); ok {
+		return v
+	} else {
+		return DefaultQuorum
+	}
+}
+
+func Quorum(n int) registry.Option {
+	return func(o *registry.Options) {
+		o.Context = context.WithValue(o.Context, contextQuorumKey{}, n)
+	}
+}
+
+// Options allow to inject a nats.Options struct for configuring
+// the nats connection.
+func NatsOptions(nopts nats.Options) registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, optionsKey{}, nopts)
+	}
+}
+
+// QueryTopic allows to set a custom nats topic on which service registries
+// query (survey) other services. All registries listen on this topic and
+// then respond to the query message.
+func QueryTopic(s string) registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, queryTopicKey{}, s)
+	}
+}
+
+// WatchTopic allows to set a custom nats topic on which registries broadcast
+// changes (e.g. when services are added, updated or removed). Since we don't
+// have a central registry service, each service typically broadcasts in a
+// determined frequency on this topic.
+func WatchTopic(s string) registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, watchTopicKey{}, s)
+	}
+}
+
+// RegisterAction allows to set the action to use when registering to nats.
+// As of now there are three different options:
+// - "create" (default) only registers if there is noone already registered under the same key.
+// - "update" only updates the registration if it already exists.
+// - "put" creates or updates a registration
+func RegisterAction(s string) registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, registerActionKey{}, s)
+	}
+}

registry/nats/nats_registry.go

@@ -0,0 +1,5 @@
+package nats
+
+var (
+	DefaultRegistry = NewNatsRegistry()
+)

registry/nats/nats_test.go

@@ -0,0 +1,95 @@
+package nats_test
+
+import (
+	"testing"
+
+	"go-micro.dev/v6/registry"
+)
+
+func TestRegister(t *testing.T) {
+	service := registry.Service{Name: "test"}
+	assertNoError(t, e.registryOne.Register(&service))
+	defer e.registryOne.Deregister(&service)
+
+	services, err := e.registryOne.ListServices()
+	assertNoError(t, err)
+	assertEqual(t, 3, len(services))
+
+	services, err = e.registryTwo.ListServices()
+	assertNoError(t, err)
+	assertEqual(t, 3, len(services))
+}
+
+func TestDeregister(t *testing.T) {
+	service1 := registry.Service{Name: "test-deregister", Version: "v1"}
+	service2 := registry.Service{Name: "test-deregister", Version: "v2"}
+
+	assertNoError(t, e.registryOne.Register(&service1))
+	services, err := e.registryOne.GetService(service1.Name)
+	assertNoError(t, err)
+	assertEqual(t, 1, len(services))
+
+	assertNoError(t, e.registryOne.Register(&service2))
+	services, err = e.registryOne.GetService(service2.Name)
+	assertNoError(t, err)
+	assertEqual(t, 2, len(services))
+
+	assertNoError(t, e.registryOne.Deregister(&service1))
+	services, err = e.registryOne.GetService(service1.Name)
+	assertNoError(t, err)
+	assertEqual(t, 1, len(services))
+
+	assertNoError(t, e.registryOne.Deregister(&service2))
+	services, err = e.registryOne.GetService(service1.Name)
+	assertNoError(t, err)
+	assertEqual(t, 0, len(services))
+}
+
+func TestGetService(t *testing.T) {
+	services, err := e.registryTwo.GetService("one")
+	assertNoError(t, err)
+	assertEqual(t, 1, len(services))
+	assertEqual(t, "one", services[0].Name)
+	assertEqual(t, 1, len(services[0].Nodes))
+}
+
+func TestGetServiceWithNoNodes(t *testing.T) {
+	services, err := e.registryOne.GetService("missing")
+	assertNoError(t, err)
+	assertEqual(t, 0, len(services))
+}
+
+func TestGetServiceFromMultipleNodes(t *testing.T) {
+	services, err := e.registryOne.GetService("two")
+	assertNoError(t, err)
+	assertEqual(t, 1, len(services))
+	assertEqual(t, "two", services[0].Name)
+	assertEqual(t, 2, len(services[0].Nodes))
+}
+
+func BenchmarkGetService(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		services, err := e.registryTwo.GetService("one")
+		assertNoError(b, err)
+		assertEqual(b, 1, len(services))
+		assertEqual(b, "one", services[0].Name)
+	}
+}
+
+func BenchmarkGetServiceWithNoNodes(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		services, err := e.registryOne.GetService("missing")
+		assertNoError(b, err)
+		assertEqual(b, 0, len(services))
+	}
+}
+
+func BenchmarkGetServiceFromMultipleNodes(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		services, err := e.registryTwo.GetService("two")
+		assertNoError(b, err)
+		assertEqual(b, 1, len(services))
+		assertEqual(b, "two", services[0].Name)
+		assertEqual(b, 2, len(services[0].Nodes))
+	}
+}

registry/nats/nats_util.go

@@ -0,0 +1,107 @@
+package nats
+
+import "go-micro.dev/v6/registry"
+
+func cp(current []*registry.Service) []*registry.Service {
+	var services []*registry.Service
+
+	for _, service := range current {
+		// copy service
+		s := new(registry.Service)
+		*s = *service
+
+		// copy nodes
+		var nodes []*registry.Node
+		for _, node := range service.Nodes {
+			n := new(registry.Node)
+			*n = *node
+			nodes = append(nodes, n)
+		}
+		s.Nodes = nodes
+
+		// copy endpoints
+		var eps []*registry.Endpoint
+		for _, ep := range service.Endpoints {
+			e := new(registry.Endpoint)
+			*e = *ep
+			eps = append(eps, e)
+		}
+		s.Endpoints = eps
+
+		// append service
+		services = append(services, s)
+	}
+
+	return services
+}
+
+func addNodes(old, neu []*registry.Node) []*registry.Node {
+	for _, n := range neu {
+		var seen bool
+		for i, o := range old {
+			if o.Id == n.Id {
+				seen = true
+				old[i] = n
+				break
+			}
+		}
+		if !seen {
+			old = append(old, n)
+		}
+	}
+	return old
+}
+
+func addServices(old, neu []*registry.Service) []*registry.Service {
+	for _, s := range neu {
+		var seen bool
+		for i, o := range old {
+			if o.Version == s.Version {
+				s.Nodes = addNodes(o.Nodes, s.Nodes)
+				seen = true
+				old[i] = s
+				break
+			}
+		}
+		if !seen {
+			old = append(old, s)
+		}
+	}
+	return old
+}
+
+func delNodes(old, del []*registry.Node) []*registry.Node {
+	var nodes []*registry.Node
+	for _, o := range old {
+		var rem bool
+		for _, n := range del {
+			if o.Id == n.Id {
+				rem = true
+				break
+			}
+		}
+		if !rem {
+			nodes = append(nodes, o)
+		}
+	}
+	return nodes
+}
+
+func delServices(old, del []*registry.Service) []*registry.Service {
+	var services []*registry.Service
+	for i, o := range old {
+		var rem bool
+		for _, s := range del {
+			if o.Version == s.Version {
+				old[i].Nodes = delNodes(o.Nodes, s.Nodes)
+				if len(old[i].Nodes) == 0 {
+					rem = true
+				}
+			}
+		}
+		if !rem {
+			services = append(services, o)
+		}
+	}
+	return services
+}

registry/nats/nats_watcher.go

@@ -0,0 +1,39 @@
+package nats
+
+import (
+	"encoding/json"
+	"time"
+
+	"github.com/nats-io/nats.go"
+	"go-micro.dev/v6/registry"
+)
+
+type natsWatcher struct {
+	sub *nats.Subscription
+	wo  registry.WatchOptions
+}
+
+func (n *natsWatcher) Next() (*registry.Result, error) {
+	var result *registry.Result
+	for {
+		m, err := n.sub.NextMsg(time.Minute)
+		if err != nil && err == nats.ErrTimeout {
+			continue
+		} else if err != nil {
+			return nil, err
+		}
+		if err := json.Unmarshal(m.Data, &result); err != nil {
+			return nil, err
+		}
+		if len(n.wo.Service) > 0 && result.Service.Name != n.wo.Service {
+			continue
+		}
+		break
+	}
+
+	return result, nil
+}
+
+func (n *natsWatcher) Stop() {
+	_ = n.sub.Unsubscribe()
+}