This commit is contained in:
@@ -8,6 +8,7 @@ from core.errors import DeviceRuntimeError
|
||||
from device.manager import DeviceManager
|
||||
from host_agent.client import HostAgentClient
|
||||
from host_agent.config import HostAgentConfig
|
||||
from host_agent.status import AgentStatusTracker
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
@@ -34,18 +35,31 @@ class HeartbeatSynchronizer:
|
||||
*,
|
||||
address: str | None = None,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
status_tracker: AgentStatusTracker | None = None,
|
||||
on_sync: Callable[[int], None] | None = None,
|
||||
) -> None:
|
||||
self.manager = manager
|
||||
self.client = client
|
||||
self.config = config
|
||||
self.address = address
|
||||
self._sleep = sleep
|
||||
self.status_tracker = status_tracker
|
||||
self.on_sync = on_sync
|
||||
|
||||
async def sync_once(self) -> HeartbeatResponse:
|
||||
return await self.client.heartbeat(
|
||||
build_device_snapshot(self.manager),
|
||||
snapshot = build_device_snapshot(self.manager)
|
||||
response = await self.client.heartbeat(
|
||||
snapshot,
|
||||
address=self.address,
|
||||
)
|
||||
if self.status_tracker is not None:
|
||||
self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot))
|
||||
if self.on_sync is not None:
|
||||
try:
|
||||
self.on_sync(len(snapshot))
|
||||
except Exception:
|
||||
pass
|
||||
return response
|
||||
|
||||
async def run(self, stop: asyncio.Event) -> None:
|
||||
self.connect_devices()
|
||||
|
||||
Reference in New Issue
Block a user