consul: Adding a simple notification mechansim

consul/notify.go

@@ -0,0 +1,42 @@
+package consul
+
+import (
+	"sync"
+)
+
+// NotifyGroup is used to allow a simple notification mechanism.
+// Channels can be marked as waiting, and when notify is invoked,
+// all the waiting channels get a message and are cleared from the
+// notify list.
+type NotifyGroup struct {
+	l      sync.Mutex
+	notify []chan struct{}
+}
+
+// Notify will do a non-blocking send to all waiting channels, and
+// clear the notify list
+func (n *NotifyGroup) Notify() {
+	n.l.Lock()
+	defer n.l.Unlock()
+	for _, ch := range n.notify {
+		select {
+		case ch <- struct{}{}:
+		default:
+		}
+	}
+	n.notify = n.notify[:0]
+}
+
+// Wait adds a channel to the notify group
+func (n *NotifyGroup) Wait(ch chan struct{}) {
+	n.l.Lock()
+	defer n.l.Unlock()
+	n.notify = append(n.notify, ch)
+}
+
+// WaitCh allocates a channel that is subscribed to notifications
+func (n *NotifyGroup) WaitCh() chan struct{} {
+	ch := make(chan struct{}, 1)
+	n.Wait(ch)
+	return ch
+}

consul/notify_test.go

@@ -0,0 +1,56 @@
+package consul
+
+import (
+	"testing"
+)
+
+func TestNotifyGroup(t *testing.T) {
+	grp := &NotifyGroup{}
+
+	ch1 := grp.WaitCh()
+	ch2 := grp.WaitCh()
+
+	select {
+	case <-ch1:
+		t.Fatalf("should block")
+	default:
+	}
+	select {
+	case <-ch2:
+		t.Fatalf("should block")
+	default:
+	}
+
+	grp.Notify()
+
+	select {
+	case <-ch1:
+	default:
+		t.Fatalf("should not block")
+	}
+	select {
+	case <-ch2:
+	default:
+		t.Fatalf("should not block")
+	}
+
+	// Should be unregistered
+	ch3 := grp.WaitCh()
+	grp.Notify()
+
+	select {
+	case <-ch1:
+		t.Fatalf("should block")
+	default:
+	}
+	select {
+	case <-ch2:
+		t.Fatalf("should block")
+	default:
+	}
+	select {
+	case <-ch3:
+	default:
+		t.Fatalf("should not block")
+	}
+}