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