dynalab-core/test/manual/peer.py

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, monotonic_ns, 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 ValueBatch, 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 = 1000

FREQUENCY_SAMPLE_SIZE = 20000

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 = ValueBatch(values=[])
            for i in range(10):
                value = ValueDescriptor(
                    signal_id=signal_id,
                    value=2.0,
                    timestamp=monotonic_ns(),
                )
                message.values.append(value)

            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", timeout_ms=5000
            )
        ],
    )

    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")