robosim.network.connection_manager

  1"""connection_manager.py - Handles the connection protocol and handshakes of all clients talking to the simulator"""
  2from typing import Callable
  3from queue import Empty
  4import atexit
  5import traceback
  6import threading
  7import socket
  8
  9from robosim.utils import logger
 10from robosim.constants import SOCKET_TIMEOUT_S
 11from .msgpack_utils import send_packet, recv_packet
 12from .channel import ObjToSimChannel
 13
 14NetworkHandlerType = Callable[[dict, ObjToSimChannel, ObjToSimChannel | None], dict | None]
 15
 16class ConnectionManager:
 17    """connection class for managing connections and callbacks to the app. App defines the rest of the protocol"""
 18    def __init__(self, host: str, port: int, n_channels: int, connect_command: str):
 19        self.host = host
 20        self.port = port
 21        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 22        self.connect_command = connect_command
 23
 24        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 25        self.server.bind((host, port))
 26        self.server.listen(1)
 27        self.server.settimeout(SOCKET_TIMEOUT_S)
 28        atexit.register(lambda: self.server.close())
 29
 30        self.channels = [ObjToSimChannel(index=i) for i in range (n_channels)] # used to communicate with robots.
 31        self.global_channel = ObjToSimChannel(index=-1) # used to communicate with the simulator, not with a robot.
 32        self.tcp_thread: threading.Thread | None = None # created in start()
 33
 34    def start(self, network_handler: NetworkHandlerType):
 35        """Starts the main thread that listens to new connections"""
 36        args = (self.server, self.channels, network_handler, self.global_channel)
 37        self.tcp_thread = threading.Thread(target=self._tcp_listener, args=args, daemon=True)
 38        self.tcp_thread.start()
 39        logger.info(f"Started TCP server on port {self.port}")
 40
 41    def _tcp_listener(self, server: socket.socket, channels: list[ObjToSimChannel], network_handler: NetworkHandlerType,
 42                      global_channel: ObjToSimChannel):
 43        """the thread that listens to the standard input for commands"""
 44        while True: # server runs forever waiting for new connection
 45            try:
 46                client, addr = server.accept()
 47                logger.info(f"Connected: '{addr[0]}:{addr[1]}'")
 48                args = (client, addr, channels, network_handler, global_channel)
 49                threading.Thread(target=self._handle_client, args=args, daemon=True).start()
 50            except socket.timeout:
 51                pass
 52            except Exception as e:
 53                logger.error(f"Error: {e}\n{''.join(traceback.format_exception(e))}")
 54
 55    def _get_first_free_channel(self, client: socket.socket, channels: list[ObjToSimChannel]) -> int:
 56        # find first free channel to assign this client the robot
 57        for i, channel in enumerate(channels):
 58            with channel.lock:
 59                if channel.client is None:
 60                    channel.client = client
 61                    return i
 62        return -1
 63
 64    def _handle_stateless_client(self, client: socket.socket, addr: tuple[str, int],
 65                                 network_handler: NetworkHandlerType,
 66                                 global_channel: ObjToSimChannel) -> bool:
 67        """We keep the TCP connection on, so the same socket can be re-used, but they are not supposed to hold any app
 68           level state. Returns true if the client needs to be closed, false if we continue (connected)"""
 69        try:
 70            while True:
 71                msg = recv_packet(client)
 72                if "cmd" not in msg:
 73                    send_packet(client, {"error": "'cmd' not in message"}, close_connection=True)
 74                    return True
 75
 76                if msg["cmd"] == self.connect_command:
 77                    # Note: this message will be replied to in _handle_client
 78                    return False
 79
 80                res = network_handler(msg, global_channel, connected_channel=None)
 81                # network_handler returned none -> message was not handled (needed connection) -> close
 82                if res is None:
 83                    return True
 84                else:
 85                    send_packet(client, res)
 86        except ConnectionError:
 87            logger.info(f"Disconnected: '{addr[0]}:{addr[1]}'")
 88        except Exception as e:
 89            logger.error(f"Error: {e}\n{''.join(traceback.format_exception(e))}")
 90        return True
 91
 92    def _handle_client(self, client: socket.socket, addr: tuple, channels: list[ObjToSimChannel],
 93                       network_handler: NetworkHandlerType, global_channel: ObjToSimChannel):
 94        # A client is handled in two steps: either it's a stateles connection or a robot connection. In the second case,
 95        # the client asks for robot control ('connect') in which we first check if the robot is already controlled
 96        # or keep it long term. Further communication with the client follows further allows stateless messages too.
 97        should_close = self._handle_stateless_client(client, addr, network_handler, global_channel)
 98        if should_close:
 99            send_packet(client, {"error": "Must connect first"}, close_connection=True)
100            return
101
102        if (channel_ix := self._get_first_free_channel(client, channels)) == -1:
103            send_packet(client, {"error": f"All {len(channels)} robots are already controlled"},
104                        close_connection=True)
105            return
106
107        # if we reach here, the user is connected and we keep this thread alive
108        send_packet(client, {"status": "connected", "id": channel_ix})
109        channel = channels[channel_ix]
110
111        try:
112            while True:
113                msg = recv_packet(client)
114                logger.log_every_s(f"Received: {msg}", "DEBUG", log_to_next_level=True)
115                # two-tier message handling: fast handler -> simulator queue via the channel (slowest)
116                # we past the message to the network handler which will route it properly (fast, global or connected)
117                if "cmd" not in msg:
118                    res = {"error": "'cmd' not in message"}
119                elif msg["cmd"] == self.connect_command:
120                    res = {"error": "robot already controlled"}
121                else:
122                    res = network_handler(msg, global_channel, connected_channel=channel)
123                send_packet(client, res)
124        except Empty as e:
125            logger.error(f"Timeout waiting for simulator response: {e}\n{''.join(traceback.format_exception(e))}")
126            raise e
127        except ConnectionError:
128            pass
129        except Exception as e:
130            logger.error(f"Error: {e}\n{''.join(traceback.format_exception(e))}")
131            if "ndarray" in str(e):
132                logger.error(f"input: {res}")
133        finally:
134            logger.info(f"Disconnected: '{addr[0]}:{addr[1]}'")
135            with channel.lock:
136                channel.client = None
137            channel.clear()
138
139    def __repr__(self):
140        return f"[Connection Manager]\n- Host: {self.host}:{self.port}\n- Channels: {len(self. channels)}"