82 lines
1.8 KiB
Go
82 lines
1.8 KiB
Go
|
package ratecounter
|
||
|
|
|
||
|
|
import (
|
||
|
|
"strconv"
|
||
|
|
"sync/atomic"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// A RateCounter is a thread-safe counter which returns the number of times
|
||
|
|
// 'Incr' has been called in the last interval
|
||
|
|
type RateCounter struct {
|
||
|
|
counter Counter
|
||
|
|
interval time.Duration
|
||
|
|
resolution int
|
||
|
|
partials []Counter
|
||
|
|
current int32
|
||
|
|
running int32
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewRateCounter Constructs a new RateCounter, for the interval provided
|
||
|
|
func NewRateCounter(intrvl time.Duration) *RateCounter {
|
||
|
|
ratecounter := &RateCounter{
|
||
|
|
interval: intrvl,
|
||
|
|
running: 0,
|
||
|
|
}
|
||
|
|
|
||
|
|
return ratecounter.WithResolution(20)
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithResolution determines the minimum resolution of this counter, default is 20
|
||
|
|
func (r *RateCounter) WithResolution(resolution int) *RateCounter {
|
||
|
|
if resolution < 1 {
|
||
|
|
panic("RateCounter resolution cannot be less than 1")
|
||
|
|
}
|
||
|
|
|
||
|
|
r.resolution = resolution
|
||
|
|
r.partials = make([]Counter, resolution)
|
||
|
|
r.current = 0
|
||
|
|
|
||
|
|
return r
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *RateCounter) run() {
|
||
|
|
if ok := atomic.CompareAndSwapInt32(&r.running, 0, 1); !ok {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
ticker := time.NewTicker(time.Duration(float64(r.interval) / float64(r.resolution)))
|
||
|
|
|
||
|
|
for range ticker.C {
|
||
|
|
current := atomic.LoadInt32(&r.current)
|
||
|
|
next := (int(current) + 1) % r.resolution
|
||
|
|
r.counter.Incr(-1 * r.partials[next].Value())
|
||
|
|
r.partials[next].Reset()
|
||
|
|
atomic.CompareAndSwapInt32(&r.current, current, int32(next))
|
||
|
|
if r.counter.Value() == 0 {
|
||
|
|
atomic.StoreInt32(&r.running, 0)
|
||
|
|
ticker.Stop()
|
||
|
|
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
|
||
|
|
// Incr Add an event into the RateCounter
|
||
|
|
func (r *RateCounter) Incr(val int64) {
|
||
|
|
r.counter.Incr(val)
|
||
|
|
r.partials[atomic.LoadInt32(&r.current)].Incr(val)
|
||
|
|
r.run()
|
||
|
|
}
|
||
|
|
|
||
|
|
// Rate Return the current number of events in the last interval
|
||
|
|
func (r *RateCounter) Rate() int64 {
|
||
|
|
return r.counter.Value()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *RateCounter) String() string {
|
||
|
|
return strconv.FormatInt(r.counter.Value(), 10)
|
||
|
|
}
|