diff --git a/cmd/cmd.go b/cmd/cmd.go index a675537e..8e13b87b 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -19,6 +19,7 @@ import ( "github.com/micro/go-micro/broker" "github.com/micro/go-micro/client" "github.com/micro/go-micro/registry" + "github.com/micro/go-micro/selector" "github.com/micro/go-micro/server" "github.com/micro/go-micro/transport" "github.com/pborman/uuid" @@ -80,6 +81,12 @@ var ( EnvVar: "MICRO_REGISTRY_ADDRESS", Usage: "Comma-separated list of registry addresses", }, + cli.StringFlag{ + Name: "selector", + EnvVar: "MICRO_SELECTOR", + Value: "selector", + Usage: "Selector used to pick nodes for querying. random, roundrobin, blacklist", + }, cli.StringFlag{ Name: "transport", EnvVar: "MICRO_TRANSPORT", @@ -137,6 +144,10 @@ var ( "consul": registry.NewRegistry, } + Selectors = map[string]func(...selector.Option) selector.Selector{ + "random": selector.NewSelector, + } + Transports = map[string]func([]string, ...transport.Option) transport.Transport{ "http": transport.NewTransport, } @@ -214,6 +225,10 @@ func Setup(c *cli.Context) error { registry.DefaultRegistry = r(strings.Split(c.String("registry_address"), ",")) } + if s, ok := Selectors[c.String("selector")]; ok { + selector.DefaultSelector = s(selector.Registry(registry.DefaultRegistry)) + } + if t, ok := Transports[c.String("transport")]; ok { transport.DefaultTransport = t(strings.Split(c.String("transport_address"), ",")) }