ffs

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-extensions',
-                '--disable-software-rasterizer',
+                '--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")
+    """Clean up resources when shutting down"""
-    global browser
+    logger.info("Application shutting down, cleaning up resources")
-    if browser:
+    await cleanup_resources(close_browser=True)
-        await browser.close()
+
-        browser = None
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)