From b409a800bdc33c6bf5e66f86edcc28ff4077a3b6 Mon Sep 17 00:00:00 2001
From: Pawel Krolikowski <pawel@uber.com>
Date: Mon, 4 Mar 2024 00:56:48 -0300
Subject: [PATCH] Fix an issue with maxSlack boundary updates.

Fixes #119
The solution is a copy from #120, but follows the testing framework
that we have - I did not want us to have a real `Sleep` in tests.

I'm not exactly thrilled by the testing setup (especially the
milliseconds) or the clock itself, but I'm not willing to totally give
up on it like #120 proposes.
I also wanted ALL implementations of the ratelimiter to be tested,
not just the currently selected.

Might follow up with some testing cleanups and/or clock migration.
---
 limiter_atomic_int64.go |  5 ++---
 ratelimit_test.go       | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/limiter_atomic_int64.go b/limiter_atomic_int64.go
index 8f2e66c..8d370b7 100644
--- a/limiter_atomic_int64.go
+++ b/limiter_atomic_int64.go
@@ -21,9 +21,8 @@
 package ratelimit // import "go.uber.org/ratelimit"
 
 import (
-	"time"
-
 	"sync/atomic"
+	"time"
 )
 
 type atomicInt64Limiter struct {
@@ -69,7 +68,7 @@ func (t *atomicInt64Limiter) Take() time.Time {
 		case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
 			// if this is our first call or t.maxSlack == 0 we need to shrink issue time to now
 			newTimeOfNextPermissionIssue = now
-		case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack):
+		case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
 			// a lot of nanoseconds passed since the last Take call
 			// we will limit max accumulated time to maxSlack
 			newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
diff --git a/ratelimit_test.go b/ratelimit_test.go
index 7268f87..ce81038 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -14,6 +14,8 @@ import (
 type testRunner interface {
 	// createLimiter builds a limiter with given options.
 	createLimiter(int, ...Option) Limiter
+	// takeOnceAfter attempts to Take at a specific time.
+	takeOnceAfter(time.Duration, Limiter)
 	// startTaking tries to Take() on passed in limiters in a loop/goroutine.
 	startTaking(rls ...Limiter)
 	// assertCountAt asserts the limiters have Taken() a number of times at the given time.
@@ -112,6 +114,16 @@ func (r *runnerImpl) startTaking(rls ...Limiter) {
 	})
 }
 
+// takeOnceAfter attempts to Take at a specific time.
+func (r *runnerImpl) takeOnceAfter(d time.Duration, rl Limiter) {
+	r.wg.Add(1)
+	r.afterFunc(d, func() {
+		rl.Take()
+		r.count.Inc()
+		r.wg.Done()
+	})
+}
+
 // assertCountAt asserts the limiters have Taken() a number of times at a given time.
 func (r *runnerImpl) assertCountAt(d time.Duration, count int) {
 	r.wg.Add(1)
@@ -269,6 +281,22 @@ func TestInitial(t *testing.T) {
 	}
 }
 
+func TestMaxSlack(t *testing.T) {
+	t.Parallel()
+	runTest(t, func(r testRunner) {
+		rl := r.createLimiter(1, WithSlack(1))
+
+		r.takeOnceAfter(time.Nanosecond, rl)
+		r.takeOnceAfter(2*time.Second+1*time.Nanosecond, rl)
+		r.takeOnceAfter(2*time.Second+2*time.Nanosecond, rl)
+		r.takeOnceAfter(2*time.Second+3*time.Nanosecond, rl)
+		r.takeOnceAfter(2*time.Second+4*time.Nanosecond, rl)
+
+		r.assertCountAt(3*time.Second, 3)
+		r.assertCountAt(10*time.Second, 5)
+	})
+}
+
 func TestSlack(t *testing.T) {
 	t.Parallel()
 	// To simulate slack, we combine two limiters.