Move to a selector package

selector/blacklist/black_list_selector.go

@@ -0,0 +1,167 @@
+package blacklist
+
+import (
+	"math/rand"
+	"sync"
+	"time"
+
+	"github.com/micro/go-micro/registry"
+	"github.com/micro/go-micro/selector"
+)
+
+type blackListNode struct {
+	age     time.Time
+	id      string
+	service string
+}
+
+type blackListSelector struct {
+	so   selector.Options
+	ttl  int64
+	exit chan bool
+	once sync.Once
+
+	sync.RWMutex
+	bl map[string]blackListNode
+}
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+
+func (r *blackListSelector) purge() {
+	now := time.Now()
+	r.Lock()
+	for k, v := range r.bl {
+		if d := v.age.Sub(now); d.Seconds() < 0 {
+			delete(r.bl, k)
+		}
+	}
+	r.Unlock()
+}
+
+func (r *blackListSelector) run() {
+	t := time.NewTicker(time.Duration(r.ttl) * time.Second)
+
+	for {
+		select {
+		case <-r.exit:
+			t.Stop()
+			return
+		case <-t.C:
+			r.purge()
+		}
+	}
+}
+
+func (r *blackListSelector) Select(service string, opts ...selector.SelectOption) (selector.Next, error) {
+	var sopts selector.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, selector.ErrNotFound
+	}
+
+	var nodes []*registry.Node
+
+	for _, service := range services {
+		for _, node := range service.Nodes {
+			nodes = append(nodes, node)
+		}
+	}
+
+	if len(nodes) == 0 {
+		return nil, selector.ErrNotFound
+	}
+
+	return func() (*registry.Node, error) {
+		var viable []*registry.Node
+
+		r.RLock()
+		for _, node := range nodes {
+			if _, ok := r.bl[node.Id]; !ok {
+				viable = append(viable, node)
+			}
+		}
+		r.RUnlock()
+
+		if len(viable) == 0 {
+			return nil, selector.ErrNoneAvailable
+		}
+
+		return viable[rand.Int()%len(viable)], nil
+	}, nil
+}
+
+func (r *blackListSelector) Mark(service string, node *registry.Node, err error) {
+	r.Lock()
+	defer r.Unlock()
+	if err == nil {
+		delete(r.bl, node.Id)
+		return
+	}
+
+	r.bl[node.Id] = blackListNode{
+		age:     time.Now().Add(time.Duration(r.ttl) * time.Second),
+		id:      node.Id,
+		service: service,
+	}
+	return
+}
+
+func (r *blackListSelector) Reset(service string) {
+	r.Lock()
+	defer r.Unlock()
+	for k, v := range r.bl {
+		if v.service == service {
+			delete(r.bl, k)
+		}
+	}
+	return
+}
+
+func (r *blackListSelector) Close() error {
+	r.once.Do(func() {
+		close(r.exit)
+	})
+	return nil
+}
+
+func NewSelector(opts ...selector.Option) selector.Selector {
+	var sopts selector.Options
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = registry.DefaultRegistry
+	}
+
+	var once sync.Once
+	bl := &blackListSelector{
+		once: once,
+		so:   sopts,
+		ttl:  60,
+		bl:   make(map[string]blackListNode),
+		exit: make(chan bool),
+	}
+
+	go bl.run()
+
+	return bl
+}

selector/blacklist/black_list_selector_test.go

@@ -0,0 +1,73 @@
+package blacklist
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/micro/go-micro/registry/mock"
+	"github.com/micro/go-micro/selector"
+)
+
+func TestBlackListSelector(t *testing.T) {
+	counts := map[string]int{}
+
+	bl := &blackListSelector{
+		so: selector.Options{
+			Registry: mock.NewRegistry(),
+		},
+		ttl:  2,
+		bl:   make(map[string]blackListNode),
+		exit: make(chan bool),
+	}
+
+	go bl.run()
+	defer bl.Close()
+
+	next, err := bl.Select("foo")
+	if err != nil {
+		t.Errorf("Unexpected error calling bl 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("BlackList Counts %v", counts)
+
+	// test blacklisting
+	for i := 0; i < 4; i++ {
+		node, err := next()
+		if err != nil {
+			t.Errorf("Expected node err, got err: %v", err)
+		}
+		bl.Mark("foo", node, errors.New("blacklist"))
+	}
+	if node, err := next(); err != selector.ErrNoneAvailable {
+		t.Errorf("Expected none available err, got node %v err %v", node, err)
+	}
+	time.Sleep(time.Second * time.Duration(bl.ttl) * 2)
+	if _, err := next(); err != nil {
+		t.Errorf("Unexpected err %v", err)
+	}
+
+	// test resetting
+	for i := 0; i < 4; i++ {
+		node, err := next()
+		if err != nil {
+			t.Errorf("Unexpected err: %v", err)
+		}
+		bl.Mark("foo", node, errors.New("blacklist"))
+	}
+	if node, err := next(); err != selector.ErrNoneAvailable {
+		t.Errorf("Expected none available err, got node %v err %v", node, err)
+	}
+	bl.Reset("foo")
+	if _, err := next(); err != nil {
+		t.Errorf("Unexpected err %v", err)
+	}
+}

selector/options.go

@@ -0,0 +1,34 @@
+package selector
+
+import (
+	"github.com/micro/go-micro/registry"
+)
+
+type Options struct {
+	Registry registry.Registry
+}
+
+type SelectOptions struct {
+	Filters []SelectFilter
+}
+
+// Option used to initialise the selector
+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
+	}
+}
+
+// Filter adds a filter function to the list of filters
+// used during the Select call.
+func Filter(fn SelectFilter) SelectOption {
+	return func(o *SelectOptions) {
+		o.Filters = append(o.Filters, fn)
+	}
+}

