feat(host-agent): shut down gracefully

apps/device-host-agent/host_agent/app.py

@@ -1,15 +1,108 @@
 from __future__ import annotations
 
+import asyncio
+from contextlib import suppress
 from dataclasses import dataclass
 
+from cloud.internal_api.models import AssignmentModel
+from device.manager import DeviceManager
+from host_agent.assignment import AssignmentExecutor
+from host_agent.client import HostAgentClient
+from host_agent.config import HostAgentConfig, load_host_agent_config
+from host_agent.execution import create_execution_factories
+from host_agent.heartbeat import HeartbeatSynchronizer
+from host_agent.lease import ActiveAssignmentRunner
+from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
+
 
 @dataclass
 class HostAgentApplication:
-    """Process shell expanded by the Host Agent implementation tasks."""
+    client: HostAgentClient
+    heartbeat: HeartbeatSynchronizer
+    processor: AssignmentProcessor
 
     def run(self) -> None:
-        return None
+        asyncio.run(self.run_async())
+
+    async def run_async(self, stop: asyncio.Event | None = None) -> None:
+        stop_requested = stop or asyncio.Event()
+        heartbeat_stop = asyncio.Event()
+        heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop))
+        active_processing: asyncio.Task[AssignmentProcessingResult] | None = None
+        try:
+            while not stop_requested.is_set():
+                assignment = await self._claim_until_stopped(stop_requested)
+                if assignment is None:
+                    continue
+                active_processing = asyncio.create_task(
+                    self.processor.process(assignment)
+                )
+                stopped = asyncio.create_task(stop_requested.wait())
+                done, _ = await asyncio.wait(
+                    {active_processing, stopped},
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if stopped in done:
+                    self.processor.request_stop()
+                else:
+                    stopped.cancel()
+                with suppress(asyncio.CancelledError):
+                    await stopped
+                await asyncio.shield(active_processing)
+                active_processing = None
+        finally:
+            self.processor.request_stop()
+            if active_processing is not None:
+                with suppress(Exception):
+                    await asyncio.shield(active_processing)
+            heartbeat_stop.set()
+            try:
+                await asyncio.gather(heartbeat_task, return_exceptions=True)
+                with suppress(Exception):
+                    await self.heartbeat.sync_once()
+            finally:
+                await self.client.aclose()
+
+    async def _claim_until_stopped(
+        self,
+        stop: asyncio.Event,
+    ) -> AssignmentModel | None:
+        claim = asyncio.create_task(self.client.claim())
+        stopped = asyncio.create_task(stop.wait())
+        try:
+            done, _ = await asyncio.wait(
+                {claim, stopped},
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        except asyncio.CancelledError:
+            claim.cancel()
+            stopped.cancel()
+            await asyncio.gather(claim, stopped, return_exceptions=True)
+            raise
+        if stopped in done:
+            claim.cancel()
+            with suppress(asyncio.CancelledError):
+                await claim
+            return None
+        stopped.cancel()
+        with suppress(asyncio.CancelledError):
+            await stopped
+        return await claim
 
 
-def create_application() -> HostAgentApplication:
-    return HostAgentApplication()
+def create_application(
+    *,
+    config: HostAgentConfig | None = None,
+    manager: DeviceManager | None = None,
+) -> HostAgentApplication:
+    resolved_config = config or load_host_agent_config()
+    resolved_manager = manager or DeviceManager()
+    client = HostAgentClient(resolved_config)
+    heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config)
+    executor = AssignmentExecutor(create_execution_factories(resolved_manager))
+    active_runner = ActiveAssignmentRunner(client, executor)
+    return HostAgentApplication(
+        client=client,
+        heartbeat=heartbeat,
+        processor=AssignmentProcessor(client, active_runner),
+    )

apps/device-host-agent/host_agent/lease.py

@@ -54,6 +54,10 @@ class ActiveAssignmentRunner:
         self.client = client
         self.executor = executor
         self._now = now or (lambda: datetime.now(UTC))
+        self._stop_requested = Event()
+
+    def request_stop(self) -> None:
+        self._stop_requested.set()
 
     async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult:
         guard = LeaseGuard()
@@ -61,7 +65,8 @@ class ActiveAssignmentRunner:
             asyncio.to_thread(
                 self.executor.execute,
                 assignment,
-                should_stop=guard.is_lost,
+                should_stop=lambda: guard.is_lost()
+                or self._stop_requested.is_set(),
             )
         )
         renewal = asyncio.create_task(

apps/device-host-agent/host_agent/processor.py

@@ -11,6 +11,8 @@ from host_agent.client import HostAgentClient
 class ActiveAssignmentExecutor(Protocol):
     async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: ...
 
+    def request_stop(self) -> None: ...
+
 
 @dataclass(frozen=True)
 class AssignmentProcessingResult:
@@ -27,6 +29,9 @@ class AssignmentProcessor:
         self.client = client
         self.active_executor = active_executor
 
+    def request_stop(self) -> None:
+        self.active_executor.request_stop()
+
     async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
         execution = await self.active_executor.run(assignment)
         status = "done" if execution.status == "done" else "failed"