feat(cloud-scheduler): renew active leases
This commit is contained in:
@@ -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.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.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.
|
- [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.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.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.
|
- [ ] 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"]
|
AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
|
||||||
TerminalTaskStatus = Literal["done", "failed"]
|
TerminalTaskStatus = Literal["done", "failed"]
|
||||||
ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
|
ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
|
||||||
|
LeaseRenewalStatus = Literal["renewed", "not_found", "conflict", "expired"]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -119,7 +120,7 @@ class CloudRepository(Protocol):
|
|||||||
host_id: str,
|
host_id: str,
|
||||||
lease_expires_at: datetime,
|
lease_expires_at: datetime,
|
||||||
now: datetime,
|
now: datetime,
|
||||||
) -> bool: ...
|
) -> LeaseRenewalStatus: ...
|
||||||
|
|
||||||
def record_task_result(
|
def record_task_result(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -333,6 +333,56 @@ class SQLAlchemyCloudRepository:
|
|||||||
session.flush()
|
session.flush()
|
||||||
return _leased_assignment_from_row(task)
|
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]:
|
def list_task_attempts(self, task_id: str) -> list[Any]:
|
||||||
with self._sessions() as session:
|
with self._sessions() as session:
|
||||||
rows = session.scalars(
|
rows = session.scalars(
|
||||||
|
|||||||
@@ -577,3 +577,175 @@ def test_expired_assignment_cannot_be_claimed(database_url: str) -> None:
|
|||||||
assert task.status == "assigned"
|
assert task.status == "assigned"
|
||||||
finally:
|
finally:
|
||||||
database.close()
|
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()
|
||||||
|
|||||||
Reference in New Issue
Block a user