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