33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
"""Claim and lease management helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .errors import LockedError
|
|
from .models import ProbeEvidence
|
|
from .store import TaskPlaneStore
|
|
|
|
|
|
class LeaseManager:
|
|
def __init__(self, store: TaskPlaneStore):
|
|
self.store = store
|
|
|
|
def claim(self, task_id: str, worker_id: str, worker_identity: str, lease_seconds: int, expected_revision: int, operation_digest: str):
|
|
return self.store.claim_task(task_id, worker_id, worker_identity, lease_seconds, expected_revision, operation_digest)
|
|
|
|
def heartbeat(self, task_id: str, claim_id: str, lease_seconds: int):
|
|
claim = self.store.heartbeat_claim(task_id, claim_id, lease_seconds)
|
|
return claim
|
|
|
|
def probe(self, claim_id: str):
|
|
claim = self.store.get_claim(claim_id)
|
|
if claim.status != "ACTIVE":
|
|
raise LockedError("claim not active")
|
|
return claim
|
|
|
|
def recover(self, task_id: str, claim_id: str, probe: ProbeEvidence):
|
|
return self.store.recover_expired_claim(
|
|
task_id,
|
|
claim_id,
|
|
probe,
|
|
)
|