chore: import upstream snapshot with attribution

selector/common_test.go

@@ -0,0 +1,47 @@
+package selector
+
+import (
+	"go-micro.dev/v6/registry"
+)
+
+var (
+	// mock data.
+	testData = map[string][]*registry.Service{
+		"foo": {
+			{
+				Name:    "foo",
+				Version: "1.0.0",
+				Nodes: []*registry.Node{
+					{
+						Id:      "foo-1.0.0-123",
+						Address: "localhost:9999",
+					},
+					{
+						Id:      "foo-1.0.0-321",
+						Address: "localhost:9999",
+					},
+				},
+			},
+			{
+				Name:    "foo",
+				Version: "1.0.1",
+				Nodes: []*registry.Node{
+					{
+						Id:      "foo-1.0.1-321",
+						Address: "localhost:6666",
+					},
+				},
+			},
+			{
+				Name:    "foo",
+				Version: "1.0.3",
+				Nodes: []*registry.Node{
+					{
+						Id:      "foo-1.0.3-345",
+						Address: "localhost:8888",
+					},
+				},
+			},
+		},
+	}
+)

selector/default.go

@@ -0,0 +1,123 @@
+package selector
+
+import (
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+
+	"go-micro.dev/v6/registry"
+	"go-micro.dev/v6/registry/cache"
+)
+
+type registrySelector struct {
+	so Options
+	rc cache.Cache
+	mu sync.RWMutex
+}
+
+func (c *registrySelector) newCache() cache.Cache {
+	opts := make([]cache.Option, 0, 1)
+
+	if c.so.Context != nil {
+		if t, ok := c.so.Context.Value("selector_ttl").(time.Duration); ok {
+			opts = append(opts, cache.WithTTL(t))
+		}
+	}
+
+	return cache.New(c.so.Registry, opts...)
+}
+
+func (c *registrySelector) Init(opts ...Option) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	for _, o := range opts {
+		o(&c.so)
+	}
+
+	c.rc.Stop()
+	c.rc = c.newCache()
+
+	return nil
+}
+
+func (c *registrySelector) Options() Options {
+	return c.so
+}
+
+func (c *registrySelector) Select(service string, opts ...SelectOption) (Next, error) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	sopts := SelectOptions{
+		Strategy: c.so.Strategy,
+	}
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	// get the service
+	// try the cache first
+	// if that fails go directly to the registry
+	services, err := c.rc.GetService(service)
+	if err != nil {
+		if errors.Is(err, registry.ErrNotFound) {
+			return nil, ErrNotFound
+		}
+
+		return nil, err
+	}
+
+	// apply the filters
+	for _, filter := range sopts.Filters {
+		services = filter(services)
+	}
+
+	// if there's nothing left, return
+	if len(services) == 0 {
+		return nil, ErrNoneAvailable
+	}
+
+	return sopts.Strategy(services), nil
+}
+
+func (c *registrySelector) Mark(service string, node *registry.Node, err error) {
+}
+
+func (c *registrySelector) Reset(service string) {
+}
+
+// Close stops the watcher and destroys the cache.
+func (c *registrySelector) Close() error {
+	c.rc.Stop()
+
+	return nil
+}
+
+func (c *registrySelector) String() string {
+	return "registry"
+}
+
+// NewSelector creates a new default selector.
+func NewSelector(opts ...Option) Selector {
+	sopts := Options{
+		Strategy: Random,
+	}
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = registry.DefaultRegistry
+	}
+
+	s := &registrySelector{
+		so: sopts,
+	}
+	s.rc = s.newCache()
+
+	return s
+}

selector/default_test.go

@@ -0,0 +1,32 @@
+package selector
+
+import (
+	"os"
+	"testing"
+
+	"go-micro.dev/v6/registry"
+)
+
+func TestRegistrySelector(t *testing.T) {
+	counts := map[string]int{}
+
+	r := registry.NewMemoryRegistry(registry.Services(testData))
+	cache := NewSelector(Registry(r))
+
+	next, err := cache.Select("foo")
+	if err != nil {
+		t.Errorf("Unexpected error calling cache select: %v", err)
+	}
+
+	for i := 0; i < 100; i++ {
+		node, err := next()
+		if err != nil {
+			t.Errorf("Expected node err, got err: %v", err)
+		}
+		counts[node.Id]++
+	}
+
+	if len(os.Getenv("IN_TRAVIS_CI")) == 0 {
+		t.Logf("Selector Counts %v", counts)
+	}
+}

