Add selector code

client/client_selector.go

@@ -1,9 +0,0 @@
-package client
-
-import (
-	"github.com/micro/go-micro/registry"
-)
-
-// Selector takes a Registry and returns a NodeSelector.
-// Used by the client to initialise a selector.
-type Selector func(registry.Registry) NodeSelector

client/node_selector.go

@@ -1,60 +0,0 @@
-package client
-
-import (
-	"math/rand"
-	"time"
-
-	"github.com/micro/go-micro/errors"
-	"github.com/micro/go-micro/registry"
-	"golang.org/x/net/context"
-)
-
-// NodeSelector is used to Retrieve a node to which a request
-// should be routed. It takes context and Request and returns a
-// single node. If a node cannot be selected it should return
-// an error. Response is called to inform the selector of the
-// response from a client call. Reset is called to zero out
-// any state.
-type NodeSelector interface {
-	Select(context.Context, Request) (*registry.Node, error)
-	Response(*registry.Node, error)
-	Reset()
-}
-
-func init() {
-	rand.Seed(time.Now().UnixNano())
-}
-
-// Built in random hashed node selector
-type nodeSelector struct {
-	r registry.Registry
-}
-
-func (n *nodeSelector) Select(ctx context.Context, req Request) (*registry.Node, error) {
-	service, err := n.r.GetService(req.Service())
-	if err != nil {
-		return nil, errors.InternalServerError("go.micro.client", err.Error())
-	}
-
-	if len(service) == 0 {
-		return nil, errors.NotFound("go.micro.client", "Service not found")
-	}
-
-	i := rand.Int()
-	j := i % len(service)
-
-	if len(service[j].Nodes) == 0 {
-		return nil, errors.NotFound("go.micro.client", "Service not found")
-	}
-
-	k := i % len(service[j].Nodes)
-	return service[j].Nodes[k], nil
-}
-
-func (n *nodeSelector) Response(node *registry.Node, err error) {
-	return
-}
-
-func (n *nodeSelector) Reset() {
-	return
-}

client/options.go

@@ -12,12 +12,14 @@ type options struct {
 	broker      broker.Broker
 	codecs      map[string]codec.NewCodec
 	registry    registry.Registry
+	selector    registry.Selector
 	transport   transport.Transport
 	wrappers    []Wrapper
-	selector    Selector
 }
 
-type callOptions struct{}
+type callOptions struct {
+	selectOptions []registry.SelectOption
+}
 
 type publishOptions struct{}
 
@@ -57,7 +59,7 @@ func Transport(t transport.Transport) Option {
 }
 
 // Select is used to select a node to route a request to
-func Select(s Selector) Option {
+func Selector(s registry.Selector) Option {
 	return func(o *options) {
 		o.selector = s
 	}
@@ -69,3 +71,11 @@ func Wrap(w Wrapper) Option {
 		o.wrappers = append(o.wrappers, w)
 	}
 }
+
+// Call Options
+
+func WithSelectOption(so registry.SelectOption) CallOption {
+	return func(o *callOptions) {
+		o.selectOptions = append(o.selectOptions, so)
+	}
+}

client/rpc_client.go

@@ -18,11 +18,9 @@ import (
 type rpcClient struct {
 	once sync.Once
 	opts options
-	sel  NodeSelector
 }
 
 func newRpcClient(opt ...Option) Client {
-	var sel NodeSelector
 	var once sync.Once
 
 	opts := options{
@@ -37,28 +35,27 @@ func newRpcClient(opt ...Option) Client {
 		opts.contentType = defaultContentType
 	}
 
-	if opts.transport == nil {
+	if opts.broker == nil {
-		opts.transport = transport.DefaultTransport
+		opts.broker = broker.DefaultBroker
 	}
 
 	if opts.registry == nil {
 		opts.registry = registry.DefaultRegistry
 	}
 
-	if opts.broker == nil {
+	if opts.selector == nil {
-		opts.broker = broker.DefaultBroker
+		opts.selector = registry.NewRandomSelector(
+			registry.SelectorRegistry(opts.registry),
+		)
 	}
 
-	if opts.selector != nil {
+	if opts.transport == nil {
-		sel = opts.selector(opts.registry)
+		opts.transport = transport.DefaultTransport
-	} else {
-		sel = &nodeSelector{opts.registry}
 	}
 
 	rc := &rpcClient{
 		once: once,
 		opts: opts,
-		sel:  sel,
 	}
 
 	c := Client(rc)
@@ -153,7 +150,17 @@ func (r *rpcClient) CallRemote(ctx context.Context, address string, request Requ
 }
 
 func (r *rpcClient) Call(ctx context.Context, request Request, response interface{}, opts ...CallOption) error {
-	node, err := r.sel.Select(ctx, request)
+	var copts callOptions
+	for _, opt := range opts {
+		opt(&copts)
+	}
+
+	next, err := r.opts.selector.Select(request.Service(), copts.selectOptions...)
+	if err != nil {
+		return err
+	}
+
+	node, err := next()
 	if err != nil {
 		return err
 	}
@@ -164,7 +171,7 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 	}
 
 	err = r.call(ctx, address, request, response)
-	r.sel.Response(node, err)
+	r.opts.selector.Mark(request.Service(), node, err)
 	return err
 }
 
@@ -173,7 +180,17 @@ func (r *rpcClient) StreamRemote(ctx context.Context, address string, request Re
 }
 
 func (r *rpcClient) Stream(ctx context.Context, request Request, responseChan interface{}, opts ...CallOption) (Streamer, error) {
-	node, err := r.sel.Select(ctx, request)
+	var copts callOptions
+	for _, opt := range opts {
+		opt(&copts)
+	}
+
+	next, err := r.opts.selector.Select(request.Service(), copts.selectOptions...)
+	if err != nil {
+		return nil, err
+	}
+
+	node, err := next()
 	if err != nil {
 		return nil, err
 	}
@@ -184,7 +201,7 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, responseChan in
 	}
 
 	stream, err := r.stream(ctx, address, request, responseChan)
-	r.sel.Response(node, err)
+	r.opts.selector.Mark(request.Service(), node, err)
 	return stream, err
 }
 

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)
+	}
+}