projects/Dockers/puppeteer-api/main.py

from fastapi import FastAPI, HTTPException, Header, BackgroundTasks
from pyppeteer import launch
import os
import asyncio
import gc
import logging
import psutil
from typing import Optional
from urllib.parse import unquote

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Get API key from environment variable
API_KEY = os.getenv('API_KEY')
if not API_KEY:
    raise ValueError("API_KEY environment variable must be set")

# Global browser instance
browser = None

async def get_browser():
    """Get or create a browser instance"""
    global browser
    if browser is None or not hasattr(browser, 'process') or browser.process is None:
        logger.info("Launching new browser instance")
        browser = await launch(
            headless=True,
            executablePath='/usr/bin/google-chrome',
            args=[
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-gpu',
                '--disable-extensions',
                '--disable-sync',
                '--disable-translate',
                '--hide-scrollbars',
                '--mute-audio',
                '--disable-background-networking',
                '--disable-background-timer-throttling',
                '--disable-backgrounding-occluded-windows',
                '--disable-breakpad',
                '--disable-client-side-phishing-detection',
                '--disable-component-extensions-with-background-pages',
                '--disable-default-apps',
                '--disable-features=TranslateUI,BlinkGenPropertyTrees',
                '--disable-hang-monitor',
                '--disable-ipc-flooding-protection',
                '--disable-popup-blocking',
                '--disable-prompt-on-repost',
                '--disable-renderer-backgrounding',
                '--force-color-profile=srgb',
                '--metrics-recording-only',
                '--no-first-run',
                '--enable-automation',
                '--password-store=basic',
                '--use-mock-keychain',
                '--js-flags=--expose-gc',
                '--memory-pressure-off',
            ]
        )
    return browser

async def cleanup_resources(page=None, close_browser=False):
    """Clean up resources to prevent memory leaks"""
    global browser
    try:
        if page:
            await page.close()

        if close_browser and browser and hasattr(browser, 'process') and browser.process is not None:
            logger.info("Closing browser instance")
            await browser.close()
            browser = None

        # Force garbage collection
        gc.collect()

        # Log memory usage
        process = psutil.Process(os.getpid())
        memory_info = process.memory_info()
        logger.info(f"Memory usage: {memory_info.rss / 1024 / 1024:.2f} MB")
    except Exception as e:
        logger.error(f"Error during cleanup: {str(e)}")

@app.head("/")
async def health_check():
    return {"status": "ok"}

@app.get("/memory")
async def memory_status(x_api_key: Optional[str] = Header(None)):
    # Validate API key
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()

    return {
        "memory_usage_mb": memory_info.rss / 1024 / 1024,
        "browser_active": browser is not None and hasattr(browser, 'process') and browser.process is not None
    }

@app.post("/cleanup")
async def force_cleanup(x_api_key: Optional[str] = Header(None)):
    # Validate API key
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

    await cleanup_resources(close_browser=True)
    return {"status": "cleanup completed"}

@app.get("/")
async def visit_url(url: str, x_api_key: Optional[str] = Header(None), background_tasks: BackgroundTasks = None):
    # Validate API key
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

    page = None
    try:
        # Decode URL if it's encoded
        decoded_url = unquote(url)
        logger.info(f"Visiting URL: {decoded_url}")

        # Get browser instance
        browser_instance = await get_browser()
        if browser_instance is None:
            raise HTTPException(status_code=500, detail="Failed to initialize browser")

        # Create new page
        page = await browser_instance.newPage()

        # Set page timeout
        await page.setDefaultNavigationTimeout(30000)  # 30 seconds

        # Set viewport
        await page.setViewport({"width": 1280, "height": 800})

        # Navigate to URL and wait for network idle
        await page.goto(decoded_url, waitUntil='networkidle2')

        # Get page content
        content = await page.content()

        # Schedule cleanup in background
        if background_tasks:
            background_tasks.add_task(cleanup_resources, page, False)
        else:
            await cleanup_resources(page, False)

        return {"status": "success", "content": content}

    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        # Ensure cleanup happens even on error
        if page:
            try:
                await cleanup_resources(page, False)
            except Exception as cleanup_error:
                logger.error(f"Error during error cleanup: {str(cleanup_error)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.on_event("shutdown")
async def shutdown_event():
    """Clean up resources when shutting down"""
    logger.info("Application shutting down, cleaning up resources")
    await cleanup_resources(close_browser=True)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)