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 math import sin, tau
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 import (
    JsonConnectorHello,
    JsonDynaLabHello,
    JsonHeartbeat,
    JsonValueBatch,
    JsonValueDescriptor,
)
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
SINE_FREQUENCY_HZ = 0.5

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_1: UUID = UUID("3f32cea3-d872-4c16-a0a2-54b57171aeb2")
signal_id_2: UUID = UUID("6578ac37-99d1-410c-b3f7-d049339919b2")


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
    sine_start_time = perf_counter()

    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

            sine_value = sin(
                tau * SINE_FREQUENCY_HZ * (perf_counter() - sine_start_time)
            )
            timestamp = monotonic_ns()
            message = JsonValueBatch(values=[])
            value1 = JsonValueDescriptor(
                signal_id=signal_id_1,
                value=2.0 * sine_value,
                timestamp=timestamp,
            )
            value2 = JsonValueDescriptor(
                signal_id=signal_id_2,
                value=sine_value,
                timestamp=timestamp,
            )
            message.values.append(value1)
            message.values.append(value2)

            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, JsonDynaLabHello):
        raise RuntimeError(f"Expected DynaLabHello, received {type(message).__name__}")

    log.debug(
        "Received hello from core instance %s",
        message.instance_id,
    )

    connector_hello = JsonConnectorHello(
        connector_uuid=uuid4(),
        protocol_version=PROTOCOL_VERSION,
        connector_name="Test connector",
        connector_version=CONNECTOR_VERSION.get_version(),
        signals=[
            SignalDescriptor(
                id=signal_id_1, name="Dummy signal 1", type="number", timeout_ms=5000
            ),
            SignalDescriptor(
                id=signal_id_2, name="Dummy signal 2", 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, JsonHeartbeat):
            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")