From 8bc7d87d4208586ee56e02a4a21e9480ef5b2153 Mon Sep 17 00:00:00 2001 From: GeiserX <9169332+GeiserX@users.noreply.github.com> Date: Sat, 7 Nov 2026 14:29:12 +0100 Subject: [PATCH 1/2] fix(proxies): retry a copy the cap reset could not start, so an offline box does not keep it stopped When a reset could not start a copy the cap had stopped, the copy kept its stopped_by_cap mark and nothing tried again until the next reset, a month later. Each heartbeat from that worker now starts such a copy again while the proxy is uncapped and the copy is reported down. --- app/main.py | 35 ++++++++++++++++++++++++++++++- docs/proxies.md | 4 +++- tests/test_proxies.py | 49 ++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 3 deletions(-) diff --git a/app/main.py b/app/main.py index 8ef60fff..fa18a011 100644 --- a/app/main.py +++ b/app/main.py @@ -4366,6 +4366,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 stopped, retried after a reset, on their way. +_proxy_cap_starting: set[tuple[int, int, str]] = set() def _running(entry: dict[str, Any]) -> bool: @@ -4392,7 +4394,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} + copies = [c for c in await database.list_proxied_copies() if c["worker_id"] == worker_id] + recorded = {(c["proxy_id"], c["slug"]) for c in copies} + held_by_cap = {(c["proxy_id"], c["slug"]) for c in copies 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 +4407,18 @@ 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"): + # A reset whose start failed leaves the copy marked as the cap's and + # down; while the proxy stays uncapped, every heartbeat tries again. + retry = { + slug + for e in items + if not _running(e) + and (proxy_id, slug := egress.container_slug(e)) in held_by_cap + and (proxy_id, worker_id, slug) not in _proxy_cap_starting + } + if retry: + _proxy_cap_starting.update((proxy_id, worker_id, slug) for slug in retry) + _spawn(_restart_capped_copies(proxy_id, worker_id, retry)) continue targets = {(worker_id, egress.container_slug(e)) for e in items if _running(e)} if crossed: @@ -4416,6 +4432,23 @@ async def _account_proxy_traffic(worker_id: int, entries: list[Any]) -> None: _spawn(_stop_capped_proxy({**proxy, "state": state}, targets)) +async def _restart_capped_copies(proxy_id: int, worker_id: int, slugs: set[str]) -> None: + """Start copies the cap stopped and a reset could not start, unless the cap is back.""" + try: + for slug in sorted(slugs): + async with _proxy_state_lock(proxy_id): + fresh = await database.get_proxy(proxy_id) + if not fresh or fresh["state"].get("capped_at"): + return + try: + await _start_proxied_copy(fresh, worker_id, slug) + await database.set_proxied_copy_capped(proxy_id, worker_id, slug, False) + except HTTPException as exc: + logger.warning("Could not restart %s through proxy %s after a reset: %s", slug, proxy_id, exc.detail) + finally: + _proxy_cap_starting.difference_update((proxy_id, worker_id, slug) for slug in slugs) + + async def _stop_capped_proxy(proxy: dict[str, Any], targets: set[tuple[int, str]]) -> None: pid = proxy["id"] unstopped = _proxy_cap_unstopped.setdefault(pid, {}) diff --git a/docs/proxies.md b/docs/proxies.md index 3790fa5b..19d261c3 100644 --- a/docs/proxies.md +++ b/docs/proxies.md @@ -227,7 +227,9 @@ the moment of the stop is stopped when it comes back. If a stop fails, the alert says which copy is still running. 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 +you stopped yourself stays stopped. If a copy cannot be started at the reset, +for example because its box is offline, every report from that box tries again +until it starts or the cap holds again. 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. diff --git a/tests/test_proxies.py b/tests/test_proxies.py index f6280c76..e8808565 100644 --- a/tests/test_proxies.py +++ b/tests/test_proxies.py @@ -322,6 +322,7 @@ def __init__(self, copies, worker_call): self.call = worker_call self.alerts: list[tuple] = [] self.pushed: list[tuple] = [] + self.mark = AsyncMock() async def heartbeat(self, worker_id, entries): spawned: list = [] @@ -343,7 +344,7 @@ async def push(*args): patch.object(main.database, "get_proxy", side_effect=get_proxy), patch.object(main.database, "update_proxy_state", side_effect=update_state), patch.object(main.database, "list_proxied_copies", new_callable=AsyncMock, return_value=self.copies), - patch.object(main.database, "set_proxied_copy_capped", new_callable=AsyncMock), + patch.object(main.database, "set_proxied_copy_capped", self.mark), patch.object(main.database, "record_alert", side_effect=record_alert), patch.object(main.database, "list_workers", new_callable=AsyncMock, return_value=[]), patch.object(main, "_proxy_worker_call", self.call), @@ -359,6 +360,7 @@ class TestCapEnforcement: def teardown_method(self): main._drop_proxy_cap_bell(1) main._proxy_cap_stopping.clear() + main._proxy_cap_starting.clear() def test_a_failed_stop_is_tried_again_on_the_next_heartbeat(self): fails = {8: [HTTPException(status_code=503, detail="Worker offline")]} @@ -406,6 +408,51 @@ def test_a_stopped_copy_is_left_alone(self): asyncio.run(rig.heartbeat(9, [_entry(5, 0, tid="t9", status="exited", tunnel_status="exited")])) assert call.call_count == 0 + def test_a_copy_a_reset_could_not_start_is_started_on_the_next_heartbeat(self): + fails = [HTTPException(status_code=503, detail="Worker offline")] + + async def worker_call(worker_id, method, path, definition, **kw): + if fails: + raise fails.pop() + return {} + + call = AsyncMock(side_effect=worker_call) + rig = _CapRig( + [ + {"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}, + {"proxy_id": 1, "worker_id": 7, "slug": "b", "stopped_by_cap": 0}, + ], + call, + ) + down = {"status": "exited", "tunnel_status": "exited"} + + async def go(): + # The first try fails, the copy stays the cap's, the next heartbeat tries again. + await rig.heartbeat(7, [_entry(0, 0, **down), _entry(0, 0, tid="tb", slug="b", **down)]) + rig.mark.assert_not_called() + await rig.heartbeat(7, [_entry(0, 0, **down), _entry(0, 0, tid="tb", slug="b", **down)]) + + asyncio.run(go()) + # Only the copy the cap stopped is started; the owner's stopped copy b is left alone. + assert [(c.args[0], c.args[2]) for c in call.call_args_list] == [ + (7, "/api/proxied/a/1/start"), + (7, "/api/proxied/a/1/start"), + ] + assert [c.args for c in rig.mark.call_args_list] == [(1, 7, "a", False)] + + def test_a_copy_the_cap_stopped_stays_down_while_the_cap_holds(self): + call = AsyncMock(return_value={}) + rig = _CapRig([{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}], call) + rig.proxy["state"] = {"used_bytes": 950_000_000, "capped_at": "x", "cap_reason": "r"} + asyncio.run(rig.heartbeat(7, [_entry(0, 0, status="exited", tunnel_status="exited")])) + assert call.call_count == 0 + + def test_a_running_copy_the_cap_stopped_is_not_started_again(self): + call = AsyncMock(return_value={}) + rig = _CapRig([{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}], call) + asyncio.run(rig.heartbeat(7, [_entry(0, 0, status="running", tunnel_status="running")])) + assert call.call_count == 0 + def test_the_bell_keeps_a_cap_that_still_holds_after_a_restart(self): capped = {"id": 3, "definition": {"name": "Oxy"}, "state": {"capped_at": "x", "cap_reason": "Oxy reached 90%"}} free = {"id": 4, "definition": {"name": "Free"}, "state": {}} From b23857aca2ddbb760c158f79be5d2d4176335efc Mon Sep 17 00:00:00 2001 From: GeiserX <9169332+GeiserX@users.noreply.github.com> Date: Sat, 7 Nov 2026 15:10:36 +0100 Subject: [PATCH 2/2] fix(proxies): let the owner's decision and a returning cap win over a pending retry A retry now rechecks the copy's cap mark before it starts (an owner start or stop clears it), keeps the mark when the cap came back during the start, and shares the in-flight set with the reset's starts so the two never double up. --- app/main.py | 25 ++++++++++++++++++++++--- tests/test_proxies.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/app/main.py b/app/main.py index fa18a011..89ceb75f 100644 --- a/app/main.py +++ b/app/main.py @@ -4328,8 +4328,10 @@ async def _release_capped_copies(proxy: dict[str, Any]) -> list[str]: _drop_proxy_cap_bell(proxy["id"]) started: list[str] = [] for copy in await database.list_proxied_copies(proxy["id"]): - if not copy.get("stopped_by_cap"): + key = (proxy["id"], copy["worker_id"], copy["slug"]) + if not copy.get("stopped_by_cap") or key in _proxy_cap_starting: continue + _proxy_cap_starting.add(key) try: await _start_proxied_copy(proxy, copy["worker_id"], copy["slug"]) await database.set_proxied_copy_capped(proxy["id"], copy["worker_id"], copy["slug"], False) @@ -4338,6 +4340,8 @@ async def _release_capped_copies(proxy: dict[str, Any]) -> list[str]: logger.warning( "Could not restart %s through proxy %s after a reset: %s", copy["slug"], proxy["id"], exc.detail ) + finally: + _proxy_cap_starting.discard(key) return started @@ -4366,7 +4370,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 stopped, retried after a reset, on their way. +# Starts of copies the cap stopped (by a reset or a heartbeat retry) on their way, +# so the two paths never send a second start for the same copy. _proxy_cap_starting: set[tuple[int, int, str]] = set() @@ -4438,13 +4443,27 @@ async def _restart_capped_copies(proxy_id: int, worker_id: int, slugs: set[str]) for slug in sorted(slugs): async with _proxy_state_lock(proxy_id): fresh = await database.get_proxy(proxy_id) + # The owner may have started or stopped it since the heartbeat; that + # clears the mark, and the owner's decision stands. + marked = 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) + ) if not fresh or fresh["state"].get("capped_at"): return + if not marked: + continue try: await _start_proxied_copy(fresh, worker_id, slug) - await database.set_proxied_copy_capped(proxy_id, worker_id, slug, False) except HTTPException as exc: logger.warning("Could not restart %s through proxy %s after a reset: %s", slug, proxy_id, exc.detail) + continue + # The cap may have come back while the start ran. Then the mark stays, so + # the next heartbeat stops the copy again and the next reset starts it. + async with _proxy_state_lock(proxy_id): + again = await database.get_proxy(proxy_id) + if again and not again["state"].get("capped_at"): + await database.set_proxied_copy_capped(proxy_id, worker_id, slug, False) finally: _proxy_cap_starting.difference_update((proxy_id, worker_id, slug) for slug in slugs) diff --git a/tests/test_proxies.py b/tests/test_proxies.py index e8808565..592aef02 100644 --- a/tests/test_proxies.py +++ b/tests/test_proxies.py @@ -440,6 +440,48 @@ async def go(): ] assert [c.args for c in rig.mark.call_args_list] == [(1, 7, "a", False)] + def test_a_retry_respects_the_owner_and_a_cap_that_came_back(self): + uncapped = {"id": 1, "definition": dict(CAPPED), "state": {}} + capped = {**uncapped, "state": {"capped_at": "x", "cap_reason": "r"}} + + def run(copies, proxies_seen): + call = AsyncMock(return_value={}) + mark = AsyncMock() + with ( + patch.object(main.database, "get_proxy", new_callable=AsyncMock, side_effect=proxies_seen), + patch.object(main.database, "list_proxied_copies", new_callable=AsyncMock, return_value=copies), + patch.object(main.database, "set_proxied_copy_capped", mark), + patch.object(main, "_proxy_worker_call", call), + ): + asyncio.run(main._restart_capped_copies(1, 7, {"a"})) + return call, mark + + # The owner started or stopped it since the heartbeat: the mark is gone, no start. + call, mark = run([{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 0}], [uncapped]) + assert call.call_count == 0 and mark.call_count == 0 + # The cap came back while the start ran: the mark stays for the next stop and reset. + call, mark = run([{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}], [uncapped, capped]) + assert call.call_count == 1 and mark.call_count == 0 + assert not main._proxy_cap_starting + + def test_a_reset_skips_a_copy_whose_start_is_already_on_its_way(self): + proxy = {"id": 1, "definition": dict(CAPPED), "state": {}} + call = AsyncMock(return_value={}) + main._proxy_cap_starting.add((1, 7, "a")) + with ( + patch.object(main.database, "clear_alerts", new_callable=AsyncMock), + patch.object( + main.database, + "list_proxied_copies", + new_callable=AsyncMock, + return_value=[{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}], + ), + patch.object(main.database, "set_proxied_copy_capped", new_callable=AsyncMock), + patch.object(main, "_proxy_worker_call", call), + ): + assert asyncio.run(main._release_capped_copies(proxy)) == [] + assert call.call_count == 0 + def test_a_copy_the_cap_stopped_stays_down_while_the_cap_holds(self): call = AsyncMock(return_value={}) rig = _CapRig([{"proxy_id": 1, "worker_id": 7, "slug": "a", "stopped_by_cap": 1}], call)