diff --git a/limiter_fixed_truncated_window.go b/limiter_fixed_truncated_window.go index c5caf74..8f56528 100644 --- a/limiter_fixed_truncated_window.go +++ b/limiter_fixed_truncated_window.go @@ -70,7 +70,7 @@ func (l *FixedTruncatedWindowRateLimiter) Dump(ctx context.Context) (r Result, e if TimeGTE(l.window.Add(l.rate.Duration()), now) { l.rateLimitReached = false - l.window = now.Truncate(l.rate.Unit) + l.window = now.Truncate(l.rate.TruncateDuration()) } c, err := l.db.Get(ctx, l.window) @@ -98,7 +98,7 @@ func (l *FixedTruncatedWindowRateLimiter) try(ctx context.Context, tokens int64) if TimeGTE(l.window.Add(l.rate.Duration()), now) { l.rateLimitReached = false - l.window = now.Truncate(l.rate.Unit) + l.window = now.Truncate(l.rate.TruncateDuration()) } ttw := l.window.Add(l.rate.Duration()).Sub(now) @@ -142,7 +142,7 @@ func (l *FixedTruncatedWindowRateLimiter) check(ctx context.Context, tokens int6 if TimeGTE(l.window.Add(l.rate.Duration()), now) { // new window so no rate Limit l.rateLimitReached = false - l.window = now.Truncate(l.rate.Unit) + l.window = now.Truncate(l.rate.TruncateDuration()) return res(0, l.capacity), nil } diff --git a/limiter_fixed_truncated_window_test.go b/limiter_fixed_truncated_window_test.go index 84c65c2..ce6829e 100644 --- a/limiter_fixed_truncated_window_test.go +++ b/limiter_fixed_truncated_window_test.go @@ -68,6 +68,93 @@ func assertFixedWindowTruncatedStepEquals( func TestNewFixedTruncatedWindowRateLimiter(t *testing.T) { tests := []testFixedWindowTruncated{ + { + + name: "window of 1 day configured with rate duration as a whole", + capacity: 3, + rate: Rate{Amount: 24, Unit: time.Hour}, + startTime: time.Date(2022, 02, 05, 4, 23, 00, 0, time.UTC), + steps: []testFixedWindowTruncatedStep{ + { + method: check, + passTime: 0, + expectedErr: nil, + expectedFreeSlots: 3, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 4 + expectedErr: nil, + expectedFreeSlots: 2, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 10 + expectedErr: nil, + expectedFreeSlots: 1, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 16 + expectedErr: nil, + expectedFreeSlots: 0, + expectedTtw: 0, + }, + { + method: try, + passTime: 0, // 22:23 + expectedErr: ErrRateLimitExceeded, + expectedFreeSlots: 0, + expectedTtw: time.Hour*1 + time.Minute*37, + }, + }, + }, + { + name: "window of 1 day configured with Unit as a whole", + capacity: 3, + rate: Rate{Amount: 1, Unit: time.Hour * 24}, + startTime: time.Date(2022, 02, 05, 4, 23, 00, 0, time.UTC), + steps: []testFixedWindowTruncatedStep{ + { + method: check, + passTime: 0, + expectedErr: nil, + expectedFreeSlots: 3, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 4 + expectedErr: nil, + expectedFreeSlots: 2, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 10 + expectedErr: nil, + expectedFreeSlots: 1, + expectedTtw: 0, + }, + { + method: try, + passTime: time.Hour * 6, // 16 + expectedErr: nil, + expectedFreeSlots: 0, + expectedTtw: 0, + }, + { + method: try, + passTime: 0, // 22:23 + expectedErr: ErrRateLimitExceeded, + expectedFreeSlots: 0, + expectedTtw: time.Hour*1 + time.Minute*37, + }, + }, + }, { name: "start of the window reaches rate limit before first tick", capacity: 2, diff --git a/rate.go b/rate.go index 1aa98bd..0a159f5 100644 --- a/rate.go +++ b/rate.go @@ -10,3 +10,13 @@ type Rate struct { func (r Rate) Duration() time.Duration { return time.Duration(r.Amount) * r.Unit } + +// TruncateDuration returns, for windows smaller than a minute, the sole unit as they scape the sexagesimal counting +// mode. Otherwise, return the product of amount and unit to produce the full rate limit window. +func (r Rate) TruncateDuration() time.Duration { + if r.Unit < time.Minute { + return r.Unit + } + + return r.Duration() +} diff --git a/tests/main.go b/tests/main.go index 61758e3..55722e2 100644 --- a/tests/main.go +++ b/tests/main.go @@ -22,10 +22,10 @@ func main() { rateLimit := pacemaker.NewFixedTruncatedWindowRateLimiter( pacemaker.FixedTruncatedWindowArgs{ - Capacity: 1200, + Capacity: 1000, Rate: pacemaker.Rate{ - Unit: time.Minute, - Amount: 3, + Unit: time.Hour, + Amount: 24, }, Clock: pacemaker.NewClock(), DB: pacemaker.NewFixedWindowRedisStorage(