From 9a9a78e7eff5b8f445586e568fda843ee6b9b8eb Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Mon, 6 Jul 2026 22:33:40 -0700 Subject: [PATCH] refactor: improve session management and concurrency in Google Search fetcher - Introduced a semaphore to limit concurrent page renders, enhancing resource management. - Updated session handling to defer browser closure until all in-flight renders are complete. - Improved comments for clarity on the behavior of concurrent fetches and session lifecycle. - Cleaned up imports in alembic environment and migration flow scripts for consistency. --- surfsense_backend/alembic/env.py | 2 +- .../platforms/google_search/fetch.py | 47 ++++++++++++-- .../scripts/check_migration_flow.py | 1 - .../google_search/test_fetch_concurrency.py | 63 +++++++++++++++++++ 4 files changed, 105 insertions(+), 8 deletions(-) create mode 100644 surfsense_backend/tests/unit/platforms/google_search/test_fetch_concurrency.py diff --git a/surfsense_backend/alembic/env.py b/surfsense_backend/alembic/env.py index f2130892a..7c524e4ba 100644 --- a/surfsense_backend/alembic/env.py +++ b/surfsense_backend/alembic/env.py @@ -5,12 +5,12 @@ import sys from logging.config import fileConfig import sqlalchemy as sa +from alembic.script import ScriptDirectory from sqlalchemy import pool from sqlalchemy.engine import Connection from sqlalchemy.ext.asyncio import async_engine_from_config from alembic import context -from alembic.script import ScriptDirectory # Ensure the app directory is in the Python path # This allows Alembic to find your models diff --git a/surfsense_backend/app/proprietary/platforms/google_search/fetch.py b/surfsense_backend/app/proprietary/platforms/google_search/fetch.py index 68cc90d26..0e290fb4b 100644 --- a/surfsense_backend/app/proprietary/platforms/google_search/fetch.py +++ b/surfsense_backend/app/proprietary/platforms/google_search/fetch.py @@ -75,7 +75,8 @@ _WALLED_ROUND_BACKOFF_S = 3.0 # The sticky IP that most recently served a real SERP; the strongest hint for # the next fetch. Re-vetted (cheap) before reuse, dropped on failure. -# ponytail: a single slot, not a pool — one render runs at a time today. +# ponytail: a single slot, not a pool — concurrent fetches share (and race) +# it; worst case a loser re-vets a fresh IP, which is the normal path anyway. _last_good_proxy: str | None = None # A usable precheck responds in <1 s; anything slower is a dead/slow sticky IP @@ -273,6 +274,23 @@ async def _in_browser_loop(coro): _sessions: dict[bool, AsyncStealthySession] = {} _session_lock = asyncio.Lock() +# How many renders may run at once. scrapling's page pool defaults to ONE +# page, and a per-fetch proxy context skips its wait-for-a-slot path, so a +# second concurrent render raised RuntimeError('Maximum page limit (1) +# reached'); the failure handler then closed the shared browser under the +# sibling render (the TargetClosedError cascade seen in production when +# several scrape runs overlap). The pool and this gate are sized together: +# the gate queues excess renders instead of tripping the pool. +_MAX_CONCURRENT_PAGES = 4 +# Only ever awaited from coroutines running on the browser loop. +_render_gate = asyncio.Semaphore(_MAX_CONCURRENT_PAGES) + +# Live renders per session, so dropping a "broken" session defers the actual +# browser close until its last in-flight render finishes — closing earlier is +# what murdered sibling renders. Browser-loop-only state. +_inflight: dict[AsyncStealthySession, int] = {} +_doomed: set[AsyncStealthySession] = set() + async def _get_session(mobile: bool) -> AsyncStealthySession: """The shared live browser session for this layout, launching it if needed.""" @@ -286,6 +304,7 @@ async def _get_session(mobile: bool) -> AsyncStealthySession: "google_search": True, "page_action": _expand_blocks, "retries": 1, # our own IP loop is the retry policy + "max_pages": _MAX_CONCURRENT_PAGES, } base = get_proxy_url() if base: @@ -305,9 +324,14 @@ async def _get_session(mobile: bool) -> AsyncStealthySession: async def _drop_session_on_loop(mobile: bool) -> None: async with _session_lock: session = _sessions.pop(mobile, None) - if session is not None: - with contextlib.suppress(Exception): # already dead; nothing to salvage - await session.close() + if session is None: + return + if _inflight.get(session, 0): + # Sibling renders are still on this browser; the last one closes it. + _doomed.add(session) + return + with contextlib.suppress(Exception): # already dead; nothing to salvage + await session.close() async def _drop_session(mobile: bool) -> None: @@ -322,8 +346,19 @@ async def close_sessions() -> None: async def _render_on_loop(url: str, proxy: str | None, mobile: bool): - session = await _get_session(mobile) - return await session.fetch(url, proxy=proxy) + async with _render_gate: + session = await _get_session(mobile) + _inflight[session] = _inflight.get(session, 0) + 1 + try: + return await session.fetch(url, proxy=proxy) + finally: + _inflight[session] -= 1 + if not _inflight[session]: + del _inflight[session] + if session in _doomed: + _doomed.discard(session) + with contextlib.suppress(Exception): + await session.close() async def _render(url: str, proxy: str | None, mobile: bool = False): diff --git a/surfsense_backend/scripts/check_migration_flow.py b/surfsense_backend/scripts/check_migration_flow.py index 18587f836..155e5ac84 100644 --- a/surfsense_backend/scripts/check_migration_flow.py +++ b/surfsense_backend/scripts/check_migration_flow.py @@ -94,7 +94,6 @@ async def set_version(version: str | None) -> None: async def assert_at_head() -> None: import asyncpg - from alembic.script import ScriptDirectory head = ScriptDirectory(str(BACKEND_DIR / "alembic")).get_current_head() diff --git a/surfsense_backend/tests/unit/platforms/google_search/test_fetch_concurrency.py b/surfsense_backend/tests/unit/platforms/google_search/test_fetch_concurrency.py new file mode 100644 index 000000000..c129f7c5f --- /dev/null +++ b/surfsense_backend/tests/unit/platforms/google_search/test_fetch_concurrency.py @@ -0,0 +1,63 @@ +"""The production regression this guards: several scrape runs render SERPs +concurrently on the shared browser. When one render failed, `_drop_session` +closed the browser immediately — under the sibling renders — so every +in-flight fetch died with TargetClosedError and the runs cascaded into +"exhausted 24 IPs". The drop must defer the actual close until the last +in-flight render on that session finishes. +""" + +import asyncio + +from app.proprietary.platforms.google_search import fetch + + +class _FakeSession: + def __init__(self): + self.release = asyncio.Event() + self.closed = 0 + + async def fetch(self, url, proxy=None): + await self.release.wait() + return "page" + + async def close(self): + self.closed += 1 + + +def test_drop_defers_close_until_inflight_renders_finish(monkeypatch): + session = _FakeSession() + + async def fake_get_session(mobile): + return session + + monkeypatch.setattr(fetch, "_get_session", fake_get_session) + monkeypatch.setitem(fetch._sessions, False, session) + + async def main(): + render = asyncio.create_task(fetch._render_on_loop("u", None, False)) + await asyncio.sleep(0) # let the render register as in-flight + assert fetch._inflight[session] == 1 + + await fetch._drop_session_on_loop(False) + assert session.closed == 0, "must not close under an in-flight render" + assert session in fetch._doomed + assert False not in fetch._sessions # next fetch relaunches + + session.release.set() + assert await render == "page" + assert session.closed == 1, "last render out closes the doomed browser" + assert session not in fetch._inflight + assert session not in fetch._doomed + + asyncio.run(main()) + + +def test_drop_closes_immediately_when_idle(): + session = _FakeSession() + + async def main(): + fetch._sessions[False] = session + await fetch._drop_session_on_loop(False) + assert session.closed == 1 + + asyncio.run(main())