fix: harden retention cleanup safety

2026-09-12 14:12:17 +08:00
parent da5c2b4503
commit b3ab1bfc52
6 changed files with 259 additions and 55 deletions
@@ -30,7 +30,7 @@ from .line_spacing import analyze_svg_line_spacing_file
 from .log_config import get_logger
 from .metadata_store import MetadataStore
 from .product_archive_service import ProductArchiveService, ProductPendingCleanupError
-from .product_archive_store import ProductArchiveStore
+from .product_archive_store import ProductArchiveStore, ProductPurgeInProgressError
 from .runner import JobRunner
 from .schemas import (
    Asset,
@@ -1073,10 +1073,12 @@ def _cleanup_scheduler_loop() -> None:

 def _start_cleanup_scheduler() -> None:
    global _cleanup_scheduler_thread
-    _run_scheduled_cleanup_once()
-    if _cleanup_scheduler_thread is not None and _cleanup_scheduler_thread.is_alive():
-        return
+    if _cleanup_scheduler_thread is not None:
+        if _cleanup_scheduler_thread.is_alive():
+            return
+        _cleanup_scheduler_thread = None
    _cleanup_scheduler_stop.clear()
+    _run_scheduled_cleanup_once()
    _cleanup_scheduler_thread = threading.Thread(
        target=_cleanup_scheduler_loop,
        name="retention-cleanup",
@@ -1088,8 +1090,10 @@ def _start_cleanup_scheduler() -> None:
 def _stop_cleanup_scheduler() -> None:
    global _cleanup_scheduler_thread
    _cleanup_scheduler_stop.set()
-    if _cleanup_scheduler_thread is not None:
-        _cleanup_scheduler_thread.join(timeout=1)
+    worker = _cleanup_scheduler_thread
+    if worker is not None:
+        worker.join(timeout=5)
+    if worker is not None and not worker.is_alive() and _cleanup_scheduler_thread is worker:
        _cleanup_scheduler_thread = None


@@ -1240,6 +1244,8 @@ def restore_product(request: Request, product_id: str) -> ProductRecord:
    _get_product_or_404(product_id)
    try:
        return product_archive_store.restore_product(product_id, now=datetime.now(timezone.utc))
+    except ProductPurgeInProgressError as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

@@ -117,25 +117,25 @@ class CleanupService:
            report.deleted_job_ids.append(item.job_id)

        for candidate in report.pending_products:
-            # Re-read logical state so a concurrent restore wins over cleanup.
-            try:
-                product = self.product_store.get_product(candidate.product_id)
-            except ValueError:
-                continue
-            if (
-                product.status != "pending_cleanup"
-                or product.purge_after is None
-                or _as_utc(product.purge_after) > now_utc
-            ):
+            product = self.product_store.claim_product_purge(candidate.product_id, now_utc)
+            if product is None:
                continue
            product_dir = self.product_store.root.parent / product.product_id
-            if product_dir.exists():
-                shutil.rmtree(product_dir)
-            self.product_store.purge_product(product.product_id, now=now_utc)
+            try:
+                self._before_product_files_delete(product.product_id)
+                if product_dir.exists():
+                    shutil.rmtree(product_dir)
+            except Exception:
+                self.product_store.release_product_purge(product.product_id)
+                raise
+            self.product_store.finalize_product_purge(product.product_id, now=now_utc)
            report.deleted_product_ids.append(product.product_id)

        return report

+    def _before_product_files_delete(self, product_id: str) -> None:
+        """Interleaving seam between the durable claim and physical deletion."""
+
    @staticmethod
    def _directory_size(root: Path) -> int:
        if not root.exists():
@@ -23,6 +23,10 @@ def _now() -> datetime:
    return datetime.now(timezone.utc)


+class ProductPurgeInProgressError(ValueError):
+    """Raised when restore loses the atomic race to physical purge."""
+
+
 class ProductArchiveStore:
    """Owns durable product metadata and logical cleanup state."""

@@ -328,43 +332,111 @@ class ProductArchiveStore:
        return self.get_product(product_id)

    def restore_product(self, product_id: str, now: datetime | None = None) -> ProductRecord:
-        self.get_product(product_id)
        restored_at = now or _now()
-        self._execute(
-            "UPDATE products SET status = ?, purge_after = NULL, updated_at = ? WHERE product_id = ?",
-            ("active", restored_at.isoformat(), product_id),
-        )
-        self._execute(
-            """UPDATE cleanup_records SET status = ?, completed_at = ?
-               WHERE product_id = ? AND status = ?""",
-            ("restored", restored_at.isoformat(), product_id, "pending_cleanup"),
-        )
+        with self._lock, self._connect() as conn:
+            conn.execute("BEGIN IMMEDIATE")
+            product = conn.execute(
+                "SELECT product_id, status FROM products WHERE product_id = ?", (product_id,)
+            ).fetchone()
+            if product is None:
+                raise ValueError(f"unknown product_id: {product_id}")
+            if product["status"] == "purged":
+                raise ValueError("purged product cannot be restored")
+            purging = conn.execute(
+                """SELECT 1 FROM cleanup_records
+                   WHERE product_id = ? AND status = 'purging' LIMIT 1""",
+                (product_id,),
+            ).fetchone()
+            if purging is not None:
+                raise ProductPurgeInProgressError("product is purging")
+            conn.execute(
+                "UPDATE products SET status = ?, purge_after = NULL, updated_at = ? WHERE product_id = ?",
+                ("active", restored_at.isoformat(), product_id),
+            )
+            conn.execute(
+                """UPDATE cleanup_records SET status = ?, completed_at = ?
+                   WHERE product_id = ? AND status = ?""",
+                ("restored", restored_at.isoformat(), product_id, "pending_cleanup"),
+            )
        return self.get_product(product_id)

    def due_product_cleanups(self, now: datetime) -> list[ProductRecord]:
        rows = self._fetchall(
-            """SELECT * FROM products WHERE status = ? AND purge_after IS NOT NULL AND purge_after <= ?
-               ORDER BY purge_after, product_id""",
+            """SELECT products.* FROM products
+               WHERE products.status = ?
+                 AND products.purge_after IS NOT NULL
+                 AND products.purge_after <= ?
+                 AND EXISTS (
+                     SELECT 1 FROM cleanup_records
+                     WHERE cleanup_records.product_id = products.product_id
+                       AND cleanup_records.status = 'pending_cleanup'
+                 )
+               ORDER BY products.purge_after, products.product_id""",
            ("pending_cleanup", now.isoformat()),
        )
        return [self._product_from_row(row) for row in rows]

+    def claim_product_purge(self, product_id: str, now: datetime) -> ProductRecord | None:
+        """Atomically win the right to delete a due product's files."""
+        with self._lock, self._connect() as conn:
+            conn.execute("BEGIN IMMEDIATE")
+            row = conn.execute(
+                "SELECT * FROM products WHERE product_id = ?", (product_id,)
+            ).fetchone()
+            if row is None:
+                return None
+            purge_after = datetime.fromisoformat(row["purge_after"]) if row["purge_after"] else None
+            if row["status"] != "pending_cleanup" or purge_after is None or purge_after > now:
+                return None
+            cleanup = conn.execute(
+                """SELECT cleanup_id FROM cleanup_records
+                   WHERE product_id = ? AND status = 'pending_cleanup'
+                   ORDER BY created_at DESC, cleanup_id DESC LIMIT 1""",
+                (product_id,),
+            ).fetchone()
+            if cleanup is None:
+                return None
+            updated = conn.execute(
+                """UPDATE cleanup_records SET status = 'purging'
+                   WHERE cleanup_id = ? AND status = 'pending_cleanup'""",
+                (cleanup["cleanup_id"],),
+            )
+            if updated.rowcount != 1:
+                return None
+            return self._product_from_row(row)
+
+    def release_product_purge(self, product_id: str) -> None:
+        """Return a failed physical deletion claim to its pending state."""
+        self._execute(
+            """UPDATE cleanup_records SET status = 'pending_cleanup'
+               WHERE product_id = ? AND status = 'purging'""",
+            (product_id,),
+        )
+
+    def finalize_product_purge(self, product_id: str, now: datetime) -> ProductRecord:
+        """Atomically mark a claimed product purged after file removal."""
+        with self._lock, self._connect() as conn:
+            conn.execute("BEGIN IMMEDIATE")
+            claim = conn.execute(
+                """SELECT 1 FROM cleanup_records
+                   WHERE product_id = ? AND status = 'purging' LIMIT 1""",
+                (product_id,),
+            ).fetchone()
+            if claim is None:
+                raise ValueError("product purge is not claimed")
+            conn.execute(
+                "UPDATE products SET status = ?, updated_at = ? WHERE product_id = ?",
+                ("purged", now.isoformat(), product_id),
+            )
+            conn.execute(
+                """UPDATE cleanup_records SET status = ?, completed_at = ?
+                   WHERE product_id = ? AND status = 'purging'""",
+                ("purged", now.isoformat(), product_id),
+            )
+        return self.get_product(product_id)
+
    def purge_product(self, product_id: str, now: datetime | None = None) -> ProductRecord:
        purged_at = now or _now()
-        product = self.get_product(product_id)
-        if (
-            product.status != "pending_cleanup"
-            or product.purge_after is None
-            or product.purge_after > purged_at
-        ):
+        if self.claim_product_purge(product_id, purged_at) is None:
            raise ValueError("product is not due for purge")
-        self._execute(
-            "UPDATE products SET status = ?, updated_at = ? WHERE product_id = ?",
-            ("purged", purged_at.isoformat(), product_id),
-        )
-        self._execute(
-            """UPDATE cleanup_records SET status = ?, completed_at = ?
-               WHERE product_id = ? AND status = ?""",
-            ("purged", purged_at.isoformat(), product_id, "pending_cleanup"),
-        )
-        return self.get_product(product_id)
+        return self.finalize_product_purge(product_id, purged_at)
@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
-"""Inspect and optionally clean stale job directories.
+"""Inspect retention candidates and optionally clean verified managed jobs.

-Default mode is a safe dry-run that reports reclaimable bytes. Pass --apply to
-actually remove unreferenced job directories.
+Default mode is a safe dry-run. Metadata-less workspaces are reported for
+manual investigation but are never removed by ``--apply``.
 """

 from __future__ import annotations
@@ -46,7 +46,11 @@ def referenced_job_ids() -> set[str]:
 def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--max-age-days", type=float, default=30)
-    parser.add_argument("--apply", action="store_true", help="Actually delete stale job directories")
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="Delete only CleanupService-verified candidates; orphan workspaces remain report-only",
+    )
    parser.add_argument("--json", type=Path, default=None, help="Write JSON report")
    args = parser.parse_args()

@@ -89,10 +93,11 @@ def main() -> None:
        if os.environ.get("CLEANUP_APPLY_ENABLED", "false").strip().lower() != "true":
            parser.error("--apply requires CLEANUP_APPLY_ENABLED=true")
        applied = cleanup.apply(datetime.now(timezone.utc))
-        freed = 0
-        for item in orphaned:
-            freed += storage.remove_job_dir(item["job_id"])
-            store.delete_job(item["job_id"])
+        freed = sum(
+            item.size_bytes
+            for item in applied.temporary_jobs
+            if item.job_id in applied.deleted_job_ids
+        )
        print(
            "freed_bytes="
            f"{freed} deleted_job_ids={applied.deleted_job_ids} "