feat(cloud-scheduler): renew active leases

packages/cloud-platform/cloud/repository.py

@@ -13,6 +13,7 @@ if TYPE_CHECKING:
 AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
 TerminalTaskStatus = Literal["done", "failed"]
 ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
+LeaseRenewalStatus = Literal["renewed", "not_found", "conflict", "expired"]
 
 
 @dataclass(frozen=True)
@@ -119,7 +120,7 @@ class CloudRepository(Protocol):
         host_id: str,
         lease_expires_at: datetime,
         now: datetime,
-    ) -> bool: ...
+    ) -> LeaseRenewalStatus: ...
 
     def record_task_result(
         self,

packages/cloud-platform/cloud/sql_repository.py

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