feat(cloud-scheduler): renew active leases

2026-07-12 17:23:07 +08:00
parent b0a5b1426f
commit 1b15a8a218
4 changed files with 225 additions and 2 deletions
@@ -333,6 +333,56 @@ class SQLAlchemyCloudRepository:
            session.flush()
            return _leased_assignment_from_row(task)

+    def renew_lease(
+        self,
+        *,
+        task_id: str,
+        attempt: int,
+        lease_id: str,
+        host_id: str,
+        lease_expires_at: datetime,
+        now: datetime,
+    ) -> str:
+        with self._sessions.begin() as session:
+            task = session.get(
+                ScheduledTaskRow,
+                task_id,
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
+            if task is None:
+                return "not_found"
+            if (
+                task.status not in {"assigned", "dispatched"}
+                or task.attempt_count != attempt
+                or task.lease_id != lease_id
+                or task.assigned_host_id != host_id
+            ):
+                return "conflict"
+            current_expiry = _parse_dt(task.lease_expires_at)
+            if current_expiry is None or current_expiry <= now:
+                return "expired"
+            if lease_expires_at <= now:
+                return "conflict"
+
+            attempt_row = session.get(
+                TaskAttemptRow,
+                (task_id, attempt),
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
+            if (
+                attempt_row is None
+                or attempt_row.status not in {"assigned", "dispatched"}
+                or attempt_row.lease_id != lease_id
+                or attempt_row.host_id != host_id
+            ):
+                return "conflict"
+
+            renewed_until = _iso(lease_expires_at)
+            task.lease_expires_at = renewed_until
+            task.updated_at = _iso(now)
+            attempt_row.lease_expires_at = renewed_until
+            return "renewed"
+
    def list_task_attempts(self, task_id: str) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(