selector/filter.go

@@ -0,0 +1,73 @@
+package selector
+
+import (
+	"go-micro.dev/v6/registry"
+)
+
+// FilterEndpoint is an endpoint based Select Filter which will
+// only return services with the endpoint specified.
+func FilterEndpoint(name string) Filter {
+	return func(old []*registry.Service) []*registry.Service {
+		var services []*registry.Service
+
+		for _, service := range old {
+			for _, ep := range service.Endpoints {
+				if ep.Name == name {
+					services = append(services, service)
+					break
+				}
+			}
+		}
+
+		return services
+	}
+}
+
+// FilterLabel is a label based Select Filter which will
+// only return services with the label specified.
+func FilterLabel(key, val string) Filter {
+	return func(old []*registry.Service) []*registry.Service {
+		var services []*registry.Service
+
+		for _, service := range old {
+			serv := new(registry.Service)
+			var nodes []*registry.Node
+
+			for _, node := range service.Nodes {
+				if node.Metadata == nil {
+					continue
+				}
+
+				if node.Metadata[key] == val {
+					nodes = append(nodes, node)
+				}
+			}
+
+			// only add service if there's some nodes
+			if len(nodes) > 0 {
+				// copy
+				*serv = *service
+				serv.Nodes = nodes
+				services = append(services, serv)
+			}
+		}
+
+		return services
+	}
+}
+
+// FilterVersion is a version based Select Filter which will
+// only return services with the version specified.
+func FilterVersion(version string) Filter {
+	return func(old []*registry.Service) []*registry.Service {
+		var services []*registry.Service
+
+		for _, service := range old {
+			if service.Version == version {
+				services = append(services, service)
+			}
+		}
+
+		return services
+	}
+}

selector/filter_test.go

