Add a test for qps.go

server/util/qps/BUILD

@@ -1,8 +1,18 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "qps",
     srcs = ["qps.go"],
     importpath = "github.com/buildbuddy-io/buildbuddy/server/util/qps",
     visibility = ["//visibility:public"],
 )
+
+go_test(
+    name = "qps_test",
+    size = "small",
+    srcs = ["qps_test.go"],
+    deps = [
+        ":qps",
+        "@com_github_stretchr_testify//require",
+    ],
+)

server/util/qps/qps.go

@@ -31,16 +31,26 @@ type Counter struct {
 	nValidBins uint64
 	window     time.Duration
 	startOnce  sync.Once
+	ticker     <-chan time.Time
 	stop       chan struct{}
 }
 
 // NewCounter returns a QPS counter using the given duration as the averaging
 // window. The caller must call Stop() on the returned counter when it is no
 // longer needed.
 func NewCounter(window time.Duration) *Counter {
+	return new(window, nil)
+}
+
+func NewCounterForTesting(window time.Duration, ticker <-chan time.Time) *Counter {
+	return new(window, ticker)
+}
+
+func new(window time.Duration, ticker <-chan time.Time) *Counter {
 	return &Counter{
 		nValidBins: 1,
 		window:     window,
+		ticker:     ticker,
 		stop:       make(chan struct{}),
 	}
 }
@@ -68,13 +78,16 @@ func (c *Counter) Get() float64 {
 }
 
 func (c *Counter) start() {
-	t := time.NewTicker(time.Duration(float64(c.window) / float64(len(c.counts))))
-	defer t.Stop()
+	if c.ticker == nil {
+		ticker := time.NewTicker(time.Duration(float64(c.window) / float64(len(c.counts))))
+		c.ticker = ticker.C
+		defer ticker.Stop()
+	}
 	for {
 		select {
 		case <-c.stop:
 			return
-		case <-t.C:
+		case <-c.ticker:
 		}
 		// Advance to the next bin, reset its current count, and mark it valid
 		// if we haven't done so already.

server/util/qps/qps_test.go

@@ -0,0 +1,71 @@
+package qps_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/buildbuddy-io/buildbuddy/server/util/qps"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQPS(t *testing.T) {
+	ticker := make(chan time.Time, 1)
+	counter := qps.NewCounterForTesting(5*time.Second, ticker)
+
+	// Test cases where the ring buffer is not full.
+	require.Equal(t, float64(0), counter.Get())
+	counter.Inc()
+	require.Equal(t, float64(12), counter.Get())
+	for i := 0; i < 59; i++ {
+		counter.Inc()
+	}
+	require.Equal(t, float64(720), counter.Get())
+	for i := 0; i < 29; i++ {
+		ticker <- time.Now()
+	}
+	// The current bin is incremented asynchronously upon receiving a tick via
+	// the ticker channel. Give it some time to run.
+	time.Sleep(10 * time.Millisecond)
+	require.Equal(t, float64(24), counter.Get())
+
+	// Test cases where the ring buffer is  full.
+	for i := 0; i < 100; i++ {
+		ticker <- time.Now()
+	}
+	time.Sleep(10 * time.Millisecond)
+	require.Equal(t, float64(0), counter.Get())
+	for i := 0; i < 61; i++ {
+		ticker <- time.Now()
+		time.Sleep(10 * time.Millisecond)
+		counter.Inc()
+	}
+	require.Equal(t, float64(12), counter.Get())
+	for i := 0; i < 60; i++ {
+		ticker <- time.Now()
+		time.Sleep(10 * time.Millisecond)
+		for j := 0; j < i; j++ {
+			counter.Inc()
+		}
+	}
+	// Sum(0, n) = n * (n-1) / 2
+	// So, sum(0, 60) = 60 * 59 / 2 = 1,770
+	// 1,770 Queries / 5s = 354 QPS.
+	require.Equal(t, float64(354), counter.Get())
+}
+
+func TestRaciness(t *testing.T) {
+	ticker := make(chan time.Time, 1)
+	counter := qps.NewCounterForTesting(5*time.Second, ticker)
+
+	for i := 1; i < 1000; i++ {
+		go func() {
+			for j := i; j < 1000; j++ {
+				counter.Inc()
+				if i%2 == 0 {
+					counter.Get()
+				}
+			}
+			ticker <- time.Now()
+		}()
+	}
+}