1 // Copyright 2015 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
19 func TestLimit(t *testing.T) {
21 t.Errorf("Limit(10) == Inf should be false")
25 func closeEnough(a, b Limit) bool {
26 return (math.Abs(float64(a)/float64(b)) - 1.0) < 1e-9
29 func TestEvery(t *testing.T) {
31 interval time.Duration
36 {1 * time.Nanosecond, Limit(1e9)},
37 {1 * time.Microsecond, Limit(1e6)},
38 {1 * time.Millisecond, Limit(1e3)},
39 {10 * time.Millisecond, Limit(100)},
40 {100 * time.Millisecond, Limit(10)},
41 {1 * time.Second, Limit(1)},
42 {2 * time.Second, Limit(0.5)},
43 {time.Duration(2.5 * float64(time.Second)), Limit(0.4)},
44 {4 * time.Second, Limit(0.25)},
45 {10 * time.Second, Limit(0.1)},
46 {time.Duration(math.MaxInt64), Limit(1e9 / float64(math.MaxInt64))},
48 for _, tc := range cases {
49 lim := Every(tc.interval)
50 if !closeEnough(lim, tc.lim) {
51 t.Errorf("Every(%v) = %v want %v", tc.interval, lim, tc.lim)
57 d = 100 * time.Millisecond
62 t1 = t0.Add(time.Duration(1) * d)
63 t2 = t0.Add(time.Duration(2) * d)
64 t3 = t0.Add(time.Duration(3) * d)
65 t4 = t0.Add(time.Duration(4) * d)
66 t5 = t0.Add(time.Duration(5) * d)
67 t9 = t0.Add(time.Duration(9) * d)
76 func run(t *testing.T, lim *Limiter, allows []allow) {
77 for i, allow := range allows {
78 ok := lim.AllowN(allow.t, allow.n)
80 t.Errorf("step %d: lim.AllowN(%v, %v) = %v want %v",
81 i, allow.t, allow.n, ok, allow.ok)
86 func TestLimiterBurst1(t *testing.T) {
87 run(t, NewLimiter(10, 1), []allow{
94 {t2, 2, false}, // burst size is 1, so n=2 always fails
100 func TestLimiterBurst3(t *testing.T) {
101 run(t, NewLimiter(10, 3), []allow{
118 func TestLimiterJumpBackwards(t *testing.T) {
119 run(t, NewLimiter(10, 3), []allow{
120 {t1, 1, true}, // start at t1
121 {t0, 1, true}, // jump back to t0, two tokens remain
125 {t1, 1, true}, // got a token
128 {t2, 1, true}, // got another token
134 // Ensure that tokensFromDuration doesn't produce
135 // rounding errors by truncating nanoseconds.
136 // See golang.org/issues/34861.
137 func TestLimiter_noTruncationErrors(t *testing.T) {
138 if !NewLimiter(0.7692307692307693, 1).Allow() {
139 t.Fatal("expected true")
143 func TestSimultaneousRequests(t *testing.T) {
154 // Very slow replenishing bucket.
155 lim := NewLimiter(limit, burst)
157 // Tries to take a token, atomically updates the counter and decreases the wait
161 if ok := lim.Allow(); ok {
162 atomic.AddUint32(&numOK, 1)
167 for i := 0; i < numRequests; i++ {
172 t.Errorf("numOK = %d, want %d", numOK, burst)
176 func TestLongRunningQPS(t *testing.T) {
178 t.Skip("skipping in short mode")
180 if runtime.GOOS == "openbsd" {
181 t.Skip("low resolution time.Sleep invalidates test (golang.org/issue/14183)")
185 // The test runs for a few seconds executing many requests and then checks
186 // that overall number of requests is reasonable.
193 lim := NewLimiter(limit, burst)
195 var wg sync.WaitGroup
197 if ok := lim.Allow(); ok {
198 atomic.AddInt32(&numOK, 1)
204 end := start.Add(5 * time.Second)
205 for time.Now().Before(end) {
209 // This will still offer ~500 requests per second, but won't consume
210 // outrageous amount of CPU.
211 time.Sleep(2 * time.Millisecond)
214 elapsed := time.Since(start)
215 ideal := burst + (limit * float64(elapsed) / float64(time.Second))
217 // We should never get more requests than allowed.
218 if want := int32(ideal + 1); numOK > want {
219 t.Errorf("numOK = %d, want %d (ideal %f)", numOK, want, ideal)
221 // We should get very close to the number of requests allowed.
222 if want := int32(0.999 * ideal); numOK < want {
223 t.Errorf("numOK = %d, want %d (ideal %f)", numOK, want, ideal)
227 type request struct {
234 // dFromDuration converts a duration to a multiple of the global constant d
235 func dFromDuration(dur time.Duration) int {
236 // Adding a millisecond to be swallowed by the integer division
237 // because we don't care about small inaccuracies
238 return int((dur + time.Millisecond) / d)
241 // dSince returns multiples of d since t0
242 func dSince(t time.Time) int {
243 return dFromDuration(t.Sub(t0))
246 func runReserve(t *testing.T, lim *Limiter, req request) *Reservation {
247 return runReserveMax(t, lim, req, InfDuration)
250 func runReserveMax(t *testing.T, lim *Limiter, req request, maxReserve time.Duration) *Reservation {
251 r := lim.reserveN(req.t, req.n, maxReserve)
252 if r.ok && (dSince(r.timeToAct) != dSince(req.act)) || r.ok != req.ok {
253 t.Errorf("lim.reserveN(t%d, %v, %v) = (t%d, %v) want (t%d, %v)",
254 dSince(req.t), req.n, maxReserve, dSince(r.timeToAct), r.ok, dSince(req.act), req.ok)
259 func TestSimpleReserve(t *testing.T) {
260 lim := NewLimiter(10, 2)
262 runReserve(t, lim, request{t0, 2, t0, true})
263 runReserve(t, lim, request{t0, 2, t2, true})
264 runReserve(t, lim, request{t3, 2, t4, true})
267 func TestMix(t *testing.T) {
268 lim := NewLimiter(10, 2)
270 runReserve(t, lim, request{t0, 3, t1, false}) // should return false because n > Burst
271 runReserve(t, lim, request{t0, 2, t0, true})
272 run(t, lim, []allow{{t1, 2, false}}) // not enought tokens - don't allow
273 runReserve(t, lim, request{t1, 2, t2, true})
274 run(t, lim, []allow{{t1, 1, false}}) // negative tokens - don't allow
275 run(t, lim, []allow{{t3, 1, true}})
278 func TestCancelInvalid(t *testing.T) {
279 lim := NewLimiter(10, 2)
281 runReserve(t, lim, request{t0, 2, t0, true})
282 r := runReserve(t, lim, request{t0, 3, t3, false})
283 r.CancelAt(t0) // should have no effect
284 runReserve(t, lim, request{t0, 2, t2, true}) // did not get extra tokens
287 func TestCancelLast(t *testing.T) {
288 lim := NewLimiter(10, 2)
290 runReserve(t, lim, request{t0, 2, t0, true})
291 r := runReserve(t, lim, request{t0, 2, t2, true})
292 r.CancelAt(t1) // got 2 tokens back
293 runReserve(t, lim, request{t1, 2, t2, true})
296 func TestCancelTooLate(t *testing.T) {
297 lim := NewLimiter(10, 2)
299 runReserve(t, lim, request{t0, 2, t0, true})
300 r := runReserve(t, lim, request{t0, 2, t2, true})
301 r.CancelAt(t3) // too late to cancel - should have no effect
302 runReserve(t, lim, request{t3, 2, t4, true})
305 func TestCancel0Tokens(t *testing.T) {
306 lim := NewLimiter(10, 2)
308 runReserve(t, lim, request{t0, 2, t0, true})
309 r := runReserve(t, lim, request{t0, 1, t1, true})
310 runReserve(t, lim, request{t0, 1, t2, true})
311 r.CancelAt(t0) // got 0 tokens back
312 runReserve(t, lim, request{t0, 1, t3, true})
315 func TestCancel1Token(t *testing.T) {
316 lim := NewLimiter(10, 2)
318 runReserve(t, lim, request{t0, 2, t0, true})
319 r := runReserve(t, lim, request{t0, 2, t2, true})
320 runReserve(t, lim, request{t0, 1, t3, true})
321 r.CancelAt(t2) // got 1 token back
322 runReserve(t, lim, request{t2, 2, t4, true})
325 func TestCancelMulti(t *testing.T) {
326 lim := NewLimiter(10, 4)
328 runReserve(t, lim, request{t0, 4, t0, true})
329 rA := runReserve(t, lim, request{t0, 3, t3, true})
330 runReserve(t, lim, request{t0, 1, t4, true})
331 rC := runReserve(t, lim, request{t0, 1, t5, true})
332 rC.CancelAt(t1) // get 1 token back
333 rA.CancelAt(t1) // get 2 tokens back, as if C was never reserved
334 runReserve(t, lim, request{t1, 3, t5, true})
337 func TestReserveJumpBack(t *testing.T) {
338 lim := NewLimiter(10, 2)
340 runReserve(t, lim, request{t1, 2, t1, true}) // start at t1
341 runReserve(t, lim, request{t0, 1, t1, true}) // should violate Limit,Burst
342 runReserve(t, lim, request{t2, 2, t3, true})
345 func TestReserveJumpBackCancel(t *testing.T) {
346 lim := NewLimiter(10, 2)
348 runReserve(t, lim, request{t1, 2, t1, true}) // start at t1
349 r := runReserve(t, lim, request{t1, 2, t3, true})
350 runReserve(t, lim, request{t1, 1, t4, true})
351 r.CancelAt(t0) // cancel at t0, get 1 token back
352 runReserve(t, lim, request{t1, 2, t4, true}) // should violate Limit,Burst
355 func TestReserveSetLimit(t *testing.T) {
356 lim := NewLimiter(5, 2)
358 runReserve(t, lim, request{t0, 2, t0, true})
359 runReserve(t, lim, request{t0, 2, t4, true})
360 lim.SetLimitAt(t2, 10)
361 runReserve(t, lim, request{t2, 1, t4, true}) // violates Limit and Burst
364 func TestReserveSetBurst(t *testing.T) {
365 lim := NewLimiter(5, 2)
367 runReserve(t, lim, request{t0, 2, t0, true})
368 runReserve(t, lim, request{t0, 2, t4, true})
369 lim.SetBurstAt(t3, 4)
370 runReserve(t, lim, request{t0, 4, t9, true}) // violates Limit and Burst
373 func TestReserveSetLimitCancel(t *testing.T) {
374 lim := NewLimiter(5, 2)
376 runReserve(t, lim, request{t0, 2, t0, true})
377 r := runReserve(t, lim, request{t0, 2, t4, true})
378 lim.SetLimitAt(t2, 10)
379 r.CancelAt(t2) // 2 tokens back
380 runReserve(t, lim, request{t2, 2, t3, true})
383 func TestReserveMax(t *testing.T) {
384 lim := NewLimiter(10, 2)
387 runReserveMax(t, lim, request{t0, 2, t0, true}, maxT)
388 runReserveMax(t, lim, request{t0, 1, t1, true}, maxT) // reserve for close future
389 runReserveMax(t, lim, request{t0, 1, t2, false}, maxT) // time to act too far in the future
396 delay int // in multiples of d
400 func runWait(t *testing.T, lim *Limiter, w wait) {
402 err := lim.WaitN(w.ctx, w.n)
403 delay := time.Now().Sub(start)
404 if (w.nilErr && err != nil) || (!w.nilErr && err == nil) || w.delay != dFromDuration(delay) {
407 errString = "<non-nil error>"
409 t.Errorf("lim.WaitN(%v, lim, %v) = %v with delay %v ; want %v with delay %v",
410 w.name, w.n, err, delay, errString, d*time.Duration(w.delay))
414 func TestWaitSimple(t *testing.T) {
415 lim := NewLimiter(10, 3)
417 ctx, cancel := context.WithCancel(context.Background())
419 runWait(t, lim, wait{"already-cancelled", ctx, 1, 0, false})
421 runWait(t, lim, wait{"exceed-burst-error", context.Background(), 4, 0, false})
423 runWait(t, lim, wait{"act-now", context.Background(), 2, 0, true})
424 runWait(t, lim, wait{"act-later", context.Background(), 3, 2, true})
427 func TestWaitCancel(t *testing.T) {
428 lim := NewLimiter(10, 3)
430 ctx, cancel := context.WithCancel(context.Background())
431 runWait(t, lim, wait{"act-now", ctx, 2, 0, true}) // after this lim.tokens = 1
436 runWait(t, lim, wait{"will-cancel", ctx, 3, 1, false})
437 // should get 3 tokens back, and have lim.tokens = 2
438 t.Logf("tokens:%v last:%v lastEvent:%v", lim.tokens, lim.last, lim.lastEvent)
439 runWait(t, lim, wait{"act-now-after-cancel", context.Background(), 2, 0, true})
442 func TestWaitTimeout(t *testing.T) {
443 lim := NewLimiter(10, 3)
445 ctx, cancel := context.WithTimeout(context.Background(), d)
447 runWait(t, lim, wait{"act-now", ctx, 2, 0, true})
448 runWait(t, lim, wait{"w-timeout-err", ctx, 3, 0, false})
451 func TestWaitInf(t *testing.T) {
452 lim := NewLimiter(Inf, 0)
454 runWait(t, lim, wait{"exceed-burst-no-error", context.Background(), 3, 0, true})
457 func BenchmarkAllowN(b *testing.B) {
458 lim := NewLimiter(Every(1*time.Second), 1)
462 b.RunParallel(func(pb *testing.PB) {
469 func BenchmarkWaitNNoDelay(b *testing.B) {
470 lim := NewLimiter(Limit(b.N), b.N)
471 ctx := context.Background()
474 for i := 0; i < b.N; i++ {