from fastapi import FastAPI, HTTPException, Header, Request from pyppeteer import launch import os import asyncio import json import re import sqlite3 import time import signal import psutil from datetime import datetime, timedelta from typing import Optional, Dict, List, Any from urllib.parse import unquote from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.triggers.cron import CronTrigger from fastapi.middleware.base import BaseHTTPMiddleware # Add imports for browser pool from asyncio import Queue, Lock, Semaphore from contextlib import asynccontextmanager 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") # Get cache expiry time from environment variable (default: 36 hours) CACHE_EXPIRY_HOURS = int(os.getenv('CACHE_EXPIRY_HOURS', '36')) # Get cleanup cron schedule from environment variable (default: every day at 3 AM) CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 3 * * *') # Define custom user agent CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36' # Browser pool configuration - Reduced for better resource management MAX_BROWSERS = int(os.getenv('MAX_BROWSERS', '3')) # Reduced from 5 to 3 BROWSER_TTL = int(os.getenv('BROWSER_TTL', '1800')) # Reduced from 3600 to 1800 seconds (30 minutes) MAX_CONCURRENT_OPERATIONS = int(os.getenv('MAX_CONCURRENT_OPERATIONS', '5')) # New: limit concurrent operations # Browser pool management browser_pool = Queue(maxsize=MAX_BROWSERS) # Add maxsize to prevent unbounded growth browser_lock = Lock() browser_creation_times = {} active_browsers = set() # Track active browsers operation_semaphore = Semaphore(MAX_CONCURRENT_OPERATIONS) # Limit concurrent operations # Rate limiting configuration RATE_LIMIT_MINUTE = int(os.getenv('RATE_LIMIT_MINUTE', '60')) # requests per minute RATE_LIMIT_WINDOW = 60 # window size in seconds class RateLimitMiddleware(BaseHTTPMiddleware): def __init__(self, app): super().__init__(app) self.requests = {} self.lock = asyncio.Lock() async def dispatch(self, request: Request, call_next): # Skip rate limiting for health check if request.url.path == "/" and request.method == "HEAD": return await call_next(request) api_key = request.headers.get("x-api-key") if not api_key: raise HTTPException(status_code=401, detail="API key required") async with self.lock: now = time.time() # Clean old requests self.requests = {k: v for k, v in self.requests.items() if now - v[-1] < RATE_LIMIT_WINDOW} # Get request times for this API key requests = self.requests.get(api_key, []) # Remove old requests outside the window requests = [t for t in requests if now - t < RATE_LIMIT_WINDOW] if len(requests) >= RATE_LIMIT_MINUTE: oldest = requests[0] wait_time = RATE_LIMIT_WINDOW - (now - oldest) raise HTTPException( status_code=429, detail=f"Rate limit exceeded. Try again in {int(wait_time)} seconds" ) requests.append(now) self.requests[api_key] = requests return await call_next(request) # Add rate limiting middleware app.add_middleware(RateLimitMiddleware) async def create_browser(): """Create a new browser instance with improved resource management""" try: browser = await launch( headless=True, executablePath='/usr/bin/google-chrome', 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', '--max_old_space_size=512', # Limit memory usage '--single-process', # Use single process to reduce resource usage '--disable-web-security', '--disable-features=VizDisplayCompositor', ], handleSIGINT=False, handleSIGTERM=False, handleSIGHUP=False, ignoreHTTPSErrors=True, autoClose=True, # Ensure browser closes automatically ) browser_creation_times[browser] = time.time() active_browsers.add(browser) print(f"Created new browser instance. Total active browsers: {len(active_browsers)}") return browser except Exception as e: print(f"Error creating browser: {e}") raise async def cleanup_browser(browser): """Safely cleanup a browser instance""" try: if browser in active_browsers: active_browsers.remove(browser) if browser in browser_creation_times: del browser_creation_times[browser] # Close all pages first pages = await browser.pages() for page in pages: try: await page.close() except Exception as e: print(f"Error closing page: {e}") # Close browser await browser.close() print(f"Browser cleaned up. Total active browsers: {len(active_browsers)}") except Exception as e: print(f"Error during browser cleanup: {e}") async def check_browser_health(): """Check browser health and recycle if needed""" while True: try: # Sleep for 2 minutes between checks (reduced from 5 minutes) await asyncio.sleep(120) async with browser_lock: # Get all browsers from the pool browsers = [] while not browser_pool.empty(): try: browsers.append(await browser_pool.get_nowait()) except asyncio.QueueEmpty: break # 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: print(f"Recycling old browser (age: {time.time() - browser_creation_times.get(browser, 0):.0f}s)") await cleanup_browser(browser) browser = await create_browser() else: # Quick health check await browser.pages() # Put back in pool if healthy if not browser_pool.full(): await browser_pool.put(browser) else: # Pool is full, cleanup this browser await cleanup_browser(browser) except Exception as e: print(f"Browser health check failed: {e}") # If unhealthy, close and create new await cleanup_browser(browser) if not browser_pool.full(): new_browser = await create_browser() await browser_pool.put(new_browser) except Exception as e: print(f"Error in browser health check: {str(e)}") async def force_cleanup_all_browsers(): """Force cleanup all browsers in emergency situations""" print("Force cleaning up all browsers...") # Clean up pool while not browser_pool.empty(): try: browser = await browser_pool.get_nowait() await cleanup_browser(browser) except asyncio.QueueEmpty: break # Clean up active browsers for browser in list(active_browsers): await cleanup_browser(browser) # Initialize browser pool @app.on_event("startup") async def init_browser_pool(): """Initialize the browser pool with some browsers""" try: for _ in range(min(2, MAX_BROWSERS)): # Start with 2 browsers instead of 3 browser = await create_browser() 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()) except Exception as e: print(f"Error initializing browser pool: {e}") @app.on_event("shutdown") async def cleanup_browser_pool(): """Clean up all browsers in the pool""" await force_cleanup_all_browsers() # Signal handlers for graceful shutdown def signal_handler(signum, frame): print(f"Received signal {signum}, shutting down gracefully...") asyncio.create_task(force_cleanup_all_browsers()) signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) # Initialize SQLite database def init_db(): global DB_PATH # Try to use the mounted volume first db_path = '/db/cache.db' db_dir = os.path.dirname(db_path) # Check if directory exists and is writable dir_writable = False if os.path.exists(db_dir): try: test_file = os.path.join(db_dir, '.write_test') with open(test_file, 'w') as f: f.write('test') os.remove(test_file) dir_writable = True except (IOError, PermissionError): print(f"Directory {db_dir} exists but is not writable") dir_writable = False # If directory doesn't exist or isn't writable, try to create it if not os.path.exists(db_dir) or not dir_writable: try: os.makedirs(db_dir, exist_ok=True) # Test if we can write to the directory test_file = os.path.join(db_dir, '.write_test') with open(test_file, 'w') as f: f.write('test') os.remove(test_file) print(f"Created directory: {db_dir}") dir_writable = True except Exception as e: print(f"Warning: Could not create or write to directory {db_dir}: {e}") # Fallback to using a local database file db_path = 'cache.db' print(f"Using local database file: {db_path}") try: conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS cache ( url TEXT, route TEXT, data TEXT, timestamp INTEGER, PRIMARY KEY (url, route) ) ''') conn.commit() conn.close() print(f"Database initialized at {db_path}") # Update the global DB_PATH DB_PATH = db_path except sqlite3.OperationalError as e: print(f"Error initializing database at {db_path}: {e}") # Fallback to using a local database file if the mounted volume has permission issues db_path = 'cache.db' print(f"Falling back to local database file: {db_path}") try: conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS cache ( url TEXT, route TEXT, data TEXT, timestamp INTEGER, PRIMARY KEY (url, route) ) ''') conn.commit() conn.close() print(f"Local database initialized at {db_path}") # Update the global DB_PATH DB_PATH = db_path except sqlite3.OperationalError as e2: print(f"Error initializing local database: {e2}") raise # Define the database path DB_PATH = '/db/cache.db' # Get cached data if it exists and is not older than the expiry time def get_cached_data(url, route): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60) # Convert hours to seconds cursor.execute( "SELECT data FROM cache WHERE url = ? AND route = ? AND timestamp > ?", (url, route, cache_expiry) ) result = cursor.fetchone() conn.close() if result: print(f"Cache hit for {url} on route {route}") return json.loads(result[0]) return None # Save data to cache def save_to_cache(url, route, data): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() timestamp = int(time.time()) # Convert data to JSON string data_json = json.dumps(data) cursor.execute( "INSERT OR REPLACE INTO cache (url, route, data, timestamp) VALUES (?, ?, ?, ?)", (url, route, data_json, timestamp) ) conn.commit() conn.close() print(f"Saved to cache: {url} on route {route}") # Function to clean up old cache entries def cleanup_old_cache_entries(): try: print(f"Running scheduled cache cleanup (entries older than {CACHE_EXPIRY_HOURS} hours, pagination: {CACHE_EXPIRY_HOURS * 31} hours)") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Calculate the timestamp for entries older than the expiry time expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60) pagination_expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60) # Get count of entries to be deleted (non-pagination) cursor.execute("SELECT COUNT(*) FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,)) count_non_pagination = cursor.fetchone()[0] # Get count of pagination entries to be deleted cursor.execute("SELECT COUNT(*) FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,)) count_pagination = cursor.fetchone()[0] # Delete old non-pagination entries cursor.execute("DELETE FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,)) # Delete old pagination entries cursor.execute("DELETE FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,)) conn.commit() conn.close() print(f"Cache cleanup completed: {count_non_pagination} non-pagination entries and {count_pagination} pagination entries removed") except Exception as e: print(f"Error during cache cleanup: {e}") # Initialize scheduler for periodic cache cleanup scheduler = BackgroundScheduler() scheduler.add_job( cleanup_old_cache_entries, CronTrigger.from_crontab(CLEANUP_CRON), id='cache_cleanup_job', replace_existing=True ) # Initialize database on startup init_db() # Start the scheduler when the application starts @app.on_event("startup") def start_scheduler(): scheduler.start() print(f"Cache cleanup scheduler started with cron: {CLEANUP_CRON}") print(f"Cache entries will expire after {CACHE_EXPIRY_HOURS} hours") # Shutdown the scheduler when the application stops @app.on_event("shutdown") def shutdown_scheduler(): scheduler.shutdown(wait=False) print("Cache cleanup scheduler stopped") async def wait_for_network_idle(page): """Wait until no network requests are in flight""" await page.waitForNetworkIdle(idleTime=500, timeout=30000) @app.head("/") async def health_check(): return {"status": "ok"} async def safe_browser_operation(url, operation_func): """Safely perform a browser operation with proper cleanup and resource limits""" async with operation_semaphore: # Limit concurrent operations async with get_browser() as browser: page = None try: # Create a new page page = await browser.newPage() # Set reasonable viewport await page.setViewport({'width': 1280, 'height': 800}) # Set user agent await page.setUserAgent(CUSTOM_USER_AGENT) # Set reasonable timeout page.setDefaultNavigationTimeout(30000) # 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: print(f"Error during browser operation: {str(e)}") raise finally: try: # Ensure page is properly closed if page: await page.close() 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)): # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "visit") if cached_result: return cached_result try: print(f"Visiting URL: {decoded_url}") # Define the operation to perform with the browser async def visit_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) if not response: print(f"Warning: No response object returned for {decoded_url}") # Get page content content = await page.content() return {"status": "success", "content": content} except Exception as e: print(f"Error during page navigation: {e}") # Try to get content anyway try: content = await page.content() return {"status": "partial", "content": content, "error": str(e)} except: raise HTTPException(status_code=500, detail=f"Failed to get page content: {str(e)}") # Perform the operation result = await safe_browser_operation(decoded_url, visit_operation) # Save to cache save_to_cache(decoded_url, "visit", result) return result except Exception as e: print(f"Error visiting URL {decoded_url}: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/seo") async def extract_seo(url: str, x_api_key: Optional[str] = Header(None)): """Extract SEO information from a website""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "seo") if cached_result: return cached_result try: print(f"Extracting SEO from: {decoded_url}") # Define the operation to perform with the browser async def seo_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) # Extract SEO information seo_data = await page.evaluate('''() => { const data = { title: document.title || '', description: '', canonical: '', h1: [], h2: [], images: 0, links: 0 }; // Get meta description const metaDescription = document.querySelector('meta[name="description"]'); if (metaDescription) { data.description = metaDescription.getAttribute('content') || ''; } // Get canonical link const canonicalLink = document.querySelector('link[rel="canonical"]'); if (canonicalLink) { data.canonical = canonicalLink.getAttribute('href') || ''; } // Get h1 tags document.querySelectorAll('h1').forEach(h1 => { const text = h1.innerText.trim(); if (text) data.h1.push(text); }); // Get h2 tags document.querySelectorAll('h2').forEach(h2 => { const text = h2.innerText.trim(); if (text) data.h2.push(text); }); // Count images data.images = document.querySelectorAll('img').length; // Count links data.links = document.querySelectorAll('a').length; return data; }''') result = { "status": "success", "url": decoded_url, "seo": seo_data } return result except Exception as e: print(f"Error during SEO extraction: {e}") return {"status": "error", "url": decoded_url, "error": str(e)} # Perform the operation result = await safe_browser_operation(decoded_url, seo_operation) # Save to cache save_to_cache(decoded_url, "seo", result) return result except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/meta") async def extract_meta_tags(url: str, x_api_key: Optional[str] = Header(None)): """Extract meta tags from a website""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "meta") if cached_result: return cached_result try: print(f"Extracting meta tags from: {decoded_url}") # Define the operation to perform with the browser async def meta_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) # Extract all meta tags meta_tags = await page.evaluate('''() => { const metas = Array.from(document.querySelectorAll('meta')); return metas.map(meta => { const attributes = {}; Array.from(meta.attributes).forEach(attr => { attributes[attr.name] = attr.value; }); return attributes; }); }''') # Extract Open Graph tags og_tags = await page.evaluate('''() => { const ogTags = {}; document.querySelectorAll('meta[property^="og:"]').forEach(tag => { const property = tag.getAttribute('property'); ogTags[property] = tag.getAttribute('content'); }); return ogTags; }''') # Extract Twitter card tags twitter_tags = await page.evaluate('''() => { const twitterTags = {}; document.querySelectorAll('meta[name^="twitter:"]').forEach(tag => { const name = tag.getAttribute('name'); twitterTags[name] = tag.getAttribute('content'); }); return twitterTags; }''') result = { "status": "success", "url": decoded_url, "meta_tags": meta_tags, "open_graph": og_tags, "twitter_card": twitter_tags, "title": await page.title() } return result except Exception as e: print(f"Error during meta tag extraction: {e}") return {"status": "error", "url": decoded_url, "error": str(e)} # Perform the operation result = await safe_browser_operation(decoded_url, meta_operation) # Save to cache save_to_cache(decoded_url, "meta", result) return result except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/cache/clear") async def clear_cache(x_api_key: Optional[str] = Header(None)): """Clear the entire cache database""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() cursor.execute("DELETE FROM cache") conn.commit() conn.close() return {"status": "success", "message": "Cache cleared successfully"} @app.get("/cache/stats") async def cache_stats(x_api_key: Optional[str] = Header(None)): """Get cache statistics""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") try: conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Get total count cursor.execute("SELECT COUNT(*) FROM cache") total_count = cursor.fetchone()[0] # Get count by route cursor.execute("SELECT route, COUNT(*) FROM cache GROUP BY route") route_counts = dict(cursor.fetchall()) # Get oldest and newest entries cursor.execute("SELECT MIN(timestamp), MAX(timestamp) FROM cache") min_time, max_time = cursor.fetchone() conn.close() return { "total_entries": total_count, "route_counts": route_counts, "oldest_entry": min_time, "newest_entry": max_time } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/status") async def system_status(x_api_key: Optional[str] = Header(None)): """Get system status and browser pool information""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") try: # Get system information process = psutil.Process() memory_info = process.memory_info() # Get browser pool information pool_size = browser_pool.qsize() active_browser_count = len(active_browsers) # Calculate browser ages browser_ages = [] for browser, creation_time in browser_creation_times.items(): age = time.time() - creation_time browser_ages.append(age) return { "system": { "cpu_percent": process.cpu_percent(), "memory_mb": memory_info.rss / 1024 / 1024, "memory_percent": process.memory_percent(), "open_files": len(process.open_files()), "connections": len(process.connections()), "threads": process.num_threads() }, "browser_pool": { "pool_size": pool_size, "active_browsers": active_browser_count, "max_browsers": MAX_BROWSERS, "browser_ttl_seconds": BROWSER_TTL, "browser_ages_seconds": browser_ages, "concurrent_operations_limit": MAX_CONCURRENT_OPERATIONS }, "timestamp": time.time() } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/emergency-cleanup") async def emergency_cleanup(x_api_key: Optional[str] = Header(None)): """Force emergency cleanup of all browsers""" # 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_all_browsers() return {"status": "success", "message": "Emergency cleanup completed"} 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 with timeout try: browser = await asyncio.wait_for(browser_pool.get(), timeout=30.0) except (asyncio.QueueEmpty, asyncio.TimeoutError): # If pool is empty or timeout, create a new browser if under the limit async with browser_lock: current_browser_count = len(active_browsers) if current_browser_count < MAX_BROWSERS: browser = await create_browser() else: # If at limit, wait for a browser to become available with timeout try: browser = await asyncio.wait_for(browser_pool.get(), timeout=60.0) except asyncio.TimeoutError: # Emergency cleanup and create new browser print("Emergency: Timeout waiting for browser, forcing cleanup") await force_cleanup_all_browsers() browser = await create_browser() yield browser except Exception as e: print(f"Error in get_browser: {e}") # Emergency cleanup if we can't get a browser await force_cleanup_all_browsers() browser = await create_browser() 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: print(f"Recycling old browser in get_browser (age: {time.time() - browser_creation_times.get(browser, 0):.0f}s)") await cleanup_browser(browser) browser = await create_browser() # Only put back if pool is not full if not browser_pool.full(): await browser_pool.put(browser) else: # Pool is full, cleanup this browser await cleanup_browser(browser) except Exception as e: print(f"Browser health check failed in get_browser: {e}") # If browser is not usable, close it and create a new one await cleanup_browser(browser) if not browser_pool.full(): new_browser = await create_browser() await browser_pool.put(new_browser) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)