|
|
|
package collection
|
|
|
|
|
|
|
|
import (
|
|
|
|
"sync"
|
|
|
|
"time"
|
|
|
|
|
|
|
|
"github.com/tal-tech/go-zero/core/timex"
|
|
|
|
)
|
|
|
|
|
|
|
|
type (
|
|
|
|
RollingWindowOption func(rollingWindow *RollingWindow)
|
|
|
|
|
|
|
|
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
|
|
|
|
}
|
|
|
|
)
|
|
|
|
|
|
|
|
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
|
|
|
|
}
|
|
|
|
|
|
|
|
func (rw *RollingWindow) Add(v float64) {
|
|
|
|
rw.lock.Lock()
|
|
|
|
defer rw.lock.Unlock()
|
|
|
|
rw.updateOffset()
|
|
|
|
rw.win.add(rw.offset, v)
|
|
|
|
}
|
|
|
|
|
|
|
|
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
|
|
|
|
} else {
|
|
|
|
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 + 1 + i)
|
|
|
|
}
|
|
|
|
|
|
|
|
rw.offset = (offset + span) % rw.size
|
|
|
|
now := timex.Now()
|
|
|
|
// align to interval time boundary
|
|
|
|
rw.lastTime = now - (now-rw.lastTime)%rw.interval
|
|
|
|
}
|
|
|
|
|
|
|
|
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()
|
|
|
|
}
|
|
|
|
|
|
|
|
func IgnoreCurrentBucket() RollingWindowOption {
|
|
|
|
return func(w *RollingWindow) {
|
|
|
|
w.ignoreCurrent = true
|
|
|
|
}
|
|
|
|
}
|