fix(restore): keep the cancel and the adoption when a stopped restore's publish ends (#360)

A publish that fails after the cancel now ends in CancelledError (after
its prune) instead of a 500, and the restored library is adopted to its
end even when a second cancel arrives. The blocking calls go through
_run_in_executor again, as the executor contract test expects.

Exit-window refusals: the comments state what was measured (a pool is
refused during threading._shutdown; Thread.start only once the
interpreter is finalizing), and the tests assert the one DEBUG line.

py/api/backup_restore.py

@@ -488,20 +488,33 @@ class BackupRestoreRoutesMixin:
         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.
+        (``_after_publish``, itself run to its end) before the cancel goes
+        on; a copy that failed is handled as usual, then the cancel goes on.
         """
-        loop = asyncio.get_running_loop()
-        future = loop.run_in_executor(None, model.publish_staged, staged)
+        future = self._in_executor_task(model.publish_staged, staged)
         cancelled = await _outlive_cancel(future)
         try:
             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)
+            if cancelled:
+                raise asyncio.CancelledError from exc
+            raise
+        except Exception as exc:
+            if cancelled:
+                raise asyncio.CancelledError from exc
             raise
         if cancelled:
             self.logger.info("The restore finished while the server was stopping")
-            await self._after_publish(model.db_path)
+            # A second cancel (the loop's own shutdown) must not cut it short.
+            adopted = asyncio.ensure_future(self._after_publish(model.db_path))
+            await _outlive_cancel(adopted)
+            if not adopted.cancelled() and adopted.exception() is not None:
+                self.logger.error(
+                    "Could not adopt the restored library",
+                    exc_info=adopted.exception(),
+                )
             raise asyncio.CancelledError
         return backup_path
 
@@ -512,14 +525,18 @@ class BackupRestoreRoutesMixin:
         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)
+        future = self._in_executor_task(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()
 
+    def _in_executor_task(self, func, *args):
+        """``_run_in_executor(func, *args)`` as a task ``_outlive_cancel`` can
+        wait for (shielded, the executor call is never cancelled)."""
+        return asyncio.ensure_future(self._run_in_executor(func, *args))
+
     async def _after_publish(self, db_path):
         """Everything after the copy replaced the library: best effort, logged.
 
@@ -957,7 +974,8 @@ class BackupRestoreRoutesMixin:
 
 
 async def _outlive_cancel(future):
