From 3c00ccb9cd03f5a80be873c3c42b0988ffeb4729 Mon Sep 17 00:00:00 2001
From: Aron Carroll <aron@replicate.com>
Date: Fri, 8 Dec 2023 12:58:29 +0000
Subject: [PATCH] Improve requests to polling endpoint

---
 middleware.ts  |  1 +
 pages/index.js | 33 +++++++++++++++++++++++++++++----
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/middleware.ts b/middleware.ts
index 5182850..55019e4 100644
--- a/middleware.ts
+++ b/middleware.ts
@@ -7,6 +7,7 @@ const ratelimit = new Ratelimit({
   // 20 requests from the same IP within a 10 second sliding window
   limiter: Ratelimit.slidingWindow(20, '10s'),
   prefix: `v2/zoo/ratelimit/${process.env.VERCEL_ENV ?? 'local'}`,
+  timeout: 500,
 });
 
 // Rate limit the /api/predictions/[id] endpoint
diff --git a/pages/index.js b/pages/index.js
index e44bbfa..b1dcf1d 100644
--- a/pages/index.js
+++ b/pages/index.js
@@ -36,7 +36,9 @@ export default function Home({ baseUrl, submissionPredictions }) {
     const response = await fetch(`/api/submissions/${seed}`, {
       method: "GET",
     });
-    submissionPredictions = await response.json();
+    if (response.ok) {
+      submissionPredictions = await response.json();
+    }
     setPredictions(submissionPredictions);
 
     // get the model names from the predictions, and update which ones are checked
@@ -151,17 +153,36 @@ export default function Home({ baseUrl, submissionPredictions }) {
       throw new Error(prediction.detail);
     }
 
+    // Add incremental backoff for polling requests.
+    const backoff = [250, 500, 500, 750, 1000, 1500, 3000, 5000, 10000, 15000, 30000];
+
     while (
       prediction.status !== "succeeded" &&
       prediction.status !== "failed"
     ) {
-      await sleep(500);
+      const jitter = random(0, 100); // Don't make all requests at the same time.
+      const delay = backoff.shift();
+      if (!delay) {
+        // We've exceeded our timeout.
+        // TODO: Better user facing messaging here.
+        break;
+      }
+
+      await sleep(delay + jitter);
       const response = await fetch("/api/predictions/" + prediction.id);
-      prediction = await response.json();
-      console.log(prediction);
+
+      // Handle Rate Limiting
+      if (response.status === 429) {
+        const reset = response.headers.get('X-Ratelimit-Reset') ?? Date.now() + 10_000;        
+        const wait = reset - Date.now();
+        await sleep(wait)
+        continue;
+      }
+
       if (response.status !== 200) {
         throw new Error(prediction.detail);
       }
+      prediction = await response.json();
     }
 
     prediction.model = model.name;
@@ -551,3 +572,7 @@ export async function getServerSideProps({ req }) {
 
   return { props: { baseUrl, submissionPredictions } };
 }
+
+function random(min, max) {
+  return Math.floor(Math.random() * (max - min + 1) + min)
+}