From 3153ca2815c5a596adec770392248573c7e448b7 Mon Sep 17 00:00:00 2001 From: Tim Hockin Date: Sun, 21 May 2017 17:12:28 -0700 Subject: [PATCH] Inject clock through flowcontrol --- .../client-go/util/flowcontrol/throttle.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/staging/src/k8s.io/client-go/util/flowcontrol/throttle.go b/staging/src/k8s.io/client-go/util/flowcontrol/throttle.go index 881a2f57d7d..c45169c40f0 100644 --- a/staging/src/k8s.io/client-go/util/flowcontrol/throttle.go +++ b/staging/src/k8s.io/client-go/util/flowcontrol/throttle.go @@ -51,6 +51,22 @@ type tokenBucketRateLimiter struct { // The maximum number of tokens in the bucket is capped at 'burst'. func NewTokenBucketRateLimiter(qps float32, burst int) RateLimiter { limiter := ratelimit.NewBucketWithRate(float64(qps), int64(burst)) + return newTokenBucketRateLimiter(limiter, qps) +} + +// An injectable, mockable clock interface. +type Clock interface { + ratelimit.Clock +} + +// NewTokenBucketRateLimiterWithClock is identical to NewTokenBucketRateLimiter +// but allows an injectable clock, for testing. +func NewTokenBucketRateLimiterWithClock(qps float32, burst int, clock Clock) RateLimiter { + limiter := ratelimit.NewBucketWithRateAndClock(float64(qps), int64(burst), clock) + return newTokenBucketRateLimiter(limiter, qps) +} + +func newTokenBucketRateLimiter(limiter *ratelimit.Bucket, qps float32) RateLimiter { return &tokenBucketRateLimiter{ limiter: limiter, qps: qps,