feat(host-agent): shut down gracefully
This commit is contained in:
@@ -1,10 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from cloud.internal_api.models import AssignmentModel
|
||||
from device.manager import DeviceManager
|
||||
from host_agent.app import HostAgentApplication, create_application
|
||||
from host_agent.config import HostAgentConfig
|
||||
|
||||
|
||||
def test_create_application_returns_host_agent_process_shell() -> None:
|
||||
application = create_application()
|
||||
def _config() -> HostAgentConfig:
|
||||
return HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
host_id="host-a",
|
||||
token="secret",
|
||||
)
|
||||
|
||||
|
||||
def _assignment() -> AssignmentModel:
|
||||
return AssignmentModel(
|
||||
task_id="task-a",
|
||||
attempt=1,
|
||||
lease_id="lease-a",
|
||||
lease_expires_at=datetime.now(UTC) + timedelta(seconds=30),
|
||||
host_id="host-a",
|
||||
device_id="device-a",
|
||||
goal="open settings",
|
||||
)
|
||||
|
||||
|
||||
def test_create_application_composes_host_agent_services(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
application = create_application(config=_config(), manager=DeviceManager())
|
||||
|
||||
assert isinstance(application, HostAgentApplication)
|
||||
assert application.run() is None
|
||||
asyncio.run(application.client.aclose())
|
||||
|
||||
|
||||
def test_shutdown_cancels_long_poll_and_sends_final_heartbeat() -> None:
|
||||
async def scenario() -> None:
|
||||
claim_started = asyncio.Event()
|
||||
claim_cancelled = asyncio.Event()
|
||||
events: list[str] = []
|
||||
|
||||
class BlockingClient:
|
||||
async def claim(self):
|
||||
claim_started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
claim_cancelled.set()
|
||||
raise
|
||||
|
||||
async def aclose(self):
|
||||
events.append("closed")
|
||||
|
||||
class RecordingHeartbeat:
|
||||
async def run(self, stop):
|
||||
await stop.wait()
|
||||
|
||||
async def sync_once(self):
|
||||
events.append("final-heartbeat")
|
||||
|
||||
class IdleProcessor:
|
||||
async def process(self, assignment):
|
||||
raise AssertionError("no assignment expected")
|
||||
|
||||
def request_stop(self):
|
||||
events.append("stop-work")
|
||||
|
||||
stop = asyncio.Event()
|
||||
application = HostAgentApplication(
|
||||
client=BlockingClient(), # type: ignore[arg-type]
|
||||
heartbeat=RecordingHeartbeat(), # type: ignore[arg-type]
|
||||
processor=IdleProcessor(), # type: ignore[arg-type]
|
||||
)
|
||||
running = asyncio.create_task(application.run_async(stop))
|
||||
await claim_started.wait()
|
||||
stop.set()
|
||||
await asyncio.wait_for(running, timeout=1)
|
||||
|
||||
assert claim_cancelled.is_set()
|
||||
assert events == ["stop-work", "final-heartbeat", "closed"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_shutdown_interrupts_active_work_before_final_heartbeat() -> None:
|
||||
async def scenario() -> None:
|
||||
processing_started = asyncio.Event()
|
||||
processing_stopped = asyncio.Event()
|
||||
events: list[str] = []
|
||||
claims = 0
|
||||
|
||||
class AssignedClient:
|
||||
async def claim(self):
|
||||
nonlocal claims
|
||||
claims += 1
|
||||
if claims == 1:
|
||||
return _assignment()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
events.append("closed")
|
||||
|
||||
class RecordingHeartbeat:
|
||||
async def run(self, stop):
|
||||
await stop.wait()
|
||||
|
||||
async def sync_once(self):
|
||||
events.append("final-heartbeat")
|
||||
|
||||
class CooperativeProcessor:
|
||||
async def process(self, assignment):
|
||||
processing_started.set()
|
||||
await processing_stopped.wait()
|
||||
events.append("work-finished")
|
||||
|
||||
def request_stop(self):
|
||||
events.append("stop-work")
|
||||
processing_stopped.set()
|
||||
|
||||
stop = asyncio.Event()
|
||||
application = HostAgentApplication(
|
||||
client=AssignedClient(), # type: ignore[arg-type]
|
||||
heartbeat=RecordingHeartbeat(), # type: ignore[arg-type]
|
||||
processor=CooperativeProcessor(), # type: ignore[arg-type]
|
||||
)
|
||||
running = asyncio.create_task(application.run_async(stop))
|
||||
await processing_started.wait()
|
||||
stop.set()
|
||||
await asyncio.wait_for(running, timeout=1)
|
||||
|
||||
assert events.index("work-finished") < events.index("final-heartbeat")
|
||||
assert events[-1] == "closed"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_final_heartbeat_failure_does_not_prevent_client_close() -> None:
|
||||
async def scenario() -> None:
|
||||
closed = False
|
||||
|
||||
class StoppedClient:
|
||||
async def claim(self):
|
||||
raise AssertionError("polling must not start")
|
||||
|
||||
async def aclose(self):
|
||||
nonlocal closed
|
||||
closed = True
|
||||
|
||||
class FailingHeartbeat:
|
||||
async def run(self, stop):
|
||||
await stop.wait()
|
||||
|
||||
async def sync_once(self):
|
||||
raise OSError("control plane unavailable")
|
||||
|
||||
class IdleProcessor:
|
||||
def request_stop(self):
|
||||
return None
|
||||
|
||||
stop = asyncio.Event()
|
||||
stop.set()
|
||||
application = HostAgentApplication(
|
||||
client=StoppedClient(), # type: ignore[arg-type]
|
||||
heartbeat=FailingHeartbeat(), # type: ignore[arg-type]
|
||||
processor=IdleProcessor(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
await application.run_async(stop)
|
||||
assert closed
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_main_task_cancellation_waits_for_active_work_shutdown() -> None:
|
||||
async def scenario() -> None:
|
||||
processing_started = asyncio.Event()
|
||||
processing_stopped = asyncio.Event()
|
||||
events: list[str] = []
|
||||
|
||||
class AssignedClient:
|
||||
async def claim(self):
|
||||
return _assignment()
|
||||
|
||||
async def aclose(self):
|
||||
events.append("closed")
|
||||
|
||||
class RecordingHeartbeat:
|
||||
async def run(self, stop):
|
||||
await stop.wait()
|
||||
|
||||
async def sync_once(self):
|
||||
events.append("final-heartbeat")
|
||||
|
||||
class CooperativeProcessor:
|
||||
async def process(self, assignment):
|
||||
processing_started.set()
|
||||
await processing_stopped.wait()
|
||||
events.append("work-finished")
|
||||
|
||||
def request_stop(self):
|
||||
processing_stopped.set()
|
||||
|
||||
application = HostAgentApplication(
|
||||
client=AssignedClient(), # type: ignore[arg-type]
|
||||
heartbeat=RecordingHeartbeat(), # type: ignore[arg-type]
|
||||
processor=CooperativeProcessor(), # type: ignore[arg-type]
|
||||
)
|
||||
running = asyncio.create_task(application.run_async())
|
||||
await processing_started.wait()
|
||||
running.cancel()
|
||||
|
||||
with suppress(asyncio.CancelledError):
|
||||
await running
|
||||
|
||||
assert events == ["work-finished", "final-heartbeat", "closed"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -126,3 +126,39 @@ def test_renewal_loop_exits_when_execution_finishes() -> None:
|
||||
assert renew_calls == 0
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_shutdown_request_stops_active_execution_cooperatively() -> None:
|
||||
async def scenario() -> None:
|
||||
execution_started = Event()
|
||||
|
||||
class CooperativeExecutor:
|
||||
def execute(self, assignment, *, should_stop=None):
|
||||
assert should_stop is not None
|
||||
execution_started.set()
|
||||
while not should_stop():
|
||||
Event().wait(0.001)
|
||||
return AssignmentExecutionResult(
|
||||
status="failed",
|
||||
failure_reason="execution interrupted",
|
||||
)
|
||||
|
||||
class RenewingClient:
|
||||
async def renew(self, assignment):
|
||||
return LeaseRenewalResponse(
|
||||
status="renewed",
|
||||
lease_expires_at=datetime.now(UTC) + timedelta(seconds=30),
|
||||
)
|
||||
|
||||
runner = ActiveAssignmentRunner(
|
||||
RenewingClient(), # type: ignore[arg-type]
|
||||
CooperativeExecutor(),
|
||||
)
|
||||
running = asyncio.create_task(runner.run(_assignment(expires_in=30)))
|
||||
assert await asyncio.to_thread(execution_started.wait, 1)
|
||||
runner.request_stop()
|
||||
|
||||
result = await asyncio.wait_for(running, timeout=1)
|
||||
assert result.failure_reason == "execution interrupted"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
Reference in New Issue
Block a user