BlueSky Network Client¶

def connect(self, callback):
    """Connect a callback to this signal.

    Idempotent: a callback that is already connected is not added again.

    Args:
        callback (Callable): Callable invoked whenever the signal is
            emitted.
    """
    if callback not in self.callbacks:
        self.callbacks.append(callback)

def disconnect(self, callback):
    """Disconnect a callback from this signal.

    A callback that is not connected is silently ignored.

    Args:
        callback (Callable): Callback to remove.
    """
    if callback in self.callbacks:
        self.callbacks.remove(callback)

def emit(self, *args, **kwargs):
    """Emit the signal to all connected callbacks.

    Iterates over a snapshot of the callback list so callbacks may connect
    or disconnect during emission. Exceptions raised by a callback are
    logged and do not stop delivery to the remaining callbacks.

    Args:
        *args (Any): Positional arguments forwarded to each callback.
        **kwargs (Any): Keyword arguments forwarded to each callback.
    """
    callbacks_snapshot = self.callbacks[
        :
    ]  # Make a copy to avoid concurrency issues
    for callback in callbacks_snapshot:
        try:
            callback(*args, **kwargs)
        except Exception as e:
            logger.warning(f"Signal {self.name}: Error in callback {callback}: {e}")
            import traceback

            traceback.print_exc()

def subscribe(self, topic: str, callback: Callable):
    """Subscribe a callback to a topic.

    Idempotent, like ``BlueSkySignal.connect``: a callback already
    subscribed to the topic is not added again.

    Args:
        topic (str): Topic name to subscribe to.
        callback (Callable): Callable invoked when data is emitted on the
            topic.
    """
    if callback not in self.subscribers[topic]:
        self.subscribers[topic].append(callback)

def emit(self, topic: str, *args, **kwargs):
    """Emit data to all subscribers of a topic.

    Exceptions raised by a callback are logged and do not stop delivery to
    the remaining callbacks. Topics without subscribers are ignored.

    Args:
        topic (str): Topic name to emit on.
        *args (Any): Positional arguments forwarded to each callback.
        **kwargs (Any): Keyword arguments forwarded to each callback.
    """
    for callback in self.subscribers.get(topic, []):
        try:
            callback(*args, **kwargs)
        except Exception as e:
            logger.warning(f"Subscriber {topic}: Error in callback {callback}: {e}")
            logger.debug(f"Subscriber {topic}: Error type: {type(e).__name__}")
            logger.debug(f"Subscriber {topic}: Args: {args}")
            logger.debug(f"Subscriber {topic}: Kwargs: {kwargs}")
            import traceback

            traceback.print_exc()

def stack(self, *cmdlines, sender_id=None):
    """Queue one or more command lines.

    Each command line is stripped of surrounding whitespace; empty lines are
    ignored. Semicolon-separated compound lines are split into individual
    commands, each queued with the given sender ID.

    Args:
        *cmdlines (str): One or more command lines to queue.
        sender_id (bytes | str | None): Identifier of the command originator, stored
            alongside each queued command.
    """
    for cmdline in cmdlines:
        cmdline = cmdline.strip()
        if cmdline:
            for line in cmdline.split(";"):
                self.cmdstack.append((line.strip(), sender_id))

def commands(self):
    """Iterate over queued commands with sender tracking.

    Pops commands from the front of the stack, exposing each one through the
    ``current`` and ``sender_id`` attributes while it is being processed.
    Both attributes are reset when the stack is exhausted.

    Yields:
        str: The next queued command.
    """
    while self.cmdstack:
        self.current, self.sender_id = self.cmdstack.popleft()
        yield self.current
    self.current = ""
    self.sender_id = None

