This commit is contained in:
@@ -1,12 +1,11 @@
|
||||
from fastapi import FastAPI, HTTPException, Header, BackgroundTasks
|
||||
from fastapi import FastAPI, HTTPException, Header
|
||||
from pyppeteer import launch
|
||||
import os
|
||||
import asyncio
|
||||
import gc
|
||||
import logging
|
||||
import psutil
|
||||
from typing import Optional
|
||||
from urllib.parse import unquote
|
||||
import logging
|
||||
import gc
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
@@ -19,103 +18,75 @@ API_KEY = os.getenv('API_KEY')
|
||||
if not API_KEY:
|
||||
raise ValueError("API_KEY environment variable must be set")
|
||||
|
||||
# Global browser instance
|
||||
# Browser pool
|
||||
browser = None
|
||||
|
||||
async def get_browser():
|
||||
"""Get or create a browser instance"""
|
||||
global browser
|
||||
if browser is None or not hasattr(browser, 'process') or browser.process is None:
|
||||
logger.info("Launching new browser instance")
|
||||
browser = await launch(
|
||||
headless=True,
|
||||
executablePath='/usr/bin/google-chrome',
|
||||
args=[
|
||||
# Browser configuration
|
||||
BROWSER_ARGS = [
|
||||
'--no-sandbox',
|
||||
'--disable-setuid-sandbox',
|
||||
'--disable-dev-shm-usage',
|
||||
'--disable-gpu',
|
||||
'--disable-extensions',
|
||||
'--disable-sync',
|
||||
'--disable-translate',
|
||||
'--hide-scrollbars',
|
||||
'--mute-audio',
|
||||
'--disable-dev-shm-usage', # Overcome limited resource problems
|
||||
'--disable-gpu', # Disable GPU hardware acceleration
|
||||
'--disable-extensions', # Disable extensions
|
||||
'--disable-audio-output', # Disable 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',
|
||||
'--enable-features=NetworkService,NetworkServiceInProcess',
|
||||
'--mute-audio',
|
||||
'--memory-pressure-off',
|
||||
]
|
||||
)
|
||||
|
||||
# Browser launch options
|
||||
BROWSER_OPTIONS = {
|
||||
'headless': True,
|
||||
'executablePath': '/usr/bin/google-chrome',
|
||||
'args': BROWSER_ARGS,
|
||||
'handleSIGINT': False,
|
||||
'handleSIGTERM': False,
|
||||
'handleSIGHUP': False,
|
||||
}
|
||||
|
||||
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")
|
||||
browser = await launch(**BROWSER_OPTIONS)
|
||||
return browser
|
||||
|
||||
async def cleanup_resources(page=None, close_browser=False):
|
||||
"""Clean up resources to prevent memory leaks"""
|
||||
global browser
|
||||
async def wait_for_network_idle(page):
|
||||
"""Wait until no network requests are in flight"""
|
||||
try:
|
||||
if page:
|
||||
await page.close()
|
||||
await page.waitForNetworkIdle(idleTime=500, timeout=30000)
|
||||
except Exception as e:
|
||||
logger.warning(f"Network idle timeout: {str(e)}")
|
||||
|
||||
if close_browser and browser and hasattr(browser, 'process') and browser.process is not None:
|
||||
logger.info("Closing browser instance")
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
"""Initialize browser on startup"""
|
||||
await get_browser()
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Close browser on shutdown"""
|
||||
global browser
|
||||
if browser:
|
||||
logger.info("Closing browser on shutdown")
|
||||
await browser.close()
|
||||
browser = None
|
||||
|
||||
# Force garbage collection
|
||||
gc.collect()
|
||||
|
||||
# Log memory usage
|
||||
process = psutil.Process(os.getpid())
|
||||
memory_info = process.memory_info()
|
||||
logger.info(f"Memory usage: {memory_info.rss / 1024 / 1024:.2f} MB")
|
||||
except Exception as 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 hasattr(browser, 'process') and browser.process is not None
|
||||
}
|
||||
|
||||
@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):
|
||||
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")
|
||||
@@ -128,47 +99,50 @@ async def visit_url(url: str, x_api_key: Optional[str] = Header(None), backgroun
|
||||
|
||||
# Get browser instance
|
||||
browser_instance = await get_browser()
|
||||
if browser_instance is None:
|
||||
raise HTTPException(status_code=500, detail="Failed to initialize browser")
|
||||
|
||||
# Create new page
|
||||
# Create new page with timeout
|
||||
page = await browser_instance.newPage()
|
||||
|
||||
# Set page timeout
|
||||
await page.setDefaultNavigationTimeout(30000) # 30 seconds
|
||||
# Set page timeout and resource limits
|
||||
await page.setDefaultNavigationTimeout(60000) # 60 seconds timeout
|
||||
await page.setRequestInterception(True)
|
||||
|
||||
# Set viewport
|
||||
await page.setViewport({"width": 1280, "height": 800})
|
||||
# Block unnecessary resources
|
||||
async def intercept_request(request):
|
||||
if request.resourceType in ['image', 'media', 'font']:
|
||||
await request.abort()
|
||||
else:
|
||||
await request.continue_()
|
||||
|
||||
page.on('request', intercept_request)
|
||||
|
||||
# Navigate to URL and wait for network idle
|
||||
await page.goto(decoded_url, waitUntil='networkidle2')
|
||||
await page.goto(decoded_url, waitUntil='networkidle2', timeout=60000)
|
||||
|
||||
# Get page content
|
||||
content = await page.content()
|
||||
|
||||
# Schedule cleanup in background
|
||||
if background_tasks:
|
||||
background_tasks.add_task(cleanup_resources, page, False)
|
||||
else:
|
||||
await cleanup_resources(page, False)
|
||||
# Close page to free up resources
|
||||
await page.close()
|
||||
|
||||
# Force garbage collection
|
||||
gc.collect()
|
||||
|
||||
return {"status": "success", "content": content}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing request: {str(e)}")
|
||||
# Ensure cleanup happens even on error
|
||||
# Ensure browser is still running
|
||||
if browser is None or not browser.isConnected():
|
||||
await get_browser()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
finally:
|
||||
# Always ensure page is closed
|
||||
if page:
|
||||
try:
|
||||
await cleanup_resources(page, False)
|
||||
except Exception as cleanup_error:
|
||||
logger.error(f"Error during error cleanup: {str(cleanup_error)}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Clean up resources when shutting down"""
|
||||
logger.info("Application shutting down, cleaning up resources")
|
||||
await cleanup_resources(close_browser=True)
|
||||
await page.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
Reference in New Issue
Block a user