Switch up the selector so it actually allows you to inform it how the node performed

client/node_selector.go

@@ -6,21 +6,40 @@ import (
 
 	"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 a list of services and selects
-// a single node. If a node cannot be selected it should return
-// an error. A list of services is provided as a service may
-// have 1 or more versions.
-type NodeSelector func(service []*registry.Service) (*registry.Node, error)
+// 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 {
+	Retrieve(context.Context, Request) (*registry.Node, error)
+	Response(*registry.Node, error)
+	Reset()
+}
+
+// Selector takes a Registry and returns a NodeSelector.
+// Used by the client to initialise a selector.
+type Selector func(registry.Registry) NodeSelector
 
 func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
 // Built in random hashed node selector
-func nodeSelector(service []*registry.Service) (*registry.Node, error) {
+type nodeSelector struct {
+	r registry.Registry
+}
+
+func (n *nodeSelector) Retrieve(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")
 	}
@@ -32,6 +51,14 @@ func nodeSelector(service []*registry.Service) (*registry.Node, error) {
 		return nil, errors.NotFound("go.micro.client", "Service not found")
 	}
 
-	n := i % len(service[j].Nodes)
-	return service[j].Nodes[n], nil
+	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

@@ -14,7 +14,7 @@ type options struct {
 	registry    registry.Registry
 	transport   transport.Transport
 	wrappers    []Wrapper
-	selector    NodeSelector
+	selector    Selector
 }
 
 // Broker to be used for pub/sub
@@ -52,8 +52,8 @@ func Transport(t transport.Transport) Option {
 	}
 }
 
-// Selector is used to select a node to route a request to
-func Selector(s NodeSelector) Option {
+// Select is used to select a node to route a request to
+func Select(s Selector) Option {
 	return func(o *options) {
 		o.selector = s
 	}

client/rpc_client.go

@@ -18,9 +18,11 @@ 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{
@@ -39,17 +41,24 @@ func newRpcClient(opt ...Option) Client {
 		opts.transport = transport.DefaultTransport
 	}
 
+	if opts.registry == nil {
+		opts.registry = registry.DefaultRegistry
+	}
+
 	if opts.broker == nil {
 		opts.broker = broker.DefaultBroker
 	}
 
-	if opts.selector == nil {
-		opts.selector = nodeSelector
+	if opts.selector != nil {
+		sel = opts.selector(opts.registry)
+	} else {
+		sel = &nodeSelector{opts.registry}
 	}
 
 	rc := &rpcClient{
 		once: once,
 		opts: opts,
+		sel:  sel,
 	}
 
 	c := Client(rc)
@@ -145,12 +154,7 @@ func (r *rpcClient) CallRemote(ctx context.Context, address string, request Requ
 
 // TODO: Call(..., opts *Options) error {
 func (r *rpcClient) Call(ctx context.Context, request Request, response interface{}) error {
-	service, err := registry.GetService(request.Service())
-	if err != nil {
-		return errors.InternalServerError("go.micro.client", err.Error())
-	}
-
-	node, err := r.opts.selector(service)
+	node, err := r.sel.Retrieve(ctx, request)
 	if err != nil {
 		return err
 	}
@@ -160,7 +164,9 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 		address = fmt.Sprintf("%s:%d", address, node.Port)
 	}
 
-	return r.call(ctx, address, request, response)
+	err = r.call(ctx, address, request, response)
+	r.sel.Response(node, err)
+	return err
 }
 
 func (r *rpcClient) StreamRemote(ctx context.Context, address string, request Request, responseChan interface{}) (Streamer, error) {
@@ -168,12 +174,7 @@ func (r *rpcClient) StreamRemote(ctx context.Context, address string, request Re
 }
 
 func (r *rpcClient) Stream(ctx context.Context, request Request, responseChan interface{}) (Streamer, error) {
-	service, err := registry.GetService(request.Service())
-	if err != nil {
-		return nil, errors.InternalServerError("go.micro.client", err.Error())
-	}
-
-	node, err := r.opts.selector(service)
+	node, err := r.sel.Retrieve(ctx, request)
 	if err != nil {
 		return nil, err
 	}
@@ -183,7 +184,9 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, responseChan in
 		address = fmt.Sprintf("%s:%d", address, node.Port)
 	}
 
-	return r.stream(ctx, address, request, responseChan)
+	stream, err := r.stream(ctx, address, request, responseChan)
+	r.sel.Response(node, err)
+	return stream, err
 }
 
 func (r *rpcClient) Publish(ctx context.Context, p Publication) error {

examples/client/node_selector/node_selector.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/micro/go-micro/client"
 	"github.com/micro/go-micro/cmd"
-	c "github.com/micro/go-micro/context"
 	"github.com/micro/go-micro/errors"
 	example "github.com/micro/go-micro/examples/server/proto/example"
 	"github.com/micro/go-micro/registry"
@@ -18,33 +17,43 @@ func init() {
 	rand.Seed(time.Now().Unix())
 }
 
-// A random node selector
-func randomSelector(s []*registry.Service) (*registry.Node, error) {
-	if len(s) == 0 {
+// Built in random hashed node selector
+type nodeSelector struct {
+	r registry.Registry
+}
+
+func (n *nodeSelector) Retrieve(ctx context.Context, req client.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(s)
+	j := i % len(service)
 
-	if len(s[j].Nodes) == 0 {
+	if len(service[j].Nodes) == 0 {
 		return nil, errors.NotFound("go.micro.client", "Service not found")
 	}
 
-	n := i % len(s[j].Nodes)
-	return s[j].Nodes[n], nil
+	k := i % len(service[j].Nodes)
+	return service[j].Nodes[k], nil
 }
 
-// Wraps the node selector so that it will log what node was selected
-func wrapSelector(fn client.NodeSelector) client.NodeSelector {
-	return func(s []*registry.Service) (*registry.Node, error) {
-		n, err := fn(s)
-		if err != nil {
-			return nil, err
-		}
-		fmt.Printf("Selected node %v\n", n)
-		return n, nil
-	}
+func (n *nodeSelector) Response(node *registry.Node, err error) {
+	return
+}
+
+func (n *nodeSelector) Reset() {
+	return
+}
+
+// Return a new random node selector
+func RandomSelector(r registry.Registry) client.NodeSelector {
+	return &nodeSelector{r}
 }
 
 func call(i int) {
@@ -53,16 +62,10 @@ func call(i int) {
 		Name: "John",
 	})
 
-	// create context with metadata
-	ctx := c.WithMetadata(context.Background(), map[string]string{
-		"X-User-Id": "john",
-		"X-From-Id": "script",
-	})
-
 	rsp := &example.Response{}
 
 	// Call service
-	if err := client.Call(ctx, req, rsp); err != nil {
+	if err := client.Call(context.Background(), req, rsp); err != nil {
 		fmt.Println("call err: ", err, rsp)
 		return
 	}
@@ -74,7 +77,7 @@ func main() {
 	cmd.Init()
 
 	client.DefaultClient = client.NewClient(
-		client.Selector(wrapSelector(randomSelector)),
+		client.Select(RandomSelector),
 	)
 
 	fmt.Println("\n--- Call example ---\n")