feat(cloud-scheduler): renew active leases

2026-07-12 17:23:07 +08:00
parent b0a5b1426f
commit 1b15a8a218
4 changed files with 225 additions and 2 deletions
@@ -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.
@@ -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,
@@ -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(
@@ -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()