From 0ea5a0fcc039c4f1f0517a56ebd8afba6ac11cfc Mon Sep 17 00:00:00 2001
From: bohdanstorozhuk <storozhuk.b.m@gmail.com>
Date: Fri, 1 Dec 2023 16:55:03 +0000
Subject: [PATCH] Fix immediate slack re-apply

---
 limiter_atomic_int64.go |  2 +-
 ratelimit_test.go       | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/limiter_atomic_int64.go b/limiter_atomic_int64.go
index 8f2e66c..6c7d838 100644
--- a/limiter_atomic_int64.go
+++ b/limiter_atomic_int64.go
@@ -69,7 +69,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..d4776ed 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -269,6 +269,21 @@ func TestInitial(t *testing.T) {
 	}
 }
 
+func TestImmediateSlackReapply(t *testing.T) {
+	t.Parallel()
+	rl := New(1, Per(time.Second), WithSlack(1))
+	rl.Take()                   // we take one immediately
+	time.Sleep(2 * time.Second) // waiting for slack to accumulate
+	start := time.Now()
+	for i := 0; i < 5; i++ {
+		rl.Take() // consume 2 immediately and then 1 every second after it
+		// sleep 1ns for now-timeOfNextPermissionIssue != int64(t.maxSlack)
+		time.Sleep(1 * time.Nanosecond)
+	}
+
+	assert.Condition(t, func() bool { return time.Since(start) > 3*time.Second }, "too fast consumption: %v", time.Since(start))
+}
+
 func TestSlack(t *testing.T) {
 	t.Parallel()
 	// To simulate slack, we combine two limiters.