forked from asynkron/protoactor-go
/
throttler.go
91 lines (77 loc) · 2.39 KB
/
throttler.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
package actor
import (
"log/slog"
"sync/atomic"
"time"
)
type ShouldThrottle func() Valve
type Valve int32
const (
Open Valve = iota
Closing
Closed
)
// NewThrottle
// This has no guarantees that the throttle opens exactly after the period, since it is reset asynchronously
// Throughput has been prioritized over exact re-opening
// throttledCallBack, This will be called with the number of events what was throttled after the period
func NewThrottle(maxEventsInPeriod int32, period time.Duration, throttledCallBack func(int32)) ShouldThrottle {
currentEvents := int32(0)
startTimer := func(duration time.Duration, back func(int32)) {
go func() {
// crete ticker to mimic sleep, we do not want to put the goroutine to sleep
// as it will schedule it out of the P making a syscall, we just want it to
// halt for the given period of time
ticker := time.NewTicker(duration)
defer ticker.Stop()
<-ticker.C // wait for the ticker to tick once
timesCalled := atomic.SwapInt32(¤tEvents, 0)
if timesCalled > maxEventsInPeriod {
throttledCallBack(timesCalled - maxEventsInPeriod)
}
}()
}
return func() Valve {
tries := atomic.AddInt32(¤tEvents, 1)
if tries == 1 {
startTimer(period, throttledCallBack)
}
if tries == maxEventsInPeriod {
return Closing
} else if tries > maxEventsInPeriod {
return Closed
} else {
return Open
}
}
}
func NewThrottleWithLogger(logger *slog.Logger, maxEventsInPeriod int32, period time.Duration, throttledCallBack func(*slog.Logger, int32)) ShouldThrottle {
currentEvents := int32(0)
startTimer := func(duration time.Duration, back func(*slog.Logger, int32)) {
go func() {
// crete ticker to mimic sleep, we do not want to put the goroutine to sleep
// as it will schedule it out of the P making a syscall, we just want it to
// halt for the given period of time
ticker := time.NewTicker(duration)
defer ticker.Stop()
<-ticker.C // wait for the ticker to tick once
timesCalled := atomic.SwapInt32(¤tEvents, 0)
if timesCalled > maxEventsInPeriod {
throttledCallBack(logger, timesCalled-maxEventsInPeriod)
}
}()
}
return func() Valve {
tries := atomic.AddInt32(¤tEvents, 1)
if tries == 1 {
startTimer(period, throttledCallBack)
}
if tries == maxEventsInPeriod {
return Closing
} else if tries > maxEventsInPeriod {
return Closed
} else {
return Open
}
}
}