make util.Until return immediately if stop is written inside the loop

pkg/util/util.go

@@ -114,17 +114,22 @@ func Forever(f func(), period time.Duration) {
 // stop channel is already closed. Pass NeverStop to Until if you
 // don't want it stop.
 func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
+	select {
+	case <-stopCh:
+		return
+	default:
+	}
+
 	for {
-		select {
-		case <-stopCh:
-			return
-		default:
-		}
 		func() {
 			defer HandleCrash()
 			f()
 		}()
-		time.Sleep(period)
+		select {
+		case <-stopCh:
+			return
+		case <-time.After(period):
+		}
 	}
 }
 

pkg/util/util_test.go

@@ -24,6 +24,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/ghodss/yaml"
 )
@@ -48,6 +49,17 @@ func TestUntil(t *testing.T) {
 	<-called
 }
 
+func TestUntilReturnsImmediately(t *testing.T) {
+	now := time.Now()
+	ch := make(chan struct{})
+	Until(func() {
+		close(ch)
+	}, 30*time.Second, ch)
+	if now.Add(25 * time.Second).Before(time.Now()) {
+		t.Errorf("Until did not return immediately when the stop chan was closed inside the func")
+	}
+}
+
 func TestHandleCrash(t *testing.T) {
 	count := 0
 	expect := 10