345 lines
8.9 KiB
Python
345 lines
8.9 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
from concurrent.futures import CancelledError as FutureCancelledError
|
|
from concurrent.futures import TimeoutError as FutureTimeoutError
|
|
from statistics import mean
|
|
from time import monotonic, perf_counter
|
|
from uuid import UUID, uuid4
|
|
|
|
from rich.logging import RichHandler
|
|
|
|
from dynalab_core.protocols.common import VersionDescriptor
|
|
from dynalab_core.protocols.constants import PROTOCOL_VERSION
|
|
from dynalab_core.protocols.json.wire import read_message, write_message
|
|
from dynalab_core.protocols.packets.data import ValueDescriptor
|
|
from dynalab_core.protocols.packets.handshake import (
|
|
ConnectorHello,
|
|
DynaLabHello,
|
|
SignalDescriptor,
|
|
)
|
|
from dynalab_core.protocols.packets.heartbeat import Heartbeat
|
|
|
|
|
|
HOST = "127.0.0.1"
|
|
PORT = 8765
|
|
|
|
# Set to None to send as quickly as possible.
|
|
# For a controlled rate, use something like 5_000.0.
|
|
TARGET_FREQUENCY_HZ: float | None = None
|
|
|
|
FREQUENCY_SAMPLE_SIZE = 500
|
|
|
|
CONNECTOR_VERSION = VersionDescriptor(
|
|
type="alpha",
|
|
major=0,
|
|
minor=0,
|
|
patch=1,
|
|
)
|
|
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
format="%(name)s %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
handlers=[
|
|
RichHandler(
|
|
rich_tracebacks=True,
|
|
show_path=False,
|
|
)
|
|
],
|
|
force=True,
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
signal_id: UUID = uuid4()
|
|
|
|
|
|
def monotonic_ms() -> int:
|
|
return round(monotonic() * 1_000)
|
|
|
|
|
|
async def send_message(
|
|
writer: asyncio.StreamWriter,
|
|
write_lock: asyncio.Lock,
|
|
message: object,
|
|
) -> None:
|
|
"""
|
|
Serialize access to the asyncio StreamWriter.
|
|
|
|
Both heartbeat replies and value messages use this function.
|
|
"""
|
|
async with write_lock:
|
|
await write_message(writer, message)
|
|
|
|
|
|
def value_sender_thread(
|
|
loop: asyncio.AbstractEventLoop,
|
|
writer: asyncio.StreamWriter,
|
|
write_lock: asyncio.Lock,
|
|
stop_event: threading.Event,
|
|
) -> None:
|
|
"""
|
|
Run the high-frequency timing loop in a dedicated OS thread.
|
|
|
|
The asyncio StreamWriter is not thread-safe, so each write is submitted
|
|
back to the asyncio event loop with run_coroutine_threadsafe().
|
|
"""
|
|
intervals: list[float] = []
|
|
previous_send_time: float | None = None
|
|
|
|
if TARGET_FREQUENCY_HZ is not None:
|
|
period = 1.0 / TARGET_FREQUENCY_HZ
|
|
next_send_time = perf_counter()
|
|
else:
|
|
period = None
|
|
next_send_time = 0.0
|
|
|
|
log.debug("Value sender thread started")
|
|
|
|
try:
|
|
while not stop_event.is_set():
|
|
if period is not None:
|
|
while True:
|
|
remaining = next_send_time - perf_counter()
|
|
|
|
if remaining <= 0:
|
|
break
|
|
|
|
# Sleep for larger remaining times, then spin for the
|
|
# final fraction of a millisecond.
|
|
if remaining > 0.001:
|
|
stop_event.wait(remaining - 0.0005)
|
|
|
|
if stop_event.is_set():
|
|
return
|
|
|
|
next_send_time += period
|
|
|
|
message = ValueDescriptor(
|
|
signal_id=signal_id,
|
|
value=2.0,
|
|
)
|
|
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
send_message(
|
|
writer,
|
|
write_lock,
|
|
message,
|
|
),
|
|
loop,
|
|
)
|
|
|
|
try:
|
|
# Waiting prevents an ever-growing queue of scheduled writes.
|
|
future.result(timeout=2.0)
|
|
except FutureTimeoutError:
|
|
log.warning("Timed out sending value message")
|
|
future.cancel()
|
|
return
|
|
except FutureCancelledError:
|
|
return
|
|
except Exception:
|
|
log.exception("Value sender failed")
|
|
return
|
|
|
|
current_send_time = perf_counter()
|
|
|
|
if previous_send_time is not None:
|
|
intervals.append(current_send_time - previous_send_time)
|
|
|
|
previous_send_time = current_send_time
|
|
|
|
if len(intervals) >= FREQUENCY_SAMPLE_SIZE:
|
|
average_interval = mean(intervals)
|
|
frequency_hz = 1.0 / average_interval
|
|
|
|
log.debug(
|
|
"Value frequency: %.3f kHz",
|
|
frequency_hz / 1_000.0,
|
|
)
|
|
|
|
intervals.clear()
|
|
|
|
if period is not None:
|
|
# Do not try to catch up by rapidly sending many old periods.
|
|
current_time = perf_counter()
|
|
|
|
if next_send_time < current_time - period:
|
|
missed_periods = int((current_time - next_send_time) / period)
|
|
next_send_time += missed_periods * period
|
|
|
|
finally:
|
|
log.debug("Value sender thread stopped")
|
|
|
|
|
|
async def perform_handshake(
|
|
reader: asyncio.StreamReader,
|
|
writer: asyncio.StreamWriter,
|
|
write_lock: asyncio.Lock,
|
|
) -> None:
|
|
try:
|
|
message = await asyncio.wait_for(
|
|
read_message(reader),
|
|
timeout=30.0,
|
|
)
|
|
except TimeoutError as exc:
|
|
raise TimeoutError("Timed out waiting for server hello") from exc
|
|
|
|
if not isinstance(message, DynaLabHello):
|
|
raise RuntimeError(f"Expected DynaLabHello, received {type(message).__name__}")
|
|
|
|
log.debug(
|
|
"Received hello from core instance %s",
|
|
message.instance_id,
|
|
)
|
|
|
|
connector_hello = ConnectorHello(
|
|
connector_uuid=uuid4(),
|
|
protocol_version=PROTOCOL_VERSION,
|
|
connector_name="Test connector",
|
|
connector_version=CONNECTOR_VERSION.get_version(),
|
|
signals=[
|
|
SignalDescriptor(
|
|
id=signal_id,
|
|
name="Dummy signal",
|
|
type="number",
|
|
)
|
|
],
|
|
)
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
send_message(
|
|
writer,
|
|
write_lock,
|
|
connector_hello,
|
|
),
|
|
timeout=5.0,
|
|
)
|
|
except TimeoutError as exc:
|
|
raise TimeoutError("Timed out sending connector hello") from exc
|
|
|
|
log.info(
|
|
"Connector hello sent for %s",
|
|
connector_hello.connector_uuid,
|
|
)
|
|
|
|
|
|
async def receive_messages(
|
|
reader: asyncio.StreamReader,
|
|
writer: asyncio.StreamWriter,
|
|
write_lock: asyncio.Lock,
|
|
) -> None:
|
|
while True:
|
|
message = await read_message(reader)
|
|
|
|
if isinstance(message, Heartbeat):
|
|
message.return_timestamp = monotonic_ms()
|
|
|
|
await send_message(
|
|
writer,
|
|
write_lock,
|
|
message,
|
|
)
|
|
|
|
log.debug(
|
|
"Returned heartbeat %s",
|
|
message.sequence,
|
|
)
|
|
else:
|
|
log.debug(
|
|
"Received unexpected message: %s",
|
|
type(message).__name__,
|
|
)
|
|
|
|
|
|
async def main() -> None:
|
|
writer: asyncio.StreamWriter | None = None
|
|
sender_thread: threading.Thread | None = None
|
|
sender_stop_event = threading.Event()
|
|
|
|
try:
|
|
reader, writer = await asyncio.open_connection(
|
|
host=HOST,
|
|
port=PORT,
|
|
)
|
|
|
|
log.info(
|
|
"Connected to DynaLab core at %s:%d",
|
|
HOST,
|
|
PORT,
|
|
)
|
|
|
|
write_lock = asyncio.Lock()
|
|
|
|
await perform_handshake(
|
|
reader,
|
|
writer,
|
|
write_lock,
|
|
)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
sender_thread = threading.Thread(
|
|
target=value_sender_thread,
|
|
name="value_sender",
|
|
args=(
|
|
loop,
|
|
writer,
|
|
write_lock,
|
|
sender_stop_event,
|
|
),
|
|
daemon=True,
|
|
)
|
|
sender_thread.start()
|
|
|
|
await receive_messages(
|
|
reader,
|
|
writer,
|
|
write_lock,
|
|
)
|
|
|
|
except TimeoutError as exc:
|
|
log.warning("%s", exc)
|
|
|
|
except asyncio.IncompleteReadError:
|
|
log.warning("DynaLab core closed the connection")
|
|
|
|
except ConnectionError as exc:
|
|
log.warning("Connection error: %s", exc)
|
|
|
|
except RuntimeError as exc:
|
|
log.warning("Protocol error: %s", exc)
|
|
|
|
finally:
|
|
sender_stop_event.set()
|
|
|
|
if sender_thread is not None:
|
|
# Do not call thread.join() directly because that would block the
|
|
# asyncio event loop while the thread may be waiting on it.
|
|
await asyncio.to_thread(
|
|
sender_thread.join,
|
|
3.0,
|
|
)
|
|
|
|
if sender_thread.is_alive():
|
|
log.warning("Value sender thread did not stop cleanly")
|
|
|
|
if writer is not None:
|
|
writer.close()
|
|
|
|
try:
|
|
await writer.wait_closed()
|
|
except ConnectionError:
|
|
pass
|
|
|
|
log.info("Disconnected from DynaLab core")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
log.info("Stopped by user")
|