Merge pull request #356 from anthonybishopric/service

Add watches for label selectors

Author: anthonybishopric

pkg/labels/applicator.go

@@ -61,4 +61,14 @@ type Applicator interface {
 
 	// Return all objects of the given type that match the given selector
 	GetMatches(selector labels.Selector, labelType Type) ([]Labeled, error)
+
+	// Watch a label selector of a given type and see updates to that set
+	// of Labeled over time. If an error occurs, the Applicator implementation
+	// is responsible for handling it and recovering gracefully.
+	//
+	// The watch may be terminated by the underlying implementation without signaling on
+	// the quit channel - this will be indicated by the closing of the result channel. For
+	// this reason, the safest way to process the result of this channel is to use
+	// it in a for loop. Alternatively, you may check for nullity of the result []Labeled
+	WatchMatches(selector labels.Selector, labelType Type, quitCh chan struct{}) chan *[]Labeled
 }

pkg/labels/consul_aggregator.go

@@ -0,0 +1,174 @@
+package labels
+
+import (
+	"sync"
+	"time"
+
+	"github.com/square/p2/pkg/kp/consulutil"
+	"github.com/square/p2/pkg/logging"
+
+	"github.com/square/p2/Godeps/_workspace/src/github.com/Sirupsen/logrus"
+	"github.com/square/p2/Godeps/_workspace/src/github.com/hashicorp/consul/api"
+	"github.com/square/p2/Godeps/_workspace/src/k8s.io/kubernetes/pkg/labels"
+)
+
+// minimum required time between retrievals of label subtrees.
+var AggregationRateCap = 10 * time.Second
+
+// linked list of watches with control channels
+type selectorWatch struct {
+	selector labels.Selector
+	resultCh chan *[]Labeled
+	canceled chan struct{}
+	next     *selectorWatch
+}
+
+func (s *selectorWatch) append(w *selectorWatch) {
+	if w == nil {
+		return
+	}
+	if s.next != nil {
+		s.next.append(w)
+	} else {
+		s.next = w
+	}
+}
+
+func (s *selectorWatch) delete(w *selectorWatch) {
+	if w == nil {
+		return
+	}
+	if s.next == w {
+		s.next = w.next
+	} else if s.next != nil {
+		s.next.delete(w)
+	}
+}
+
+type consulAggregator struct {
+	logger         logging.Logger
+	labelType      Type
+	watcherLock    sync.Mutex
+	path           string
+	kv             consulutil.ConsulLister
+	watchers       *selectorWatch
+	labeledCache   []Labeled // cached contents of the label subtree
+	aggregatorQuit chan struct{}
+}
+
+func NewConsulAggregator(labelType Type, kv consulutil.ConsulLister, logger logging.Logger) *consulAggregator {
+	return &consulAggregator{
+		kv:             kv,
+		logger:         logger,
+		labelType:      labelType,
+		path:           typePath(labelType),
+		aggregatorQuit: make(chan struct{}),
+	}
+}
+
+// Add a new selector to the aggregator. New values on the output channel may not appear
+// right away.
+func (c *consulAggregator) Watch(selector labels.Selector, quitCh chan struct{}) chan *[]Labeled {
+	resCh := make(chan *[]Labeled)
+	select {
+	case <-c.aggregatorQuit:
+		c.logger.WithField("selector", selector.String()).Warnln("New selector added after aggregator was closed")
+		close(resCh)
+		return resCh
+	default:
+	}
+	c.watcherLock.Lock()
+	defer c.watcherLock.Unlock()
+	watch := &selectorWatch{
+		selector: selector,
+		resultCh: resCh,
+		canceled: make(chan struct{}),
+	}
+	if c.watchers == nil {
+		c.watchers = watch
+	} else {
+		c.watchers.append(watch)
+	}
+	go func() {
+		select {
+		case <-quitCh:
+		case <-c.aggregatorQuit:
+		}
+		c.removeWatch(watch)
+	}()
+	return watch.resultCh
+}
+
+func (c *consulAggregator) removeWatch(watch *selectorWatch) {
+	c.watcherLock.Lock()
+	defer c.watcherLock.Unlock()
+	close(watch.canceled)
+	close(watch.resultCh)
+	if c.watchers == watch {
+		c.watchers = c.watchers.next
+	} else {
+		c.watchers.delete(watch)
+	}
+}
+
+func (c *consulAggregator) Quit() {
+	close(c.aggregatorQuit)
+}
+
+// Aggregate does the labor of querying Consul for all labels under a given type,
+// applying each watcher's label selector to the results and sending those results on each
+// watcher's output channel respectively.
+// Aggregate will loop forever, constantly sending matches to each watcher
+// until Quit() has been invoked.
+func (c *consulAggregator) Aggregate() {
+	outPairs := make(chan api.KVPairs)
+	done := make(chan struct{})
+	outErrors := make(chan error)
+	go consulutil.WatchPrefix(c.path, c.kv, outPairs, done, outErrors)
+	for {
+		loopTime := time.After(AggregationRateCap)
+		select {
+		case <-c.aggregatorQuit:
+			return
+		case pairs := <-outPairs:
+			// Convert watch result to []Labeled
+			c.labeledCache = make([]Labeled, len(pairs))
+			for i, kvp := range pairs {
+				val, err := convertKVPToLabeled(kvp)
+				if err != nil {
+					c.logger.WithErrorAndFields(err, logrus.Fields{
+						"key":   kvp.Key,
+						"value": string(kvp.Value),
+					}).Errorln("Invalid key encountered, skipping this value")
+					continue
+				}
+				c.labeledCache[i] = val
+			}
+
+			// Iterate over each watcher and send the []Labeled
+			// that match the watcher's selector to the watcher's out channel.
+			c.watcherLock.Lock()
+			watcher := c.watchers
+			for watcher != nil {
+				matches := []Labeled{}
+				for _, labeled := range c.labeledCache {
+					if watcher.selector.Matches(labeled.Labels) {
+						matches = append(matches, labeled)
+					}
+				}
+				select {
+				case watcher.resultCh <- &matches:
+				case <-watcher.canceled:
+				}
+				watcher = watcher.next
+			}
+			c.watcherLock.Unlock()
+		}
+		select {
+		case <-c.aggregatorQuit:
+			return
+		case <-loopTime:
+		}
+
+	}
+}

