feat(host-agent): shut down gracefully

2026-07-12 19:47:50 +08:00
parent 91f08509a7
commit b1303569e3
6 changed files with 360 additions and 9 deletions
@@ -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),
+    )