Drain msgd work before shutdown and stop new effect claims

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

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

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

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.

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:

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*.

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)

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

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