Add selector code

registry/mock_registry.go

@@ -0,0 +1,62 @@
+package registry
+
+type mockRegistry struct{}
+
+func (m *mockRegistry) GetService(service string) ([]*Service, error) {
+	return []*Service{
+		{
+			Name:    "foo",
+			Version: "1.0.0",
+			Nodes: []*Node{
+				{
+					Id:      "foo-1.0.0-123",
+					Address: "localhost",
+					Port:    9999,
+				},
+				{
+					Id:      "foo-1.0.0-321",
+					Address: "localhost",
+					Port:    9999,
+				},
+			},
+		},
+		{
+			Name:    "foo",
+			Version: "1.0.1",
+			Nodes: []*Node{
+				{
+					Id:      "foo-1.0.1-321",
+					Address: "localhost",
+					Port:    6666,
+				},
+			},
+		},
+		{
+			Name:    "foo",
+			Version: "1.0.3",
+			Nodes: []*Node{
+				{
+					Id:      "foo-1.0.3-345",
+					Address: "localhost",
+					Port:    8888,
+				},
+			},
+		},
+	}, nil
+}
+
+func (m *mockRegistry) ListServices() ([]*Service, error) {
+	return []*Service{}, nil
+}
+
+func (m *mockRegistry) Register(s *Service) error {
+	return nil
+}
+
+func (m *mockRegistry) Deregister(s *Service) error {
+	return nil
+}
+
+func (m *mockRegistry) Watch() (Watcher, error) {
+	return nil, nil
+}

registry/random_selector.go

@@ -0,0 +1,79 @@
+package registry
+
+import (
+	"math/rand"
+	"time"
+)
+
+type randomSelector struct {
+	so SelectorOptions
+}
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+
+func (r *randomSelector) Select(service string, opts ...SelectOption) (SelectNext, error) {
+	var sopts SelectOptions
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	// get the service
+	services, err := r.so.Registry.GetService(service)
+	if err != nil {
+		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, ErrNotFound
+	}
+
+	var nodes []*Node
+
+	for _, service := range services {
+		for _, node := range service.Nodes {
+			nodes = append(nodes, node)
+		}
+	}
+
+	if len(nodes) == 0 {
+		return nil, ErrNotFound
+	}
+
+	return func() (*Node, error) {
+		return nodes[rand.Int()%len(nodes)], nil
+	}, nil
+}
+
+func (r *randomSelector) Mark(service string, node *Node, err error) {
+	return
+}
+
+func (r *randomSelector) Reset(service string) {
+	return
+}
+
+func (r *randomSelector) Close() error {
+	return nil
+}
+
+func NewRandomSelector(opts ...SelectorOption) Selector {
+	var sopts SelectorOptions
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = DefaultRegistry
+	}
+
+	return &randomSelector{sopts}
+}

registry/random_selector_test.go

@@ -0,0 +1,30 @@
+package registry
+
+import (
+	"testing"
+)
+
+func TestRandomSelector(t *testing.T) {
+	counts := map[string]int{}
+
+	rr := &randomSelector{
+		so: SelectorOptions{
+			Registry: &mockRegistry{},
+		},
+	}
+
+	next, err := rr.Select("foo")
+	if err != nil {
+		t.Errorf("Unexpected error calling rr 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]++
+	}
+
+	t.Logf("Random Counts %v", counts)
+}

registry/registry.go

@@ -1,5 +1,9 @@
 package registry
 
+import (
+	"errors"
+)
+
 type Registry interface {
 	Register(*Service) error
 	Deregister(*Service) error
@@ -14,6 +18,8 @@ type Option func(*options)
 
 var (
 	DefaultRegistry = newConsulRegistry([]string{})
+
+	ErrNotFound = errors.New("not found")
 )
 
 func NewRegistry(addrs []string, opt ...Option) Registry {

registry/round_robin_selector.go

@@ -0,0 +1,80 @@
+package registry
+
+import (
+	"sync"
+)
+
+type roundRobinSelector struct {
+	so SelectorOptions
+}
+
+func (r *roundRobinSelector) Select(service string, opts ...SelectOption) (SelectNext, error) {
+	var sopts SelectOptions
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	// get the service
+	services, err := r.so.Registry.GetService(service)
+	if err != nil {
+		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, ErrNotFound
+	}
+
+	var nodes []*Node
+
+	for _, service := range services {
+		for _, node := range service.Nodes {
+			nodes = append(nodes, node)
+		}
+	}
+
+	if len(nodes) == 0 {
+		return nil, ErrNotFound
+	}
+
+	var i int
+	var mtx sync.Mutex
+
+	return func() (*Node, error) {
+		mtx.Lock()
+		defer mtx.Unlock()
+		i++
+		return nodes[i%len(nodes)], nil
+	}, nil
+}
+
+func (r *roundRobinSelector) Mark(service string, node *Node, err error) {
+	return
+}
+
+func (r *roundRobinSelector) Reset(service string) {
+	return
+}
+
+func (r *roundRobinSelector) Close() error {
+	return nil
+}
+
+func NewRoundRobinSelector(opts ...SelectorOption) Selector {
+	var sopts SelectorOptions
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = DefaultRegistry
+	}
+
+	return &roundRobinSelector{sopts}
+}

registry/round_robin_selector_test.go

@@ -0,0 +1,30 @@
+package registry
+
+import (
+	"testing"
+)
+
+func TestRoundRobinSelector(t *testing.T) {
+	counts := map[string]int{}
+
+	rr := &roundRobinSelector{
+		so: SelectorOptions{
+			Registry: &mockRegistry{},
+		},
+	}
+
+	next, err := rr.Select("foo")
+	if err != nil {
+		t.Errorf("Unexpected error calling rr 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]++
+	}
+
+	t.Logf("Round Robin Counts %v", counts)
+}

registry/selector.go

@@ -0,0 +1,44 @@
+package 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 {
+	Select(service string, opts ...SelectOption) (SelectNext, error)
+	Mark(service string, node *Node, err error)
+	Reset(service string)
+	Close() error
+}
+
+// SelectNext is a function that returns the next node
+// based on the selector's algorithm
+type SelectNext func() (*Node, error)
+
+type SelectorOptions struct {
+	Registry Registry
+}
+
+type SelectOptions struct {
+	Filters []func([]*Service) []*Service
+}
+
+// Option used to initialise the selector
+type SelectorOption func(*SelectorOptions)
+
+// Option used when making a select call
+type SelectOption func(*SelectOptions)
+
+// SelectorRegistry sets the registry used by the selector
+func SelectorRegistry(r Registry) SelectorOption {
+	return func(o *SelectorOptions) {
+		o.Registry = r
+	}
+}
+
+// SelectFilter adds a filter function to the list of filters
+// used during the Select call.
+func SelectFilter(fn func([]*Service) []*Service) SelectOption {
+	return func(o *SelectOptions) {
+		o.Filters = append(o.Filters, fn)
+	}
+}