59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import TYPE_CHECKING
|
|
|
|
from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse
|
|
from device.manager import DeviceManager
|
|
from host_agent.client import HostAgentClient
|
|
from host_agent.config import HostAgentConfig
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
|
|
def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]:
|
|
return [
|
|
DeviceSnapshotModel(
|
|
device_id=device.id,
|
|
driver_type=device.driver_type,
|
|
status=device.status,
|
|
capability_tags=list(device.capability_tags),
|
|
)
|
|
for device in sorted(manager.list_devices(), key=lambda item: item.id)
|
|
]
|
|
|
|
|
|
class HeartbeatSynchronizer:
|
|
def __init__(
|
|
self,
|
|
manager: DeviceManager,
|
|
client: HostAgentClient,
|
|
config: HostAgentConfig,
|
|
*,
|
|
address: str | None = None,
|
|
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
|
) -> None:
|
|
self.manager = manager
|
|
self.client = client
|
|
self.config = config
|
|
self.address = address
|
|
self._sleep = sleep
|
|
|
|
async def sync_once(self) -> HeartbeatResponse:
|
|
return await self.client.heartbeat(
|
|
build_device_snapshot(self.manager),
|
|
address=self.address,
|
|
)
|
|
|
|
async def run(self, stop: asyncio.Event) -> None:
|
|
while not stop.is_set():
|
|
await self.sync_once()
|
|
try:
|
|
await asyncio.wait_for(
|
|
stop.wait(),
|
|
timeout=self.config.heartbeat_interval_seconds,
|
|
)
|
|
except TimeoutError:
|
|
continue
|