fix(backup): re-review round for restore repair, apply and backup cleanup

- R1: when re-keying merges an intact body into a holder whose own body was
  damaged, the rewritten holder leaves the damaged list, so the clear pass
  no longer deletes it and unlinks the images just moved to it.
- R2: a losing duplicate apply removes a stage only when the store still
  held it, never one another apply's job has taken.
- R3: a restore that times out waiting for requests keeps its review only
  while it is unexpired and its slot is free; otherwise the stage goes and
  the job says to try again.
- R4: a backup download cancelled while its snapshot is written removes
  the snapshot once the worker finishes.

database/workflow_documents.py

@@ -545,13 +545,17 @@ def _inline_copies(conn: sqlite3.Connection, wanted: Dict[bytes, int]):
     return found
 
 
-def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> None:
-    """Give an intact body its true key, or merge it into the document that has it."""
+def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> Optional[int]:
+    """Give an intact body its true key, or merge it into the document that has it.
+
+    Returns the holder's id when its own body was damaged and has just been
+    rewritten from ``raw`` (it is repaired, never to be cleared), else None.
+    """
     digest = hashlib.sha256(raw).digest()
     holder = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
     if holder is None or holder[0] == document_id:
         conn.execute(REKEY_SQL, (digest, len(raw), document_id))
-        return
+        return None
     for column in DOCUMENT_ID_COLUMNS:
         # ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data.
         conn.execute(
@@ -563,6 +567,8 @@ def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> None:
         conn.execute(
             REPAIR_SQL, (len(raw), zlib.compress(raw, COMPRESSION_LEVEL), holder[0])
         )
+        return holder[0]
+    return None
 
 
 def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
@@ -596,8 +602,14 @@ def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
     if not damaged and not intact:
         return result
     with conn:
+        healed = []
         for document_id, raw in intact.items():
-            _rekey(conn, document_id, raw)
+            holder = _rekey(conn, document_id, raw)
+            if holder is not None and holder in damaged.values():
+                # Rewritten from the merged body: repaired, not to be cleared
+                # (it now also holds the merged images' references).
+                damaged = {k: v for k, v in damaged.items() if v != holder}
+                healed.append(holder)
         copies = _inline_copies(conn, damaged) if damaged else {}
         for document_id, text in copies.items():
             raw = text.encode("utf-8")
@@ -625,7 +637,7 @@ def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
                 f"DELETE FROM workflow_documents WHERE id IN ({marks})",  # nosec B608
                 chunk,
             )
-    for document_id in list(intact) + list(copies):
+    for document_id in list(intact) + healed + list(copies):
         logger.warning(f"Workflow document {document_id} was damaged; repaired")
     if lost:
         logger.warning(
@@ -633,6 +645,8 @@ def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
             f"{len(images)} image(s) show no workflow"
         )
     result.update(
-        repaired=len(intact) + len(copies), cleared=len(lost), images=len(images)
+        repaired=len(intact) + len(healed) + len(copies),
+        cleared=len(lost),
+        images=len(images),
     )
     return result

py/api/backup_restore.py

@@ -273,8 +273,17 @@ class BackupRestoreRoutesMixin:
 
     async def _make_backup_snapshot(self, model):
         """``(path, warning)``, or the error response (inside the hold)."""
+        loop = asyncio.get_running_loop()
+        future = loop.run_in_executor(None, write_backup_snapshot, model)
         try:
