diff --git a/deploy/msgd-hosting.service b/deploy/msgd-hosting.service index 0b81a887..b11d3717 100644 --- a/deploy/msgd-hosting.service +++ b/deploy/msgd-hosting.service @@ -11,6 +11,10 @@ WorkingDirectory=/ ExecStart=/usr/bin/msgd --config-dir /etc/msgd-hosting hosting --listen 127.0.0.1 --port 8043 Environment=PGSERVICEFILE=/etc/msgd-hosting/pg_service.conf Environment=PGPASSFILE=/etc/msgd-hosting/pgpass +# Signal the parent first so in-flight work can drain before killing children. +KillSignal=SIGTERM +KillMode=mixed +TimeoutStopSec=90 Restart=on-failure UMask=0077 NoNewPrivileges=true diff --git a/deploy/msgd-worker.service b/deploy/msgd-worker.service index 181aa52b..4c655c5e 100644 --- a/deploy/msgd-worker.service +++ b/deploy/msgd-worker.service @@ -8,6 +8,10 @@ User=msgd Group=msgd WorkingDirectory=/var/lib/msgd ExecStart=/usr/bin/msgd worker +# Signal the parent first so in-flight work can drain before killing children. +KillSignal=SIGTERM +KillMode=mixed +TimeoutStopSec=90 Restart=on-failure RestartSec=3 UMask=0077 diff --git a/deploy/msgd.service b/deploy/msgd.service index 26690709..066e4860 100644 --- a/deploy/msgd.service +++ b/deploy/msgd.service @@ -8,6 +8,10 @@ User=msgd Group=msgd WorkingDirectory=/var/lib/msgd ExecStart=/usr/bin/msgd serve +# Signal the parent first so in-flight work can drain before killing children. +KillSignal=SIGTERM +KillMode=mixed +TimeoutStopSec=90 Restart=on-failure RestartSec=3 UMask=0077 diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index 1019aa21..5819b57a 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -119,3 +119,19 @@ msgd --config-dir /new/etc/msgd restore /secure-backup/service.zip --data-dir /n 旧服务、解释器、数据和配置先保留。旧写入冻结后完成最终保护快照和增量核对；新 Root 的身份映射及旧内容归属按[导入流程](NEW_ROOT_LEGACY_IMPORT_RUNBOOK.md)核验，不能把旧权限字段当成新授权。只有本机 Root/基础 CA、依赖、隔离、日志链和实际功能验收通过才切代理 upstream；`nginx -t` 成功后 reload，并记录真实外部回读、服务版本和配置摘要，再启用已验收的后台服务。 切流后尚无新业务写入时，可停止新服务并恢复旧 upstream/服务。已有新写入时，先冻结新 writer/worker并保护新 PostgreSQL/Git/CAS，核对数据差异后执行已验证的数据回迁或维持维护状态，不能直接退回旧 SQLite 丢弃新写入。新旧 Root 不同，回滚不会让新凭据自动被旧版接受。 + +### Graceful service stop + +`msgd serve` and `msgd hosting` stop accepting connections on SIGTERM and +allow in-flight HTTP work up to 60 seconds before cancellation. Their lifespan +cleanup closes storage after draining. The worker stops starting maintenance +and claiming effects, allows its current stage up to 60 seconds, then cancels +it and closes storage and Valkey clients. Cancellation does not acknowledge +an effect: a running lease remains durable and expires to `uncertain`, since +an external side effect may already have happened. + +The three service units use `KillMode=mixed`: SIGTERM goes to the main process +first so sandbox children are not terminated ahead of their parent. +`TimeoutStopSec=90` bounds the entire shutdown, including cleanup; after this +budget systemd kills remaining processes. This is a bounded drain, not an +exactly-once guarantee for external effects. diff --git a/src/msg/daemon.py b/src/msg/daemon.py index b5da35c7..feae80ed 100644 --- a/src/msg/daemon.py +++ b/src/msg/daemon.py @@ -37,27 +37,37 @@ def load_application(directory): return Application(load_settings(directory)) +SHUTDOWN_GRACE_SECONDS = 60 + + async def worker_loop(app, *, once=False): from msg.storage.valkey_bus import ValkeyOutboxSignal from msg.workers.effects import EffectWorker from msg.workers.maintenance import run_maintenance - await app.load() - worker = EffectWorker(app) - wakeup = ( - ValkeyOutboxSignal(app.settings.server.valkey_url) - if app.settings.server.valkey_url - else None - ) stop = asyncio.Event() loop = asyncio.get_running_loop() + installed_signals = [] for sig in (signal.SIGINT, signal.SIGTERM): try: loop.add_signal_handler(sig, stop.set) + installed_signals.append(sig) except NotImplementedError, RuntimeError: pass - next_cleanup = 0.0 - try: + wakeup = None + work = stopped = None + + async def run(): + nonlocal wakeup + await app.load() + worker = EffectWorker(app) + worker.stopping = stop.is_set + wakeup = ( + ValkeyOutboxSignal(app.settings.server.valkey_url) + if app.settings.server.valkey_url + else None + ) + next_cleanup = 0.0 while not stop.is_set(): executor = getattr(app, 'executor', None) if executor is not None and executor.recovery_drill_active(): @@ -69,30 +79,61 @@ async def worker_loop(app, *, once=False): pass continue if loop.time() >= next_cleanup: - await run_maintenance(app, 'cleanup_expired', scheduled=True) - await run_maintenance(app, 'collect_garbage', scheduled=True) - await run_maintenance(app, 'deliver_due_todos', scheduled=True) + for action in ('cleanup_expired', 'collect_garbage', 'deliver_due_todos'): + if stop.is_set(): + return None + await run_maintenance(app, action, scheduled=True) next_cleanup = loop.time() + 60 + # A signal during maintenance must not lead to claiming another job. + if stop.is_set(): + return None processed = await worker.run_once() if once: return {'processed': processed} - if not processed: + if not processed and not stop.is_set(): if wakeup is not None: try: await wakeup.wait_for_pending(1) + continue except Exception: # Pub/Sub is only a hint; durable jobs remain in PostgreSQL. - try: - await asyncio.wait_for(stop.wait(), 1) - except TimeoutError: - pass - else: - try: - await asyncio.wait_for(stop.wait(), 1) - except TimeoutError: pass + try: + await asyncio.wait_for(stop.wait(), 1) + except TimeoutError: + pass + return None + + try: + work = asyncio.create_task(run()) + stopped = asyncio.create_task(stop.wait()) + done, _ = await asyncio.wait((work, stopped), return_when=asyncio.FIRST_COMPLETED) + if work in done: + return await work + try: + # Finish the current stage, never start a new claim. Cancellation on + # expiry preserves the running lease; it is not an acknowledgement. + return await asyncio.wait_for(asyncio.shield(work), SHUTDOWN_GRACE_SECONDS) + except TimeoutError: + work.cancel() + await asyncio.gather(work, return_exceptions=True) + return None finally: - await app.close() + for task in (work, stopped): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather( + *(task for task in (work, stopped) if task is not None), return_exceptions=True + ) + try: + if wakeup is not None: + await wakeup.close() + finally: + try: + await app.close() + finally: + for sig in installed_signals: + loop.remove_signal_handler(sig) def parser(): @@ -415,7 +456,14 @@ async def lifespan(asgi): asgi = hosting_app(app) asgi.router.lifespan_context = lifespan - uvicorn.run(asgi, host=args.listen, port=args.port, access_log=False, ws='none') + uvicorn.run( + asgi, + host=args.listen, + port=args.port, + access_log=False, + ws='none', + timeout_graceful_shutdown=SHUTDOWN_GRACE_SECONDS, + ) return 0 app = load_application(args.config_dir) if args.command == 'backup': @@ -446,6 +494,7 @@ async def save(): port=app.settings.port, access_log=False, ws='none', + timeout_graceful_shutdown=SHUTDOWN_GRACE_SECONDS, ) return 0 except (Failure, OSError, ValueError) as exc: diff --git a/src/msg/storage/postgres.py b/src/msg/storage/postgres.py index e865160b..c94b934d 100644 --- a/src/msg/storage/postgres.py +++ b/src/msg/storage/postgres.py @@ -674,7 +674,9 @@ def blocking(): ) async def close(self): - return None + close = getattr(self.signal, 'close', None) + if close is not None: + await close() def backup(self, destination: Path): """Write a pg_dump custom-format archive at *destination*. diff --git a/src/msg/storage/valkey_bus.py b/src/msg/storage/valkey_bus.py index b2a3f359..5c6666f6 100644 --- a/src/msg/storage/valkey_bus.py +++ b/src/msg/storage/valkey_bus.py @@ -43,3 +43,6 @@ def _wait_for_pending(self, timeout: float) -> bool: if message and message.get('type') == 'message': return True return False + + async def close(self) -> None: + await asyncio.to_thread(self.client.close) diff --git a/src/msg/transports/http_routes.py b/src/msg/transports/http_routes.py index df0bbd10..0a353d8f 100644 --- a/src/msg/transports/http_routes.py +++ b/src/msg/transports/http_routes.py @@ -736,10 +736,12 @@ def create_app(service): @asynccontextmanager async def lifespan(app): - if not service._loaded: - await service.load() - yield - await service.close() + try: + if not service._loaded: + await service.load() + yield + finally: + await service.close() async def dispatch(request: Request): nonlocal graphql_adapter, short_codes diff --git a/src/msg/workers/effects.py b/src/msg/workers/effects.py index 4218da31..e19ea1fd 100644 --- a/src/msg/workers/effects.py +++ b/src/msg/workers/effects.py @@ -122,6 +122,7 @@ def __init__( lease_seconds=120, ): self.app = app + self.stopping = lambda: False if tool_runner is None: from msg.workers.sandbox import BubblewrapRunner @@ -142,6 +143,8 @@ async def _claim(self): from msg.extensions.tools import tool_concurrency async with self.app.metadata.transaction(write=True) as tx: + if self.stopping(): + return None, False self.app.runtime_generation.require_current(tx) if quarantine_active(tx): return None, False @@ -180,6 +183,8 @@ async def _claim(self): return job, False if running.get(ref.id, 0) >= limits[ref.id, ref.revision]: continue + if self.stopping(): + return None, False job = replace( job, state='running', @@ -540,6 +545,8 @@ async def _webhook(self, job): await self._finish(job, 'done' if state == 'delivered' else 'failed', state) async def run_once(self): + if self.stopping(): + return False await self.app.executor.require_current_runtime() if self.app.executor.recovery_drill_active(): return False diff --git a/tests/test_graceful_shutdown.py b/tests/test_graceful_shutdown.py new file mode 100644 index 00000000..941f4a40 --- /dev/null +++ b/tests/test_graceful_shutdown.py @@ -0,0 +1,256 @@ +"""Exercise actual SIGTERM delivery while a worker stage is in flight.""" + +import asyncio +import os +import signal +from types import SimpleNamespace + +import pytest + +from msg import daemon +from msg.workers import effects, maintenance + + +@pytest.mark.asyncio +@pytest.mark.parametrize('stage', ['load', 'maintenance', 'job']) +async def test_sigterm_drains_current_stage_without_new_claim(monkeypatch, stage): + entered = asyncio.Event() + release = asyncio.Event() + completed = [] + claims = [] + + async def pause(name): + if name == stage: + entered.set() + await release.wait() + completed.append(name) + + class App: + settings = SimpleNamespace(server=SimpleNamespace(valkey_url=None)) + closed = False + + async def load(self): + await pause('load') + + async def close(self): + self.closed = True + + class Worker: + def __init__(self, app): + pass + + async def run_once(self): + claims.append('claim') + await pause('job') + return True + + async def maintain(app, action, **kwargs): + await pause('maintenance') + + monkeypatch.setattr(effects, 'EffectWorker', Worker) + monkeypatch.setattr(maintenance, 'run_maintenance', maintain) + app = App() + task = asyncio.create_task(daemon.worker_loop(app)) + try: + await asyncio.wait_for(entered.wait(), 2) + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.sleep(0.02) + assert not task.done() + assert not app.closed + release.set() + await asyncio.wait_for(task, 2) + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert app.closed + assert completed.count(stage) == 1 + assert claims == (['claim'] if stage == 'job' else []) + + +@pytest.mark.asyncio +async def test_sigterm_budget_cancels_job_without_reporting_completion(monkeypatch): + entered = asyncio.Event() + states = [] + + class App: + settings = SimpleNamespace(server=SimpleNamespace(valkey_url=None)) + closed = False + + async def load(self): + pass + + async def close(self): + self.closed = True + + class Worker: + def __init__(self, app): + pass + + async def run_once(self): + states.append('running') + entered.set() + try: + await asyncio.Event().wait() + states.append('done') + except asyncio.CancelledError: + states.append('cancelled') + raise + + async def maintain(*args, **kwargs): + pass + + monkeypatch.setattr(effects, 'EffectWorker', Worker) + monkeypatch.setattr(maintenance, 'run_maintenance', maintain) + monkeypatch.setattr(daemon, 'SHUTDOWN_GRACE_SECONDS', 0.02) + app = App() + task = asyncio.create_task(daemon.worker_loop(app)) + await asyncio.wait_for(entered.wait(), 2) + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.wait_for(task, 2) + assert states == ['running', 'cancelled'] + assert app.closed + + +@pytest.mark.asyncio +async def test_cancelled_effect_keeps_lease_and_expires_uncertain(installed, monkeypatch): + from dataclasses import replace + from datetime import timedelta + + from msg.core.models import EffectJob, Principal + + app, _ = installed + entered = asyncio.Event() + now = app.clock() + job = EffectJob( + id='job_shutdown_lease', + event_id='e_shutdown_lease', + kind='maintenance', + dedupe_key='shutdown-lease', + principal=Principal( + actor=None, + subject=None, + credential_id=None, + method='anonymous', + certificates=(), + ceiling=(), + ), + operation='system.maintenance', + arguments={'action': 'rebuild_search'}, + state='pending', + attempts=0, + next_attempt_at=now, + lease_until=None, + ) + async with app.metadata.transaction(write=True) as tx: + await tx.enqueue(job) + + async def maintain(app, action, **kwargs): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(maintenance, 'run_maintenance', maintain) + worker = effects.EffectWorker(app) + task = asyncio.create_task(worker.run_once()) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + async with app.metadata.transaction(write=True) as tx: + current = await tx.job(job.id) + assert current.state == 'running' + assert current.attempts == 1 and current.lease_until > now + assert tx.setting('job_status:' + job.id) is None + await tx.save_job(replace(current, lease_until=now - timedelta(seconds=1))) + claimed, execute = await worker._claim() + assert claimed.id == job.id and not execute + async with app.metadata.transaction(write=False) as tx: + current = await tx.job(job.id) + assert current.state == 'uncertain' + assert current.attempts == 1 + assert tx.setting('job_status:' + job.id) == {'code': 'expired_execution_lease'} + + +@pytest.mark.asyncio +async def test_stop_prevents_pending_effect_claim(installed): + app, _ = installed + worker = effects.EffectWorker(app) + worker.stopping = lambda: True + assert await worker._claim() == (None, False) + assert not await worker.run_once() + + +@pytest.mark.asyncio +async def test_http_sigterm_drains_response_before_storage_close(installed, monkeypatch): + import socket + from contextlib import nullcontext + + import httpx + import uvicorn + from starlette.responses import PlainTextResponse + from starlette.routing import Route + + from msg.transports.http import create_app + + app, _ = installed + entered = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + original_close = app.close + + async def close(): + await original_close() + closed.set() + + async def slow(request): + entered.set() + await release.wait() + assert not closed.is_set() + return PlainTextResponse('finished') + + monkeypatch.setattr(app, 'close', close) + asgi = create_app(app) + asgi.routes.insert(0, Route('/shutdown-test', slow)) + server = uvicorn.Server( + uvicorn.Config( + asgi, + access_log=False, + ws='none', + log_level='error', + timeout_graceful_shutdown=daemon.SHUTDOWN_GRACE_SECONDS, + ) + ) + # Use asyncio's OS signal dispatch rather than Uvicorn's post-run re-raise, + # which would terminate the pytest process after a successful shutdown. + monkeypatch.setattr(server, 'capture_signals', nullcontext) + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGTERM, server.handle_exit, signal.SIGTERM, None) + listener = socket.socket() + listener.bind(('127.0.0.1', 0)) + listener.listen() + port = listener.getsockname()[1] + serving = asyncio.create_task(server.serve(sockets=[listener])) + try: + async with httpx.AsyncClient(timeout=3) as client: + for _ in range(100): + if server.started: + break + await asyncio.sleep(0.01) + response = asyncio.create_task(client.get(f'http://127.0.0.1:{port}/shutdown-test')) + await asyncio.wait_for(entered.wait(), 2) + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.sleep(0.25) + assert not closed.is_set() + with pytest.raises(httpx.ConnectError): + await client.get(f'http://127.0.0.1:{port}/shutdown-test') + release.set() + assert (await response).text == 'finished' + await asyncio.wait_for(serving, 2) + assert closed.is_set() + finally: + release.set() + server.should_exit = True + await asyncio.wait_for(serving, 3) + listener.close() + loop.remove_signal_handler(signal.SIGTERM)