blob: 83d4059c2af00df15dec05e9c950f6c7f73bf6f5 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause
package logger
import (
"time"
)
// tokenBucket is a simple token bucket style rate limiter.
// It's similar in function to golang.org/x/time/rate.Limiter, which we
// can't use because:
// - It doesn't give access to the number of accumulated tokens, which we
// need for implementing hysteresis;
// - It doesn't let us provide our own time function, which we need for
// implementing proper unit tests.
//
// rate.Limiter is also much more complex than necessary, but that wouldn't
// be enough to disqualify it on its own.
//
// Unlike rate.Limiter, this token bucket does not attempt to
// do any locking of its own. Don't try to access it reentrantly.
// That's fine inside this types/logger package because we already have
// locking at a higher level.
type tokenBucket struct {
remaining int
max int
tick time.Duration
t time.Time
}
func newTokenBucket(tick time.Duration, max int, now time.Time) *tokenBucket {
return &tokenBucket{max, max, tick, now}
}
func (tb *tokenBucket) Get() bool {
if tb.remaining > 0 {
tb.remaining--
return true
}
return false
}
func (tb *tokenBucket) Refund(n int) {
b := tb.remaining + n
if b > tb.max {
tb.remaining = tb.max
} else {
tb.remaining = b
}
}
func (tb *tokenBucket) AdvanceTo(t time.Time) {
diff := t.Sub(tb.t)
// only use up whole ticks. The remainder will be used up
// next time.
ticks := int(diff / tb.tick)
tb.t = tb.t.Add(time.Duration(ticks) * tb.tick)
tb.Refund(ticks)
}
|
