Next phase of selector

selector/internal/blacklist/blacklist.go

@@ -0,0 +1,164 @@
+package blacklist
+
+import (
+	"math/rand"
+	"sync"
+	"time"
+
+	"github.com/micro/go-micro/registry"
+)
+
+type blackListNode struct {
+	age     time.Time
+	id      string
+	service string
+	count   int
+}
+
+type BlackList struct {
+	ttl  int
+	exit chan bool
+
+	sync.RWMutex
+	bl map[string]blackListNode
+}
+
+var (
+	// number of times we see an error before blacklisting
+	count = 3
+
+	// the ttl to blacklist for
+	ttl = 30
+)
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+
+func (r *BlackList) 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 *BlackList) 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 *BlackList) Filter(services []*registry.Service) ([]*registry.Service, error) {
+	var viableServices []*registry.Service
+
+	r.RLock()
+
+	for _, service := range services {
+		var viableNodes []*registry.Node
+
+		for _, node := range service.Nodes {
+			n, ok := r.bl[node.Id]
+			if !ok {
+				// blacklist miss so add it
+				viableNodes = append(viableNodes, node)
+				continue
+			}
+
+			// got some blacklist info
+			// skip the node if it exceeds count
+			if n.count >= count {
+				continue
+			}
+
+			// doesn't exceed count, still viable
+			viableNodes = append(viableNodes, node)
+		}
+
+		if len(viableNodes) == 0 {
+			continue
+		}
+
+		viableService := new(registry.Service)
+		*viableService = *service
+		viableService.Nodes = viableNodes
+		viableServices = append(viableServices, viableService)
+	}
+
+	r.RUnlock()
+
+	return viableServices, nil
+}
+
+func (r *BlackList) Mark(service string, node *registry.Node, err error) {
+	r.Lock()
+	defer r.Unlock()
+
+	// reset when error is nil
+	// basically closing the circuit
+	if err == nil {
+		delete(r.bl, node.Id)
+		return
+	}
+
+	n, ok := r.bl[node.Id]
+	if !ok {
+		n = blackListNode{
+			id:      node.Id,
+			service: service,
+		}
+	}
+
+	// mark it
+	n.count++
+
+	// set age to ttl seconds in future
+	n.age = time.Now().Add(time.Duration(r.ttl) * time.Second)
+
+	// save
+	r.bl[node.Id] = n
+}
+
+func (r *BlackList) Reset(service string) {
+	r.Lock()
+	defer r.Unlock()
+
+	for k, v := range r.bl {
+		// delete every node that matches the service
+		if v.service == service {
+			delete(r.bl, k)
+		}
+	}
+}
+
+func (r *BlackList) Close() error {
+	select {
+	case <-r.exit:
+		return nil
+	default:
+		close(r.exit)
+	}
+	return nil
+}
+
+func New() *BlackList {
+	bl := &BlackList{
+		ttl:  ttl,
+		bl:   make(map[string]blackListNode),
+		exit: make(chan bool),
+	}
+
+	go bl.run()
+	return bl
+}

selector/internal/blacklist/blacklist_test.go

@@ -0,0 +1,107 @@
+package blacklist
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/micro/go-micro/registry"
+)
+
+func TestBlackList(t *testing.T) {
+	bl := &BlackList{
+		ttl:  1,
+		bl:   make(map[string]blackListNode),
+		exit: make(chan bool),
+	}
+
+	go bl.run()
+	defer bl.Close()
+
+	services := []*registry.Service{
+		&registry.Service{
+			Name: "foo",
+			Nodes: []*registry.Node{
+				&registry.Node{
+					Id:      "foo-1",
+					Address: "localhost",
+					Port:    10001,
+				},
+				&registry.Node{
+					Id:      "foo-2",
+					Address: "localhost",
+					Port:    10002,
+				},
+				&registry.Node{
+					Id:      "foo-3",
+					Address: "localhost",
+					Port:    10002,
+				},
+			},
+		},
+	}
+
+	// check nothing is filtered on clean run
+	filterTest := func() {
+		for i := 0; i < 3; i++ {
+			srvs, err := bl.Filter(services)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(srvs) != len(services) {
+				t.Fatal("nodes were filtered when they shouldn't be")
+			}
+
+			for _, node := range srvs[0].Nodes {
+				var seen bool
+				for _, n := range srvs[0].Nodes {
+					if n.Id == node.Id {
+						seen = true
+						break
+					}
+				}
+				if !seen {
+					t.Fatalf("Missing node %s", node.Id)
+				}
+			}
+		}
+	}
+
+	// run filter test
+	filterTest()
+
+	blacklistTest := func() {
+		// test blacklisting
+		// mark until failure
+		for i := 0; i < count+1; i++ {
+			for _, node := range services[0].Nodes {
+				bl.Mark("foo", node, errors.New("blacklist"))
+			}
+		}
+
+		filtered, err := bl.Filter(services)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(filtered) > 0 {
+			t.Fatalf("Expected zero nodes got %+v", filtered)
+		}
+	}
+
+	// sleep the ttl duration
+	time.Sleep(time.Second * time.Duration(bl.ttl) * 2)
+
+	// now run filterTest again
+	filterTest()
+
+	// run the blacklist test
+	blacklistTest()
+
+	// reset
+	bl.Reset("foo")
+
+	// check again
+	filterTest()
+}