From 1b15a8a2180428724c7dd41f7bfb6c7727df61c0 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 17:23:07 +0800 Subject: [PATCH] feat(cloud-scheduler): renew active leases --- .../cloud-control-plane-integration/tasks.md | 2 +- packages/cloud-platform/cloud/repository.py | 3 +- .../cloud-platform/cloud/sql_repository.py | 50 +++++ tests/test_cloud_repository_contract.py | 172 ++++++++++++++++++ 4 files changed, 225 insertions(+), 2 deletions(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index f182bc3..467066d 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -20,7 +20,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. - [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. +- [x] 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. - [ ] 3.7 Add PostgreSQL concurrency tests proving one assignment/claim winner and SQLite tests documenting single-control-plane behavior. diff --git a/packages/cloud-platform/cloud/repository.py b/packages/cloud-platform/cloud/repository.py index cce943b..8b92889 100644 --- a/packages/cloud-platform/cloud/repository.py +++ b/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, diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index 7a3f773..6e866a6 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/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( diff --git a/tests/test_cloud_repository_contract.py b/tests/test_cloud_repository_contract.py index fc5dd5a..f5e8248 100644 --- a/tests/test_cloud_repository_contract.py +++ b/tests/test_cloud_repository_contract.py @@ -577,3 +577,175 @@ def test_expired_assignment_cannot_be_claimed(database_url: str) -> None: assert task.status == "assigned" finally: database.close() + + +def test_active_lease_renews_for_owning_host(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("renew-host") + device_id = _unique_id("renew-device") + task_id = _unique_id("renew-task") + now = datetime(2026, 7, 12, 8, 0, tzinfo=UTC) + initial_expiry = now + timedelta(minutes=1) + renewed_expiry = now + timedelta(minutes=2) + + 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="renew 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="renew-lease", + lease_expires_at=initial_expiry, + now=now, + ) + database.repository.claim_assignment(host_id=host_id, now=now) + + status = database.repository.renew_lease( + task_id=task_id, + attempt=1, + lease_id="renew-lease", + host_id=host_id, + lease_expires_at=renewed_expiry, + now=now + timedelta(seconds=30), + ) + + assert status == "renewed" + task = database.repository.get_task(task_id) + assert task is not None + assert task.lease_expires_at == renewed_expiry + attempts = database.repository.list_task_attempts(task_id) + assert attempts[0].lease_expires_at == renewed_expiry + finally: + database.close() + + +@pytest.mark.parametrize( + ("attempt", "lease_id", "host_id"), + [ + (2, "lease-current", "owner"), + (1, "lease-stale", "owner"), + (1, "lease-current", "foreign"), + ], +) +def test_stale_or_foreign_lease_renewal_conflicts( + database_url: str, + attempt: int, + lease_id: str, + host_id: str, +) -> None: + database = CloudDatabase(database_url) + owner_id = _unique_id("renew-owner") + device_id = _unique_id("renew-conflict-device") + task_id = _unique_id("renew-conflict-task") + now = datetime(2026, 7, 12, 9, 0, tzinfo=UTC) + initial_expiry = now + timedelta(minutes=1) + + try: + database.repository.upsert_host(owner_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + owner_id, + [_device(device_id, owner_id)], + ) + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="do not renew", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + database.repository.assign_task( + task_id=task_id, + host_id=owner_id, + device_id=device_id, + lease_id="lease-current", + lease_expires_at=initial_expiry, + now=now, + ) + + status = database.repository.renew_lease( + task_id=task_id, + attempt=attempt, + lease_id=lease_id, + host_id=owner_id if host_id == "owner" else _unique_id("foreign"), + lease_expires_at=now + timedelta(minutes=2), + now=now + timedelta(seconds=30), + ) + + assert status == "conflict" + task = database.repository.get_task(task_id) + assert task is not None + assert task.lease_expires_at == initial_expiry + finally: + database.close() + + +def test_expired_or_missing_lease_cannot_be_renewed(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("expired-renew-host") + device_id = _unique_id("expired-renew-device") + task_id = _unique_id("expired-renew-task") + now = datetime(2026, 7, 12, 10, 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="expired renewal", + 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-renew-lease", + lease_expires_at=now + timedelta(seconds=1), + now=now, + ) + + assert ( + database.repository.renew_lease( + task_id=task_id, + attempt=1, + lease_id="expired-renew-lease", + host_id=host_id, + lease_expires_at=now + timedelta(minutes=2), + now=now + timedelta(seconds=2), + ) + == "expired" + ) + assert ( + database.repository.renew_lease( + task_id=_unique_id("missing-task"), + attempt=1, + lease_id="missing-lease", + host_id=host_id, + lease_expires_at=now + timedelta(minutes=2), + now=now, + ) + == "not_found" + ) + finally: + database.close()