Merge pull request #6314 from bparees/tokenbucket

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)

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)
 }

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")
+	}
+}