fa273bcc26
- Add Python toki_socket package (base_client, communicator, tcp/ws client/server) - Add protobuf message definitions for Python - Add go_python.go cross-test implementation and python_go_client - Add agent-task/python_impl plan and code review docs - Update .gitignore for Python cache files with recursive patterns
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from toki_socket.communicator import Communicator
|
|
from toki_socket.packets.message_common_pb2 import HeartBeat
|
|
|
|
|
|
class _HbTimer:
|
|
def __init__(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
|
|
self._task = asyncio.create_task(self._run(delay, callback))
|
|
|
|
async def _run(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
|
|
await asyncio.sleep(delay)
|
|
await callback()
|
|
|
|
def cancel(self) -> None:
|
|
self._task.cancel()
|
|
|
|
|
|
class BaseClient:
|
|
def __init__(
|
|
self,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
do_close: Callable[[], Awaitable[None]],
|
|
) -> None:
|
|
self.communicator = Communicator()
|
|
self._interval_sec = interval_sec
|
|
self._wait_sec = wait_sec
|
|
self._do_close = do_close
|
|
self._close_called = False
|
|
self._close_lock = asyncio.Lock()
|
|
self._disconnect_listeners: list[Callable[[Any], Any]] = []
|
|
self._hb_timer: _HbTimer | None = None
|
|
self._waiting_hb_response = False
|
|
|
|
def _init_base(self) -> None:
|
|
self._close_lock = asyncio.Lock()
|
|
|
|
async def close(self) -> None:
|
|
async with self._close_lock:
|
|
if self._close_called:
|
|
return
|
|
self._close_called = True
|
|
|
|
self.communicator._shutdown()
|
|
self._stop_heartbeat()
|
|
await self._do_close()
|
|
await self._notify_disconnected()
|
|
|
|
def add_disconnect_listener(self, fn: Callable[[Any], Any]) -> None:
|
|
self._disconnect_listeners.append(fn)
|
|
|
|
def remove_disconnect_listeners(self) -> None:
|
|
self._disconnect_listeners.clear()
|
|
|
|
async def _notify_disconnected(self) -> None:
|
|
listeners = list(self._disconnect_listeners)
|
|
self._disconnect_listeners.clear()
|
|
for listener in listeners:
|
|
result = listener(self)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
|
|
async def _on_disconnected(self) -> None:
|
|
await self.close()
|
|
|
|
async def _send_heartbeat(self) -> None:
|
|
if not self.communicator.is_alive() or self._interval_sec <= 0:
|
|
return
|
|
self._stop_heartbeat()
|
|
self._hb_timer = _HbTimer(self._interval_sec, self._heartbeat_interval_elapsed)
|
|
|
|
async def _heartbeat_interval_elapsed(self) -> None:
|
|
if not self.communicator.is_alive():
|
|
return
|
|
self._waiting_hb_response = True
|
|
with suppress(Exception):
|
|
await self.communicator.send(HeartBeat())
|
|
self._stop_heartbeat()
|
|
if self._wait_sec > 0:
|
|
self._hb_timer = _HbTimer(self._wait_sec, self._heartbeat_wait_elapsed)
|
|
|
|
async def _heartbeat_wait_elapsed(self) -> None:
|
|
if self.communicator.is_alive():
|
|
await self._on_disconnected()
|
|
|
|
async def _on_heartbeat(self) -> None:
|
|
if self._waiting_hb_response:
|
|
self._waiting_hb_response = False
|
|
return
|
|
with suppress(Exception):
|
|
await self.communicator.send(HeartBeat())
|
|
|
|
def _stop_heartbeat(self) -> None:
|
|
if self._hb_timer is not None:
|
|
self._hb_timer.cancel()
|
|
self._hb_timer = None
|
|
|