From 70a46d18d3743f645a0e355e4ccac460a80c9c67 Mon Sep 17 00:00:00 2001 From: GeiserX <9169332+GeiserX@users.noreply.github.com> Date: Sat, 7 Nov 2026 14:37:52 +0100 Subject: [PATCH] fix(proxies): start a copy the cap released once its box is back When a cap reset or a cap change lifted the stop, each copy the cap had stopped was started once. If its box was offline at that moment, the start failed and nothing tried again: the copy stayed stopped, marked as stopped by the cap, until the next reset. Each heartbeat now starts the copies of that box that are still marked and not running, once the proxy is no longer capped. The same mark is set when a stop lands, the cap is lifted meanwhile, and the start back up fails. --- app/main.py | 45 +++++++++++++++++++++- docs/proxies.md | 1 + tests/test_proxies.py | 89 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/app/main.py b/app/main.py index 8ef60fff..87d9ed53 100644 --- a/app/main.py +++ b/app/main.py @@ -4336,7 +4336,10 @@ async def _release_capped_copies(proxy: dict[str, Any]) -> list[str]: started.append(f"{copy['slug']} on worker {copy['worker_id']}") except HTTPException as exc: logger.warning( - "Could not restart %s through proxy %s after a reset: %s", copy["slug"], proxy["id"], exc.detail + "Could not restart %s through proxy %s after a reset (its box's next heartbeat tries again): %s", + copy["slug"], + proxy["id"], + exc.detail, ) return started @@ -4366,6 +4369,8 @@ async def _run_proxy_cap_resets() -> None: _proxy_cap_unstopped: dict[int, dict[tuple[int, str], str]] = {} # Stops already on their way, so a slow worker is not sent a second one. _proxy_cap_stopping: set[tuple[int, int, str]] = set() +# Starts of copies the cap released, already on their way. +_proxy_cap_starting: set[tuple[int, int, str]] = set() def _running(entry: dict[str, Any]) -> bool: @@ -4392,7 +4397,9 @@ async def _account_proxy_traffic(worker_id: int, entries: list[Any]) -> None: return # Only traffic of copies the UI recorded on this worker counts, so a worker # cannot spend, or cap, a proxy it runs nothing through. - recorded = {(c["proxy_id"], c["slug"]) for c in await database.list_proxied_copies() if c["worker_id"] == worker_id} + mine = [c for c in await database.list_proxied_copies() if c["worker_id"] == worker_id] + recorded = {(c["proxy_id"], c["slug"]) for c in mine} + held = {(c["proxy_id"], c["slug"]) for c in mine if c.get("stopped_by_cap")} now = time.time() for proxy_id, items in by_proxy.items(): counted = [e for e in items if (proxy_id, egress.container_slug(e)) in recorded] @@ -4403,6 +4410,14 @@ async def _account_proxy_traffic(worker_id: int, entries: list[Any]) -> None: state, crossed = proxies.account(proxy["state"], proxy["definition"], worker_id, counted, now) await database.update_proxy_state(proxy_id, state) if not state.get("capped_at"): + # The cap was lifted, but a copy it stopped could not be started + # then (the box was offline, say). This box is answering now. + for e in items: + slug = egress.container_slug(e) + key = (proxy_id, worker_id, slug) + if (proxy_id, slug) in held and not _running(e) and key not in _proxy_cap_starting: + _proxy_cap_starting.add(key) + _spawn(_start_released_copy(proxy_id, worker_id, slug)) continue targets = {(worker_id, egress.container_slug(e)) for e in items if _running(e)} if crossed: @@ -4416,6 +4431,30 @@ async def _account_proxy_traffic(worker_id: int, entries: list[Any]) -> None: _spawn(_stop_capped_proxy({**proxy, "state": state}, targets)) +async def _start_released_copy(proxy_id: int, worker_id: int, slug: str) -> None: + """Start one copy the cap stopped, after its release could not start it.""" + try: + async with _proxy_state_lock(proxy_id): + proxy = await database.get_proxy(proxy_id) + held = any( + c["worker_id"] == worker_id and c["slug"] == slug and c.get("stopped_by_cap") + for c in await database.list_proxied_copies(proxy_id) + ) + # Capped again, or the owner started or stopped it by hand since. + if not proxy or proxy["state"].get("capped_at") or not held: + return + await _start_proxied_copy(proxy, worker_id, slug) + async with _proxy_state_lock(proxy_id): + fresh = await database.get_proxy(proxy_id) + if fresh and not fresh["state"].get("capped_at"): + await database.set_proxied_copy_capped(proxy_id, worker_id, slug, False) + logger.info("Started %s through proxy %s, which the cap had stopped", slug, proxy_id) + except HTTPException as exc: + logger.warning("Could not start %s again through proxy %s: %s", slug, proxy_id, exc.detail) + finally: + _proxy_cap_starting.discard((proxy_id, worker_id, slug)) + + async def _stop_capped_proxy(proxy: dict[str, Any], targets: set[tuple[int, str]]) -> None: pid = proxy["id"] unstopped = _proxy_cap_unstopped.setdefault(pid, {}) @@ -4440,6 +4479,8 @@ async def _stop_capped_proxy(proxy: dict[str, Any], targets: set[tuple[int, str] try: await _start_proxied_copy(fresh, worker_id, slug) except HTTPException as exc: + # Marked, so the box's next heartbeat starts it. + await database.set_proxied_copy_capped(pid, worker_id, slug, True) logger.warning("Could not start %s again through proxy %s: %s", slug, pid, exc.detail) except HTTPException as exc: if exc.status_code == 404 and any( diff --git a/docs/proxies.md b/docs/proxies.md index 3790fa5b..3bc3edd1 100644 --- a/docs/proxies.md +++ b/docs/proxies.md @@ -230,6 +230,7 @@ A reset, on the day or by hand, starts only the copies the cap stopped; a copy you stopped yourself stays stopped. Turning a cap on starts a fresh period, so traffic from before it existed does not count. Turning it off, or raising it above what this period has used, lifts the stop and starts those copies again. +A copy whose box was offline then is started when that box next reports. Workers report every 60 seconds, so the copies can run up to a minute past the threshold. That is why the default is 90%, not 100%. diff --git a/tests/test_proxies.py b/tests/test_proxies.py index f6280c76..d2dcecf1 100644 --- a/tests/test_proxies.py +++ b/tests/test_proxies.py @@ -797,6 +797,95 @@ async def go(): assert proxy["state"]["used_bytes"] == 950_001_100 +async def _heartbeat_awaiting_spawns(worker_id, entries): + spawned: list = [] + with patch.object(main, "_spawn", side_effect=spawned.append): + await main._account_proxy_traffic(worker_id, entries) + while spawned: + await spawned.pop(0) + + +def _capped_copy(pid): + return database.update_proxy_state( + pid, {"used_bytes": 950_000_000, "capped_at": "2026-11-04T01:00:00+00:00", "cap_reason": "P hit"} + ) + + +class TestReleasedCopiesThatDidNotStart: + """A lifted cap whose restart failed: the copy's box starts it when it next reports.""" + + def test_the_boxs_next_heartbeat_starts_it(self, db): + offline = AsyncMock(side_effect=HTTPException(status_code=503, detail="Worker offline")) + online = AsyncMock(return_value={}) + + async def go(): + pid = (await main.api_create_proxy(_Req(_socks(cap_gb=1)), {}))["id"] + await database.save_proxied_copy(pid, 7, "honeygain") + await database.set_proxied_copy_capped(pid, 7, "honeygain", True) + await _capped_copy(pid) + with patch.object(main, "_proxy_to_worker", offline): + await main.api_update_proxy(_Req({"cap_gb": ""}), pid, {}) + after_release = (await database.list_proxied_copies(pid))[0]["stopped_by_cap"] + hb = [_entry(0, 0, slug="honeygain", proxy_id=pid, status="exited", tunnel_status="exited")] + with patch.object(main, "_proxy_to_worker", online): + await _heartbeat_awaiting_spawns(7, hb) + return pid, after_release, (await database.list_proxied_copies(pid))[0]["stopped_by_cap"] + + pid, after_release, after_heartbeat = _run(go) + assert after_release == 1, "the failed start keeps the copy marked" + assert online.call_args.args[:3] == (7, "POST", f"/api/proxied/honeygain/{pid}/start") + assert after_heartbeat == 0 + assert not main._proxy_cap_starting + + @pytest.mark.parametrize( + ("capped", "marked", "status"), + [ + pytest.param(True, True, "exited", id="capped-again"), + pytest.param(False, False, "exited", id="started-or-stopped-by-hand"), + pytest.param(False, True, "running", id="already-running"), + ], + ) + def test_nothing_is_started_when_it_should_not_be(self, db, capped, marked, status): + call = AsyncMock(return_value={}) + + async def go(): + pid = (await main.api_create_proxy(_Req(_socks(cap_gb=1)), {}))["id"] + await database.save_proxied_copy(pid, 7, "honeygain") + await database.set_proxied_copy_capped(pid, 7, "honeygain", marked) + if capped: + await _capped_copy(pid) + hb = [_entry(0, 0, slug="honeygain", proxy_id=pid, status=status, tunnel_status=status)] + with patch.object(main, "_proxy_to_worker", call): + await _heartbeat_awaiting_spawns(7, hb) + return pid + + pid = _run(go) + starts = [c for c in call.call_args_list if c.args[2] == f"/api/proxied/honeygain/{pid}/start"] + assert starts == [] + + def test_a_restart_that_fails_after_a_cap_lifted_mid_stop_is_marked(self, db): + """The stop landed, the cap was lifted meanwhile, and the start back up failed.""" + + async def worker_call(worker_id, method, path, definition, **kw): + if path.endswith("/start"): + raise HTTPException(status_code=503, detail="Worker offline") + return {"was_running": True} + + async def go(): + pid = (await main.api_create_proxy(_Req(_socks(cap_gb=1)), {}))["id"] + await database.save_proxied_copy(pid, 7, "honeygain") + proxy = await database.get_proxy(pid) + with ( + patch.object(main, "_proxy_worker_call", side_effect=worker_call), + patch.object(main, "_spawn", side_effect=_discard_spawn), + ): + await main._stop_capped_proxy(proxy, {(7, "honeygain")}) + return (await database.list_proxied_copies(pid))[0]["stopped_by_cap"] + + assert _run(go) == 1 + main._drop_proxy_cap_bell(1) + + class TestValidationEcho: def test_a_422_on_a_proxy_route_does_not_echo_the_body(self): from contextlib import asynccontextmanager