forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
rollingwindow.go
150 lines (127 loc) · 3.06 KB
/
rollingwindow.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package collection
import (
"sync"
"time"
"github.com/zeromicro/go-zero/core/timex"
)
type (
// RollingWindowOption let callers customize the RollingWindow.
RollingWindowOption func(rollingWindow *RollingWindow)
// RollingWindow defines a rolling window to calculate the events in buckets with time interval.
RollingWindow struct {
lock sync.RWMutex
size int
win *window
interval time.Duration
offset int
ignoreCurrent bool
lastTime time.Duration // start time of the last bucket
}
)
// NewRollingWindow returns a RollingWindow that with size buckets and time interval,
// use opts to customize the RollingWindow.
func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
if size < 1 {
panic("size must be greater than 0")
}
w := &RollingWindow{
size: size,
win: newWindow(size),
interval: interval,
lastTime: timex.Now(),
}
for _, opt := range opts {
opt(w)
}
return w
}
// Add adds value to current bucket.
func (rw *RollingWindow) Add(v float64) {
rw.lock.Lock()
defer rw.lock.Unlock()
rw.updateOffset()
rw.win.add(rw.offset, v)
}
// Reduce runs fn on all buckets, ignore current bucket if ignoreCurrent was set.
func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
rw.lock.RLock()
defer rw.lock.RUnlock()
var diff int
span := rw.span()
// ignore current bucket, because of partial data
if span == 0 && rw.ignoreCurrent {
diff = rw.size - 1
} else {
diff = rw.size - span
}
if diff > 0 {
offset := (rw.offset + span + 1) % rw.size
rw.win.reduce(offset, diff, fn)
}
}
func (rw *RollingWindow) span() int {
offset := int(timex.Since(rw.lastTime) / rw.interval)
if 0 <= offset && offset < rw.size {
return offset
}
return rw.size
}
func (rw *RollingWindow) updateOffset() {
span := rw.span()
if span <= 0 {
return
}
offset := rw.offset
// reset expired buckets
for i := 0; i < span; i++ {
rw.win.resetBucket((offset + i + 1) % rw.size)
}
rw.offset = (offset + span) % rw.size
now := timex.Now()
// align to interval time boundary
rw.lastTime = now - (now-rw.lastTime)%rw.interval
}
// Bucket defines the bucket that holds sum and num of additions.
type Bucket struct {
Sum float64
Count int64
}
func (b *Bucket) add(v float64) {
b.Sum += v
b.Count++
}
func (b *Bucket) reset() {
b.Sum = 0
b.Count = 0
}
type window struct {
buckets []*Bucket
size int
}
func newWindow(size int) *window {
buckets := make([]*Bucket, size)
for i := 0; i < size; i++ {
buckets[i] = new(Bucket)
}
return &window{
buckets: buckets,
size: size,
}
}
func (w *window) add(offset int, v float64) {
w.buckets[offset%w.size].add(v)
}
func (w *window) reduce(start, count int, fn func(b *Bucket)) {
for i := 0; i < count; i++ {
fn(w.buckets[(start+i)%w.size])
}
}
func (w *window) resetBucket(offset int) {
w.buckets[offset%w.size].reset()
}
// IgnoreCurrentBucket lets the Reduce call ignore current bucket.
func IgnoreCurrentBucket() RollingWindowOption {
return func(w *RollingWindow) {
w.ignoreCurrent = true
}
}