-    """Await an executor ``future`` to its end, even through cancels (#360).
+    """Await ``future`` (an executor call's task) 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

tests/test_exit_window_threads.py

@@ -4,9 +4,10 @@ When the interpreter exits, ``threading._shutdown`` first shuts every
 ``ThreadPoolExecutor`` (``submit`` then raises "cannot schedule new futures
 after interpreter shutdown") and joins the non-daemon threads; only then do
 the ``atexit`` callbacks (the package teardown) run. Our daemon threads keep
-working meanwhile, and a late ``Thread.start`` can be refused with "can't
-create new thread at interpreter shutdown". Such a refusal is logged once at
-DEBUG, with no traceback, and leaves nothing pending.
+working meanwhile. ``Thread.start`` is refused ("can't create new thread at
+interpreter shutdown") only later, once the interpreter is finalizing; the
+unit tests cover that case, the child process the pool case. Such a refusal
+is logged once at DEBUG, with no traceback, and leaves nothing pending.
 
 ``threading.Thread`` is replaced only in the module under test (a copy of
 the ``threading`` namespace), never process-wide. The child-process test runs
@@ -64,6 +65,7 @@ class _Capture(logging.Handler):
 
 
 def _capture(testcase, name):
+    parallel._noted.clear()  # "once" is per process: start each test afresh
     logger = logging.getLogger(name)
     handler = _Capture()
     previous = logger.level
@@ -78,6 +80,15 @@ def _capture(testcase, name):
     return handler.records
 
 
+def _assert_one_debug_line(testcase, records):
+    """Exactly one refusal line, at DEBUG and without a traceback."""
+    refusals = [r for r in records if "interpreter is exiting" in r.getMessage()]
+    testcase.assertEqual(len(refusals), 1, [r.getMessage() for r in records])
+    testcase.assertEqual(refusals[0].levelno, logging.DEBUG)
+    testcase.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
+    testcase.assertTrue(all(r.exc_info is None for r in records))
+
+
 class TestScheduleRefusedAtShutdown(unittest.TestCase):
     def setUp(self):
         self.handler = image_monitor.ImageGenerationHandler(
@@ -91,8 +102,7 @@ class TestScheduleRefusedAtShutdown(unittest.TestCase):
             self.assertFalse(self.handler.schedule("/out/a.png"))
             self.assertFalse(self.handler.schedule("/out/b.png"))
         self.assertEqual(self.handler.pending_count(), 0)
-        self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
-        self.assertTrue(all(r.exc_info is None for r in records))
+        _assert_one_debug_line(self, records)
 
     def test_another_start_failure_still_raises_and_leaves_nothing_pending(self):
         class Failing(threading.Thread):
@@ -132,8 +142,7 @@ class TestPreviewQueueRefusedAtShutdown(unittest.TestCase):
         self.assertEqual(len(thumbnailer._queue), 0)
         self.assertEqual(thumbnailer._queued, set())
         self.assertTrue(thumbnailer.wait_idle(0))
-        self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
-        self.assertTrue(all(r.exc_info is None for r in records))
+        _assert_one_debug_line(self, records)
 
 
 class TestMapParallelRefusedAtShutdown(unittest.TestCase):
@@ -143,8 +152,7 @@ class TestMapParallelRefusedAtShutdown(unittest.TestCase):
             self.assertEqual(
                 parallel.map_parallel(lambda n: n * n, [1, 2, 3], 4), [1, 4, 9]
             )
-        self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
-        self.assertTrue(all(r.exc_info is None for r in records))
+        _assert_one_debug_line(self, records)
 
     def test_another_pool_error_still_raises(self):
         class Broken(ThreadPoolExecutor):

tests/test_restore_shutdown.py

@@ -13,11 +13,13 @@ held.
 import asyncio
 import os
 import shutil
+import sqlite3
 import threading
 from unittest.mock import patch
 
 from database import restore_stage
 from database.models import PromptModel
+from py.api import admin as admin_mod
 from py.api.jobs import RESTORE_JOB_KIND
 
 try:
@@ -155,6 +157,57 @@ class TestLegacyRestoreAtShutdown(ShutdownRestoreCase):
         self.assertEqual(self.uploads(), [])
         self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND))
 
+    async def test_a_publish_that_fails_after_the_cancel_still_cancels(self):
+        self.save("old")
+
+        def refuse(model, staged):
+            raise sqlite3.OperationalError("refused")
+
+        gate = self.gate(refuse)
+        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)
+            task.cancel()
+            await self.assert_still_pending(task)
+            gate.release.set()
+            # Not answered as a 500: the cancel goes on.
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.wait_for(task, WAIT)
+        self.assertEqual(_texts(self.db_path), ["old"])
+        self.assertEqual(self.stages(), [])
+        self.assertEqual(self.uploads(), [])
+        self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND))
+
+    async def test_a_second_cancel_still_adopts_the_restored_library(self):
+        self.save("old")
+        old_db = self.api.db
+        publish = self.gate(PromptModel.publish_staged)
+        prune = self.gate(admin_mod._prune_safety_backups)
+        with (
+            self.fake_upload(["restored"]),
+            patch.object(PromptModel, "publish_staged", publish.method()),
+            patch.object(admin_mod, "_prune_safety_backups", prune),
+        ):
+            task = asyncio.ensure_future(self.api.restore_database(None))
+            await self.wait_thread(publish.entered)
+            task.cancel()
+            await self.assert_still_pending(task)
+            publish.release.set()
+            await self.wait_thread(prune.entered)  # adopting the library
+            task.cancel()  # the loop's own shutdown cancels again
+            await self.assert_still_pending(task)
+            prune.release.set()
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.wait_for(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.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND))
+
 
 class TestReviewedRestoreAtShutdown(ShutdownRestoreCase):
     async def test_cancelled_job_waits_for_the_publish(self):

utils/parallel.py

@@ -7,12 +7,14 @@ the GIL while decoding, resizing and encoding, so threads still scale close
 to linearly for both jobs.
 
 Interpreter exit (#375): ``threading._shutdown`` shuts every thread pool
-before it joins the non-daemon threads and long before the package teardown
-runs, while our daemon threads keep working. A pool then refuses new work
-("cannot schedule new futures after interpreter shutdown") and a late
-``Thread.start`` can be refused ("can't create new thread at interpreter
-shutdown"). ``refused_at_shutdown`` recognises both; ``note_refusal`` logs
-one DEBUG line per place, with no traceback.
+(``submit`` then raises "cannot schedule new futures after interpreter
+shutdown") before it joins the non-daemon threads and long before the package
+teardown runs, while our daemon threads keep working. Measured on 3.12 and
+3.14, ``Thread.start`` still succeeds in that window; it is refused ("can't
+create new thread at interpreter shutdown") only once the interpreter is
+marked finalizing, after the ``atexit`` callbacks. ``refused_at_shutdown``
+recognises both; ``note_refusal`` logs one DEBUG line per place, with no
+traceback.
 """
 
 import logging