diff --git a/Dockers/puppeteer-api/Dockerfile b/Dockers/puppeteer-api/Dockerfile index 18c43c0..1503a96 100644 --- a/Dockers/puppeteer-api/Dockerfile +++ b/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 diff --git a/Dockers/puppeteer-api/main.py b/Dockers/puppeteer-api/main.py index 49403a0..a26d390 100644 --- a/Dockers/puppeteer-api/main.py +++ b/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)):