def connect(
    self, hostname="localhost", recv_port=11000, send_port=11001, protocol="tcp"
):
    """Connect to a BlueSky server following the ZMQ pattern.

    Creates the ZMQ context, SUB (data) and XPUB (command) sockets, tunes
    buffer and shutdown options, connects both sockets, registers them with a
    poller, and subscribes to messages targeted at this node so the server
    can discover it. On failure all partially-created resources are released
    via ``close()``.

    Args:
        hostname (str): BlueSky server hostname or IP address.
        recv_port (int): Port for receiving simulation data.
        send_port (int): Port for sending commands and events.
        protocol (str): ZMQ transport protocol (e.g. "tcp").

    Returns:
        bool: True if the connection was established, False otherwise.
    """
    logger.info(f"Connecting to {hostname}:{recv_port}/{send_port}...")

    try:
        # Create ZMQ context and sockets (following ZMQ pattern)
        self.zmq_context = zmq.Context()
        self.sock_recv = self.zmq_context.socket(zmq.SUB)
        self.sock_send = self.zmq_context.socket(zmq.XPUB)
        self.poller = zmq.Poller()

        # Tune buffers and shutdown behaviour BEFORE connecting (HWM/LINGER
        # must be set before connect to take effect):
        #  - RCVHWM/SNDHWM absorb bursty traffic instead of dropping it.
        #  - LINGER=0 makes close()/reconnect return immediately rather than
        #    blocking on undeliverable queued messages, so a disconnect can
        #    never hang the timer/command threads.
        self.sock_recv.setsockopt(zmq.RCVHWM, RECV_HWM)
        self.sock_recv.setsockopt(zmq.LINGER, 0)
        self.sock_send.setsockopt(zmq.SNDHWM, SEND_HWM)
        self.sock_send.setsockopt(zmq.LINGER, 0)

        # Connect sockets
        recv_addr = f"{protocol}://{hostname}:{recv_port}"
        send_addr = f"{protocol}://{hostname}:{send_port}"

        self.sock_recv.connect(recv_addr)
        self.sock_send.connect(send_addr)

        # Register with poller
        self.poller.register(self.sock_recv, zmq.POLLIN)
        self.poller.register(self.sock_send, zmq.POLLIN)

        # CRITICAL: Register this node by subscribing to targeted messages
        # This is what makes the node discoverable to the server!
        self._subscribe("", "", self.node_id)

        self.connected = True
        self.running = True

        logger.info(
            f"Connected to BlueSky server host at {hostname} with node_id={safe_decode(self.node_id)}"
        )
        return True

    except Exception as e:
        logger.error(f"Connection failed: {e}")
        self.close()
        return False

def close(self):
    """Close all connections following ZMQ pattern."""
    logger.debug("Closing connections...")

    self.running = False
    self.connected = False

    # Take the socket lock so teardown cannot race a receive()/send() in
    # another thread (the timer thread may still be mid-poll when a command
    # thread or shutdown triggers close()). RLock keeps this safe even if a
    # close path is re-entered.
    with self._sock_lock:
        self._close_locked()

def update(self):
    """Update function - call periodically to receive and process data."""
    if not self.running or not self.connected:
        return False

    try:
        return self.receive(timeout=0)
    except Exception as e:
        logger.error(f"Error in update: {e}")
        return False

def receive(self, timeout=0):
    """Receive and process incoming messages (following ZMQ recv pattern).

    Socket I/O (poll + recv) is done under ``_sock_lock`` so it cannot race
    a concurrent send() from a command thread. Every currently-available
    message is drained in one pass — a single recv per tick would let a
    fast-forward burst back up for many ticks (stale map, growing latency).
    Handler dispatch runs *outside* the lock: handlers can re-enter the
    client (e.g. node discovery triggers a REQUEST send), and keeping the
    lock hold short means command sends are never blocked by serialisation.
    """
    if not self.running or not self.poller:
        return False

    # (is_data_socket, msg) pairs collected under the lock, dispatched after.
    collected: list[tuple[bool, list]] = []

    try:
        with self._sock_lock:
            if not self.poller:
                return False

            events = dict(self.poller.poll(timeout))

            for sock, event in events.items():
                if event != zmq.POLLIN:
                    continue

                is_data_socket = sock is self.sock_recv

                # Drain everything queued on this socket right now, bounded
                # so a flood can't pin the timer thread indefinitely.
                for _ in range(MAX_DRAIN_PER_SOCKET):
                    try:
                        msg = sock.recv_multipart(zmq.DONTWAIT)
                    except zmq.Again:
                        break
                    if msg:
                        collected.append((is_data_socket, msg))

    except zmq.Again:
        # No messages available (expected with timeout=0)
        return True
    except Exception as e:
        if "ENOTSOCK" in str(e):
            logger.debug("Socket closed during receive")
            return False
        else:
            logger.error(f"Error receiving: {e}")
            return False

    # Dispatch outside the socket lock.
    for is_data_socket, msg in collected:
        if is_data_socket:
            self._process_data_message(msg)
        else:
            self._process_subscription_message(msg)

    return True

