Add util/math to dskit

math/math.go

@@ -0,0 +1,33 @@
+package math
+
+// Max returns the maximum of two ints
+func Max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
+
+// Min returns the minimum of two ints
+func Min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// Max64 returns the maximum of two int64s
+func Max64(a, b int64) int64 {
+	if a > b {
+		return a
+	}
+	return b
+}
+
+// Min64 returns the minimum of two int64s
+func Min64(a, b int64) int64 {
+	if a < b {
+		return a
+	}
+	return b
+}

math/rate.go

@@ -0,0 +1,59 @@
+package math
+
+import (
+	"sync"
+	"time"
+
+	"go.uber.org/atomic"
+)
+
+// EwmaRate tracks an exponentially weighted moving average of a per-second rate.
+type EwmaRate struct {
+	newEvents atomic.Int64
+
+	alpha    float64
+	interval time.Duration
+
+	mutex    sync.RWMutex
+	lastRate float64
+	init     bool
+}
+
+func NewEWMARate(alpha float64, interval time.Duration) *EwmaRate {
+	return &EwmaRate{
+		alpha:    alpha,
+		interval: interval,
+	}
+}
+
+// Rate returns the per-second rate.
+func (r *EwmaRate) Rate() float64 {
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+	return r.lastRate
+}
+
+// Tick assumes to be called every r.interval.
+func (r *EwmaRate) Tick() {
+	newEvents := r.newEvents.Swap(0)
+	instantRate := float64(newEvents) / r.interval.Seconds()
+
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if r.init {
+		r.lastRate += r.alpha * (instantRate - r.lastRate)
+	} else {
+		r.init = true
+		r.lastRate = instantRate
+	}
+}
+
+// Inc counts one event.
+func (r *EwmaRate) Inc() {
+	r.newEvents.Inc()
+}
+
+func (r *EwmaRate) Add(delta int64) {
+	r.newEvents.Add(delta)
+}

math/rate_test.go

@@ -0,0 +1,47 @@
+package math
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRate(t *testing.T) {
+	ticks := []struct {
+		events int
+		want   float64
+	}{
+		{60, 1},
+		{30, 0.9},
+		{0, 0.72},
+		{60, 0.776},
+		{0, 0.6208},
+		{0, 0.49664},
+		{0, 0.397312},
+		{0, 0.3178496},
+		{0, 0.25427968},
+		{0, 0.203423744},
+		{0, 0.1627389952},
+	}
+	r := NewEWMARate(0.2, time.Minute)
+
+	for _, tick := range ticks {
+		for e := 0; e < tick.events; e++ {
+			r.Inc()
+		}
+		r.Tick()
+		// We cannot do double comparison, because double operations on different
+		// platforms may actually produce results that differ slightly.
+		// There are multiple issues about this in Go's github, eg: 18354 or 20319.
+		require.InDelta(t, tick.want, r.Rate(), 0.0000000001, "unexpected rate")
+	}
+
+	r = NewEWMARate(0.2, time.Minute)
+
+	for _, tick := range ticks {
+		r.Add(int64(tick.events))
+		r.Tick()
+		require.InDelta(t, tick.want, r.Rate(), 0.0000000001, "unexpected rate")
+	}
+}