From 1dc12f813da1280e572a65eb9750234a6c54bad3 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 3 Feb 2026 21:46:59 +0000
Subject: [PATCH] Speed up screenshotting

---
 og-screenshot/src/screenshot.ts | 148 +++++++++++++++++++++++---------
 1 file changed, 109 insertions(+), 39 deletions(-)

diff --git a/og-screenshot/src/screenshot.ts b/og-screenshot/src/screenshot.ts
index d202816..f097aa5 100644
--- a/og-screenshot/src/screenshot.ts
+++ b/og-screenshot/src/screenshot.ts
@@ -1,60 +1,124 @@
-import { chromium, type Browser, type BrowserContext } from 'playwright';
+import { chromium, type Browser, type Page } from 'playwright';
+import { existsSync, readdirSync } from 'fs';
 
 const VIEWPORT = { width: 1200, height: 630 };
 const NAVIGATION_TIMEOUT = 15_000;
 const TILE_BUFFER_MS = 500;
-const MAX_CONCURRENT = 2;
+const POOL_SIZE = 3;
 
 let browser: Browser | null = null;
-let concurrency = 0;
-const queue: Array<{ resolve: () => void }> = [];
+const pagePool: Page[] = [];
+let warmingUp = false;
+
+/**
+ * Detect if a GPU is available for hardware-accelerated rendering.
+ * Checks for DRI devices on Linux which indicate GPU + driver availability.
+ */
+function detectGpu(): boolean {
+  try {
+    // Check for DRI render nodes (Linux)
+    if (existsSync('/dev/dri')) {
+      const devices = readdirSync('/dev/dri');
+      // Look for renderD* (render nodes) or card* (display cards)
+      const hasGpu = devices.some((d) => d.startsWith('renderD') || d.startsWith('card'));
+      if (hasGpu) {
+        console.log(`GPU detected: /dev/dri contains ${devices.join(', ')}`);
+        return true;
+      }
+    }
+  } catch {
+    // Ignore errors - fall back to software rendering
+  }
+  console.log('No GPU detected, using SwiftShader software rendering');
+  return false;
+}
+
+const hasGpu = detectGpu();
+
+function getBrowserArgs(): string[] {
+  const baseArgs = ['--no-sandbox', '--disable-dev-shm-usage'];
+
+  if (hasGpu) {
+    // Use hardware GPU acceleration
+    return [...baseArgs, '--enable-gpu', '--enable-webgl'];
+  } else {
+    // Fall back to SwiftShader software rendering
+    return [...baseArgs, '--use-gl=angle', '--use-angle=swiftshader'];
+  }
+}
 
 async function ensureBrowser(): Promise<Browser> {
   if (!browser || !browser.isConnected()) {
-    browser = await chromium.launch({
-      args: [
-        '--no-sandbox',
-        '--disable-dev-shm-usage',
-        '--use-gl=angle',
-        '--use-angle=swiftshader',
-      ],
-    });
+    const args = getBrowserArgs();
+    console.log(`Launching browser with args: ${args.join(' ')}`);
+    browser = await chromium.launch({ args });
   }
   return browser;
 }
 
-async function acquireSlot(): Promise<void> {
-  if (concurrency < MAX_CONCURRENT) {
-    concurrency++;
-    return;
-  }
-  return new Promise<void>((resolve) => {
-    queue.push({ resolve });
+async function createPage(): Promise<Page> {
+  const instance = await ensureBrowser();
+  const context = await instance.newContext({
+    viewport: VIEWPORT,
+    deviceScaleFactor: 1,
   });
+  return context.newPage();
 }
 
-function releaseSlot(): void {
-  concurrency--;
-  const next = queue.shift();
-  if (next) {
-    concurrency++;
-    next.resolve();
+async function warmPool(): Promise<void> {
+  if (warmingUp) return;
+  warmingUp = true;
+
+  try {
+    const instance = await ensureBrowser();
+    while (pagePool.length < POOL_SIZE && instance.isConnected()) {
+      const page = await createPage();
+      pagePool.push(page);
+    }
+  } finally {
+    warmingUp = false;
+  }
+}
+
+async function acquirePage(): Promise<Page> {
+  // Try to get a page from the pool
+  const page = pagePool.shift();
+  if (page && !page.isClosed()) {
+    // Refill pool in background
+    warmPool().catch(() => {});
+    return page;
+  }
+
+  // No pooled page available, create one
+  const newPage = await createPage();
+  // Start warming pool in background
+  warmPool().catch(() => {});
+  return newPage;
+}
+
+async function releasePage(page: Page): Promise<void> {
+  try {
+    if (page.isClosed()) return;
+
+    // Reset page state for reuse
+    await page.goto('about:blank', { timeout: 5000 }).catch(() => {});
+
+    if (!page.isClosed() && pagePool.length < POOL_SIZE) {
+      pagePool.push(page);
+    } else {
+      await page.context().close().catch(() => {});
+    }
+  } catch {
+    await page.context().close().catch(() => {});
   }
 }
 
 export async function takeScreenshot(url: string): Promise<Buffer> {
-  await acquireSlot();
-  let context: BrowserContext | null = null;
+  const page = await acquirePage();
 
   try {
-    const instance = await ensureBrowser();
-    context = await instance.newContext({
-      viewport: VIEWPORT,
-      deviceScaleFactor: 1,
-    });
-    const page = await context.newPage();
-
-    await page.goto(url, { waitUntil: 'networkidle', timeout: NAVIGATION_TIMEOUT });
+    // Use domcontentloaded instead of networkidle - let __og_ready handle readiness
+    await page.goto(url, { waitUntil: 'domcontentloaded', timeout: NAVIGATION_TIMEOUT });
 
     // Wait for the frontend to signal readiness
     try {
@@ -71,16 +135,22 @@ export async function takeScreenshot(url: string): Promise<Buffer> {
     const screenshot = await page.screenshot({ type: 'png' });
     return Buffer.from(screenshot);
   } finally {
-    if (context) {
-      await context.close().catch(() => {});
-    }
-    releaseSlot();
+    await releasePage(page);
   }
 }
 
 export async function closeBrowser(): Promise<void> {
+  // Close all pooled pages
+  for (const page of pagePool) {
+    await page.context().close().catch(() => {});
+  }
+  pagePool.length = 0;
+
   if (browser) {
     await browser.close().catch(() => {});
     browser = null;
   }
 }
+
+// Pre-warm the pool on module load
+warmPool().catch(() => {});