Merge pull request #65 from replicate/rate-limit

Add rate limiting for API endpoints

middleware.ts

@@ -0,0 +1,36 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { Ratelimit } from '@upstash/ratelimit';
+import { kv } from '@vercel/kv';
+
+const ratelimit = new Ratelimit({
+  redis: kv,
+  // 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
+export const config = {
+  matcher: ['/api/predictions/:path+'],
+};
+
+export default async function middleware(request: NextRequest) {
+  if (!process.env.VERCEL_ENV || !process.env.KV_REST_API_URL || !process.env.KV_REST_API_URL) {
+    console.warn('Skipping ratelimiting middleware');
+    return NextResponse.next();
+  }
+
+  const ip = request.ip ?? '127.0.0.1';
+  const { success, limit, remaining, reset } = await ratelimit.limit(ip);
+  const headers = {
+    'X-Ratelimit-Hit': String(!success),
+    'X-Ratelimit-Limit': String(limit),
+    'X-Ratelimit-Remaining': String(remaining),
+    'X-Ratelimit-Reset': String(reset),
+  }
+
+  return success
+    ? NextResponse.next({headers})
+    : NextResponse.json({}, { status: 429, headers });
+}

package.json

@@ -14,7 +14,9 @@
     "@heroicons/react": "^2.0.18",
     "@prisma/client": "^4.14.0",
     "@supabase/supabase-js": "^2.24.0",
+    "@upstash/ratelimit": "^1.0.0",
     "@vercel/analytics": "^1.0.1",
+    "@vercel/kv": "^1.0.1",
     "@vercel/og": "^0.5.4",
     "@vercel/postgres": "^0.3.0",
     "eslint": "8.29.0",

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)
+}