cleanup

Dockers/puppeteer-api/main.py

@@ -14,10 +14,12 @@ from urllib.parse import unquote
 from apscheduler.schedulers.background import BackgroundScheduler
 from apscheduler.triggers.cron import CronTrigger
 from fastapi.middleware.base import BaseHTTPMiddleware
+from app.config import BROWSER_INSTANCE_TIMEOUT_MINUTES
 
 # Add imports for browser pool
 from asyncio import Queue, Lock, Semaphore
 from contextlib import asynccontextmanager
+from app.utils.browser_utils import force_cleanup_old_pages
 
 app = FastAPI()
 
@@ -44,7 +46,9 @@ MAX_CONCURRENT_OPERATIONS = int(os.getenv('MAX_CONCURRENT_OPERATIONS', '5'))  #
 browser_pool = Queue(maxsize=MAX_BROWSERS)  # Add maxsize to prevent unbounded growth
 browser_lock = Lock()
 browser_creation_times = {}
+page_creation_times = {}  # Track page creation times for force cleanup
 active_browsers = set()  # Track active browsers
+active_pages = set()  # Track active pages for force cleanup
 operation_semaphore = Semaphore(MAX_CONCURRENT_OPERATIONS)  # Limit concurrent operations
 playwright_instance = None  # Global playwright instance
 
@@ -152,10 +156,14 @@ async def cleanup_browser(browser):
         if browser in browser_creation_times:
             del browser_creation_times[browser]
 
-        # Close all pages first
+        # Close all pages first and clean up page tracking
         pages = browser.contexts[0].pages if browser.contexts else []
         for page in pages:
             try:
+                if page in active_pages:
+                    active_pages.remove(page)
+                if page in page_creation_times:
+                    del page_creation_times[page]
                 await page.close()
             except Exception as e:
                 print(f"Error closing page: {e}")
@@ -173,6 +181,9 @@ async def check_browser_health():
             # Sleep for 2 minutes between checks (reduced from 5 minutes)
             await asyncio.sleep(120)
 
+            # First, force cleanup old instances
+            await force_cleanup_old_instances()
+
             async with browser_lock:
                 # Get all browsers from the pool
                 browsers = []
@@ -223,6 +234,58 @@ async def check_browser_health():
             print(f"Error in browser health check: {e}")
             await asyncio.sleep(60)  # Wait before retrying
 
+async def force_cleanup_old_instances():
+    """Force cleanup old browser and page instances based on timeout"""
+    print(f"Checking for old browser/page instances (timeout: {BROWSER_INSTANCE_TIMEOUT_MINUTES} minutes)...")
+
+    current_time = time.time()
+    timeout_seconds = BROWSER_INSTANCE_TIMEOUT_MINUTES * 60
+    cleaned_browsers = 0
+    cleaned_pages = 0
+
+    async with browser_lock:
+        # Clean up old browsers
+        browsers_to_cleanup = []
+        for browser in list(active_browsers):
+            creation_time = browser_creation_times.get(browser, 0)
+            if current_time - creation_time > timeout_seconds:
+                browsers_to_cleanup.append(browser)
+                print(f"Marking browser for cleanup (age: {(current_time - creation_time)/60:.1f} minutes)")
+
+        for browser in browsers_to_cleanup:
+            try:
+                await cleanup_browser(browser)
+                cleaned_browsers += 1
+            except Exception as e:
+                print(f"Error cleaning up old browser: {e}")
+
+        # Clean up old pages (this is a fallback for pages that might not be properly tracked)
+        for browser in list(active_browsers):
+            try:
+                if browser.contexts:
+                    for context in browser.contexts:
+                        for page in context.pages:
+                            if page in page_creation_times:
+                                creation_time = page_creation_times[page]
+                                if current_time - creation_time > timeout_seconds:
+                                    try:
+                                        if page in active_pages:
+                                            active_pages.remove(page)
+                                        if page in page_creation_times:
+                                            del page_creation_times[page]
+                                        await page.close()
+                                        cleaned_pages += 1
+                                        print(f"Force closed old page (age: {(current_time - creation_time)/60:.1f} minutes)")
+                                    except Exception as e:
+                                        print(f"Error closing old page: {e}")
+            except Exception as e:
+                print(f"Error checking pages in browser: {e}")
+
+    # Also cleanup pages from browser_utils module
+    await force_cleanup_old_pages()
+
+    print(f"Force cleanup completed: {cleaned_browsers} browsers, {cleaned_pages} pages cleaned")
+
 async def force_cleanup_all_browsers():
     """Force cleanup all browsers - useful for emergency situations"""
     print("Force cleaning up all browsers...")
@@ -439,6 +502,10 @@ async def safe_browser_operation(url, operation_func):
 
             page = await context.new_page()
 
+            # Track page creation time for force cleanup
+            page_creation_times[page] = time.time()
+            active_pages.add(page)
+
             # Set up request interception for better performance
             await page.route("**/*", lambda route: route.abort()
                            if route.request.resource_type in ['image', 'stylesheet', 'font', 'media']
@@ -455,6 +522,11 @@ async def safe_browser_operation(url, operation_func):
             # Cleanup
             if page:
                 try:
+                    # Remove from tracking
+                    if page in active_pages:
+                        active_pages.remove(page)
+                    if page in page_creation_times:
+                        del page_creation_times[page]
                     await page.close()
                 except:
                     pass
@@ -757,6 +829,7 @@ async def system_status(x_api_key: Optional[str] = Header(None)):
         # Get browser pool information
         pool_size = browser_pool.qsize()
         active_browser_count = len(active_browsers)
+        active_page_count = len(active_pages)
 
         # Get cache statistics
         conn = sqlite3.connect('/db/cache.db')
@@ -778,7 +851,9 @@ async def system_status(x_api_key: Optional[str] = Header(None)):
                 "pool_size": pool_size,
                 "max_browsers": MAX_BROWSERS,
                 "active_browsers": active_browser_count,
-                "browser_ttl_seconds": BROWSER_TTL
+                "active_pages": active_page_count,
+                "browser_ttl_seconds": BROWSER_TTL,
+                "instance_timeout_minutes": BROWSER_INSTANCE_TIMEOUT_MINUTES
             },
             "cache": {
                 "total_entries": cache_count,
@@ -805,6 +880,22 @@ async def emergency_cleanup(x_api_key: Optional[str] = Header(None)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/force-cleanup-old")
+async def force_cleanup_old(x_api_key: Optional[str] = Header(None)):
+    """Force cleanup old browser and page instances based on timeout"""
+    # Validate API key
+    if not x_api_key or x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    try:
+        await force_cleanup_old_instances()
+        return {
+            "status": "success",
+            "message": f"Force cleanup of old instances completed (timeout: {BROWSER_INSTANCE_TIMEOUT_MINUTES} minutes)"
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @asynccontextmanager
 async def get_browser():
     """Context manager for getting a browser from the pool"""