diff --git a/Dockers/puppeteer-api/main.py b/Dockers/puppeteer-api/main.py index c7d251e..6c7b710 100644 --- a/Dockers/puppeteer-api/main.py +++ b/Dockers/puppeteer-api/main.py @@ -8,18 +8,22 @@ 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 browser.isConnected(): logger.info("Launching new browser instance") @@ -31,58 +35,136 @@ async def get_browser(): '--disable-setuid-sandbox', '--disable-dev-shm-usage', '--disable-gpu', - '--single-process', - '--disable-software-rasterizer', + '--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): +async def cleanup_resources(page=None, close_browser=False): + """Clean up resources to prevent memory leaks""" try: if page: - try: - await page.evaluate("window.gc()") # Trigger JS garbage collection - except: - pass # Ignore if GC is not exposed await page.close() + + if close_browser and browser and browser.isConnected(): + logger.info("Closing browser instance") + await browser.close() + global browser + browser = None + + # Force garbage collection gc.collect() + + # Log memory usage process = psutil.Process(os.getpid()) - logger.info(f"Memory usage: {process.memory_info().rss / 1024 / 1024:.2f} MB") + memory_info = process.memory_info() + logger.info(f"Memory usage: {memory_info.rss / 1024 / 1024:.2f} MB") except Exception as e: - logger.error(f"Cleanup error: {str(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 browser.isConnected() if browser else False + } + +@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() - context = await browser_instance.createIncognitoBrowserContext() - page = await context.newPage() + + # 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) + background_tasks.add_task(cleanup_resources, page, False) else: - await cleanup_resources(page) + await cleanup_resources(page, False) + return {"status": "success", "content": content} + except Exception as e: - logger.error(f"Error: {str(e)}") + logger.error(f"Error processing request: {str(e)}") + # Ensure cleanup happens even on error if page: - await cleanup_resources(page) + await cleanup_resources(page, False) raise HTTPException(status_code=500, detail=str(e)) @app.on_event("shutdown") async def shutdown_event(): - logger.info("Shutting down, cleaning up browser instance") - global browser - if browser: - await browser.close() - browser = None + """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)