From 12de164a4fac6ca6ad03e4ce2e28663ccc669461 Mon Sep 17 00:00:00 2001
From: Alex Gleason <alex@alexgleason.me>
Date: Sat, 25 Jan 2025 13:36:49 -0600
Subject: [PATCH] Add a custom RateLimiter implementation

---
 src/utils/ratelimiter/MemoryRateLimiter.ts | 77 ++++++++++++++++++++++
 src/utils/ratelimiter/RateLimitError.ts    | 10 +++
 src/utils/ratelimiter/types.ts             | 12 ++++
 3 files changed, 99 insertions(+)
 create mode 100644 src/utils/ratelimiter/MemoryRateLimiter.ts
 create mode 100644 src/utils/ratelimiter/RateLimitError.ts
 create mode 100644 src/utils/ratelimiter/types.ts

diff --git a/src/utils/ratelimiter/MemoryRateLimiter.ts b/src/utils/ratelimiter/MemoryRateLimiter.ts
new file mode 100644
index 00000000..b3f14d81
--- /dev/null
+++ b/src/utils/ratelimiter/MemoryRateLimiter.ts
@@ -0,0 +1,77 @@
+import { RateLimitError } from './RateLimitError.ts';
+import { RateLimiter, RateLimiterClient } from './types.ts';
+
+interface MemoryRateLimiterOpts {
+  limit: number;
+  window: number;
+}
+
+export class MemoryRateLimiter implements RateLimiter {
+  private iid: number;
+
+  private previous = new Map<string, RateLimiterClient>();
+  private current = new Map<string, RateLimiterClient>();
+
+  constructor(private opts: MemoryRateLimiterOpts) {
+    this.iid = setInterval(() => {
+      this.previous = this.current;
+      this.current = new Map();
+    }, opts.window);
+  }
+
+  get limit(): number {
+    return this.opts.limit;
+  }
+
+  get window(): number {
+    return this.opts.window;
+  }
+
+  client(key: string): RateLimiterClient {
+    const curr = this.current.get(key);
+    const prev = this.previous.get(key);
+
+    if (curr) {
+      return curr;
+    }
+
+    if (prev) {
+      this.current.set(key, prev);
+      this.previous.delete(key);
+      return prev;
+    }
+
+    const next = new MemoryRateLimiterClient(this);
+    this.current.set(key, next);
+    return next;
+  }
+
+  [Symbol.dispose](): void {
+    clearInterval(this.iid);
+  }
+}
+
+class MemoryRateLimiterClient implements RateLimiterClient {
+  private _hits: number = 0;
+  readonly resetAt: Date;
+
+  constructor(private limiter: MemoryRateLimiter) {
+    this.resetAt = new Date(Date.now() + limiter.window);
+  }
+
+  get hits(): number {
+    return this._hits;
+  }
+
+  get remaining(): number {
+    return this.limiter.limit - this.hits;
+  }
+
+  hit(n: number = 1): void {
+    this._hits += n;
+
+    if (this.remaining < 0) {
+      throw new RateLimitError(this.limiter, this);
+    }
+  }
+}
diff --git a/src/utils/ratelimiter/RateLimitError.ts b/src/utils/ratelimiter/RateLimitError.ts
new file mode 100644
index 00000000..ce21af72
--- /dev/null
+++ b/src/utils/ratelimiter/RateLimitError.ts
@@ -0,0 +1,10 @@
+import { RateLimiter, RateLimiterClient } from './types.ts';
+
+export class RateLimitError extends Error {
+  constructor(
+    readonly limiter: RateLimiter,
+    readonly client: RateLimiterClient,
+  ) {
+    super('Rate limit exceeded');
+  }
+}
diff --git a/src/utils/ratelimiter/types.ts b/src/utils/ratelimiter/types.ts
new file mode 100644
index 00000000..c1a6b2f0
--- /dev/null
+++ b/src/utils/ratelimiter/types.ts
@@ -0,0 +1,12 @@
+export interface RateLimiter extends Disposable {
+  readonly limit: number;
+  readonly window: number;
+  client(key: string): RateLimiterClient;
+}
+
+export interface RateLimiterClient {
+  readonly hits: number;
+  readonly resetAt: Date;
+  readonly remaining: number;
+  hit(n?: number): void;
+}