@@ -0,0 +1,239 @@
+package selector
+
+import (
+	"testing"
+
+	"go-micro.dev/v6/registry"
+)
+
+func TestFilterEndpoint(t *testing.T) {
+	testData := []struct {
+		services []*registry.Service
+		endpoint string
+		count    int
+	}{
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+					Endpoints: []*registry.Endpoint{
+						{
+							Name: "Foo.Bar",
+						},
+					},
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+					Endpoints: []*registry.Endpoint{
+						{
+							Name: "Baz.Bar",
+						},
+					},
+				},
+			},
+			endpoint: "Foo.Bar",
+			count:    1,
+		},
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+					Endpoints: []*registry.Endpoint{
+						{
+							Name: "Foo.Bar",
+						},
+					},
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+					Endpoints: []*registry.Endpoint{
+						{
+							Name: "Foo.Bar",
+						},
+					},
+				},
+			},
+			endpoint: "Bar.Baz",
+			count:    0,
+		},
+	}
+
+	for _, data := range testData {
+		filter := FilterEndpoint(data.endpoint)
+		services := filter(data.services)
+
+		if len(services) != data.count {
+			t.Fatalf("Expected %d services, got %d", data.count, len(services))
+		}
+
+		for _, service := range services {
+			var seen bool
+
+			for _, ep := range service.Endpoints {
+				if ep.Name == data.endpoint {
+					seen = true
+					break
+				}
+			}
+
+			if !seen && data.count > 0 {
+				t.Fatalf("Expected %d services but seen is %t; result %+v", data.count, seen, services)
+			}
+		}
+	}
+}
+
+func TestFilterLabel(t *testing.T) {
+	testData := []struct {
+		services []*registry.Service
+		label    [2]string
+		count    int
+	}{
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+					Nodes: []*registry.Node{
+						{
+							Id:      "test-1",
+							Address: "localhost",
+							Metadata: map[string]string{
+								"foo": "bar",
+							},
+						},
+					},
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+					Nodes: []*registry.Node{
+						{
+							Id:      "test-2",
+							Address: "localhost",
+							Metadata: map[string]string{
+								"foo": "baz",
+							},
+						},
+					},
+				},
+			},
+			label: [2]string{"foo", "bar"},
+			count: 1,
+		},
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+					Nodes: []*registry.Node{
+						{
+							Id:      "test-1",
+							Address: "localhost",
+						},
+					},
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+					Nodes: []*registry.Node{
+						{
+							Id:      "test-2",
+							Address: "localhost",
+						},
+					},
+				},
+			},
+			label: [2]string{"foo", "bar"},
+			count: 0,
+		},
+	}
+
+	for _, data := range testData {
+		filter := FilterLabel(data.label[0], data.label[1])
+		services := filter(data.services)
+
+		if len(services) != data.count {
+			t.Fatalf("Expected %d services, got %d", data.count, len(services))
+		}
+
+		for _, service := range services {
+			var seen bool
+
+			for _, node := range service.Nodes {
+				if node.Metadata[data.label[0]] != data.label[1] {
+					t.Fatalf("Expected %s=%s but got %s=%s for service %+v node %+v",
+						data.label[0], data.label[1], data.label[0], node.Metadata[data.label[0]], service, node)
+				}
+				seen = true
+			}
+
+			if !seen {
+				t.Fatalf("Expected node for %s=%s but saw none; results %+v", data.label[0], data.label[1], service)
+			}
+		}
+	}
+}
+
+func TestFilterVersion(t *testing.T) {
+	testData := []struct {
+		services []*registry.Service
+		version  string
+		count    int
+	}{
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+				},
+			},
+			version: "1.0.0",
+			count:   1,
+		},
+		{
+			services: []*registry.Service{
+				{
+					Name:    "test",
+					Version: "1.0.0",
+				},
+				{
+					Name:    "test",
+					Version: "1.1.0",
+				},
+			},
+			version: "2.0.0",
+			count:   0,
+		},
+	}
+
+	for _, data := range testData {
+		filter := FilterVersion(data.version)
+		services := filter(data.services)
+
+		if len(services) != data.count {
+			t.Fatalf("Expected %d services, got %d", data.count, len(services))
+		}
+
+		var seen bool
+
+		for _, service := range services {
+			if service.Version != data.version {
+				t.Fatalf("Expected version %s, got %s", data.version, service.Version)
+			}
+			seen = true
+		}
+
+		if !seen && data.count > 0 {
+			t.Fatalf("Expected %d services but seen is %t; result %+v", data.count, seen, services)
+		}
+	}
+}

selector/options.go

@@ -0,0 +1,70 @@
+package selector
+
+import (
+	"context"
+
+	"go-micro.dev/v6/logger"
+	"go-micro.dev/v6/registry"
+)
+
+type Options struct {
+	Registry registry.Registry
+	Strategy Strategy
+
+	// Other options for implementations of the interface
+	// can be stored in a context
+	Context context.Context
+	// Logger is the underline logger
+	Logger logger.Logger
+}
+
+type SelectOptions struct {
+
+	// Other options for implementations of the interface
+	// can be stored in a context
+	Context  context.Context
+	Strategy Strategy
+
+	Filters []Filter
+}
+
+type Option func(*Options)
+
+// SelectOption used when making a select call.
+type SelectOption func(*SelectOptions)
+
+// Registry sets the registry used by the selector.
+func Registry(r registry.Registry) Option {
+	return func(o *Options) {
+		o.Registry = r
+	}
+}
+
+// SetStrategy sets the default strategy for the selector.
+func SetStrategy(fn Strategy) Option {
+	return func(o *Options) {
+		o.Strategy = fn
+	}
+}
+
+// WithFilter adds a filter function to the list of filters
+// used during the Select call.
+func WithFilter(fn ...Filter) SelectOption {
+	return func(o *SelectOptions) {
+		o.Filters = append(o.Filters, fn...)
+	}
+}
+
+// Strategy sets the selector strategy.
+func WithStrategy(fn Strategy) SelectOption {
+	return func(o *SelectOptions) {
+		o.Strategy = fn
+	}
+}
+
+// WithLogger sets the underline logger.
+func WithLogger(l logger.Logger) Option {
+	return func(o *Options) {
+		o.Logger = l
+	}
+}

