diff --git a/backend/app/sync/coordinator.py b/backend/app/sync/coordinator.py index 837188d..b55413b 100644 --- a/backend/app/sync/coordinator.py +++ b/backend/app/sync/coordinator.py @@ -202,7 +202,21 @@ class SyncCoordinator: await self._release(account_id) async def _heartbeat_once(self) -> None: - payload = json.dumps({"owned": len(self._runners), "ts": time.time()}) + """Publish the worker's load for peer rebalancing + observability (§2).""" + tiers = {"hot": 0, "warm": 0, "cold": 0} + sockets_open = 0 + for runner in self._runners.values(): + tiers[runner.tier] = tiers.get(runner.tier, 0) + 1 + if runner.tier == "hot": + sockets_open += 1 # one persistent IDLE connection per hot account + payload = json.dumps( + { + "owned": len(self._runners), + "sockets_open": sockets_open, + "tiers": tiers, + "ts": time.time(), + } + ) await self._redis.set(leases.worker_key(self.worker_id), payload, ex=self.heartbeat_ttl) async def _alive_worker_count(self) -> int: diff --git a/backend/app/sync/supervisor.py b/backend/app/sync/supervisor.py deleted file mode 100644 index 4bb5ec8..0000000 --- a/backend/app/sync/supervisor.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Live-sync entrypoint (refactor §2, step 3). - -The sync worker's ``on_startup`` starts a :class:`SyncCoordinator`, which claims -accounts from the ``sync:live`` registry via Redis leases and runs one -:class:`~app.sync.runner.AccountRunner` per owned account. ``on_shutdown`` stops -it, releasing every lease so peers can take the accounts over immediately. -""" - -from __future__ import annotations - -import logging -from typing import Any - -from app.sync.coordinator import SyncCoordinator - -logger = logging.getLogger(__name__) - - -async def start_supervisor(ctx: dict[str, Any]) -> None: - coordinator = SyncCoordinator() - await coordinator.start() - ctx["coordinator"] = coordinator - - -async def stop_supervisor(ctx: dict[str, Any]) -> None: - coordinator: SyncCoordinator | None = ctx.get("coordinator") - if coordinator is not None: - await coordinator.stop() - ctx["coordinator"] = None diff --git a/backend/tests/sync/test_acceptance.py b/backend/tests/sync/test_acceptance.py index a661c9b..993c9b8 100644 --- a/backend/tests/sync/test_acceptance.py +++ b/backend/tests/sync/test_acceptance.py @@ -297,6 +297,35 @@ async def test_event_driven_teardown(connections, make_accounts) -> None: # typ await c.stop() +# --- acceptance 6: graceful shutdown ----------------------------------------- + + +async def test_acceptance_6_graceful_shutdown(connections, make_accounts) -> None: # type: ignore[no-untyped-def] + """A graceful stop releases every lease so a peer adopts the accounts fast. + + Leases are long-lived, so adoption within a couple of seconds proves the + accounts were handed over by explicit release — not by waiting out the TTL. + """ + ids = set(await make_accounts(6)) + await _seed(ids) # hot => persistent IDLE + c1 = _fast(SyncCoordinator("w1")) + c2 = _fast(SyncCoordinator("w2")) + c1.lease_ttl = c2.lease_ttl = 30 # adoption must come from release, not expiry + await c1.start() + await _until(lambda: c1.owned == ids) # sole owner + await c2.start() # nothing to claim yet — all leased by w1 + try: + await c1.stop() # graceful: stops claimer, cancels runners, releases leases + assert c1.owned == set() + assert await get_redis().exists(leases.worker_key("w1")) == 0 # heartbeat gone + # w2 adopts every account well within the 30s TTL, one connection each. + await _until(lambda: c2.owned == ids, timeout_s=4.0) + await _until(lambda: all(connections[a] == 1 for a in ids)) + finally: + await c2.stop() + await _until(lambda: all(connections[a] == 0 for a in ids)) + + # --- §5 activity tiers ------------------------------------------------------- diff --git a/backend/workers/run.py b/backend/workers/run.py index 2eac061..9d4a4d5 100644 --- a/backend/workers/run.py +++ b/backend/workers/run.py @@ -1,12 +1,16 @@ -"""Worker entrypoint (refactor §1). +"""Worker entrypoint (refactor §1, §7). One image, two roles selected by ``WORKER_ROLE``: - ``jobs`` — arq consumer for the ``interactive`` and ``bulk`` queues (both in one process, per-queue concurrency/timeout). -- ``sync`` — runs the account-sync supervisor (long-lived per-account loops). +- ``sync`` — runs the lease-owned account-sync coordinator (§2) directly. - ``all`` — both, for single-process development (the default). +Graceful shutdown (SIGINT/SIGTERM): stop the coordinator first — it stops the +claimer, cancels the runners, and releases every lease so a peer adopts the +accounts immediately — then drain the arq consumers. Target < 10 s. + Run with: ``python -m workers.run``. """ @@ -21,77 +25,61 @@ from arq.worker import Worker from app.config import get_settings from app.queue import QUEUE_BULK, QUEUE_INTERACTIVE, redis_settings -from workers.worker import ( - BULK_FUNCTIONS, - INTERACTIVE_FUNCTIONS, - sync_on_shutdown, - sync_on_startup, -) +from app.sync.coordinator import SyncCoordinator +from workers.worker import BULK_FUNCTIONS, INTERACTIVE_FUNCTIONS logger = logging.getLogger(__name__) -# arq requires a non-empty function set; the sync worker consumes a queue nothing -# is ever enqueued to, so it just keeps the event loop alive for the supervisor. -_SYNC_CONTROL_QUEUE = "sync:control" +_SHUTDOWN_DRAIN_S = 10.0 # §7: release leases + drain within this budget -async def _noop(ctx: dict[str, Any]) -> None: - return None - - -def _worker( - functions: list[Any], - queue_name: str, - max_jobs: int, - job_timeout: int, - on_startup: Any = None, - on_shutdown: Any = None, -) -> Worker: +def _worker(functions: list[Any], queue_name: str, max_jobs: int, job_timeout: int) -> Worker: return Worker( functions=functions, queue_name=queue_name, redis_settings=redis_settings(), max_jobs=max_jobs, job_timeout=job_timeout, - on_startup=on_startup, - on_shutdown=on_shutdown, handle_signals=False, # the runner owns signal handling (graceful shutdown) ) -def build_workers(role: str) -> list[Worker]: - """Construct the arq Worker instances for the given role.""" +def build_arq_workers(role: str) -> list[Worker]: + """Construct the arq Worker instances for the given role (jobs only).""" s = get_settings() - workers: list[Worker] = [] - if role in ("sync", "all"): - workers.append( - _worker([_noop], _SYNC_CONTROL_QUEUE, 1, 3600, sync_on_startup, sync_on_shutdown) - ) - if role in ("jobs", "all"): - workers.append( - _worker(INTERACTIVE_FUNCTIONS, QUEUE_INTERACTIVE, s.interactive_concurrency, 120) - ) - workers.append(_worker(BULK_FUNCTIONS, QUEUE_BULK, s.bulk_concurrency, 3600)) - return workers + if role not in ("jobs", "all"): + return [] + return [ + _worker(INTERACTIVE_FUNCTIONS, QUEUE_INTERACTIVE, s.interactive_concurrency, 120), + _worker(BULK_FUNCTIONS, QUEUE_BULK, s.bulk_concurrency, 3600), + ] async def _run() -> None: logging.basicConfig(level=get_settings().log_level) role = get_settings().worker_role - workers = build_workers(role) - logger.info("worker starting: role=%s (%d arq worker(s))", role, len(workers)) - tasks = [asyncio.create_task(w.async_run()) for w in workers] + coordinator = SyncCoordinator() if role in ("sync", "all") else None + if coordinator is not None: + await coordinator.start() + arq_tasks = [asyncio.create_task(w.async_run()) for w in build_arq_workers(role)] + logger.info("worker starting: role=%s (%d arq worker(s))", role, len(arq_tasks)) + stop = asyncio.Event() loop = asyncio.get_running_loop() for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, stop.set) - await stop.wait() - logger.info("worker stopping: cancelling %d worker(s)", len(tasks)) - for task in tasks: - task.cancel() # arq runs each worker's on_shutdown in its finally - await asyncio.gather(*tasks, return_exceptions=True) + + logger.info("worker stopping: releasing leases + draining") + if coordinator is not None: + try: + await asyncio.wait_for(coordinator.stop(), timeout=_SHUTDOWN_DRAIN_S) + except TimeoutError: + logger.warning("coordinator stop exceeded %.0fs drain budget", _SHUTDOWN_DRAIN_S) + for task in arq_tasks: + task.cancel() + await asyncio.gather(*arq_tasks, return_exceptions=True) from app.db import dispose_engine from app.queue import close_arq_pool diff --git a/backend/workers/worker.py b/backend/workers/worker.py index 59d2ec9..c354964 100644 --- a/backend/workers/worker.py +++ b/backend/workers/worker.py @@ -1,8 +1,9 @@ -"""Worker job functions and the sync-worker lifecycle (refactor §1, §6). +"""Worker job callables and queue function-sets (refactor §1, §6). The process is launched via ``workers.run`` (``WORKER_ROLE`` decides which arq workers run). This module only defines the job callables and the two queue function-sets; the runner wires them onto the ``interactive`` / ``bulk`` queues. +The sync role runs the coordinator directly (see ``workers.run``), not via arq. """ from __future__ import annotations @@ -46,19 +47,3 @@ INTERACTIVE_FUNCTIONS = [ # Bulk work (§6): long timeout, low concurrency, isolated from interactive jobs. BULK_FUNCTIONS = [sync_bootstrap_account] - - -async def sync_on_startup(ctx: dict[str, Any]) -> None: - """Sync-worker startup: launch the account-sync supervisor.""" - from app.sync.supervisor import start_supervisor - - await start_supervisor(ctx) - logger.info("sync worker: supervisor started") - - -async def sync_on_shutdown(ctx: dict[str, Any]) -> None: - """Sync-worker shutdown: stop the supervisor (process cleanup is in the runner).""" - from app.sync.supervisor import stop_supervisor - - await stop_supervisor(ctx) - logger.info("sync worker: supervisor stopped")