Fixed leaky bucket ResetTime invalid on first request

mailgun · Dec 16, 2020 · 79420e9 · 79420e9
1 parent 4d87ed8
commit 79420e9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/algorithms.go b/algorithms.go
@@ -313,7 +313,7 @@ func leakyBucket(s Store, c Cache, r *RateLimitReq) (resp *RateLimitResp, err er
 		Status:    Status_UNDER_LIMIT,
 		Limit:     r.Limit,
 		Remaining: r.Limit - r.Hits,
-		ResetTime: duration / r.Limit,
+		ResetTime: now + duration / r.Limit,
 	}
 
 	// Client could be requesting that we start with the bucket OVER_LIMIT

diff --git a/functional_test.go b/functional_test.go
@@ -191,6 +191,7 @@ func TestLeakyBucket(t *testing.T) {
 		},
 	}
 
+	now := clock.Now()
 	for i, test := range tests {
 		resp, err := client.GetRateLimits(context.Background(), &guber.GetRateLimitsReq{
 			Requests: []*guber.RateLimitReq{
@@ -212,7 +213,7 @@ func TestLeakyBucket(t *testing.T) {
 		assert.Equal(t, test.Status, rl.Status, i)
 		assert.Equal(t, test.Remaining, rl.Remaining, i)
 		assert.Equal(t, int64(5), rl.Limit, i)
-		assert.True(t, rl.ResetTime != 0)
+		assert.True(t, rl.ResetTime > now.Unix())
 		clock.Advance(test.Sleep)
 	}
 }