89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from fastapi import FastAPI, HTTPException, Header, BackgroundTasks
|
|
from pyppeteer import launch
|
|
import os
|
|
import asyncio
|
|
import gc
|
|
import logging
|
|
import psutil
|
|
from typing import Optional
|
|
from urllib.parse import unquote
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = FastAPI()
|
|
|
|
API_KEY = os.getenv('API_KEY')
|
|
if not API_KEY:
|
|
raise ValueError("API_KEY environment variable must be set")
|
|
|
|
browser = None
|
|
|
|
async def get_browser():
|
|
global browser
|
|
if browser is None or not browser.isConnected():
|
|
logger.info("Launching new browser instance")
|
|
browser = await launch(
|
|
headless=True,
|
|
executablePath='/usr/bin/google-chrome',
|
|
args=[
|
|
'--no-sandbox',
|
|
'--disable-setuid-sandbox',
|
|
'--disable-dev-shm-usage',
|
|
'--disable-gpu',
|
|
'--single-process',
|
|
'--disable-software-rasterizer',
|
|
'--js-flags=--expose-gc',
|
|
'--memory-pressure-off',
|
|
]
|
|
)
|
|
return browser
|
|
|
|
async def cleanup_resources(page=None):
|
|
try:
|
|
if page:
|
|
try:
|
|
await page.evaluate("window.gc()") # Trigger JS garbage collection
|
|
except:
|
|
pass # Ignore if GC is not exposed
|
|
await page.close()
|
|
gc.collect()
|
|
process = psutil.Process(os.getpid())
|
|
logger.info(f"Memory usage: {process.memory_info().rss / 1024 / 1024:.2f} MB")
|
|
except Exception as e:
|
|
logger.error(f"Cleanup error: {str(e)}")
|
|
|
|
@app.get("/")
|
|
async def visit_url(url: str, x_api_key: Optional[str] = Header(None), background_tasks: BackgroundTasks = None):
|
|
if not x_api_key or x_api_key != API_KEY:
|
|
raise HTTPException(status_code=401, detail="Invalid API key")
|
|
|
|
page = None
|
|
try:
|
|
decoded_url = unquote(url)
|
|
logger.info(f"Visiting URL: {decoded_url}")
|
|
browser_instance = await get_browser()
|
|
context = await browser_instance.createIncognitoBrowserContext()
|
|
page = await context.newPage()
|
|
await page.setViewport({"width": 1280, "height": 800})
|
|
await page.goto(decoded_url, waitUntil='networkidle2')
|
|
content = await page.content()
|
|
if background_tasks:
|
|
background_tasks.add_task(cleanup_resources, page)
|
|
else:
|
|
await cleanup_resources(page)
|
|
return {"status": "success", "content": content}
|
|
except Exception as e:
|
|
logger.error(f"Error: {str(e)}")
|
|
if page:
|
|
await cleanup_resources(page)
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
logger.info("Shutting down, cleaning up browser instance")
|
|
global browser
|
|
if browser:
|
|
await browser.close()
|
|
browser = None
|