131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
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 driver.registry import build_driver_factory
|
|
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
|
|
from storage.device_config import DeviceConfigStore
|
|
|
|
|
|
@dataclass
|
|
class HostAgentApplication:
|
|
client: HostAgentClient
|
|
heartbeat: HeartbeatSynchronizer
|
|
processor: AssignmentProcessor
|
|
|
|
def run(self) -> 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(
|
|
*,
|
|
config: HostAgentConfig | None = None,
|
|
manager: DeviceManager | None = None,
|
|
device_config_store: DeviceConfigStore | None = None,
|
|
) -> HostAgentApplication:
|
|
resolved_config = config or load_host_agent_config()
|
|
resolved_manager = manager or _configured_device_manager(device_config_store)
|
|
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),
|
|
)
|
|
|
|
|
|
def _configured_device_manager(
|
|
config_store: DeviceConfigStore | None,
|
|
) -> DeviceManager:
|
|
manager = DeviceManager()
|
|
store = config_store or DeviceConfigStore()
|
|
for device in store.list():
|
|
manager.register_device(
|
|
device["device_id"],
|
|
build_driver_factory(
|
|
device["driver_type"],
|
|
device["connection_info"],
|
|
),
|
|
name=device["name"],
|
|
driver_type=device["driver_type"],
|
|
connection_info=device["connection_info"],
|
|
)
|
|
return manager
|