diff --git a/database/workflow_documents.py b/database/workflow_documents.py index 3c8434a2..d4795eb3 100644 --- a/database/workflow_documents.py +++ b/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 diff --git a/py/api/backup_restore.py b/py/api/backup_restore.py index 493904dd..f3b11c3b 100644 --- a/py/api/backup_restore.py +++ b/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) diff --git a/py/api/restore_plans.py b/py/api/restore_plans.py index 7181c423..13d08c1c 100644 --- a/py/api/restore_plans.py +++ b/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).""" diff --git a/tests/test_restore_review_fixes.py b/tests/test_restore_review_fixes.py index f625e307..d6f2f327 100644 --- a/tests/test_restore_review_fixes.py +++ b/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