1import { assert, it } from "@effect/vitest"
2import { Duration, Effect, type Layer } from "effect"
3import { RateLimiter } from "effect/persistence"
4import { TestClock } from "effect/testing"
5
6type OnExceeded = "fail" | "delay"
7
8const options = (key: string, onExceeded?: OnExceeded) =>
9 ({ algorithm: "token-bucket", window: "5 minutes", limit: 5, key, onExceeded }) as const
10
11const exceeded = (error: RateLimiter.RateLimiterError) =>
12 error.reason._tag === "RateLimitExceeded" ? error.reason : assert.fail("Expected RateLimitExceeded")
13
14/**
15 * Consumes `tokens`, runs `idle`, and asserts the bucket is back at capacity
16 * with a fresh refill interval.
17 */
18export const restartsInterval = Effect.fnUntraced(function*<E, R>(
19 key: string,
20 onExceeded: OnExceeded,
21 tokens: number,
22 idle: Effect.Effect<void, E, R>
23) {
24 const limiter = yield* RateLimiter.make
25 const opts = options(key, onExceeded)
26 yield* limiter.consume({ ...opts, tokens })
27 yield* idle
28 const result = yield* limiter.consume({ ...opts, tokens })
29
30 assert.strictEqual(result.remaining, 5 - tokens)
31 assert.deepStrictEqual(result.delay, Duration.zero)
32 assert.strictEqual(Duration.toMillis(result.resetAfter), tokens * 60_000)
33
34 const excess = limiter.consume({ ...opts, tokens: 6 - tokens })
35 if (onExceeded === "fail") {
36 const error = exceeded(yield* Effect.flip(excess))
37 assert.strictEqual(Duration.toMillis(error.retryAfter), 60_000)
38 } else {
39 const delayed = yield* excess
40 assert.strictEqual(delayed.remaining, -1)
41 assert.strictEqual(Duration.toMillis(delayed.delay), 60_000)
42 assert.strictEqual(Duration.toMillis(delayed.resetAfter), 360_000)
43 }
44})
45
46export const consumeFractionalCosts = Effect.fnUntraced(function*(key: string) {
47 const limiter = yield* RateLimiter.make
48 return yield* Effect.forEach(
49 [0.1, 0.1, 0.4, 0.15, 0.15, 0.1],
50 (tokens) => limiter.consume({ ...options(key), tokens })
51 )
52})
53
54export const suite = (
55 name: string,
56 layer: Layer.Layer<RateLimiter.RateLimiterStore, unknown>
57) => {
58 it.layer(layer, { timeout: "30 seconds", concurrent: false })(`RateLimiter token-bucket timing (${name})`, (it) => {
59 for (const onExceeded of ["fail", "delay"] as const) {
60 for (
61 const [elapsed, tokens, expected, resetAfter] of [
62 [59_000, 1, 1_000, 301_000],
63 [59_500, 1, 500, 300_500],
64 [59_000, 2, 61_000, 361_000],
65 [119_000, 2, 1_000, 301_000]
66 ] as const
67 ) {
68 it.effect(`${onExceeded}: ${tokens} tokens after ${elapsed}ms waits ${expected}ms`, () =>
69 Effect.gen(function*() {
70 const limiter = yield* RateLimiter.make
71 const opts = options(`timing-${onExceeded}-${elapsed}-${tokens}`, onExceeded)
72 yield* limiter.consume({ ...opts, tokens: 5 })
73 yield* TestClock.adjust(elapsed)
74
75 if (onExceeded === "fail") {
76 const error = exceeded(yield* Effect.flip(limiter.consume({ ...opts, tokens })))
77 assert.strictEqual(Duration.toMillis(error.retryAfter), expected)
78 assert.strictEqual(error.remaining, 0)
79 } else {
80 const result = yield* limiter.consume({ ...opts, tokens })
81 assert.strictEqual(Duration.toMillis(result.delay), expected)
82 assert.strictEqual(result.remaining, Math.floor(elapsed / 60_000) - tokens)
83 assert.strictEqual(Duration.toMillis(result.resetAfter), resetAfter)
84 }
85 }))
86 }
87
88 it.effect(`${onExceeded}: resetAfter accounts for elapsed time on an allowed request`, () =>
89 Effect.gen(function*() {
90 const limiter = yield* RateLimiter.make
91 const opts = options(`timing-allowed-${onExceeded}`, onExceeded)
92 const initial = yield* limiter.consume({ ...opts, tokens: 4 })
93 assert.strictEqual(Duration.toMillis(initial.resetAfter), 240_000)
94 yield* TestClock.adjust("59 seconds")
95 const result = yield* limiter.consume(opts)
96
97 assert.deepStrictEqual(result.delay, Duration.zero)
98 assert.strictEqual(result.remaining, 0)
99 assert.strictEqual(Duration.toMillis(result.resetAfter), 241_000)
100
101 yield* TestClock.adjust("241 seconds")
102 const full = yield* limiter.consume({ ...opts, tokens: 0 })
103 assert.strictEqual(full.remaining, 5)
104 assert.deepStrictEqual(full.resetAfter, Duration.zero)
105 }))
106
107 for (const [tokens, idle] of [[1, "61 seconds"], [5, "359 seconds"]] as const) {
108 it.effect(`${onExceeded}: restarts the refill interval after ${idle} idle`, () =>
109 restartsInterval(`timing-idle-${onExceeded}-${tokens}`, onExceeded, tokens, TestClock.adjust(idle)))
110 }
111 }
112
113 it.effect("restarts the interval after a zero-token call without a whole-token refill", () =>
114 Effect.gen(function*() {
115 const limiter = yield* RateLimiter.make
116 const opts = options("timing-zero-tokens")
117 const initial = yield* limiter.consume({ ...opts, tokens: 0 })
118 assert.strictEqual(initial.remaining, 5)
119 assert.deepStrictEqual(initial.resetAfter, Duration.zero)
120
121 yield* TestClock.adjust("30 seconds")
122 const result = yield* limiter.consume(opts)
123
124 assert.strictEqual(result.remaining, 4)
125 assert.deepStrictEqual(result.delay, Duration.zero)
126 assert.strictEqual(Duration.toMillis(result.resetAfter), 60_000)
127 }))
128
129 for (const allowOverflow of [false, true]) {
130 it.effect(`preserves signed fractional token counts with allowOverflow=${allowOverflow}`, () =>
131 Effect.gen(function*() {
132 const store = yield* RateLimiter.RateLimiterStore
133 const opts = {
134 key: `fractional-count-${allowOverflow}`,
135 limit: 5,
136 refillRate: Duration.minutes(1),
137 allowOverflow
138 }
139 const initial = yield* store.tokenBucket({ ...opts, tokens: 4.5 })
140 const excess = yield* store.tokenBucket({ ...opts, tokens: 1 })
141 const stored = yield* store.tokenBucket({ ...opts, tokens: 0 })
142
143 assert.deepStrictEqual([initial, excess, stored], [
144 [0.5, 0],
145 [-0.5, 0],
146 [allowOverflow ? -0.5 : 0.5, 0]
147 ])
148 }))
149 }
150
151 it.effect("preserves accumulated fractional balances and their reset timing", () =>
152 Effect.gen(function*() {
153 const results = yield* consumeFractionalCosts("fractional-accumulation")
154
155 assert.deepStrictEqual(
156 results.map((result) => ({
157 remaining: result.remaining,
158 delay: Duration.toMillis(result.delay),
159 resetAfter: Duration.toMillis(result.resetAfter)
160 })),
161 [
162 { remaining: 4.9, delay: 0, resetAfter: 60_000 },
163 { remaining: 4.800000000000001, delay: 0, resetAfter: 60_000 },
164 { remaining: 4.4, delay: 0, resetAfter: 60_000 },
165 { remaining: 4.25, delay: 0, resetAfter: 60_000 },
166 { remaining: 4.1, delay: 0, resetAfter: 60_000 },
167 { remaining: 3.9999999999999996, delay: 0, resetAfter: 120_000 }
168 ]
169 )
170 }))
171
172 it.effect("rejects a request with a fractional token deficit", () =>
173 Effect.gen(function*() {
174 const limiter = yield* RateLimiter.make
175 const opts = options("fractional-deficit", "fail")
176 yield* limiter.consume({ ...opts, tokens: 4.5 })
177
178 const error = exceeded(yield* Effect.flip(limiter.consume(opts)))
179 assert.strictEqual(error.remaining, 0)
180 assert.strictEqual(error.limit, 5)
181 }))
182
183 for (const onExceeded of ["fail", "delay"] as const) {
184 it.effect(`${onExceeded}: fractional costs reset at the next whole-token refill`, () =>
185 Effect.gen(function*() {
186 const limiter = yield* RateLimiter.make
187 const opts = options(`fractional-reset-${onExceeded}`, onExceeded)
188 const initial = yield* limiter.consume({ ...opts, tokens: 0.5 })
189 yield* TestClock.adjust("45 seconds")
190 const partial = yield* limiter.consume({ ...opts, tokens: 0 })
191 yield* TestClock.adjust("15 seconds")
192 const full = yield* limiter.consume({ ...opts, tokens: 0 })
193
194 assert.strictEqual(full.remaining, 5)
195 assert.deepStrictEqual(
196 [initial, partial, full].map((result) => Duration.toMillis(result.resetAfter)),
197 [60_000, 15_000, 0]
198 )
199 }))
200 }
201
202 it.effect("fractional deficits remain rejected until the whole-token refill", () =>
203 Effect.gen(function*() {
204 const limiter = yield* RateLimiter.make
205 const opts = options("fractional-retry-boundary", "fail")
206 yield* limiter.consume({ ...opts, tokens: 4.5 })
207 yield* TestClock.adjust("45 seconds")
208 const first = exceeded(yield* Effect.flip(limiter.consume(opts)))
209 yield* TestClock.adjust(14_999)
210 const beforeRefill = exceeded(yield* Effect.flip(limiter.consume(opts)))
211 yield* TestClock.adjust(1)
212 const allowed = yield* limiter.consume(opts)
213
214 assert.deepStrictEqual(
215 [first, beforeRefill].map((error) => Duration.toMillis(error.retryAfter)),
216 [15_000, 1]
217 )
218 assert.strictEqual(allowed.remaining, 0.5)
219 assert.deepStrictEqual(allowed.delay, Duration.zero)
220 }))
221
222 it.effect("fractional delay reservations wait for enough whole-token refills", () =>
223 Effect.gen(function*() {
224 const limiter = yield* RateLimiter.make
225 const opts = options("fractional-reservation-boundary", "delay")
226 yield* limiter.consume({ ...opts, tokens: 4.5 })
227 yield* TestClock.adjust("45 seconds")
228 const first = yield* limiter.consume(opts)
229 const second = yield* limiter.consume({ ...opts, tokens: 1.25 })
230 yield* TestClock.adjust(74_999)
231 const beforeRefill = yield* limiter.consume({ ...opts, tokens: 0 })
232 yield* TestClock.adjust(1)
233 const available = yield* limiter.consume({ ...opts, tokens: 0 })
234
235 assert.deepStrictEqual(
236 [first, second, beforeRefill, available].map((result) => ({
237 remaining: result.remaining,
238 delay: Duration.toMillis(result.delay),
239 resetAfter: Duration.toMillis(result.resetAfter)
240 })),
241 [
242 { remaining: -0.5, delay: 15_000, resetAfter: 315_000 },
243 { remaining: -1.75, delay: 75_000, resetAfter: 375_000 },
244 { remaining: -0.75, delay: 1, resetAfter: 300_001 },
245 { remaining: 0.25, delay: 0, resetAfter: 300_000 }
246 ]
247 )
248 }))
249
250 it.effect("preserves fractional elapsed milliseconds in the store tuple and timing metadata", () =>
251 Effect.gen(function*() {
252 const store = yield* RateLimiter.RateLimiterStore
253 const opts = {
254 key: "timing-fractional",
255 limit: 3,
256 refillRate: Duration.millis(1_000 / 3),
257 allowOverflow: true
258 }
259 const initial = yield* store.tokenBucket({ ...opts, tokens: 3 })
260 assert.deepStrictEqual(initial, [0, 0])
261 yield* TestClock.adjust(500)
262 const [remaining, elapsedMillis] = yield* store.tokenBucket({ ...opts, tokens: 1 })
263
264 assert.strictEqual(remaining, 0)
265 assert.closeTo(elapsedMillis, 1_000 / 6, 0.000_001)
266
267 const limiter = yield* RateLimiter.make
268 const result = yield* limiter.consume({
269 algorithm: "token-bucket",
270 onExceeded: "delay",
271 window: "1 second",
272 limit: 3,
273 key: opts.key
274 })
275 assert.strictEqual(result.remaining, -1)
276 assert.closeTo(Duration.toMillis(result.delay), 1_000 / 6, 0.000_01)
277 assert.closeTo(Duration.toMillis(result.resetAfter), 7_000 / 6, 0.000_01)
278 }))
279
280 it.effect("failed retries do not reserve tokens or move the refill boundary", () =>
281 Effect.gen(function*() {
282 const limiter = yield* RateLimiter.make
283 const opts = options("timing-fail-boundary", "fail")
284 yield* limiter.consume({ ...opts, tokens: 5 })
285 yield* Effect.flip(limiter.consume(opts))
286 yield* TestClock.adjust("59 seconds")
287 yield* Effect.flip(limiter.consume(opts))
288 yield* Effect.flip(limiter.consume(opts))
289 yield* TestClock.adjust("1 second")
290
291 const result = yield* limiter.consume(opts)
292 assert.deepStrictEqual(result.delay, Duration.zero)
293 assert.strictEqual(result.remaining, 0)
294 assert.strictEqual(Duration.toMillis(result.resetAfter), 300_000)
295 }))
296
297 it.effect("accumulated delay reservations retain the partial refill interval", () =>
298 Effect.gen(function*() {
299 const limiter = yield* RateLimiter.make
300 const opts = options("timing-delay-reservations", "delay")
301 yield* limiter.consume({ ...opts, tokens: 5 })
302 yield* TestClock.adjust("59 seconds")
303 yield* limiter.consume(opts)
304 const result = yield* limiter.consume(opts)
305
306 assert.strictEqual(Duration.toMillis(result.delay), 61_000)
307 assert.strictEqual(result.remaining, -2)
308 assert.strictEqual(Duration.toMillis(result.resetAfter), 361_000)
309 }))
310 })
311}