From 801c1c642fe443680914380adce41e41d1cf5902 Mon Sep 17 00:00:00 2001
From: Will Hoskings <resynth1943@outlook.com>
Date: Wed, 11 Mar 2020 22:06:00 +0000
Subject: [PATCH] LOOK BUT DON'T TOUCH: Ratelimits

---
 module/asyncutil.ts               |  3 +++
 module/discord-request-manager.ts | 29 +++++++++++++--------
 module/ratelimiter.ts             | 42 +++++++++++++++++++++++++++++--
 module/routed-request-manager.ts  | 21 ++++++++++++++++
 4 files changed, 83 insertions(+), 12 deletions(-)
 create mode 100644 module/asyncutil.ts

diff --git a/module/asyncutil.ts b/module/asyncutil.ts
new file mode 100644
index 000000000..b0f3c64f5
--- /dev/null
+++ b/module/asyncutil.ts
@@ -0,0 +1,3 @@
+export function sleep (timeout: number) {
+    return new Promise(resolve => setTimeout(resolve, timeout));
+}
diff --git a/module/discord-request-manager.ts b/module/discord-request-manager.ts
index 61b9322ae..e6291fb21 100644
--- a/module/discord-request-manager.ts
+++ b/module/discord-request-manager.ts
@@ -1,19 +1,27 @@
 import Client from "../module/client.ts"
 import { RequestMethod } from "../types/fetch.ts"
+import { Ratelimiter } from './ratelimiter';
 
 // type RequestBody = string | Blob | ArrayBufferView | ArrayBuffer | FormData | URLSearchParams | null | undefined
 
 export default class DiscordRequestManager {
+  public ratelimiter = new Ratelimiter();
+
   constructor(public client: Client) {
     this.client = client
   }
 
   async get(url: string, body?: unknown) {
-    const response = await fetch(this.resolveURL(url), {
-      headers: this.getDiscordHeaders(),
-      body: body ? JSON.stringify(body) : undefined
-    })
-    return await response.json()
+    return this.runMethod(RequestMethod.Get, url, body);
+  }
+
+  protected async addBucket (headers: Headers) {
+    this.ratelimiter.addBucket(headers.get('X-RateLimit-Bucket')!, {
+      retryAfter: parseInt(headers.get('X-RateLimit-Retry-After')!),
+      limit: parseInt(headers.get('X-RateLimit-Limit')!),
+      remaining: parseInt(headers.get('X-RateLimit-Remaining')!),
+      reset: parseInt(headers.get('X-RateLimit-Reset')!)
+    });
   }
 
   async post(url: string, body?: unknown) {
@@ -32,16 +40,17 @@ export default class DiscordRequestManager {
     return this.runMethod(RequestMethod.Put, url, body)
   }
 
-  async runMethod(method: RequestMethod, url: string, body?: unknown) {
-    const response = await fetch(this.resolveURL(url), {
+  protected async baseCreateRequestForMethod (method: RequestMethod, url: string, body?: unknown) {
+    return fetch(this.resolveURL(url), {
       method,
       headers: this.getDiscordHeaders(),
       body: body ? JSON.stringify(body) : undefined
     })
+  }
 
-    const json = await response.json()
-
-    return json
+  async runMethod(method: RequestMethod, url: string, body?: unknown) {
+    const response = await this.baseCreateRequestForMethod(method, url, body);
+    return response.json();
   }
 
   // A hook for the RouteAwareRequestManager to override URLs.
diff --git a/module/ratelimiter.ts b/module/ratelimiter.ts
index c7eb410b3..debec21fb 100644
--- a/module/ratelimiter.ts
+++ b/module/ratelimiter.ts
@@ -1,3 +1,41 @@
+import { sleep } from './asyncutil';
+
+export interface Ratelimit {
+    retryAfter: number;
+    limit: number;
+    remaining: number;
+    reset: number;
+}
+
 export class Ratelimiter {
-    
-}
\ No newline at end of file
+    public buckets: Record<string, Ratelimit> = {};
+
+    awaitRatelimit (ratelimit: Ratelimit): Promise<unknown> {
+        if (ratelimit.remaining === 0) {
+            return sleep(ratelimit.retryAfter);
+        }
+
+        return Promise.resolve();
+    }
+
+    addBucket (bucket: string, ratelimit: Ratelimit) {
+        if (this.buckets[bucket]) {
+            return;
+        }
+
+        // Otherwise, add this ratelimit to the registry.
+        this.buckets[bucket] = ratelimit;
+    }
+
+    async awaitBucket (bucket: string) {
+        if (this.buckets[bucket]) {
+            // POSSIBLE MEMORY LEAK: Some buckets might never get cleaned up.
+            await this.awaitRatelimit(this.buckets[bucket]);
+
+            // IIRC, we avoid `delete` so v8 doesn't deoptimize this?
+            this.buckets[bucket] = undefined as any;
+        }
+
+        return Promise.resolve();
+    }
+}
diff --git a/module/routed-request-manager.ts b/module/routed-request-manager.ts
index da228f90c..969a3bc62 100644
--- a/module/routed-request-manager.ts
+++ b/module/routed-request-manager.ts
@@ -2,8 +2,13 @@ import DiscordRequestManager from "./discord-request-manager.ts"
 import Client from "./client.ts"
 import { resolveURLs } from "./url.ts"
 import { baseEndpoints } from "../constants/discord.ts"
+import { Ratelimit, Ratelimiter } from './ratelimiter';
+import { RequestMethod } from '../types/fetch';
 
 export class RouteAwareDiscordRequestManager extends DiscordRequestManager {
+  protected currentRatelimit?: Ratelimit;
+  public ratelimiter = new Ratelimiter();
+
   constructor(public client: Client, public routeName: string) {
     super(client)
   }
@@ -11,6 +16,22 @@ export class RouteAwareDiscordRequestManager extends DiscordRequestManager {
   protected resolveURL(url: string) {
     return resolveURLs(baseEndpoints.BASE_URL, this.routeName, url)
   }
+
+  async runMethod (method: RequestMethod, url: string, body?: unknown) {
+    if (this.currentRatelimit) {
+      await this.ratelimiter.awaitRatelimit(this.currentRatelimit);
+    }
+
+    const response = await this.baseCreateRequestForMethod(method, url, body);
+
+    // Capture the ratelimit from this request in our cute little store.
+
+    return response.json();
+  }
+
+  protected createRatelimitFromRequest (request: Request) {
+    
+  }
 }
 
 export class RoutedDiscordRequestManager {