@stimulcross/rate-limiter 0.0.1

package/tests/integration/limiters/token-bucket.limiter.spec.ts

@@ -0,0 +1,363 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { type Clock, RateLimiterDestroyedError, RateLimitErrorCode } from '../../../src/index.js';
+import { TokenBucketLimiter, type TokenBucketState } from '../../../src/limiters/token-bucket/index.js';
+import { InMemoryStateStore } from '../../../src/runtime/in-memory-state-store.js';
+
+describe('TokenBucketLimiter (Integration)', () => {
+	let clock: Clock;
+	let store: InMemoryStateStore<TokenBucketState>;
+
+	beforeEach(() => {
+		vi.useFakeTimers();
+		vi.setSystemTime(10_000);
+
+		clock = { now: () => Date.now() };
+		store = new InMemoryStateStore(clock);
+	});
+
+	afterEach(() => {
+		vi.clearAllTimers();
+		vi.useRealTimers();
+	});
+
+	describe('Reject mode (immediate decisions)', () => {
+		it('should allow requests within the capacity immediately', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 2,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			const res1 = await limiter.run(async () => 'A');
+			const res2 = await limiter.run(async () => 'B');
+
+			expect(res1).toBe('A');
+			expect(res2).toBe('B');
+		});
+
+		it('should reject requests that exceed the capacity if behavior is "reject"', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await limiter.run(() => 'A');
+
+			const promise = limiter.run(() => 'B');
+			await expect(promise).rejects.toMatchObject({ code: RateLimitErrorCode.LimitExceeded });
+		});
+
+		it('should consume multiple tokens if cost > 1', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 5,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await expect(limiter.run(() => 'A', { cost: 3 })).resolves.toBe('A');
+			await expect(limiter.run(() => 'B', { cost: 3 })).rejects.toMatchObject({
+				code: RateLimitErrorCode.LimitExceeded,
+			});
+		});
+
+		it('should allow new requests after tokens are refilled', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await limiter.run(() => 'A');
+			await expect(limiter.run(() => 'B')).rejects.toMatchObject({ code: RateLimitErrorCode.LimitExceeded });
+
+			// 1 token refill
+			await vi.advanceTimersByTimeAsync(100);
+
+			await expect(limiter.run(() => 'C')).resolves.toBe('C');
+		});
+	});
+
+	describe('Enqueue mode (queueing & scheduling)', () => {
+		it('should delay request and execute it when enough tokens are refilled', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			void limiter.run(() => 'A');
+
+			const pBSpy = vi.fn().mockReturnValue('B');
+			const pB = limiter.run<unknown>(pBSpy);
+
+			expect(pBSpy).not.toHaveBeenCalled();
+
+			await vi.advanceTimersByTimeAsync(100);
+
+			const result = await pB;
+			expect(pBSpy).toHaveBeenCalledOnce();
+			expect(result).toBe('B');
+		});
+
+		it('should maintain order and timings for multiple queued requests', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			void limiter.run(() => {});
+
+			const results: string[] = [];
+
+			// delay 100ms
+			const pA = limiter.run(() => results.push('A'));
+			// delay 200ms
+			const pB = limiter.run(() => results.push('B'));
+			// delay 300ms
+			const pC = limiter.run(() => results.push('C'));
+
+			await vi.advanceTimersByTimeAsync(100);
+			await pA;
+			expect(results).toEqual(['A']);
+
+			await vi.advanceTimersByTimeAsync(100);
+			await pB;
+			expect(results).toEqual(['A', 'B']);
+
+			await vi.advanceTimersByTimeAsync(100);
+			await pC;
+			expect(results).toEqual(['A', 'B', 'C']);
+		});
+
+		it('should reject with QueueOverflow if executor queue exceeds maxSize', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				queue: { capacity: 1 },
+				clock,
+				store,
+			});
+
+			const pA = limiter.run(() => 'A');
+			const pB = limiter.run(() => 'B');
+
+			await vi.advanceTimersByTimeAsync(0);
+
+			// overflow
+			const pC = limiter.run(() => 'C');
+
+			await expect(pC).rejects.toMatchObject({ code: RateLimitErrorCode.QueueOverflow });
+
+			await vi.advanceTimersByTimeAsync(0);
+			await expect(pA).resolves.toBe('A');
+
+			await vi.advanceTimersByTimeAsync(100);
+			await expect(pB).resolves.toBe('B');
+		});
+
+		it('should execute queued tasks based on priority order, not just chronological', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 2,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			void limiter.run(() => 'A', { priority: 1 });
+			void limiter.run(() => 'B', { priority: 1 });
+
+			const order: string[] = [];
+
+			// delay 100ms
+			void limiter.run(() => order.push('Lowest'), { priority: 1 });
+			// delay 200мс
+			void limiter.run(() => order.push('Highest'), { priority: 5 });
+
+			await vi.advanceTimersByTimeAsync(100);
+			expect(order).toEqual(['Highest']);
+
+			await vi.advanceTimersByTimeAsync(100);
+			expect(order).toEqual(['Highest', 'Lowest']);
+		});
+
+		it('should enqueue the task and reject it with Expired when TTL is reached', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				queue: { maxWaitMs: 150 },
+				clock,
+				store,
+			});
+
+			const pA = limiter.run(() => 'A');
+			// delay 100ms
+			const pB = limiter.run(() => 'B');
+			// delay 200ms, expires after 150ms
+			const spyC = vi.fn().mockReturnValue('C');
+			const pC = limiter.run(spyC);
+			pC.catch(() => {});
+
+			await expect(pA).resolves.toBe('A');
+
+			await vi.advanceTimersByTimeAsync(100);
+			await expect(pB).resolves.toBe('B');
+			expect(spyC).not.toHaveBeenCalled();
+
+			await vi.advanceTimersByTimeAsync(50);
+			await expect(pC).rejects.toMatchObject({ code: RateLimitErrorCode.Expired });
+			expect(spyC).not.toHaveBeenCalled();
+		});
+
+		it('should free up the canceled ticket for new requests, while keeping already queued requests at their scheduled time', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			const pA = limiter.run(() => 'A');
+
+			const controllerB = new AbortController();
+			// delay 100ms
+			const pB = limiter.run(() => 'B', { signal: controllerB.signal });
+
+			const spyC = vi.fn().mockReturnValue('C');
+			// delay 200ms
+			const pC = limiter.run(spyC);
+
+			await vi.advanceTimersByTimeAsync(0);
+
+			// canceling B
+			// the 200ms ticket should be freed up, the consumed token should be returned to the bucket
+			controllerB.abort();
+			await expect(pB).rejects.toMatchObject({ code: RateLimitErrorCode.Cancelled });
+
+			const spyD = vi.fn().mockReturnValue('D');
+			// delay 200ms (it was freed up by canceling B)
+			const pD = limiter.run(spyD);
+
+			await expect(pA).resolves.toBe('A');
+			expect(spyC).not.toHaveBeenCalled();
+			expect(spyD).not.toHaveBeenCalled();
+
+			// t = 10_100
+			// C takes the 100ms ticket created by B
+			await vi.advanceTimersByTimeAsync(100);
+
+			await expect(pC).resolves.toBe('C');
+			expect(spyC).toHaveBeenCalledOnce();
+			expect(spyD).not.toHaveBeenCalled();
+
+			// t = 10_200
+			// D takes the 200ms ticket created by itself
+			await vi.advanceTimersByTimeAsync(100);
+
+			await expect(pD).resolves.toBe('D');
+			expect(spyD).toHaveBeenCalledOnce();
+		});
+	});
+
+	describe('Runtime overrides', () => {
+		it('should allow overriding limitBehavior per task', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await limiter.run(() => 'A');
+			const pB = limiter.run(() => 'B', { limitBehavior: 'enqueue' });
+
+			await vi.advanceTimersByTimeAsync(100);
+			await expect(pB).resolves.toBe('B');
+		});
+
+		it('should override max wait time for a specific task and expire it independently', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'enqueue',
+				capacity: 1,
+				refillRate: 10,
+				queue: { maxWaitMs: 5000 },
+				clock,
+				store,
+			});
+
+			void limiter.run(() => 'A');
+
+			// delay 100ms
+			const pB = limiter.run(() => 'B');
+
+			// delay 200ms
+			// expires after 150ms, so it should be rejected
+			const pC = limiter.run(() => 'C', { maxWaitMs: 150 });
+			pC.catch(() => {});
+
+			// delay 300ms
+			const pD = limiter.run(() => 'D', { maxWaitMs: 400 });
+
+			// t=10_100
+			await vi.advanceTimersByTimeAsync(100);
+			await expect(pB).resolves.toBe('B');
+
+			// t=10_150
+			await vi.advanceTimersByTimeAsync(50);
+			await expect(pC).rejects.toMatchObject({ code: RateLimitErrorCode.Expired });
+
+			// t=10_300
+			await vi.advanceTimersByTimeAsync(150);
+			await expect(pD).resolves.toBe('D');
+		});
+	});
+
+	describe('State lifecycle', () => {
+		it('should reset limits and clear queued tasks on clear()', async () => {
+			const limiter = new TokenBucketLimiter({
+				limitBehavior: 'reject',
+				capacity: 1,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await limiter.run(() => 'A');
+
+			await limiter.clear();
+
+			await expect(limiter.run(() => 'B')).resolves.toBe('B');
+		});
+
+		it('should throw RateLimiterDestroyedError after destroy() is called', async () => {
+			const limiter = new TokenBucketLimiter({
+				capacity: 5,
+				refillRate: 10,
+				clock,
+				store,
+			});
+
+			await limiter.destroy();
+
+			await expect(limiter.run(() => 'A')).rejects.toThrow(RateLimiterDestroyedError);
+		});
+	});
+});

