diff --git a/pkg/util/util.go b/pkg/util/util.go index 519203f5971..e8234e1e8b0 100644 --- a/pkg/util/util.go +++ b/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): + } } } diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go index 99711691916..6ea25d851e8 100644 --- a/pkg/util/util_test.go +++ b/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