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
@@ -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)