def send(self, topic: str, data="", to_group=""):
    """Send data to a topic."""
    if not self.running or not self.sock_send:
        return False

    try:
        btopic = asbytestr(topic)
        bto_group = asbytestr(to_group or "")

        header = bto_group.ljust(IDLEN, b"*") + btopic + self.node_id
        payload = msgpack.packb(data, use_bin_type=True)

        # Serialise socket access against the receive loop / other senders;
        # ZMQ sockets are not thread-safe (see _sock_lock). DONTWAIT means a
        # momentarily-full send buffer raises zmq.Again instead of blocking
        # the caller (a Socket.IO command thread) on the socket.
        with self._sock_lock:
            if not self.sock_send:
                return False
            self.sock_send.send_multipart([header, payload], zmq.DONTWAIT)

        return True

    except zmq.Again:
        logger.warning(f"Send buffer full, dropping '{topic}' command")
        return False
    except Exception as e:
        logger.error(f"Error sending: {e}")
        return False

def subscribe(self, topic: str, callback: Callable, actonly=False):
    """Subscribe to a topic with callback."""
    self.subscriber.subscribe(topic, callback)
    # Subscribe on network level with actonly support
    self._subscribe(topic, GROUPID_DEFAULT, "", actonly=actonly)

def actnode(self, newact=None):
    """Set or get the active simulation node."""
    if newact:
        if newact not in self.nodes:
            logger.error(
                f"Error selecting active node (unknown node): {safe_decode(newact)}"
            )
            return None

        if self.act_id is None:
            # First time selecting active node - disconnect auto-selection
            self.node_added.disconnect(self.actnode)

        # Update subscriptions for new active node
        if newact != self.act_id:
            for topic, groupset in self.acttopics.items():
                for to_group in groupset:
                    if self.act_id:
                        self._unsubscribe(topic, self.act_id, to_group)
                    self._subscribe(topic, newact, to_group)

            self.act_id = newact
            self.actnode_changed.emit(newact)

    return self.act_id

def addnodes(self, count=1, server_id=None):
    """Tell server to add nodes."""
    target_server = server_id or (
        self.act_id[:-1] + seqidx2id(0) if self.act_id else self.server_id
    )
    return self.send("ADDNODES", {"count": count}, target_server)

def delnode(self, node_id):
    """Tell the owning server to terminate a single simulation node.

    The DELNODE message carries the raw node id and is addressed to the
    server that spawned the node (same group, sequence index 0). Servers
    without DELNODE support ignore unknown topics, so this is safe to
    send to older BlueSky versions.
    """
    target_server = node_id[:-1] + seqidx2id(0)
    return self.send("DELNODE", node_id, target_server)

def on_node_added_request_data(self, node_id):
    """When a new node is announced, request the initial/current state of all
    subscribed shared states."""
    logger.info("A new node has been added! request topics")

    # TODO: fix request
    # Request all BlueSky topics we want to receive add #STACK
    # topics = ['RESET', 'REQUEST', 'PLOT', 'SHOWDIALOG', 'SIMINFO',
    #          'SIMSETTINGS', 'TRAILS', 'ROUTEDATA', 'ACDATA', 'DEFWPT',
    #          'POLY', 'STACKCMDS']

    topics = ["POLY", "STACKCMDS"]

    logger.debug(
        f"Requesting topics {topics} from all nodes (triggered by new node {safe_decode(node_id)})"
    )
    self.send("REQUEST", topics)
    self.send("REQUEST", topics)