package/tests/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "../tsconfig.json",
+	"include": ["./", "../src"]
+}

package/tests/unit/policies/composite.policy.spec.ts

@@ -0,0 +1,244 @@
+import { describe, expect, it } from 'vitest';
+import { InvalidCostError } from '../../../src/index.js';
+import { CompositePolicy } from '../../../src/limiters/composite.policy.js';
+import { FixedWindowPolicy } from '../../../src/limiters/fixed-window/fixed-window.policy.js';
+import { type FixedWindowState, type FixedWindowStatus } from '../../../src/limiters/fixed-window/index.js';
+
+const createFixedWindowCompositePolicy = (
+	policies: FixedWindowPolicy[],
+): CompositePolicy<FixedWindowState, FixedWindowStatus, FixedWindowPolicy> =>
+	new CompositePolicy<FixedWindowState, FixedWindowStatus, FixedWindowPolicy>(policies);
+
+describe('CompositePolicy', () => {
+	describe('Initialization', () => {
+		it('should initialize with provided policies and expose them', () => {
+			const policies = [new FixedWindowPolicy(10, 1000), new FixedWindowPolicy(5, 5000)];
+			const composite = createFixedWindowCompositePolicy(policies);
+
+			expect(composite.policies).toHaveLength(2);
+			expect(composite.policies).toStrictEqual(policies);
+		});
+
+		it('should aggregate initial states from all underlying policies', () => {
+			const composite = createFixedWindowCompositePolicy([
+				new FixedWindowPolicy(10, 1000),
+				new FixedWindowPolicy(5, 5000),
+			]);
+
+			const states = composite.getInitialState(0);
+
+			expect(states).toHaveLength(2);
+			expect(states[0]).toEqual({ windowStart: 0, used: 0, reserved: 0 });
+			expect(states[1]).toEqual({ windowStart: 0, used: 0, reserved: 0 });
+		});
+	});
+
+	describe('Get Status', () => {
+		const policies = [new FixedWindowPolicy(10, 1000), new FixedWindowPolicy(5, 5000)];
+		const composite = createFixedWindowCompositePolicy(policies);
+
+		it('should return correct status for empty state', () => {
+			const state = composite.getInitialState(1000);
+			const status = composite.getStatus(state, 500);
+
+			expect(status.map(inf => inf.windowStart)).toEqual([0, 0]);
+			expect(status.map(inf => inf.windowEnd)).toEqual([1000, 5000]);
+			expect(status.map(inf => inf.used)).toEqual([0, 0]);
+			expect(status.map(inf => inf.reserved)).toEqual([0, 0]);
+			expect(status.map(inf => inf.remaining)).toEqual([10, 5]);
+			expect(status.map(inf => inf.nextAvailableAt)).toEqual([0, 0]);
+			expect(status.map(inf => inf.resetAt)).toEqual([1000, 5000]);
+		});
+
+		it('should return correct timing when queue is partially filled', () => {
+			const state = [
+				{ windowStart: 0, used: 10, reserved: 5 },
+				{ windowStart: 0, used: 5, reserved: 3 },
+			];
+			const status = composite.getStatus(state, 500);
+
+			expect(status.map(i => i.nextAvailableAt)).toEqual([1000, 5000]);
+			expect(status.map(i => i.resetAt)).toEqual([2000, 10_000]);
+		});
+
+		it('should return correct timing when multiple windows are blocked', () => {
+			const state = [
+				{ windowStart: 0, used: 10, reserved: 25 },
+				{ windowStart: 0, used: 5, reserved: 13 },
+			];
+			const status = composite.getStatus(state, 500);
+
+			expect(status.map(i => i.nextAvailableAt)).toEqual([3000, 15_000]);
+			expect(status.map(i => i.resetAt)).toEqual([4000, 20_000]);
+		});
+	});
+
+	describe('Evaluation', () => {
+		describe('Allow', () => {
+			it('should allow and update states when all policies allow', () => {
+				const composite = createFixedWindowCompositePolicy([
+					new FixedWindowPolicy(10, 1000),
+					new FixedWindowPolicy(5, 5000),
+				]);
+				const states = composite.getInitialState(0);
+
+				const result = composite.evaluate(states, 500, 1);
+
+				expect(result.decision).toEqual({ kind: 'allow' });
+				expect(result.nextState[0].used).toBe(1);
+				expect(result.nextState[1].used).toBe(1);
+			});
+		});
+
+		describe('Deny and Revert', () => {
+			it('should deny and return maximum retryAt when all policies deny', () => {
+				const composite = createFixedWindowCompositePolicy([
+					new FixedWindowPolicy(1, 1000),
+					new FixedWindowPolicy(1, 5000),
+				]);
+				const states = [
+					{ windowStart: 0, used: 1, reserved: 0 },
+					{ windowStart: 0, used: 1, reserved: 0 },
+				];
+
+				const result = composite.evaluate(states, 500, 1);
+
+				expect(result.decision).toEqual({ kind: 'deny', retryAt: 5000 });
+			});
+
+			it('should deny, return maximum retryAt, and revert states for policies that allowed', () => {
+				const composite = createFixedWindowCompositePolicy([
+					new FixedWindowPolicy(10, 1000),
+					new FixedWindowPolicy(1, 5000),
+				]);
+				const states = [
+					{ windowStart: 0, used: 0, reserved: 0 },
+					{ windowStart: 0, used: 1, reserved: 0 },
+				];
+
+				const result = composite.evaluate(states, 500, 1);
+
+				expect(result.decision).toEqual({ kind: 'deny', retryAt: 5000 });
+				expect(result.nextState[0].used).toBe(0);
+				expect(result.nextState[1].used).toBe(1);
+			});
+		});
+
+		describe('Delay and Mixed Decisions', () => {
+			it('should delay and return maximum runAt when policies delay and none deny', () => {
+				const composite = createFixedWindowCompositePolicy([
+					new FixedWindowPolicy(1, 1000, 5),
+					new FixedWindowPolicy(1, 2000, 5),
+				]);
+				const states = [
+					{ windowStart: 0, used: 1, reserved: 0 },
+					{ windowStart: 0, used: 1, reserved: 0 },
+				];
+
+				const result = composite.evaluate(states, 500, 1, true);
+
+				expect(result.decision).toEqual({ kind: 'delay', runAt: 2000 });
+				expect(result.nextState[0].reserved).toBe(1);
+				expect(result.nextState[1].reserved).toBe(1);
+			});
+
+			it('should prioritize deny over delay and revert both delayed and allowed policies', () => {
+				const composite = createFixedWindowCompositePolicy([
+					new FixedWindowPolicy(10, 1000),
+					new FixedWindowPolicy(1, 2000, 5),
+					new FixedWindowPolicy(1, 5000, 0),
+				]);
+				const states = [
+					{ windowStart: 0, used: 0, reserved: 0 },
+					{ windowStart: 0, used: 1, reserved: 0 },
+					{ windowStart: 0, used: 1, reserved: 0 },
+				];
+
+				const result = composite.evaluate(states, 500, 1, true);
+
+				expect(result.decision).toEqual({ kind: 'deny', retryAt: 5000 });
+				expect(result.nextState[0].used).toBe(0);
+				expect(result.nextState[1].reserved).toBe(0);
+				expect(result.nextState[2].used).toBe(1);
+			});
+		});
+
+		describe('Edge Cases', () => {
+			it('should default cost to 1 if omitted', () => {
+				const composite = createFixedWindowCompositePolicy([new FixedWindowPolicy(10, 1000)]);
+				const states = composite.getInitialState(0);
+
+				const result = composite.evaluate(states, 500);
+
+				expect(result.nextState[0].used).toBe(1);
+			});
+
+			it('should throw when evaluate cost validation fails', () => {
+				const composite = createFixedWindowCompositePolicy([new FixedWindowPolicy(10, 1000)]);
+				const states = composite.getInitialState(0);
+
+				expect(() => composite.evaluate(states, 500, -5)).toThrow(InvalidCostError);
+				expect(() => composite.evaluate(states, 500, 1.5)).toThrow(InvalidCostError);
+			});
+		});
+	});
+
+	describe('Revert', () => {
+		it('should correctly proxy revert calls to all underlying policies within the same window', () => {
+			const composite = createFixedWindowCompositePolicy([
+				new FixedWindowPolicy(10, 1000),
+				new FixedWindowPolicy(5, 5000),
+			]);
+			const states = [
+				{ windowStart: 0, used: 5, reserved: 0 },
+				{ windowStart: 0, used: 3, reserved: 0 },
+			];
+
+			const revertedStates = composite.revert(states, 2, 500);
+
+			expect(revertedStates[0]).toEqual({ windowStart: 0, used: 3, reserved: 0 });
+			expect(revertedStates[1]).toEqual({ windowStart: 0, used: 1, reserved: 0 });
+		});
+
+		it('should correctly revert single unit cost', () => {
+			const composite = createFixedWindowCompositePolicy([new FixedWindowPolicy(10, 1000)]);
+			const states = [{ windowStart: 0, used: 5, reserved: 0 }];
+
+			const revertedStates = composite.revert(states, 1, 500);
+
+			expect(revertedStates[0]).toEqual({ windowStart: 0, used: 4, reserved: 0 });
+		});
+
+		it('should sync state for each policy independently based on their window sizes before reverting', () => {
+			const composite = createFixedWindowCompositePolicy([
+				new FixedWindowPolicy(10, 1000),
+				new FixedWindowPolicy(5, 5000),
+			]);
+			const states = [
+				{ windowStart: 0, used: 10, reserved: 5 },
+				{ windowStart: 0, used: 5, reserved: 3 },
+			];
+
+			const revertedStates = composite.revert(states, 2, 1500);
+
+			expect(revertedStates[0]).toEqual({ windowStart: 1000, used: 3, reserved: 0 });
+			expect(revertedStates[1]).toEqual({ windowStart: 0, used: 5, reserved: 1 });
+		});
+
+		it('should handle revert when time jumps past multiple windows differently for each policy', () => {
+			const composite = createFixedWindowCompositePolicy([
+				new FixedWindowPolicy(10, 1000),
+				new FixedWindowPolicy(5, 5000),
+			]);
+			const states = [
+				{ windowStart: 0, used: 10, reserved: 5 },
+				{ windowStart: 0, used: 5, reserved: 5 },
+			];
+
+			const revertedStates = composite.revert(states, 2, 6000);
+
+			expect(revertedStates[0]).toEqual({ windowStart: 6000, used: 0, reserved: 0 });
+			expect(revertedStates[1]).toEqual({ windowStart: 5000, used: 3, reserved: 0 });
+		});
+	});
+});