fix(restore): a restore cut off by shutdown waits for its worker thread (#360)

A cancelled restore handler or job no longer removes the stage or the
upload, or ends its exclusive registry admission, while the executor
thread still stages or publishes. The cancel waits for the worker; a
publish that finished is adopted before the cancel goes on, and a stage
finished after the cancel is removed.

py/api/backup_restore.py

@@ -483,13 +483,42 @@ class BackupRestoreRoutesMixin:
             self.logger.info(f"Removed {len(removed)} old safety backups")
 
     async def _publish(self, model, staged):
-        """Publish ``staged``; on a page-copy failure prune and re-raise."""
+        """Publish ``staged``; on a page-copy failure prune and re-raise.
+
+        A cancel (server stop) waits for the publishing thread (#360): the
+        caller's cleanup and its registry admission must outlive it. When the
+        copy then succeeded the library was replaced, so it is adopted
+        (``_after_publish``) before the cancel goes on.
+        """
+        loop = asyncio.get_running_loop()
+        future = loop.run_in_executor(None, model.publish_staged, staged)
+        cancelled = await _outlive_cancel(future)
         try:
-            return await self._run_in_executor(model.publish_staged, staged)
+            backup_path = future.result()
         except sqlite3.OperationalError as exc:
             if getattr(exc, "_prompt_manager_restore_backup_created", False) is True:
                 await self._prune_after(model.db_path, failed=True)
             raise
+        if cancelled:
+            self.logger.info("The restore finished while the server was stopping")
+            await self._after_publish(model.db_path)
+            raise asyncio.CancelledError
+        return backup_path
+
+    async def _stage_to_end(self, model, temp_path):
+        """Stage ``temp_path``; a cancel waits for the worker (#360).
+
+        The upload stays until the staging thread is done with it, and a
+        stage it finished after the cancel is removed before the cancel goes
+        on.
+        """
+        loop = asyncio.get_running_loop()
+        future = loop.run_in_executor(None, model.stage_restore, temp_path)
+        if await _outlive_cancel(future):
+            if not future.cancelled() and future.exception() is None:
+                await self._discard_held([_Orphan(future.result())])
+            raise asyncio.CancelledError
+        return future.result()
 
     async def _after_publish(self, db_path):
         """Everything after the copy replaced the library: best effort, logged.
@@ -609,7 +638,7 @@ class BackupRestoreRoutesMixin:
                     {"success": False, "error": REQUESTS_BUSY}, status=409
                 )
             model = self.db.model
-            staged = await self._run_in_executor(model.stage_restore, temp_path)
+            staged = await self._stage_to_end(model, temp_path)
             try:
                 backup_path = await self._publish(model, staged)
             finally:
@@ -927,6 +956,29 @@ class BackupRestoreRoutesMixin:
         job.update(progress=100, status_line=line)
 
 
+async def _outlive_cancel(future):
+    """Await an executor ``future`` to its end, even through cancels (#360).
+
+    A cancelled task cannot stop the worker thread, so the files it uses and
+    the registry admission it runs under must outlive it: the cancel is held
+    until the worker returned. Returns whether the task was cancelled
+    meanwhile; the caller reads the outcome from ``future`` and re-raises
+    ``CancelledError``. A future cancelled itself (an executor shut down) is
+    not a cancel of the task: ``future.result()`` raises it. Never catches
+    ``GeneratorExit``: a closed loop still closes the coroutine.
+    """
+    cancelled = False
+    while not future.done():
+        try:
+            await asyncio.shield(future)
+        except asyncio.CancelledError:
+            if not future.cancelled():
+                cancelled = True
+        except Exception:
+            pass  # the caller reads it from the future
+    return cancelled
+
+
 class _Orphan:
     """A stage with no store entry, in the shape ``_discard_held`` takes."""
 

tests/test_restore_shutdown.py

@@ -0,0 +1,180 @@
+"""A restore cut off by a server stop waits for its worker thread (#360).
+
+A cancelled handler (or job) task cannot stop the executor thread that
+stages or publishes. Its cleanup (the stage, the upload) and the registry's
+``restore`` admission must outlive that thread, and the live library ends as
+the old or the restored one, never partial, with the safety backup kept.
+
+Every positive wait is on an event. The one bounded wait (``SETTLE``) is a
+negative check: the cancelled task must still be pending while its worker is
+held.
+"""
+
+import asyncio
+import os
+import shutil
+import threading
+from unittest.mock import patch
+
+from database import restore_stage
+from database.models import PromptModel
+from py.api.jobs import RESTORE_JOB_KIND
+
+try:
+    from tests.test_restore_plans import RestoreTestCase, _texts
+except ImportError:  # discovered with tests/ as the top-level directory
+    from test_restore_plans import RestoreTestCase, _texts
+
+WAIT = 60  # a ceiling only a hung test reaches
+SETTLE = 0.3  # the negative check's bound: the old code's cleanup takes ms
+
+
+class _Gate:
+    """Hold one blocking call on its worker thread until released."""
+
+    def __init__(self, real):
+        self.real = real
+        self.entered = threading.Event()
+        self.release = threading.Event()
+        self.done = threading.Event()
+
+    def method(self):
+        """A function to patch onto the class (bound like the real one)."""
+
+        def held(*args):
+            return self(*args)
+
+        return held
+
+    def __call__(self, *args):
+        self.entered.set()
+        self.release.wait(WAIT)
+        try:
+            return self.real(*args)
+        finally:
+            self.done.set()
+
+
+class ShutdownRestoreCase(RestoreTestCase):
+    def gate(self, real):
+        gate = _Gate(real)
+        # A failed check never leaves the worker held until WAIT.
+        self._gates = getattr(self, "_gates", []) + [gate]
+        self.addCleanup(gate.release.set)
+        return gate
+
+    def stages(self):
+        """The stage folders (the upload's prefix starts the same way)."""
+        return [
+            name
+            for name in super().stages()
+            if os.path.isdir(os.path.join(self.tmpdir, name))
+        ]
+
+    def uploads(self):
+        """The upload files (not the -wal/-shm its readers may add)."""
+        return sorted(
+            name
+            for name in os.listdir(self.tmpdir)
+            if name.startswith(restore_stage.UPLOAD_PREFIX) and name.endswith(".tmp")
+        )
+
+    def fake_upload(self, texts):
+        """Make ``_receive_restore_upload`` hand over a prepared upload."""
+        self.upload(texts, "source.db")
+        source = os.path.join(self.tmpdir, "uploads", "source.db")
+        target = os.path.join(self.tmpdir, restore_stage.UPLOAD_PREFIX + "x.tmp")
+        shutil.copyfile(source, target)
+
+        async def receive(request, uploads):
+            uploads.append(target)
+            return target, None
+
+        return patch.object(self.api, "_receive_restore_upload", receive)
+
+    async def wait_thread(self, event):
+        loop = asyncio.get_running_loop()
+        self.assertTrue(await loop.run_in_executor(None, event.wait, WAIT))
+
+    async def assert_still_pending(self, task):
+        done, _ = await asyncio.wait({task}, timeout=SETTLE)
+        if done:
+            for gate in getattr(self, "_gates", ()):
+                gate.release.set()  # let the worker end before teardown
+            self.fail("the cancelled task ended while its worker still ran")
+
+
+class TestLegacyRestoreAtShutdown(ShutdownRestoreCase):
+    async def test_cancel_during_publish_waits_for_the_worker(self):
+        self.save("old")
+        old_db = self.api.db
+        gate = self.gate(PromptModel.publish_staged)
+        with (
+            self.fake_upload(["restored"]),
+            patch.object(PromptModel, "publish_staged", gate.method()),
+        ):
+            task = asyncio.ensure_future(self.api.restore_database(None))
+            await self.wait_thread(gate.entered)
+            self.assertEqual(len(self.stages()), 1)
+            task.cancel()
+            await self.assert_still_pending(task)
+            # The worker still has its stage and upload, and runs alone.
+            self.assertEqual(len(self.stages()), 1)
+            self.assertEqual(len(self.uploads()), 1)
+            self.assertIsNotNone(self.api._jobs.active(RESTORE_JOB_KIND))
+            gate.release.set()
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.wait_for(task, WAIT)
+        self.assertTrue(gate.done.is_set())
+        self.assertEqual(_texts(self.db_path), ["restored"])
+        self.assertIsNot(self.api.db, old_db)  # the restored library adopted
+        self.assertEqual(len(self.backups()), 1)  # the safety backup is kept
+        self.assertEqual(self.stages(), [])
+        self.assertEqual(self.uploads(), [])
+        self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND))
+
+    async def test_cancel_during_staging_waits_and_changes_nothing(self):
+        self.save("old")
+        gate = self.gate(PromptModel.stage_restore)
+        with (
+            self.fake_upload(["restored"]),
+            patch.object(PromptModel, "stage_restore", gate.method()),
+        ):
+            task = asyncio.ensure_future(self.api.restore_database(None))
+            await self.wait_thread(gate.entered)
+            task.cancel()
+            await self.assert_still_pending(task)
+            self.assertEqual(len(self.uploads()), 1)  # still being read
+            self.assertIsNotNone(self.api._jobs.active(RESTORE_JOB_KIND))
+            gate.release.set()
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.wait_for(task, WAIT)
+        self.assertEqual(_texts(self.db_path), ["old"])
+        self.assertEqual(self.backups(), [])
+        self.assertEqual(self.stages(), [])  # the finished stage is removed
+        self.assertEqual(self.uploads(), [])
+        self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND))
+
+
+class TestReviewedRestoreAtShutdown(ShutdownRestoreCase):
+    async def test_cancelled_job_waits_for_the_publish(self):
+        self.save("old")
+        old_db = self.api.db
+        body = await self.preview(self.upload(["restored"]))
+        gate = self.gate(PromptModel.publish_staged)
+        with patch.object(PromptModel, "publish_staged", gate.method()):
+            started = await self.apply(body["token"])
+            job = self.api._jobs.get(started["job"]["id"])
+            await self.wait_thread(gate.entered)
+            job.task.cancel()
+            await self.assert_still_pending(job.task)
+            self.assertEqual(len(self.stages()), 1)
+            self.assertTrue(job.running)  # its exclusive admission holds
+            gate.release.set()
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.wait_for(job.task, WAIT)
+        self.assertEqual(_texts(self.db_path), ["restored"])
+        self.assertIsNot(self.api.db, old_db)
+        self.assertEqual(len(self.backups()), 1)
+        self.assertEqual(self.stages(), [])
+        self.assertFalse(job.running)