adjust puppeteer

Dockers/puppeteer-api/Dockerfile

@@ -44,6 +44,11 @@ RUN useradd -m puppeteer && chown -R puppeteer:puppeteer /app
 # Create the database directory and set permissions
 RUN mkdir -p /db && chmod 777 /db && chown -R puppeteer:puppeteer /app
 
+# Add system limits configuration
+RUN echo "* soft nofile 65535" >> /etc/security/limits.conf && \
+    echo "* hard nofile 65535" >> /etc/security/limits.conf && \
+    echo "session required pam_limits.so" >> /etc/pam.d/common-session
+
 USER puppeteer
 
 # Expose port

Dockers/puppeteer-api/main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, Header
+from fastapi import FastAPI, HTTPException, Header, Request
 from pyppeteer import launch
 import os
 import asyncio
@@ -11,6 +11,11 @@ from typing import Optional, Dict, List, Any
 from urllib.parse import unquote
 from apscheduler.schedulers.background import BackgroundScheduler
 from apscheduler.triggers.cron import CronTrigger
+from fastapi.middleware.base import BaseHTTPMiddleware
+
+# Add imports for browser pool
+from asyncio import Queue, Lock
+from contextlib import asynccontextmanager
 
 app = FastAPI()
 
@@ -28,6 +33,121 @@ CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 3 * * *')
 # Define custom user agent
 CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36'
 
+# Browser pool configuration
+MAX_BROWSERS = int(os.getenv('MAX_BROWSERS', '5'))  # Maximum number of browser instances
+browser_pool = Queue()
+browser_lock = Lock()
+
+# Rate limiting configuration
+RATE_LIMIT_MINUTE = int(os.getenv('RATE_LIMIT_MINUTE', '60'))  # requests per minute
+RATE_LIMIT_WINDOW = 60  # window size in seconds
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app):
+        super().__init__(app)
+        self.requests = {}
+        self.lock = asyncio.Lock()
+
+    async def dispatch(self, request: Request, call_next):
+        # Skip rate limiting for health check
+        if request.url.path == "/" and request.method == "HEAD":
+            return await call_next(request)
+
+        api_key = request.headers.get("x-api-key")
+        if not api_key:
+            raise HTTPException(status_code=401, detail="API key required")
+
+        async with self.lock:
+            now = time.time()
+            # Clean old requests
+            self.requests = {k: v for k, v in self.requests.items()
+                           if now - v[-1] < RATE_LIMIT_WINDOW}
+
+            # Get request times for this API key
+            requests = self.requests.get(api_key, [])
+            # Remove old requests outside the window
+            requests = [t for t in requests if now - t < RATE_LIMIT_WINDOW]
+
+            if len(requests) >= RATE_LIMIT_MINUTE:
+                oldest = requests[0]
+                wait_time = RATE_LIMIT_WINDOW - (now - oldest)
+                raise HTTPException(
+                    status_code=429,
+                    detail=f"Rate limit exceeded. Try again in {int(wait_time)} seconds"
+                )
+
+            requests.append(now)
+            self.requests[api_key] = requests
+
+        return await call_next(request)
+
+# Add rate limiting middleware
+app.add_middleware(RateLimitMiddleware)
+
+async def create_browser():
+    """Create a new browser instance"""
+    return await launch(
+        headless=True,
+        executablePath='/usr/bin/google-chrome',
+        args=['--no-sandbox', '--disable-setuid-sandbox'],
+        handleSIGINT=False,
+        handleSIGTERM=False,
+        handleSIGHUP=False
+    )
+
+@asynccontextmanager
+async def get_browser():
+    """Get a browser from the pool or create a new one if needed"""
+    browser = None
+    try:
+        # Try to get a browser from the pool
+        try:
+            browser = await browser_pool.get_nowait()
+        except asyncio.QueueEmpty:
+            # If pool is empty, create a new browser if under the limit
+            async with browser_lock:
+                if browser_pool.qsize() + 1 <= MAX_BROWSERS:
+                    browser = await create_browser()
+                else:
+                    # If at limit, wait for a browser to become available
+                    browser = await browser_pool.get()
+
+        yield browser
+    finally:
+        # Return browser to pool if it's still viable
+        if browser:
+            try:
+                # Quick check if browser is still usable
+                await browser.pages()
+                await browser_pool.put(browser)
+            except Exception:
+                # If browser is not usable, close it and create a new one
+                try:
+                    await browser.close()
+                except:
+                    pass
+                browser = await create_browser()
+                await browser_pool.put(browser)
+
+# Initialize browser pool
+@app.on_event("startup")
+async def init_browser_pool():
+    """Initialize the browser pool with some browsers"""
+    for _ in range(min(3, MAX_BROWSERS)):  # Start with 3 browsers or MAX_BROWSERS, whichever is smaller
+        browser = await create_browser()
+        await browser_pool.put(browser)
+    print(f"Browser pool initialized with {browser_pool.qsize()} browsers")
+
+@app.on_event("shutdown")
+async def cleanup_browser_pool():
+    """Clean up all browsers in the pool"""
+    while not browser_pool.empty():
+        try:
+            browser = await browser_pool.get_nowait()
+            await browser.close()
+        except:
+            pass
+
 # Initialize SQLite database
 def init_db():
     global DB_PATH
@@ -213,40 +333,29 @@ async def health_check():
 
 async def safe_browser_operation(url, operation_func):
     """Safely perform browser operations with proper cleanup"""
-    browser = None
-    try:
-        browser = await launch(
-            headless=True,
-            executablePath='/usr/bin/google-chrome',
-            args=['--no-sandbox', '--disable-setuid-sandbox'],
-            handleSIGINT=False,
-            handleSIGTERM=False,
-            handleSIGHUP=False
-        )
+    async with get_browser() as browser:
+        try:
+            # Create new page with timeout
+            page = await browser.newPage()
+            page.setDefaultNavigationTimeout(30000)
 
-        # Create new page with timeout
-        page = await browser.newPage()
-        page.setDefaultNavigationTimeout(30000)
+            # Set custom user agent
+            await page.setUserAgent(CUSTOM_USER_AGENT)
 
-        # Set custom user agent
-        await page.setUserAgent(CUSTOM_USER_AGENT)
+            # Call the operation function that uses the page
+            result = await operation_func(page)
 
-        # Call the operation function that uses the page
-        result = await operation_func(page)
+            # Explicitly close the page
+            await page.close()
 
-        # Explicitly close the page
-        await page.close()
-
-        return result
-
-    finally:
-        # Ensure browser is closed properly
-        if browser:
+            return result
+        except Exception as e:
+            # If there's an error, close the page and re-raise
             try:
-                await browser.close()
-            except Exception as e:
-                print(f"Error closing browser: {e}")
-                # We don't re-raise here to avoid masking the original error
+                await page.close()
+            except:
+                pass
+            raise e
 
 @app.get("/")
 async def visit_url(url: str, x_api_key: Optional[str] = Header(None)):