From 10d3946f1e449ccbfb91a3bd36a211e22022db12 Mon Sep 17 00:00:00 2001
From: Lukasz Raczylo <lukasz@raczylo.com>
Date: Tue, 3 Nov 2026 00:00:24 +0000
Subject: [PATCH] fix(rps): stop dropping requests recorded during a sample

The sampler read the counter and then reset it in two steps. Requests recorded between the steps were lost. The sampler now swaps the counter to zero in one atomic step.
---
 rps_tracker.go             | 17 +++++++++------
 rps_tracker_sample_test.go | 43 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 6 deletions(-)
 create mode 100644 rps_tracker_sample_test.go

diff --git a/rps_tracker.go b/rps_tracker.go
index 0c13c97..4167455 100644
--- a/rps_tracker.go
+++ b/rps_tracker.go
@@ -57,17 +57,23 @@ func (r *RPSTracker) Shutdown() {
 
 // sample calculates RPS since last sample
 func (r *RPSTracker) sample() {
-	now := time.Now()
-	nowNano := now.UnixNano()
+	r.sampleCount()
+}
+
+// sampleCount performs one sample and returns the requests it consumed.
+func (r *RPSTracker) sampleCount() int64 {
+	nowNano := time.Now().UnixNano()
 
-	currentCount := r.lastCount.Load()
 	lastSampleNano := r.lastSampleTime.Load()
 
 	if lastSampleNano == 0 {
 		r.lastSampleTime.Store(nowNano)
-		return
+		return 0
 	}
 
+	// Swap, not Load+Store(0): requests recorded between the two would be lost.
+	currentCount := r.lastCount.Swap(0)
+
 	elapsed := float64(nowNano-lastSampleNano) / float64(time.Second)
 	if elapsed > 0 {
 		rps := float64(currentCount) / elapsed
@@ -75,9 +81,8 @@ func (r *RPSTracker) sample() {
 		r.currentRPS.Store(uint64(rps * 100))
 	}
 
-	// Reset for next sample
-	r.lastCount.Store(0)
 	r.lastSampleTime.Store(nowNano)
+	return currentCount
 }
 
 // GetCurrentRPS returns the current requests per second
diff --git a/rps_tracker_sample_test.go b/rps_tracker_sample_test.go
new file mode 100644
index 0000000..c1e140b
--- /dev/null
+++ b/rps_tracker_sample_test.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Pins theorems/Theorems/RpsSample.lean: no recorded request is dropped by sample().
+func TestRPSTracker_SampleKeepsConcurrentRequests(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	tr := NewRPSTracker(ctx)
+	tr.Shutdown() // stop the background sampler; sample() is driven by the test
+
+	const writers, perWriter = 8, 20000
+	var wg sync.WaitGroup
+	for i := 0; i < writers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < perWriter; j++ {
+				tr.RecordRequest()
+			}
+		}()
+	}
+	var sampled int64
+	done := make(chan struct{})
+	go func() { wg.Wait(); close(done) }()
+loop:
+	for {
+		select {
+		case <-done:
+			break loop
+		default:
+			sampled += tr.sampleCount()
+		}
+	}
+	sampled += tr.sampleCount()
+	assert.Equal(t, int64(writers*perWriter), sampled)
+}