From 18b969f5a70f0fc41b6615ce6d49a4f2befc241c Mon Sep 17 00:00:00 2001
From: Neal Patel <neal@ineal.me>
Date: Tue, 2 Mar 2021 21:11:10 -0800
Subject: [PATCH] Fix leaky bucket algorithm returning remaining more than
 limit

---
 algorithms.go      |  8 ++++----
 functional_test.go | 12 +++++++++---
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/algorithms.go b/algorithms.go
index 1b0dfbdd..6c7bd115 100644
--- a/algorithms.go
+++ b/algorithms.go
@@ -235,15 +235,15 @@ func leakyBucket(s Store, c Cache, r *RateLimitReq) (resp *RateLimitResp, err er
 		elapsed := now - b.UpdatedAt
 		leak := float64(elapsed) / rate
 
-		if int64(b.Remaining) > b.Limit {
-			b.Remaining = float64(b.Limit)
-		}
-
 		if int64(leak) > 0 {
 			b.Remaining += leak
 			b.UpdatedAt = now
 		}
 
+		if int64(b.Remaining) > b.Limit {
+			b.Remaining = float64(b.Limit)
+		}
+
 		rl := &RateLimitResp{
 			Limit:     b.Limit,
 			Remaining: int64(b.Remaining),
diff --git a/functional_test.go b/functional_test.go
index 46c1c8f9..ab75b8b7 100644
--- a/functional_test.go
+++ b/functional_test.go
@@ -308,11 +308,17 @@ func TestLeakyBucket(t *testing.T) {
 			Hits:      0,
 			Remaining: 1,
 			Status:    guber.Status_UNDER_LIMIT,
+			Sleep:     clock.Second * 60,
+		},
+		{
+			Name:      "should max out the limit",
+			Hits:      0,
+			Remaining: 10,
+			Status:    guber.Status_UNDER_LIMIT,
 			Sleep:     clock.Second,
 		},
 	}
 
-	now := clock.Now()
 	for _, test := range tests {
 		t.Run(test.Name, func(t *testing.T) {
 			resp, err := client.GetRateLimits(context.Background(), &guber.GetRateLimitsReq{
@@ -327,15 +333,15 @@ func TestLeakyBucket(t *testing.T) {
 					},
 				},
 			})
-			clock.Freeze(clock.Now())
 			require.NoError(t, err)
+			require.Len(t, resp.Responses, 1)
 
 			rl := resp.Responses[0]
 
 			assert.Equal(t, test.Status, rl.Status)
 			assert.Equal(t, test.Remaining, rl.Remaining)
 			assert.Equal(t, int64(10), rl.Limit)
-			assert.True(t, rl.ResetTime > now.Unix())
+			assert.Equal(t, clock.Now().Unix()+3, rl.ResetTime/1000)
 			clock.Advance(test.Sleep)
 		})
 	}