From 6c14836d817aecfbaab4507bbfbc0f4d62d63f87 Mon Sep 17 00:00:00 2001
From: Kyagara <kevinsouzarocha@outlook.com>
Date: Fri, 12 Jan 2024 10:19:10 -0300
Subject: [PATCH] comments, readme

---
 README.md           | 1 +
 equinox_test.go     | 2 +-
 ratelimit/bucket.go | 2 ++
 ratelimit/limits.go | 3 +++
 4 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fa1c344..6bfa45f 100644
--- a/README.md
+++ b/README.md
@@ -119,6 +119,7 @@ func main() {
 
 ## Todo
 
+- Create a rate limit interface
 - Add Redis store for rate limit, using a lua script
 - Maybe the context usage throughout the project could be improved
 - Maybe add more options to customize the rate limiter
diff --git a/equinox_test.go b/equinox_test.go
index d02977b..587f027 100644
--- a/equinox_test.go
+++ b/equinox_test.go
@@ -113,7 +113,7 @@ func TestRateLimitWithMock(t *testing.T) {
 	// One request left, it should be blocked as it would exceed the rate limit
 
 	// A deadline is set, however, the bucket refill would take longer than the ctx deadline, waiting would exceed that deadline
-	// It also shouldn't Reserve() a token from the bucket
+	// It also shouldn't Reserve() a request from the bucket
 	ctx := context.Background()
 	ctx, c := context.WithTimeout(ctx, 2*time.Second)
 	defer c()
diff --git a/ratelimit/bucket.go b/ratelimit/bucket.go
index d1518d6..89741d0 100644
--- a/ratelimit/bucket.go
+++ b/ratelimit/bucket.go
@@ -38,6 +38,7 @@ func NewBucket(interval time.Duration, intervalOverhead time.Duration, baseLimit
 	}
 }
 
+// Checks if the 'next' reset is in the past, and if so, resets the bucket tokens and sets the next reset.
 func (b *Bucket) check() {
 	now := time.Now()
 	if b.next.Before(now) {
@@ -46,6 +47,7 @@ func (b *Bucket) check() {
 	}
 }
 
+// Increments the number of tokens in the bucket and returns if the bucket is rate limited.
 func (b *Bucket) isRateLimited() bool {
 	b.check()
 	if b.limit == 0 {
diff --git a/ratelimit/limits.go b/ratelimit/limits.go
index e9749b6..11b7b4b 100644
--- a/ratelimit/limits.go
+++ b/ratelimit/limits.go
@@ -61,11 +61,14 @@ func (l *Limit) checkBuckets(ctx context.Context, logger zerolog.Logger, route s
 				Str("limit_type", l.limitType).
 				Object("bucket", bucket).
 				Msg("Rate limited")
+
 			err := WaitN(ctx, bucket.next, time.Until(bucket.next))
 			if err != nil {
 				bucket.mutex.Unlock()
 				return err
 			}
+
+			// 'next' reset is now in the past, so reset the bucket
 			bucket.check()
 			bucket.tokens++
 		}