Merge pull request #2661 from smarterclayton/allow_loops_to_exit

pkg/util/util.go

@@ -51,7 +51,19 @@ func HandleCrash() {
 
 // Forever loops forever running f every d.  Catches any panics, and keeps going.
 func Forever(f func(), period time.Duration) {
+	Until(f, period, nil)
+}
+
+// Until loops until stop channel is closed, running f every d.
+// Catches any panics, and keeps going. f may not be invoked if
+// stop channel is already closed.
+func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
 	for {
+		select {
+		case <-stopCh:
+			return
+		default:
+		}
 		func() {
 			defer HandleCrash()
 			f()

pkg/util/util_test.go

@@ -24,6 +24,26 @@ import (
 	"github.com/ghodss/yaml"
 )
 
+func TestUntil(t *testing.T) {
+	ch := make(chan struct{})
+	close(ch)
+	Until(func() {
+		t.Fatal("should not have been invoked")
+	}, 0, ch)
+
+	ch = make(chan struct{})
+	called := make(chan struct{})
+	go func() {
+		Until(func() {
+			called <- struct{}{}
+		}, 0, ch)
+		close(called)
+	}()
+	<-called
+	close(ch)
+	<-called
+}
+
 func TestHandleCrash(t *testing.T) {
 	count := 0
 	expect := 10