From b0a5b1426f424592c1c1829ca0b6581015aa4c4d Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 17:20:51 +0800 Subject: [PATCH] feat(cloud-scheduler): claim host assignments --- .../cloud-control-plane-integration/tasks.md | 2 +- .../cloud-platform/cloud/sql_repository.py | 43 ++++++ tests/test_cloud_repository_contract.py | 140 ++++++++++++++++++ 3 files changed, 184 insertions(+), 1 deletion(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index 6301c59..f182bc3 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -19,7 +19,7 @@ - [x] 3.1 Extend scheduled-task persistence with attempt count, lease id/expiry, terminal result, failure reason, and auditable attempt records. - [x] 3.2 Implement atomic queued-task assignment and device reservation while excluding devices with active assignments even when snapshots report idle. -- [ ] 3.3 Implement owning-host claim that atomically transitions one assigned attempt to dispatched under its active lease. +- [x] 3.3 Implement owning-host claim that atomically transitions one assigned attempt to dispatched under its active lease. - [ ] 3.4 Implement lease renewal with host/task/attempt ownership validation and conflict responses for stale leases. - [ ] 3.5 Implement idempotent terminal result recording and reservation release for active leases. - [ ] 3.6 Implement expired-lease requeue/failure behavior with bounded attempts and preserved attempt history. diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index 7b775e4..7a3f773 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/packages/cloud-platform/cloud/sql_repository.py @@ -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( diff --git a/tests/test_cloud_repository_contract.py b/tests/test_cloud_repository_contract.py index cbbae36..fc5dd5a 100644 --- a/tests/test_cloud_repository_contract.py +++ b/tests/test_cloud_repository_contract.py @@ -437,3 +437,143 @@ def test_expired_assignment_no_longer_reserves_device(database_url: str) -> None ) finally: database.close() + + +def test_owning_host_claims_one_active_assignment(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("claim-host") + device_id = _unique_id("claim-device") + task_id = _unique_id("claim-task") + now = datetime(2026, 7, 12, 5, 0, tzinfo=UTC) + + try: + database.repository.upsert_host(host_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="claim me", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + database.repository.assign_task( + task_id=task_id, + host_id=host_id, + device_id=device_id, + lease_id="claim-lease", + lease_expires_at=now + timedelta(minutes=1), + now=now, + ) + + assignment = database.repository.claim_assignment( + host_id=host_id, + now=now + timedelta(seconds=1), + ) + + assert assignment is not None + assert assignment.task_id == task_id + assert assignment.lease_id == "claim-lease" + task = database.repository.get_task(task_id) + assert task is not None + assert task.status == "dispatched" + attempts = database.repository.list_task_attempts(task_id) + assert len(attempts) == 1 + assert attempts[0].status == "dispatched" + assert database.repository.claim_assignment(host_id=host_id, now=now) is None + finally: + database.close() + + +def test_foreign_host_cannot_claim_assignment(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("owner-host") + device_id = _unique_id("owner-device") + task_id = _unique_id("owner-task") + now = datetime(2026, 7, 12, 6, 0, tzinfo=UTC) + + try: + database.repository.upsert_host(host_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="owner only", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + database.repository.assign_task( + task_id=task_id, + host_id=host_id, + device_id=device_id, + lease_id="owner-lease", + lease_expires_at=now + timedelta(minutes=1), + now=now, + ) + + assert ( + database.repository.claim_assignment( + host_id=_unique_id("foreign-host"), + now=now, + ) + is None + ) + task = database.repository.get_task(task_id) + assert task is not None + assert task.status == "assigned" + finally: + database.close() + + +def test_expired_assignment_cannot_be_claimed(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("expired-claim-host") + device_id = _unique_id("expired-claim-device") + task_id = _unique_id("expired-claim-task") + now = datetime(2026, 7, 12, 7, 0, tzinfo=UTC) + + try: + database.repository.upsert_host(host_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="too late", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + database.repository.assign_task( + task_id=task_id, + host_id=host_id, + device_id=device_id, + lease_id="expired-claim-lease", + lease_expires_at=now + timedelta(seconds=1), + now=now, + ) + + assert ( + database.repository.claim_assignment( + host_id=host_id, + now=now + timedelta(seconds=2), + ) + is None + ) + task = database.repository.get_task(task_id) + assert task is not None + assert task.status == "assigned" + finally: + database.close()