selector/selector.go

@@ -0,0 +1,43 @@
+// Package selector is a way to pick a list of service nodes
+package selector
+
+import (
+	"errors"
+
+	"go-micro.dev/v6/registry"
+)
+
+// Selector builds on the registry as a mechanism to pick nodes
+// and mark their status. This allows host pools and other things
+// to be built using various algorithms.
+type Selector interface {
+	Init(opts ...Option) error
+	Options() Options
+	// Select returns a function which should return the next node
+	Select(service string, opts ...SelectOption) (Next, error)
+	// Mark sets the success/error against a node
+	Mark(service string, node *registry.Node, err error)
+	// Reset returns state back to zero for a service
+	Reset(service string)
+	// Close renders the selector unusable
+	Close() error
+	// Name of the selector
+	String() string
+}
+
+// Next is a function that returns the next node
+// based on the selector's strategy.
+type Next func() (*registry.Node, error)
+
+// Filter is used to filter a service during the selection process.
+type Filter func([]*registry.Service) []*registry.Service
+
+// Strategy is a selection strategy e.g random, round robin.
+type Strategy func([]*registry.Service) Next
+
+var (
+	DefaultSelector = NewSelector()
+
+	ErrNotFound      = errors.New("not found")
+	ErrNoneAvailable = errors.New("none available")
+)

selector/strategy.go

@@ -0,0 +1,51 @@
+package selector
+
+import (
+	"math/rand"
+	"sync"
+
+	"go-micro.dev/v6/registry"
+)
+
+// Random is a random strategy algorithm for node selection.
+func Random(services []*registry.Service) Next {
+	nodes := make([]*registry.Node, 0, len(services))
+
+	for _, service := range services {
+		nodes = append(nodes, service.Nodes...)
+	}
+
+	return func() (*registry.Node, error) {
+		if len(nodes) == 0 {
+			return nil, ErrNoneAvailable
+		}
+
+		i := rand.Int() % len(nodes)
+		return nodes[i], nil
+	}
+}
+
+// RoundRobin is a roundrobin strategy algorithm for node selection.
+func RoundRobin(services []*registry.Service) Next {
+	nodes := make([]*registry.Node, 0, len(services))
+
+	for _, service := range services {
+		nodes = append(nodes, service.Nodes...)
+	}
+
+	var i = rand.Int()
+	var mtx sync.Mutex
+
+	return func() (*registry.Node, error) {
+		if len(nodes) == 0 {
+			return nil, ErrNoneAvailable
+		}
+
+		mtx.Lock()
+		node := nodes[i%len(nodes)]
+		i++
+		mtx.Unlock()
+
+		return node, nil
+	}
+}

selector/strategy_test.go

@@ -0,0 +1,58 @@
+package selector
+
+import (
+	"os"
+	"testing"
+
+	"go-micro.dev/v6/registry"
+)
+
+func TestStrategies(t *testing.T) {
+	testData := []*registry.Service{
+		{
+			Name:    "test1",
+			Version: "latest",
+			Nodes: []*registry.Node{
+				{
+					Id:      "test1-1",
+					Address: "10.0.0.1:1001",
+				},
+				{
+					Id:      "test1-2",
+					Address: "10.0.0.2:1002",
+				},
+			},
+		},
+		{
+			Name:    "test1",
+			Version: "default",
+			Nodes: []*registry.Node{
+				{
+					Id:      "test1-3",
+					Address: "10.0.0.3:1003",
+				},
+				{
+					Id:      "test1-4",
+					Address: "10.0.0.4:1004",
+				},
+			},
+		},
+	}
+
+	for name, strategy := range map[string]Strategy{"random": Random, "roundrobin": RoundRobin} {
+		next := strategy(testData)
+		counts := make(map[string]int)
+
+		for i := 0; i < 100; i++ {
+			node, err := next()
+			if err != nil {
+				t.Fatal(err)
+			}
+			counts[node.Id]++
+		}
+
+		if len(os.Getenv("IN_TRAVIS_CI")) == 0 {
+			t.Logf("%s: %+v\n", name, counts)
+		}
+	}
+}