More

screenshot/src/server.ts

@@ -1,4 +1,4 @@
-import express, { type Request } from 'express';
+import express, { type Request, type Response } from 'express';
 import { ScreenshotCache } from './cache.js';
 import { takeScreenshot, checkWebGL, closeBrowser, initialize } from './screenshot.js';
 import { buildScreenshotRequest, ValidationError } from './validation.js';
@@ -27,25 +27,63 @@ app.set('trust proxy', true);
 let activeScreenshots = 0;
 let lastRateLimitPrune = 0;
 const rateLimitBuckets = new Map<string, { count: number; resetAt: number }>();
+type ReleaseScreenshotSlot = () => void;
+type PendingScreenshotSlot = {
+  resolve: (release: ReleaseScreenshotSlot | null) => void;
+  cleanup: () => void;
+};
+const screenshotSlotQueue: PendingScreenshotSlot[] = [];
 
 function parsePositiveIntEnv(name: string, fallback: number): number {
   const value = Number.parseInt(process.env[name] || '', 10);
   return Number.isFinite(value) && value > 0 ? value : fallback;
 }
 
-function acquireScreenshotSlot(): (() => void) | null {
-  if (activeScreenshots >= SCREENSHOT_CONCURRENCY) {
-    return null;
-  }
+function grantScreenshotSlot(): ReleaseScreenshotSlot {
   activeScreenshots += 1;
   let released = false;
   return () => {
     if (released) return;
     released = true;
     activeScreenshots = Math.max(0, activeScreenshots - 1);
+    drainScreenshotSlotQueue();
   };
 }
 
+function drainScreenshotSlotQueue(): void {
+  while (activeScreenshots < SCREENSHOT_CONCURRENCY && screenshotSlotQueue.length > 0) {
+    const pending = screenshotSlotQueue.shift();
+    if (!pending) return;
+    pending.cleanup();
+    pending.resolve(grantScreenshotSlot());
+  }
+}
+
+function acquireScreenshotSlot(res: Response): Promise<ReleaseScreenshotSlot | null> {
+  if (activeScreenshots < SCREENSHOT_CONCURRENCY) {
+    return Promise.resolve(grantScreenshotSlot());
+  }
+
+  return new Promise((resolve) => {
+    let pending: PendingScreenshotSlot;
+    const onClose = () => {
+      if (res.writableEnded) return;
+      pending.cleanup();
+      const index = screenshotSlotQueue.indexOf(pending);
+      if (index !== -1) screenshotSlotQueue.splice(index, 1);
+      resolve(null);
+    };
+
+    pending = {
+      resolve,
+      cleanup: () => res.off('close', onClose),
+    };
+    res.on('close', onClose);
+    screenshotSlotQueue.push(pending);
+    console.log(`Queued screenshot request; queue length: ${screenshotSlotQueue.length}`);
+  });
+}
+
 function rateLimitKey(req: Request): string {
   const forwardedFor = req.get('x-forwarded-for')?.split(',')[0]?.trim();
   return forwardedFor || req.ip || req.socket.remoteAddress || 'unknown';
@@ -117,9 +155,8 @@ app.get('/screenshot', async (req, res) => {
       return;
     }
 
-    releaseSlot = acquireScreenshotSlot();
+    releaseSlot = await acquireScreenshotSlot(res);
     if (!releaseSlot) {
-      res.status(503).json({ error: 'Screenshot service busy' });
       return;
     }