forked from DataDog/datadog-agent
/
log_limit.go
67 lines (57 loc) · 1.41 KB
/
log_limit.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
55
56
57
58
59
60
61
62
63
64
65
66
67
package util
import (
"sync/atomic"
"time"
)
// LogLimit is a utility that can be used to avoid logging noisily
type LogLimit struct {
// n is the times remaining that the LogLimit will return true for ShouldLog.
// we repeatedly add 1 to it.
n int32
// exit and ticker must be different channels
// becaues Stopping a ticker will not close the ticker channel,
// and we will otherwise leak memory
ticker *time.Ticker
exit chan struct{}
}
// NewLogLimit creates a LogLimit where shouldLog will return
// true the first N times it is called, and will return true once every
// interval thereafter.
func NewLogLimit(n int, interval time.Duration) *LogLimit {
l := &LogLimit{
n: int32(n),
ticker: time.NewTicker(interval),
exit: make(chan struct{}),
}
go l.resetLoop()
return l
}
// ShouldLog returns true if the caller should log
func (l *LogLimit) ShouldLog() bool {
n := atomic.LoadInt32(&l.n)
if n > 0 {
atomic.CompareAndSwapInt32(&l.n, n, n-1)
return true
}
return false
}
// Close will stop the underlying ticker
func (l *LogLimit) Close() {
l.ticker.Stop()
close(l.exit)
}
func (l *LogLimit) resetLoop() {
for {
select {
case <-l.ticker.C:
l.resetCounter()
case <-l.exit:
return
}
}
}
func (l *LogLimit) resetCounter() {
// c.n == 0, it means we have gotten through the first few logs, and after ticker.T we should
// do another log
atomic.CompareAndSwapInt32(&l.n, 0, 1)
}