feat(cloud-scheduler): reserve devices atomically

2026-07-12 17:19:10 +08:00
parent 0a9392b7ea
commit ac7734c7dc
6 changed files with 301 additions and 19 deletions
@@ -207,6 +207,89 @@ class SQLAlchemyCloudRepository:
            row = session.get(PluginRow, name)
        return _plugin_from_row(row) if row else None

+    def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
+        with self._sessions() as session:
+            device_ids = session.scalars(
+                select(ScheduledTaskRow.assigned_device_id).where(
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.assigned_device_id.is_not(None),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+            ).all()
+        return {device_id for device_id in device_ids if device_id is not None}
+
+    def assign_task(
+        self,
+        *,
+        task_id: str,
+        host_id: str,
+        device_id: str,
+        lease_id: str,
+        lease_expires_at: datetime,
+        now: datetime,
+    ) -> Any | None:
+        with self._sessions.begin() as session:
+            task_statement = select(ScheduledTaskRow).where(
+                ScheduledTaskRow.id == task_id
+            )
+            device_statement = select(PooledDeviceRow).where(
+                PooledDeviceRow.host_id == host_id,
+                PooledDeviceRow.device_id == device_id,
+            )
+            if self.engine.dialect.name == "postgresql":
+                task_statement = task_statement.with_for_update(skip_locked=True)
+                device_statement = device_statement.with_for_update()
+
+            task = session.scalars(task_statement).first()
+            if task is None or task.status != "queued":
+                return None
+            device = session.scalars(device_statement).first()
+            if device is None or device.status != "idle":
+                return None
+
+            active_reservation = session.scalar(
+                select(ScheduledTaskRow.id)
+                .where(
+                    ScheduledTaskRow.id != task_id,
+                    ScheduledTaskRow.assigned_device_id == device_id,
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+                .limit(1)
+            )
+            if active_reservation is not None:
+                return None
+
+            attempt = task.attempt_count + 1
+            task.status = "assigned"
+            task.assigned_host_id = host_id
+            task.assigned_device_id = device_id
+            task.attempt_count = attempt
+            task.lease_id = lease_id
+            task.lease_expires_at = _iso(lease_expires_at)
+            task.failure_reason = None
+            task.result_json = None
+            task.updated_at = _iso(now)
+            session.add(
+                TaskAttemptRow(
+                    task_id=task.id,
+                    attempt=attempt,
+                    lease_id=lease_id,
+                    host_id=host_id,
+                    device_id=device_id,
+                    status="assigned",
+                    lease_expires_at=_iso(lease_expires_at),
+                    created_at=_iso(now),
+                    completed_at=None,
+                    failure_reason=None,
+                    result_json=None,
+                )
+            )
+            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(
@@ -311,6 +394,29 @@ def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
    )


+def _leased_assignment_from_row(row: ScheduledTaskRow) -> Any:
+    from cloud.repository import LeasedAssignment
+
+    lease_expires_at = _parse_dt(row.lease_expires_at)
+    if (
+        row.lease_id is None
+        or lease_expires_at is None
+        or row.assigned_host_id is None
+        or row.assigned_device_id is None
+    ):
+        raise ValueError(f"task {row.id!r} does not contain a complete lease")
+    return LeasedAssignment(
+        task_id=row.id,
+        attempt=row.attempt_count,
+        lease_id=row.lease_id,
+        lease_expires_at=lease_expires_at,
+        host_id=row.assigned_host_id,
+        device_id=row.assigned_device_id,
+        goal=row.goal,
+        workflow_definition_id=row.workflow_definition_id,
+    )
+
+
 def _parse_json_object(value: str | None) -> dict[str, Any] | None:
    if value is None:
        return None