diff --git a/pkg/apiserver/handlers_test.go b/pkg/apiserver/handlers_test.go index 80c63c2698..40f0a9a72c 100644 --- a/pkg/apiserver/handlers_test.go +++ b/pkg/apiserver/handlers_test.go @@ -36,6 +36,7 @@ type fakeRL bool func (fakeRL) Stop() {} func (f fakeRL) CanAccept() bool { return bool(f) } +func (f fakeRL) Accept() {} func expectHTTP(url string, code int, t *testing.T) { r, err := http.Get(url) diff --git a/pkg/util/throttle.go b/pkg/util/throttle.go index 0b337c5e3c..c961811faa 100644 --- a/pkg/util/throttle.go +++ b/pkg/util/throttle.go @@ -24,6 +24,8 @@ import ( type RateLimiter interface { // CanAccept returns true if the rate is below the limit, false otherwise CanAccept() bool + // Accept returns once a token becomes available. + Accept() // Stop stops the rate limiter, subsequent calls to CanAccept will return false Stop() } @@ -73,6 +75,11 @@ func (t *tickRateLimiter) CanAccept() bool { } } +// Accept will block until a token becomes available +func (t *tickRateLimiter) Accept() { + <-t.tokens +} + func (t *tickRateLimiter) Stop() { close(t.stop) } diff --git a/pkg/util/throttle_test.go b/pkg/util/throttle_test.go index 328cf232cb..e290fe1336 100644 --- a/pkg/util/throttle_test.go +++ b/pkg/util/throttle_test.go @@ -60,3 +60,17 @@ func TestOverBurst(t *testing.T) { r.step() } } + +func TestThrottle(t *testing.T) { + r := NewTokenBucketRateLimiter(10, 5) + + // Should consume 5 tokens immediately, then + // the remaining 11 should take at least 1 second (0.1s each) + expectedFinish := time.Now().Add(time.Second * 1) + for i := 0; i < 16; i++ { + r.Accept() + } + if time.Now().Before(expectedFinish) { + t.Error("rate limit was not respected, finished too early") + } +}