pkg/labels/consul_aggregator_test.go

@@ -0,0 +1,158 @@
+package labels
+
+import (
+	"testing"
+	"time"
+
+	"github.com/square/p2/pkg/logging"
+
+	. "github.com/square/p2/Godeps/_workspace/src/github.com/anthonybishopric/gotcha"
+	"github.com/square/p2/Godeps/_workspace/src/k8s.io/kubernetes/pkg/labels"
+)
+
+func alterAggregationTime(dur time.Duration) {
+	AggregationRateCap = dur
+}
+
+func fakeLabeledPods() map[string][]byte {
+	return map[string][]byte{
+		objectPath(POD, "maroono"):  []byte(`{"color": "red", "deployment": "production"}`),
+		objectPath(POD, "emeralda"): []byte(`{"color": "green", "deployment": "canary"}`),
+		objectPath(POD, "slashi"):   []byte(`{"color": "red", "deployment": "canary"}`),
+	}
+}
+
+// Check that two clients can share an aggregator
+func TestTwoClients(t *testing.T) {
+	alterAggregationTime(time.Millisecond)
+
+	fakeKV := &fakeLabelStore{fakeLabeledPods(), nil}
+	aggreg := NewConsulAggregator(POD, fakeKV, logging.DefaultLogger)
+	go aggreg.Aggregate()
+	defer aggreg.Quit()
+
+	quitCh := make(chan struct{})
+	labeledChannel1 := aggreg.Watch(labels.Everything().Add("color", labels.EqualsOperator, []string{"green"}), quitCh)
+	labeledChannel2 := aggreg.Watch(labels.Everything().Add("deployment", labels.EqualsOperator, []string{"canary"}), quitCh)
+
+	var checked string
+	for i := 0; i < 2; i++ {
+		select {
+		case <-time.After(time.Second):
+			t.Fatal("Should not have taken a second to get results")
+		case labeledPtr := <-labeledChannel1:
+			Assert(t).IsNotNil(labeledPtr, "ptr should not have been nil")
+			labeled := *labeledPtr
+			Assert(t).AreNotEqual("green", checked, "Should not have already checked the green selector result")
+			checked = "green" // ensure that both sides get checked
+			Assert(t).AreEqual(1, len(labeled), "Should have received one result from the color watch")
+			Assert(t).AreEqual("emeralda", labeled[0].ID, "should have received the emerald app")
+		case labeledPtr := <-labeledChannel2:
+			Assert(t).IsNotNil(labeledPtr, "ptr should not have been nil")
+			labeled := *labeledPtr
+			Assert(t).AreNotEqual("canary", checked, "Should not have already checked the canary selector result")
+			checked = "canary" // ensure that both sides get checked
+			Assert(t).AreEqual(2, len(labeled), "Should have received two results from the canary watch")
+			emeraldaIndex := 0
+			slashiIndex := 1
+			if labeled[0].ID == "slashi" { // order doesn't matter
+				emeraldaIndex, slashiIndex = slashiIndex, emeraldaIndex
+			}
+			Assert(t).AreEqual("emeralda", labeled[emeraldaIndex].ID, "should have received the emerald app")
+			Assert(t).AreEqual("slashi", labeled[slashiIndex].ID, "should have received the slashi app")
+		}
+	}
+}
+
+func TestQuitAggregateAfterResults(t *testing.T) {
+	alterAggregationTime(time.Millisecond)
+
+	fakeKV := &fakeLabelStore{fakeLabeledPods(), nil}
+	aggreg := NewConsulAggregator(POD, fakeKV, logging.DefaultLogger)
+	go aggreg.Aggregate()
+
+	quitCh := make(chan struct{})
+	res := aggreg.Watch(labels.Everything().Add("color", labels.EqualsOperator, []string{"green"}), quitCh)
+
+	// Quit now. We expect that the aggregator will close the res channels
+	aggreg.Quit()
+	success := make(chan struct{})
+	go func() {
+		for _ = range res {
+		}
+		success <- struct{}{}
+	}()
+
+	select {
+	case <-success:
+	case <-time.After(time.Second):
+		t.Fatal("Should not be waiting or processing results after a second")
+	}
+}
+
+func TestQuitAggregateBeforeResults(t *testing.T) {
+	alterAggregationTime(time.Millisecond)
+
+	// this channel prevents the List from returning, so the aggregator
+	// must quit prior to entering the loop
+	trigger := make(chan struct{})
+	fakeKV := &fakeLabelStore{fakeLabeledPods(), trigger}
+	aggreg := NewConsulAggregator(POD, fakeKV, logging.DefaultLogger)
+	go aggreg.Aggregate()
+
+	quitCh := make(chan struct{})
+	res := aggreg.Watch(labels.Everything().Add("color", labels.EqualsOperator, []string{"green"}), quitCh)
+
+	// Quit now. We expect that the aggregator will close the res channels
+	aggreg.Quit()
+
+	select {
+	case labeled := <-res:
+		Assert(t).IsNotNil(labeled, "Should not have received any results")
+	case <-time.After(time.Second):
+		t.Fatal("Should still be waiting or processing results after a second")
+	}
+}
+
+func TestQuitIndividualWatch(t *testing.T) {
+	alterAggregationTime(time.Millisecond)
+
+	fakeKV := &fakeLabelStore{fakeLabeledPods(), nil}
+	aggreg := NewConsulAggregator(POD, fakeKV, logging.DefaultLogger)
+	go aggreg.Aggregate()
+
+	quitCh1 := make(chan struct{})
+	labeledChannel1 := aggreg.Watch(labels.Everything().Add("color", labels.EqualsOperator, []string{"green"}), quitCh1)
+
+	quitCh2 := make(chan struct{})
+	labeledChannel2 := aggreg.Watch(labels.Everything().Add("deployment", labels.EqualsOperator, []string{"production"}), quitCh2)
+
+	close(quitCh1) // this should not interrupt the flow of messages to the second channel
+
+	// iterate twice to show that we are not waiting on other now-closed channels
+	for i := 0; i < 2; i++ {
+		select {
+		case <-time.After(time.Second):
+			t.Fatalf("Should not have taken a second to get results on iteration %v", i)
+		case labeledPtr := <-labeledChannel2:
+			Assert(t).IsNotNil(labeledPtr, "ptr should not have been nil")
+			labeled := *labeledPtr
+			Assert(t).AreEqual(1, len(labeled), "Should have one result with a production deployment")
+			Assert(t).AreEqual("maroono", labeled[0].ID, "Should have received maroono as the one production deployment")
+		}
+	}
+
+	// drain the first channel to show that it was closed. We do this
+	// in a loop since it is possible that a value was sent on the channel
+	success := make(chan struct{})
+	go func() {
+		for _ = range labeledChannel1 {
+		}
+		success <- struct{}{}
+	}()
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("Should not have taken a second to see the closed label channel")
+	case <-success:
+	}
+}