try to avoid memory leak

Dockers/puppeteer-api/main.py

@@ -35,8 +35,10 @@ CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.3
 
 # Browser pool configuration
 MAX_BROWSERS = int(os.getenv('MAX_BROWSERS', '5'))  # Maximum number of browser instances
+BROWSER_TTL = int(os.getenv('BROWSER_TTL', '3600'))  # Time to live for browser instances in seconds
 browser_pool = Queue()
 browser_lock = Lock()
+browser_creation_times = {}
 
 # Rate limiting configuration
 RATE_LIMIT_MINUTE = int(os.getenv('RATE_LIMIT_MINUTE', '60'))  # requests per minute
@@ -86,48 +88,80 @@ app.add_middleware(RateLimitMiddleware)
 
 async def create_browser():
     """Create a new browser instance"""
-    return await launch(
+    browser = await launch(
         headless=True,
         executablePath='/usr/bin/google-chrome',
-        args=['--no-sandbox', '--disable-setuid-sandbox'],
+        args=[
+            '--no-sandbox',
+            '--disable-setuid-sandbox',
+            '--disable-dev-shm-usage',
+            '--disable-accelerated-2d-canvas',
+            '--disable-gpu',
+            '--disable-extensions',
+            '--disable-sync',
+            '--disable-background-networking',
+            '--disable-default-apps',
+            '--disable-translate',
+            '--disable-background-timer-throttling',
+            '--disable-backgrounding-occluded-windows',
+            '--disable-client-side-phishing-detection',
+            '--disable-features=site-per-process',
+            '--disable-hang-monitor',
+            '--disable-ipc-flooding-protection',
+            '--disable-popup-blocking',
+            '--disable-prompt-on-repost',
+            '--disable-renderer-backgrounding',
+            '--memory-pressure-off',
+            '--no-first-run',
+            '--safebrowsing-disable-auto-update',
+        ],
         handleSIGINT=False,
         handleSIGTERM=False,
-        handleSIGHUP=False
+        handleSIGHUP=False,
+        ignoreHTTPSErrors=True
     )
+    browser_creation_times[browser] = time.time()
+    return browser
 
-@asynccontextmanager
-async def get_browser():
-    """Get a browser from the pool or create a new one if needed"""
-    browser = None
+async def check_browser_health():
+    """Check browser health and recycle if needed"""
+    while True:
         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
+            # Sleep for 5 minutes between checks
+            await asyncio.sleep(300)
+
             async with browser_lock:
-                if browser_pool.qsize() + 1 <= MAX_BROWSERS:
+                # Get all browsers from the pool
+                browsers = []
+                while not browser_pool.empty():
+                    browsers.append(await browser_pool.get())
+
+                # Check each browser
+                for browser in browsers:
+                    try:
+                        # Check if browser is too old
+                        if time.time() - browser_creation_times.get(browser, 0) > BROWSER_TTL:
+                            await browser.close()
+                            del browser_creation_times[browser]
                             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
+                            # Quick health check
                             await browser.pages()
+
+                        # Put back in pool if healthy
                         await browser_pool.put(browser)
                     except Exception:
-                # If browser is not usable, close it and create a new one
+                        # If unhealthy, close and create new
                         try:
                             await browser.close()
                         except:
                             pass
-                browser = await create_browser()
-                await browser_pool.put(browser)
+                        if browser in browser_creation_times:
+                            del browser_creation_times[browser]
+                        new_browser = await create_browser()
+                        await browser_pool.put(new_browser)
+        except Exception as e:
+            print(f"Error in browser health check: {str(e)}")
 
 # Initialize browser pool
 @app.on_event("startup")
@@ -138,6 +172,9 @@ async def init_browser_pool():
         await browser_pool.put(browser)
     print(f"Browser pool initialized with {browser_pool.qsize()} browsers")
 
+    # Start browser health check task
+    asyncio.create_task(check_browser_health())
+
 @app.on_event("shutdown")
 async def cleanup_browser_pool():
     """Clean up all browsers in the pool"""
@@ -145,6 +182,8 @@ async def cleanup_browser_pool():
         try:
             browser = await browser_pool.get_nowait()
             await browser.close()
+            if browser in browser_creation_times:
+                del browser_creation_times[browser]
         except:
             pass
 
@@ -332,30 +371,47 @@ async def health_check():
     return {"status": "ok"}
 
 async def safe_browser_operation(url, operation_func):
-    """Safely perform browser operations with proper cleanup"""
+    """Safely perform a browser operation with proper cleanup"""
     async with get_browser() as browser:
         try:
-            # Create new page with timeout
+            # Create a new page
             page = await browser.newPage()
-            page.setDefaultNavigationTimeout(30000)
 
-            # Set custom user agent
+            # Set reasonable viewport
+            await page.setViewport({'width': 1280, 'height': 800})
+
+            # Set user agent
             await page.setUserAgent(CUSTOM_USER_AGENT)
 
-            # Call the operation function that uses the page
-            result = await operation_func(page)
+            # Set reasonable timeout
+            page.setDefaultNavigationTimeout(30000)
 
-            # Explicitly close the page
-            await page.close()
+            # Enable request interception to block unnecessary resources
+            await page.setRequestInterception(True)
+
+            async def intercept(request):
+                # Block unnecessary resource types
+                if request.resourceType in ['image', 'media', 'font', 'stylesheet']:
+                    await request.abort()
+                else:
+                    await request.continue_()
+
+            page.on('request', lambda req: asyncio.ensure_future(intercept(req)))
+
+            # Perform the operation
+            result = await operation_func(page)
 
             return result
         except Exception as e:
-            # If there's an error, close the page and re-raise
+            print(f"Error during browser operation: {str(e)}")
+            raise
+        finally:
             try:
+                # Ensure page is properly closed
+                if 'page' in locals():
                     await page.close()
-            except:
-                pass
-            raise e
+            except Exception as e:
+                print(f"Error closing page: {str(e)}")
 
 @app.get("/")
 async def visit_url(url: str, x_api_key: Optional[str] = Header(None)):
@@ -1071,6 +1127,48 @@ async def detect_pagination(url: str, x_api_key: Optional[str] = Header(None)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@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()
+                # Check if browser is too old
+                if time.time() - browser_creation_times.get(browser, 0) > BROWSER_TTL:
+                    await browser.close()
+                    if browser in browser_creation_times:
+                        del browser_creation_times[browser]
+                    browser = await create_browser()
+                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
+                if browser in browser_creation_times:
+                    del browser_creation_times[browser]
+                browser = await create_browser()
+                await browser_pool.put(browser)
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)