feat(cloud-scheduler): claim host assignments

2026-07-12 17:20:51 +08:00
parent ac7734c7dc
commit b0a5b1426f
3 changed files with 184 additions and 1 deletions
@@ -290,6 +290,49 @@ class SQLAlchemyCloudRepository:
            session.flush()
            return _leased_assignment_from_row(task)

+    def claim_assignment(
+        self,
+        *,
+        host_id: str,
+        now: datetime,
+    ) -> Any | None:
+        with self._sessions.begin() as session:
+            statement = (
+                select(ScheduledTaskRow)
+                .where(
+                    ScheduledTaskRow.status == "assigned",
+                    ScheduledTaskRow.assigned_host_id == host_id,
+                    ScheduledTaskRow.lease_id.is_not(None),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+                .order_by(ScheduledTaskRow.created_at, ScheduledTaskRow.id)
+                .limit(1)
+            )
+            if self.engine.dialect.name == "postgresql":
+                statement = statement.with_for_update(skip_locked=True)
+
+            task = session.scalars(statement).first()
+            if task is None:
+                return None
+            attempt = session.get(
+                TaskAttemptRow,
+                (task.id, task.attempt_count),
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
+            if (
+                attempt is None
+                or attempt.status != "assigned"
+                or attempt.lease_id != task.lease_id
+            ):
+                return None
+
+            task.status = "dispatched"
+            task.updated_at = _iso(now)
+            attempt.status = "dispatched"
+            session.flush()
+            return _leased_assignment_from_row(task)
+
    def list_task_attempts(self, task_id: str) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(