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.
2026-11-14 13:05:09 +00:00 · 2026-09-22 10:40:49 -06:00
parent 6e681af513
commit d239964f0b
8 changed files with 267 additions and 20 deletions
@@ -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()
@@ -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()]