Merge pull request #1578 from MODSetter/ci_mvp

refactor: improve session management and concurrency in Google Search

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

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):

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()

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())