From d239964f0b2ad539b2886347efb1c463e30dabcd Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Tue, 22 Sep 2026 10:40:49 -0600 Subject: [PATCH] fix(serve): apply staged service reloads when sessions end, with a max staging age A staged tag was only retried on the refresh tick or a job completion, so a service with steady remote sessions never swapped in. Session end, expiry and eviction now trigger a single-flight apply, the tick applies after the pull, and serve.services_staged_max_age swaps an overdue tag in even while busy. --- docs/reference/configuration/services.md | 3 +- tests/core/test_services_reload.py | 16 +++ .../remote/unit/test_serve_service_refresh.py | 57 +++++++++ tests/remote/unit/test_session_store.py | 121 ++++++++++++++++++ unshackle/commands/serve.py | 5 +- unshackle/core/api/download_manager.py | 54 +++++--- unshackle/core/api/session_store.py | 14 ++ unshackle/core/services.py | 17 ++- 8 files changed, 267 insertions(+), 20 deletions(-) create mode 100644 tests/remote/unit/test_serve_service_refresh.py diff --git a/docs/reference/configuration/services.md b/docs/reference/configuration/services.md index 227fa15f..3778fa73 100644 --- a/docs/reference/configuration/services.md +++ b/docs/reference/configuration/services.md @@ -125,7 +125,8 @@ is the [REST API](../../dev/rest-api/index.md) section. These are the config key | `max_sessions` | int or null | `100` | Maximum concurrent remote sessions; at the cap the server evicts the oldest. Set it to `null` or `0` for no limit, which is also what `/api/dashboard/status` then reports, so a dashboard can tell "no cap" from a cap that happens to sit at the default. | | `history_limit` | int | `100` | How many finished jobs to retain in history. | | `compression_level` | int | `1` | gzip level for responses. | -| `services_refresh_interval` | int (s) | `0` | How often the server pulls the git-backed service repositories in `directories.services` and hot-reloads the services that changed. `0` turns it off. A service with a running or queued job, or a live remote session, swaps to the new code as soon as its last job and session finish. | +| `services_refresh_interval` | int (s) | `0` | How often the server pulls the git-backed service repositories in `directories.services` and hot-reloads the services that changed. `0` turns it off. A service with a running or queued job, or a live remote session, is staged and swaps to the new code as soon as its last job and session finish. | +| `services_staged_max_age` | int (s) | `0` | How long a staged service waits for its jobs and sessions before the server swaps it in anyway. `0` waits forever. Set it on a server whose popular services never have a quiet moment. A session that is live at the swap keeps the old code until it ends; a running job keeps its old code too. | | `global_speed_limit` | str | *(unlimited)* | Server-wide download speed cap, e.g. `10M`, `1.5G` or plain bytes/sec (same format as `speed_limit`). One shared budget across all concurrent jobs; the server ignores per-job speed limits while it is set. | | `cdm_overrides` | list or bool | *(unset)* | Allowed per-request CDM overrides: a list of permitted device names, or `true` for any. Unset rejects every override. | | `allow_job_credentials` | bool | `false` | Permit clients to supply credentials per job. | diff --git a/tests/core/test_services_reload.py b/tests/core/test_services_reload.py index f9096326..05c4b108 100644 --- a/tests/core/test_services_reload.py +++ b/tests/core/test_services_reload.py @@ -105,3 +105,19 @@ def test_apply_pending_never_reports_a_failed_reimport_as_applied( services.PENDING.add("FOO") assert services.apply_pending(busy=set()) == ["FOO"] assert services.LOADED_COMMITS["FOO"] == "new-head" + + +def test_overdue_staged_service_applies_while_busy(service_dir: Path, monkeypatch: pytest.MonkeyPatch): + write_service(service_dir, 2) + monkeypatch.setattr(services, "repo_specs", lambda: ["example/repo"]) + monkeypatch.setattr(services, "refresh_repo", lambda spec: (service_dir, ["~FOO"])) + monkeypatch.setattr(services, "staged_max_age", lambda: 60.0) + + services.refresh_and_reload(busy={"FOO"}) + assert services.apply_pending(busy={"FOO"}) == [] + assert services.MODULES["FOO"].VERSION == 1 + + services.PENDING_SINCE["FOO"] -= 61 + assert services.apply_pending(busy={"FOO"}) == ["FOO"] + assert services.MODULES["FOO"].VERSION == 2 + assert services.PENDING == set() and "FOO" not in services.PENDING_SINCE diff --git a/tests/remote/unit/test_serve_service_refresh.py b/tests/remote/unit/test_serve_service_refresh.py new file mode 100644 index 00000000..2941ff67 --- /dev/null +++ b/tests/remote/unit/test_serve_service_refresh.py @@ -0,0 +1,57 @@ +"""Pins the serve refresh loop: staged reloads apply after the pull, not only on the next tick.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from aiohttp import web + +from unshackle.commands.serve import _install_service_refresh +from unshackle.core import services +from unshackle.core.api import download_manager +from unshackle.core.config import config + +pytestmark = pytest.mark.unit + + +async def test_refresh_tick_applies_tags_staged_during_the_pull(monkeypatch: pytest.MonkeyPatch) -> None: + """A tag whose last session ended while the pull ran is staged although idle; the same tick swaps it in.""" + pulled = asyncio.Event() + seen: list[set[str]] = [] + + def fake_refresh(busy: set[str]) -> list[dict[str, Any]]: + # The pull read the busy set before the session ended, so the tag lands in PENDING while idle. + services.PENDING.add("EXAMPLE") + pulled.set() + return [] + + def fake_apply(busy: set[str]) -> list[str]: + seen.append(set(services.PENDING)) + ready = sorted(services.PENDING - busy) + services.PENDING.difference_update(ready) + return ready + + monkeypatch.setattr(config, "serve", {"services_refresh_interval": 1}) + monkeypatch.setattr(services, "PENDING", set()) + monkeypatch.setattr(services, "repo_specs", lambda: ["x"]) + monkeypatch.setattr(services, "log_load_issues", lambda: None) + monkeypatch.setattr(services, "record_loaded_commits", lambda: None) + monkeypatch.setattr(services, "refresh_and_reload", fake_refresh) + monkeypatch.setattr(services, "apply_pending", fake_apply) + monkeypatch.setattr(download_manager, "busy_services", set) + monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None) + monkeypatch.setattr(download_manager, "_reload_task", None) + + app = web.Application() + _install_service_refresh(app) + await app.on_startup[-1](app) + try: + await asyncio.wait_for(pulled.wait(), timeout=3) + await asyncio.sleep(0.2) + finally: + await app.on_cleanup[-1](app) + + assert {"EXAMPLE"} in seen + assert services.PENDING == set() diff --git a/tests/remote/unit/test_session_store.py b/tests/remote/unit/test_session_store.py index 32dc11a3..4de3a695 100644 --- a/tests/remote/unit/test_session_store.py +++ b/tests/remote/unit/test_session_store.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import time import pytest @@ -142,3 +143,123 @@ async def test_get_expiry_removes_cache_dir(store: SessionStore, monkeypatch: py entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100) assert await store.get(entry.session_id) is None assert removed == ["_sessions/k/s/EXAMPLE"] + + +async def test_session_end_applies_staged_service_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None: + """A staged reload swaps in as soon as the last session on that tag ends, not on the next refresh tick.""" + from unshackle.core import services + from unshackle.core.api import download_manager + + applied: list[set[str]] = [] + + def fake_apply(busy: set[str]) -> list[str]: + applied.append(busy) + return sorted(services.PENDING - busy) + + monkeypatch.setattr(services, "PENDING", {"EXAMPLE"}) + monkeypatch.setattr(services, "apply_pending", fake_apply) + monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None) + + entry = await store.create("EXAMPLE", _FakeService()) + assert await store.delete(entry.session_id) + await asyncio.sleep(0.05) + assert applied == [set()] + + +def _stage_reload(monkeypatch: pytest.MonkeyPatch, pending: set[str], delay: float = 0.0) -> list[set[str]]: + """Stage ``pending`` for reload and record every apply_pending call; returns the call log.""" + from unshackle.core import services + from unshackle.core.api import download_manager + + applied: list[set[str]] = [] + + def fake_apply(busy: set[str]) -> list[str]: + if delay: + time.sleep(delay) + applied.append(busy) + return sorted(services.PENDING - busy) + + monkeypatch.setattr(services, "PENDING", pending) + monkeypatch.setattr(services, "apply_pending", fake_apply) + monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None) + monkeypatch.setattr(download_manager, "_reload_task", None) + return applied + + +async def test_session_end_on_unstaged_tag_skips_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None: + applied = _stage_reload(monkeypatch, {"OTHER"}) + + entry = await store.create("EXAMPLE", _FakeService()) + assert await store.delete(entry.session_id) + await asyncio.sleep(0.05) + assert applied == [] + + +async def test_concurrent_session_ends_reload_once(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None: + """Session ends that overlap an in-flight reload share it instead of each re-importing the tag.""" + applied = _stage_reload(monkeypatch, {"EXAMPLE"}, delay=0.05) + + ids = [(await store.create("EXAMPLE", _FakeService())).session_id for _ in range(10)] + await asyncio.gather(*(store.delete(sid) for sid in ids)) + await asyncio.sleep(0.3) + assert len(applied) == 1 + + +async def test_eviction_applies_staged_service_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None: + applied = _stage_reload(monkeypatch, {"EXAMPLE"}) + monkeypatch.setattr(type(store), "max_sessions", property(lambda _: 1)) + + await store.create("EXAMPLE", _FakeService(), session_id="a") + await store.create("OTHER", _FakeService(), session_id="b") + await asyncio.sleep(0.05) + assert len(applied) == 1 + + +async def test_inline_expiry_applies_staged_service_reload( + store: SessionStore, monkeypatch: pytest.MonkeyPatch +) -> None: + from datetime import datetime, timedelta, timezone + + applied = _stage_reload(monkeypatch, {"EXAMPLE"}) + + entry = await store.create("EXAMPLE", _FakeService()) + entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100) + assert await store.get(entry.session_id) is None + await asyncio.sleep(0.05) + assert applied == [set()] + + +async def test_cleanup_expired_applies_staged_service_reload_once( + store: SessionStore, monkeypatch: pytest.MonkeyPatch +) -> None: + """A cleanup sweep that drops several sessions on one tag schedules one reload for that tag.""" + from datetime import datetime, timedelta, timezone + + applied = _stage_reload(monkeypatch, {"EXAMPLE"}) + stale = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100) + for sid in ("a", "b"): + (await store.create("EXAMPLE", _FakeService(), session_id=sid)).last_accessed = stale + + assert await store.cleanup_expired() == 2 + await asyncio.sleep(0.05) + assert applied == [set()] + + +async def test_live_session_on_same_tag_keeps_it_busy(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None: + """Ending one of two sessions on a staged tag reports the tag busy; ending the last one frees it.""" + from unshackle.core.api import download_manager + from unshackle.core.api import session_store as session_store_module + + applied = _stage_reload(monkeypatch, {"EXAMPLE"}) + monkeypatch.setattr(session_store_module, "session_store", store) + monkeypatch.setattr(download_manager, "download_manager", None) + + first = await store.create("EXAMPLE", _FakeService()) + second = await store.create("EXAMPLE", _FakeService()) + assert await store.delete(first.session_id) + await asyncio.sleep(0.05) + assert applied == [{"EXAMPLE"}] + + assert await store.delete(second.session_id) + await asyncio.sleep(0.05) + assert applied == [{"EXAMPLE"}, set()] diff --git a/unshackle/commands/serve.py b/unshackle/commands/serve.py index 64493877..4e8d8a57 100644 --- a/unshackle/commands/serve.py +++ b/unshackle/commands/serve.py @@ -52,7 +52,7 @@ def _install_service_refresh(app: web.Application) -> None: return async def loop() -> None: - from unshackle.core.api.download_manager import busy_services + from unshackle.core.api.download_manager import busy_services, schedule_pending_reload while True: await asyncio.sleep(interval) @@ -70,6 +70,9 @@ def _install_service_refresh(app: web.Application) -> None: for err in r["load_errors"]: log.error(f"Service reload failed: {err}") publish_refresh_events(repos) + # The busy set was read before the pull. A tag whose last session ended during the pull + # is staged although idle, and its session-end trigger ran before it was staged. + schedule_pending_reload() except Exception: log.exception("Service refresh failed") diff --git a/unshackle/core/api/download_manager.py b/unshackle/core/api/download_manager.py index 5990aa0b..97f55eb9 100644 --- a/unshackle/core/api/download_manager.py +++ b/unshackle/core/api/download_manager.py @@ -747,21 +747,7 @@ class DownloadQueueManager: def schedule_pending_reload(self) -> None: """Swap in staged service updates whose last job has finished, off the event loop.""" - from unshackle.core import services - - if not services.PENDING: - return - - async def run() -> None: - try: - applied = await asyncio.to_thread(services.apply_pending, busy_services()) - if applied: - log.info(f"Services reloaded after job completion: {', '.join(applied)}") - publish_service_event("applied", applied) - except Exception: - log.exception("Applying pending service reloads failed") - - asyncio.get_running_loop().create_task(run()) + schedule_pending_reload() @staticmethod def put(queue: asyncio.Queue, item: Optional[Dict[str, Any]]) -> None: @@ -1210,3 +1196,41 @@ def busy_services() -> set[str]: if stats.mode != "remote_only": busy |= get_download_manager().busy_services() return busy + + +_reload_task: Optional[asyncio.Task[None]] = None + + +def schedule_pending_reload(tag: Optional[str] = None) -> None: + """Swap in staged service updates whose last job or remote session has ended, off the event loop. + + Runs on every job completion and remote session removal, because a busy service on a popular + tag can otherwise stay staged across every periodic refresh tick. With ``tag``, the reload only + runs when ``PENDING`` holds that tag. + """ + global _reload_task + from unshackle.core import services + + if not services.PENDING or (tag is not None and tag not in services.PENDING): + return + # apply_pending computes its ready set before it takes RELOAD_LOCK, so concurrent + # session ends would each re-import the same tag; one in-flight task covers them all. + if _reload_task is not None and not _reload_task.done(): + return + + async def run() -> None: + try: + applied = await asyncio.to_thread(services.apply_pending, busy_services()) + if applied: + log.info(f"Services reloaded after their jobs and sessions ended: {', '.join(applied)}") + publish_service_event("applied", applied) + except Exception: + log.exception("Applying pending service reloads failed") + + _reload_task = asyncio.get_running_loop().create_task(run()) + + def clear(_: asyncio.Task[None]) -> None: + global _reload_task + _reload_task = None + + _reload_task.add_done_callback(clear) diff --git a/unshackle/core/api/session_store.py b/unshackle/core/api/session_store.py index efde333b..df8fb575 100644 --- a/unshackle/core/api/session_store.py +++ b/unshackle/core/api/session_store.py @@ -95,6 +95,13 @@ class SessionEntry: } +def _schedule_pending_reload(service_tag: str) -> None: + """Apply staged service reloads now that a remote session ended; the task runs after the store lock releases.""" + from unshackle.core.api.download_manager import schedule_pending_reload + + schedule_pending_reload(service_tag) + + def _publish(action: str, entry: SessionEntry, reason: Optional[str] = None) -> None: data = {"action": action, **entry.summary()} if reason: @@ -145,6 +152,7 @@ class SessionStore: evicted.input_bridge.cancel() self.cleanup_cache_dir(evicted.cache_tag) _publish("delete", evicted, "evicted") + _schedule_pending_reload(evicted.service_tag) session_id = session_id or str(uuid.uuid4()) entry = SessionEntry( @@ -173,6 +181,7 @@ class SessionStore: log.info(f"Session {sanitize_log(session_id)} expired (elapsed={elapsed:.0f}s, ttl={self.ttl}s)") self.cleanup_cache_dir(entry.cache_tag) _publish("delete", self._sessions.pop(session_id), "expired") + _schedule_pending_reload(entry.service_tag) return None entry.touch() @@ -196,6 +205,7 @@ class SessionStore: self.cleanup_cache_dir(entry.cache_tag) _publish("delete", entry, "closed") log.info(f"Deleted session {sanitize_log(session_id)}") + _schedule_pending_reload(entry.service_tag) return True return False @@ -211,12 +221,16 @@ class SessionStore: expired.append(sid) elif elapsed > self.ttl: expired.append(sid) + ended_tags: set[str] = set() for sid in expired: entry = self._sessions.pop(sid) if entry.input_bridge: entry.input_bridge.cancel() self.cleanup_cache_dir(entry.cache_tag) _publish("delete", entry, "expired") + ended_tags.add(entry.service_tag) + for tag in ended_tags: + _schedule_pending_reload(tag) if expired: log.info(f"Cleaned up {len(expired)} expired sessions") return len(expired) diff --git a/unshackle/core/services.py b/unshackle/core/services.py index 846c2e78..2c935c3d 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -272,13 +272,24 @@ def failed_tags(errors: Iterable[str]) -> set[str]: return {err.split(":", 1)[0] for err in errors if ":" in err} +def staged_max_age() -> float: + """Seconds a staged tag may wait for its jobs and sessions before it swaps in anyway; 0 waits forever.""" + return float((config.serve or {}).get("services_staged_max_age", 0) or 0) + + def apply_pending(busy: set[str] | None = None) -> list[str]: """Reload every staged tag whose service is no longer busy; returns the tags now running. - A tag whose re-import failed is not one of them: the old module keeps serving, so calling - it applied would tell the dashboard that code is live when it is not. + A tag staged for longer than ``serve.services_staged_max_age`` reloads even while busy: + its live sessions keep the old class objects they already hold, and a popular tag would + otherwise never see a quiet moment. A tag whose re-import failed is not counted as + running: the old module keeps serving, so calling it applied would tell the dashboard + that code is live when it is not. """ - ready = sorted(PENDING - (busy or set())) + max_age = staged_max_age() + now = time.time() + overdue = {tag for tag, since in PENDING_SINCE.items() if max_age and now - since >= max_age} + ready = sorted((PENDING - (busy or set())) | (PENDING & overdue)) if not ready: return [] errors = reload_services(ready)