Skip to content

Commit

Permalink
comments, readme
Browse files Browse the repository at this point in the history
  • Loading branch information
Kyagara committed Jan 12, 2024
1 parent 41b4b34 commit 6c14836
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 1 deletion.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion equinox_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 2 additions & 0 deletions ratelimit/bucket.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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 {
Expand Down
3 changes: 3 additions & 0 deletions ratelimit/limits.go
Original file line number Diff line number Diff line change
Expand Up @@ -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++
}
Expand Down

0 comments on commit 6c14836

Please sign in to comment.