fix: truncated window on rates for units greater than seconds

sonirico · Nov 23, 2022 · f1a5bcd · f1a5bcd
1 parent 33b09c7
commit f1a5bcd
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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(