118 lines
2.3 KiB
Go
118 lines
2.3 KiB
Go
package counter
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type bucket struct {
|
|
val int64
|
|
next *bucket
|
|
}
|
|
|
|
func (b *bucket) Add(val int64) {
|
|
b.val += val
|
|
}
|
|
|
|
func (b *bucket) Value() int64 {
|
|
return b.val
|
|
}
|
|
|
|
func (b *bucket) Reset() {
|
|
b.val = 0
|
|
}
|
|
|
|
var _ Counter = new(rollingCounter)
|
|
|
|
type rollingCounter struct {
|
|
mu sync.RWMutex
|
|
buckets []bucket
|
|
bucketTime int64
|
|
lastAccess int64
|
|
cur *bucket
|
|
}
|
|
|
|
// NewRolling creates a new window. windowTime is the time covering the entire
|
|
// window. windowBuckets is the number of buckets the window is divided into.
|
|
// An example: a 10 second window with 10 buckets will have 10 buckets covering
|
|
// 1 second each.
|
|
func NewRolling(window time.Duration, winBucket int) Counter {
|
|
buckets := make([]bucket, winBucket)
|
|
bucket := &buckets[0]
|
|
for i := 1; i < winBucket; i++ {
|
|
bucket.next = &buckets[i]
|
|
bucket = bucket.next
|
|
}
|
|
bucket.next = &buckets[0]
|
|
bucketTime := time.Duration(window.Nanoseconds() / int64(winBucket))
|
|
return &rollingCounter{
|
|
cur: &buckets[0],
|
|
buckets: buckets,
|
|
bucketTime: int64(bucketTime),
|
|
lastAccess: time.Now().UnixNano(),
|
|
}
|
|
}
|
|
|
|
// Add increments the counter by value and return new value.
|
|
func (r *rollingCounter) Add(val int64) {
|
|
r.mu.Lock()
|
|
r.lastBucket().Add(val)
|
|
r.mu.Unlock()
|
|
}
|
|
|
|
// Value get the counter value.
|
|
func (r *rollingCounter) Value() (sum int64) {
|
|
now := time.Now().UnixNano()
|
|
r.mu.RLock()
|
|
b := r.cur
|
|
i := r.elapsed(now)
|
|
for j := 0; j < len(r.buckets); j++ {
|
|
// skip all future reset bucket.
|
|
if i > 0 {
|
|
i--
|
|
} else {
|
|
sum += b.Value()
|
|
}
|
|
b = b.next
|
|
}
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
|
|
// Reset reset the counter.
|
|
func (r *rollingCounter) Reset() {
|
|
r.mu.Lock()
|
|
for i := range r.buckets {
|
|
r.buckets[i].Reset()
|
|
}
|
|
r.mu.Unlock()
|
|
}
|
|
|
|
func (r *rollingCounter) elapsed(now int64) (i int) {
|
|
var e int64
|
|
if e = now - r.lastAccess; e <= r.bucketTime {
|
|
return
|
|
}
|
|
if i = int(e / r.bucketTime); i > len(r.buckets) {
|
|
i = len(r.buckets)
|
|
}
|
|
return
|
|
}
|
|
|
|
func (r *rollingCounter) lastBucket() (b *bucket) {
|
|
now := time.Now().UnixNano()
|
|
b = r.cur
|
|
// reset the buckets between now and number of buckets ago. If
|
|
// that is more that the existing buckets, reset all.
|
|
if i := r.elapsed(now); i > 0 {
|
|
r.lastAccess = now
|
|
for ; i > 0; i-- {
|
|
// replace the next used bucket.
|
|
b = b.next
|
|
b.Reset()
|
|
}
|
|
}
|
|
r.cur = b
|
|
return
|
|
}
|