selector/random/random.go

@@ -0,0 +1,7 @@
+package random
+
+import "github.com/micro/go-micro/selector"
+
+func NewSelector(opts ...selector.Option) selector.Selector {
+	return selector.NewSelector(opts...)
+}

selector/random_selector.go

@@ -0,0 +1,89 @@
+package selector
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/micro/go-micro/registry"
+)
+
+type randomSelector struct {
+	so Options
+}
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+
+func (r *randomSelector) Select(service string, opts ...SelectOption) (Next, 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 []*registry.Node
+
+	for _, service := range services {
+		for _, node := range service.Nodes {
+			nodes = append(nodes, node)
+		}
+	}
+
+	if len(nodes) == 0 {
+		return nil, ErrNotFound
+	}
+
+	return func() (*registry.Node, error) {
+		i := rand.Int()
+		j := i % len(services)
+
+		if len(services[j].Nodes) == 0 {
+			return nil, ErrNotFound
+		}
+
+		k := i % len(services[j].Nodes)
+		return services[j].Nodes[k], nil
+	}, nil
+}
+
+func (r *randomSelector) Mark(service string, node *registry.Node, err error) {
+	return
+}
+
+func (r *randomSelector) Reset(service string) {
+	return
+}
+
+func (r *randomSelector) Close() error {
+	return nil
+}
+
+func newRandomSelector(opts ...Option) Selector {
+	var sopts Options
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = registry.DefaultRegistry
+	}
+
+	return &randomSelector{sopts}
+}

selector/random_selector_test.go

@@ -0,0 +1,32 @@
+package selector
+
+import (
+	"testing"
+
+	"github.com/micro/go-micro/registry/mock"
+)
+
+func TestRandomSelector(t *testing.T) {
+	counts := map[string]int{}
+
+	bl := &randomSelector{
+		so: Options{
+			Registry: mock.NewRegistry(),
+		},
+	}
+
+	next, err := bl.Select("foo")
+	if err != nil {
+		t.Errorf("Unexpected error calling random 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)
+}

selector/roundrobin/round_robin_selector.go

@@ -0,0 +1,83 @@
+package roundrobin
+
+import (
+	"sync"
+
+	"github.com/micro/go-micro/registry"
+	"github.com/micro/go-micro/selector"
+)
+
+type roundRobinSelector struct {
+	so selector.Options
+}
+
+func (r *roundRobinSelector) Select(service string, opts ...selector.SelectOption) (selector.Next, error) {
+	var sopts selector.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, selector.ErrNotFound
+	}
+
+	var nodes []*registry.Node
+
+	for _, service := range services {
+		for _, node := range service.Nodes {
+			nodes = append(nodes, node)
+		}
+	}
+
+	if len(nodes) == 0 {
+		return nil, selector.ErrNotFound
+	}
+
+	var i int
+	var mtx sync.Mutex
+
+	return func() (*registry.Node, error) {
+		mtx.Lock()
+		defer mtx.Unlock()
+		i++
+		return nodes[i%len(nodes)], nil
+	}, nil
+}
+
+func (r *roundRobinSelector) Mark(service string, node *registry.Node, err error) {
+	return
+}
+
+func (r *roundRobinSelector) Reset(service string) {
+	return
+}
+
+func (r *roundRobinSelector) Close() error {
+	return nil
+}
+
+func NewRoundRobinSelector(opts ...selector.Option) selector.Selector {
+	var sopts selector.Options
+
+	for _, opt := range opts {
+		opt(&sopts)
+	}
+
+	if sopts.Registry == nil {
+		sopts.Registry = registry.DefaultRegistry
+	}
+
+	return &roundRobinSelector{sopts}
+}

selector/roundrobin/round_robin_selector_test.go

@@ -0,0 +1,33 @@
+package roundrobin
+
+import (
+	"testing"
+
+	"github.com/micro/go-micro/registry/mock"
+	"github.com/micro/go-micro/selector"
+)
+
+func TestRoundRobinSelector(t *testing.T) {
+	counts := map[string]int{}
+
+	rr := &roundRobinSelector{
+		so: selector.Options{
+			Registry: mock.NewRegistry(),
+		},
+	}
+
+	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)
+}

selector/selector.go

@@ -0,0 +1,93 @@
+/*
+The Selector package provides a way to algorithmically filter and return
+nodes required by the client or any other system. Selector's implemented
+by Micro build on the registry but it's of optional use. One could
+provide a static Selector that has a fixed pool.
+
+	func (r *randomSelector) Select(service string, opts ...SelectOption) (Next, 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 []*registry.Node
+
+		for _, service := range services {
+			for _, node := range service.Nodes {
+				nodes = append(nodes, node)
+			}
+		}
+
+		if len(nodes) == 0 {
+			return nil, ErrNotFound
+		}
+
+		return func() (*registry.Node, error) {
+			i := rand.Int()
+			j := i % len(services)
+
+			if len(services[j].Nodes) == 0 {
+				return nil, ErrNotFound
+			}
+
+			k := i % len(services[j].Nodes)
+			return services[j].Nodes[k], nil
+		}, nil
+	}
+
+
+*/
+package selector
+
+import (
+	"errors"
+	"github.com/micro/go-micro/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 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
+}
+
+// Next is a function that returns the next node
+// based on the selector's algorithm
+type Next func() (*registry.Node, error)
+
+// SelectFilter is used to filter a service during the selection process
+type SelectFilter func([]*registry.Service) []*registry.Service
+
+var (
+	DefaultSelector = newRandomSelector()
+
+	ErrNotFound      = errors.New("not found")
+	ErrNoneAvailable = errors.New("none available")
+)
+
+func NewSelector(opts ...Option) Selector {
+	return newRandomSelector(opts...)
+}