From 9c4d62edd97bcec0393c45164ef1ac1646af5a24 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Tue, 4 Feb 2014 12:21:30 -0800
Subject: [PATCH] consul: Adding a simple notification mechansim

---
 consul/notify.go      | 42 ++++++++++++++++++++++++++++++++
 consul/notify_test.go | 56 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 consul/notify.go
 create mode 100644 consul/notify_test.go

diff --git a/consul/notify.go b/consul/notify.go
new file mode 100644
index 0000000000..88da09242c
--- /dev/null
+++ b/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
+}
diff --git a/consul/notify_test.go b/consul/notify_test.go
new file mode 100644
index 0000000000..4c3be55901
--- /dev/null
+++ b/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")
+	}
+}