56 lines
1.1 KiB
Go
56 lines
1.1 KiB
Go
package utils
|
|
|
|
import (
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type RollingStats struct {
|
|
buckets []atomic.Int64
|
|
bucketIndex atomic.Int32
|
|
numBuckets int32
|
|
bucketSize time.Duration
|
|
}
|
|
|
|
func NewRollingStats(numBuckets int32, bucketSize time.Duration) *RollingStats {
|
|
rs := &RollingStats{
|
|
buckets: make([]atomic.Int64, numBuckets),
|
|
numBuckets: numBuckets,
|
|
bucketSize: bucketSize,
|
|
}
|
|
go rs.rotateBuckets()
|
|
return rs
|
|
}
|
|
|
|
func (rs *RollingStats) rotateBuckets() {
|
|
ticker := time.NewTicker(rs.bucketSize)
|
|
defer ticker.Stop()
|
|
for range ticker.C {
|
|
newIdx := (rs.bucketIndex.Load() + 1) % int32(rs.numBuckets)
|
|
rs.bucketIndex.Store(newIdx)
|
|
rs.buckets[newIdx].Store(0)
|
|
}
|
|
}
|
|
|
|
func (rs *RollingStats) Increment() {
|
|
idx := rs.bucketIndex.Load()
|
|
rs.buckets[idx].Add(1)
|
|
}
|
|
|
|
func (rs *RollingStats) GetRPS() float64 {
|
|
var total int64
|
|
for i := 0; i < int(rs.numBuckets); i++ {
|
|
total += rs.buckets[i].Load()
|
|
}
|
|
|
|
windowSeconds := float64(rs.numBuckets) * rs.bucketSize.Seconds()
|
|
return float64(total) / windowSeconds
|
|
}
|
|
|
|
func (rs *RollingStats) GetTotal() int64 {
|
|
var total int64
|
|
for i := 0; i < int(rs.numBuckets); i++ {
|
|
total += rs.buckets[i].Load()
|
|
}
|
|
return total
|
|
}
|