-            temp, warning = await self._run_in_executor(write_backup_snapshot, model)
+            temp, warning = await asyncio.shield(future)
+        except asyncio.CancelledError:
+            # The worker cannot be stopped: remove what it writes once it is
+            # done (#335 re-review R4).
+            future.add_done_callback(
+                lambda done: loop.create_task(self._remove_finished_snapshot(done))
+            )
+            raise
         except SafetyBackupFailed as failed:
             self.logger.error(f"Backup failed: {failed} {failed.reason or ''}")
             return web.json_response(
@@ -294,6 +303,15 @@ class BackupRestoreRoutesMixin:
             self.logger.warning(f"Backup sent with a failed check: {warning}")
         return temp, warning
 
+    async def _remove_finished_snapshot(self, future):
+        """Remove the snapshot a cancelled download's worker finished."""
+        if future.cancelled() or future.exception() is not None:
+            return  # nothing was left on disk
+        temp, _warning = future.result()
+        await self._run_in_executor(
+            remove_quietly, *(temp + suffix for suffix in SIDECARS)
+        )
+
     async def _send_snapshot(self, request, temp, warning):
         """Stream the snapshot (outside the hold), then close and remove it."""
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -508,6 +526,16 @@ class BackupRestoreRoutesMixin:
             except Exception:
                 self.logger.exception("Could not remove a restore stage")
 
+    async def _discard_popped(self, token):
+        """Drop ``token`` and remove its stage, only if the store still held it.
+
+        A stage another apply's job has taken is not this request's to remove
+        (#335 re-review R2).
+        """
+        popped = self._restore_plans.pop(token)
+        if popped is not None:
+            await self._discard_held([popped])
+
     async def _drop_expired_restores(self):
         """Drop expired previews from the store (on the loop), then their files."""
         await self._discard_held(self._restore_plans.pop_expired())
@@ -786,7 +814,7 @@ class BackupRestoreRoutesMixin:
             return refused
         db = self.db
         if held.db_identity != self._database_identity(db):
-            await self._discard_held([self._restore_plans.pop(token) or held])
+            await self._discard_popped(token)
             return _failure(
                 "The database was replaced since the review. Review again.",
                 409,
@@ -801,7 +829,7 @@ class BackupRestoreRoutesMixin:
             return busy_response(busy)
         if fingerprint != held.fingerprint:
             # It can never match again: free the stage for a new review.
-            await self._discard_held([self._restore_plans.pop(token) or held])
+            await self._discard_popped(token)
             return _failure(PLAN_CHANGED, 409, "plan-changed")
         # The token may have been used, expired or discarded meanwhile.
         held, refused = self._peek_restore(token)
@@ -840,9 +868,9 @@ class BackupRestoreRoutesMixin:
         try:
             job.update(status_line="Waiting for other requests to finish")
             if not await self._jobs.drain_leases():
-                if token is not None:
-                    # Nothing ran: the reviewed stage stays usable (#335 M2).
-                    self._restore_plans.give_back(token, held)
+                # Nothing ran: the reviewed stage stays usable (#335 M2),
+                # unless it expired meanwhile or its slot was reused.
+                if token is not None and self._restore_plans.give_back(token, held):
                     keep_stage = True
                     self._schedule_restore_expiry()
                     return _fail_job(job, "busy", REVIEW_KEPT)

py/api/restore_plans.py

@@ -105,10 +105,17 @@ class RestorePlanStore:
             self._used.pop(next(iter(self._used)))
         return held
 
-    def give_back(self, token, held) -> None:
-        """Hold a taken stage again under its token (its job did not apply it)."""
+    def give_back(self, token, held) -> bool:
+        """Hold a taken stage again under its token (its job did not apply it).
+
+        False, holding nothing, when the review expired meanwhile or its slot
+        was reused (the store is full): the caller then discards the stage.
+        """
+        if held.expires_at <= self._clock() or self.full():
+            return False
         self._used.pop(token, None)
         self._held[token] = held
+        return True
 
     def entry(self, token) -> Optional[HeldRestore]:
         """The held entry, expired or not (for the response that issued it)."""

tests/test_restore_review_fixes.py

@@ -450,3 +450,130 @@ class TestFingerprintTransaction(RestoreTestCase):
         self.assertTrue(conn.in_transaction)
         conn.rollback()
         self.assertEqual([p["text"] for p in self.api.db.search_prompts()], [])
+
+
+class TestRereviewFixes(RestoreTestCase):
+    """PR #335 re-review: R1 to R4."""
+
+    def test_r1_a_merge_into_a_damaged_holder_keeps_both_images(self):
+        rekey = TestRekeyDocuments.library
+        path, text = rekey(self, lambda digest, length: (b"\0" * 32, length))
+        raw = text.encode()
+        with closing(sqlite3.connect(path)) as conn:
+            holder = conn.execute(
+                "INSERT INTO workflow_documents (sha256, byte_length, data) "
+                "VALUES (?,?,?)",
+                (hashlib.sha256(raw).digest(), len(raw), b"garbage-not-zlib"),
+            ).lastrowid
+            prompt_id = conn.execute("SELECT id FROM prompts").fetchone()[0]
+            conn.execute(
+                "INSERT INTO generated_images (prompt_id, image_path, filename, "
+                "workflow_document_id) VALUES (?, '/x/b.png', 'b.png', ?)",
+                (prompt_id, holder),
+            )
+            conn.commit()
+        staged = self.api.db.model.stage_restore(path)
+        try:
+            self.assertEqual(
+                staged.documents, {"repaired": 2, "cleared": 0, "images": 0}
+            )
+            self.assertEqual(
+                TestRekeyDocuments.staged_texts(self, staged), [text, text]
+            )
+            with closing(sqlite3.connect(staged.path)) as conn:
+                self.assertEqual(
+                    conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone(),
+                    (1,),
+                )
+        finally:
+            staged.discard()
+
+    async def test_r2_a_losing_duplicate_apply_leaves_the_taken_stage(self):
+        self.save("live")
+        body = await self.preview(self.upload(["restored"]))
+        store = self.api._restore_plans
+        held = store.peek(body["token"])
+        entered, release = threading.Event(), threading.Event()
+        real = self.api._live_fingerprint
+
+        def gated(db):
+            entered.set()
+            release.wait(WAIT)
+            return real(db)
+
+        with patch.object(self.api, "_live_fingerprint", side_effect=gated):
+            second = asyncio.ensure_future(
+                self.client.post(
+                    "/prompt_manager/restore/apply", json={"token": body["token"]}
+                )
+            )
+            loop = asyncio.get_running_loop()
+            self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
+            # Meanwhile the first apply's job took the stage, and a write landed.
+            store.take(body["token"], "first-job")
+            self.save("a run count")
+            release.set()
+            resp = await second
+        self.assertEqual(resp.status, 409)
+        self.assertTrue(os.path.isdir(held.staged.directory))
+        await self._discard_quietly(held)
+
+    async def _discard_quietly(self, held):
+        await asyncio.get_running_loop().run_in_executor(None, held.staged.discard)
+
+    def test_r3_give_back_refuses_a_full_store_or_an_expired_review(self):
+        clock = [1000.0]
+        store = backup_restore.RestorePlanStore(clock=lambda: clock[0], max_stages=1)
+        first = store.issue(_Stage(), "f", "id", {})
+        held = store.take(first, "job")
+        other = store.issue(_Stage(), "f", "id", {})
+        self.assertFalse(store.give_back(first, held))  # the slot was reused
+        self.assertEqual(len(store), 1)
+        store.pop(other)
+        clock[0] = held.expires_at
+        self.assertFalse(store.give_back(first, held))  # expired meanwhile
+        clock[0] -= 1
+        self.assertTrue(store.give_back(first, held))
+        self.assertIs(store.peek(first), held)
+
+    async def test_r4_a_cancelled_backup_removes_its_snapshot(self):
+        from aiohttp.test_utils import make_mocked_request
+
+        self.save("keep")
+        entered, release, done = (threading.Event() for _ in range(3))
+        real = backup_restore.write_backup_snapshot
+
+        def slow(model):
+            entered.set()
+            release.wait(WAIT)
+            try:
+                return real(model)
+            finally:
+                done.set()
+
+        with patch.object(backup_restore, "write_backup_snapshot", slow):
+            task = asyncio.ensure_future(
+                self.api._stream_backup(
+                    make_mocked_request("GET", "/prompt_manager/backup")
+                )
+            )
+            loop = asyncio.get_running_loop()
+            self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+            release.set()
+            self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT))
+
+            async def snapshot_gone():
+                while any(n.startswith(".pm-safety-") for n in os.listdir(self.tmpdir)):
+                    await asyncio.sleep(0)
+
+            # A ceiling only a regression reaches (the file never goes).
+            await asyncio.wait_for(snapshot_gone(), WAIT)
+        self.assertIsNone(self.api._jobs.active("backup"))
+
+
+class _Stage:
+    def discard(self):
+        pass