util: add a nonblocking throttle variant

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

util/throttle/throttle.go

@@ -8,6 +8,17 @@ import (
 // Throttle wraps a function so that internal function does not get called
 // more frequently than the specified duration.
 func Throttle(d time.Duration, f func()) func() {
+	return throttle(d, f, true)
+}
+
+// ThrottleAfter wraps a function so that internal function does not get called
+// more frequently than the specified duration. The delay is added after function
+// has been called.
+func ThrottleAfter(d time.Duration, f func()) func() {
+	return throttle(d, f, false)
+}
+
+func throttle(d time.Duration, f func(), wait bool) func() {
 	var next, running bool
 	var mu sync.Mutex
 	return func() {
@@ -25,12 +36,21 @@ func Throttle(d time.Duration, f func()) func() {
 						mu.Unlock()
 						return
 					}
+					if !wait {
+						next = false
+					}
 					mu.Unlock()
-					time.Sleep(d)
+
-					mu.Lock()
+					if wait {
-					next = false
+						time.Sleep(d)
-					mu.Unlock()
+						mu.Lock()
-					f()
+						next = false
+						mu.Unlock()
+						f()
+					} else {
+						f()
+						time.Sleep(d)
+					}
 				}
 			}()
 		}

util/throttle/throttle_test.go

@@ -54,3 +54,25 @@ func TestThrottle(t *testing.T) {
 	}
 
 }
+
+func TestThrottleAfter(t *testing.T) {
+	t.Parallel()
+
+	var i int64
+	f := func() {
+		atomic.AddInt64(&i, 1)
+	}
+
+	f = ThrottleAfter(100*time.Millisecond, f)
+
+	f()
+
+	time.Sleep(10 * time.Millisecond)
+	require.Equal(t, int64(1), atomic.LoadInt64(&i))
+	f()
+	time.Sleep(10 * time.Millisecond)
+	require.Equal(t, int64(1), atomic.LoadInt64(&i))
+
+	time.Sleep(200 * time.Millisecond)
+	require.Equal(t, int64(2), atomic.LoadInt64(&i))
+}