/
ratelimit.go
54 lines (47 loc) · 1.36 KB
/
ratelimit.go
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
package log
import (
"time"
"github.com/signalfx/golib/v3/eventcounter"
)
// A RateLimitedLogger will disable itself after a Limit number of log messages in a period.
// This can be useful if you want to prevent your logger from killing the thing it is logging to
// such as disk or network.
type RateLimitedLogger struct {
EventCounter eventcounter.EventCounter
Limit int64
Logger Logger
LimitLogger Logger
Now func() time.Time
}
func (r *RateLimitedLogger) now() time.Time {
if r.Now == nil {
return time.Now()
}
return r.Now()
}
// Log kvs to the wrapped Logger if the limit hasn't been reached
func (r *RateLimitedLogger) Log(kvs ...interface{}) {
now := r.now()
if r.EventCounter.Event(now) > r.Limit {
if r.LimitLogger != nil {
// Note: Log here messes up "caller" :/
r.LimitLogger.Log(kvs...)
}
return
}
r.Logger.Log(kvs...)
}
// Disabled returns true if this logger is over its limit or if the wrapped logger is
// disabled.
func (r *RateLimitedLogger) Disabled() bool {
now := r.now()
return r.EventCounter.Events(now, 0) > r.Limit || IsDisabled(r.Logger)
}
// NewOnePerSecond returns a *RateLimitedLogger that allows one message per second
func NewOnePerSecond(log Logger) *RateLimitedLogger {
return &RateLimitedLogger{
EventCounter: eventcounter.New(time.Now(), time.Second),
Limit: 1,
Logger: log,
}
}