diff --git a/app/main.py b/app/main.py index 4a04ac6a..90df1fa9 100644 --- a/app/main.py +++ b/app/main.py @@ -4346,8 +4346,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) @@ -4356,6 +4358,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 @@ -4384,6 +4388,9 @@ 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 (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() def _running(entry: dict[str, Any]) -> bool: @@ -4410,7 +4417,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] @@ -4421,6 +4430,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: @@ -4434,6 +4455,37 @@ 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) + # 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) + 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) + + 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 d967745c..49811474 100644 --- a/docs/proxies.md +++ b/docs/proxies.md @@ -228,7 +228,9 @@ 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. The reset day is checked every hour and -each time the app starts. Turning a cap on starts a fresh period, so +each time the app starts. 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..592aef02 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,93 @@ 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_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) + 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": {}}