Implement rate estimation.

estimator/estimator.go

@@ -0,0 +1,54 @@
+package estimator
+
+import (
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type Estimator struct {
+	interval time.Duration
+	count    uint32
+
+	mu   sync.Mutex
+	rate uint32
+	time time.Time
+}
+
+func New(interval time.Duration) *Estimator {
+	return &Estimator{
+		interval: interval,
+		time:     time.Now(),
+	}
+}
+
+func (e *Estimator) swap(now time.Time) {
+	interval := now.Sub(e.time)
+	count := atomic.SwapUint32(&e.count, 0)
+	if interval < time.Millisecond {
+		e.rate = 0
+	} else {
+		e.rate = uint32(uint64(count*1000) / uint64(interval/time.Millisecond))
+	}
+	e.time = now
+}
+
+func (e *Estimator) Add(count uint32) {
+	atomic.AddUint32(&e.count, count)
+}
+
+func (e *Estimator) estimate(now time.Time) uint32 {
+	if now.Sub(e.time) > e.interval {
+		e.swap(now)
+	}
+
+	return e.rate
+}
+
+func (e *Estimator) Estimate() uint32 {
+	now := time.Now()
+
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	return e.estimate(now)
+}

estimator/estimator_test.go

@@ -0,0 +1,21 @@
+package estimator
+
+import (
+	"testing"
+	"time"
+)
+
+func TestEstimator(t *testing.T) {
+	now := time.Now()
+	e := New(time.Second)
+
+	e.estimate(now)
+	e.Add(42)
+	e.Add(128)
+	e.estimate(now.Add(time.Second))
+	rate := e.estimate(now.Add(time.Second + time.Millisecond))
+
+	if rate != 42+128 {
+		t.Errorf("Expected %v, got %v", 42+128, rate)
+	}
+}