robosim.network

14def send_packet(sock: socket.socket, data: dict, close_connection: bool=False, add_timestmap: bool=True):
15    """pack the message in messagepack format + length-prefixed for TCP stream then send it through the socket"""
16    assert isinstance(data, dict), (data, type(data))
17    if add_timestmap:
18        assert "timestamp" not in data.keys(), data
19        data["timestamp"] = datetime.now().isoformat()
20    try:
21        packed_data = msgpack.packb(data)
22    except TypeError as e:
23        _find_bad(data)
24        raise e
25    packet = len(packed_data).to_bytes(4, "big") + packed_data # <SIZE-4b><packed_data>
26    logger.trace(f"Sending {len(packet)} bytes through the wire")
27    sock.sendall(packet)
28
29    if close_connection:
30        sock.close()

32def recv_packet(sock: socket.socket) -> dict:
33    """reads a packet from the server. First 4 bytes for len, then the rest of the message"""
34    packet_len = int.from_bytes(_recvall(sock, 4), "big")
35    packet: dict = msgpack.unpackb(_recvall(sock, packet_len), raw=False)
36    assert isinstance(packet, dict), f"packet is not dictionary: {type(packet)}"
37    return packet

 8@dataclass
 9class ObjToSimChannel:
10    """Channel of two queues handling the communication between the main thread (sim) and the tcp handler of a robot"""
11    index: int # the global index of the channel. Should be unique across the simulator.
12    tcp2sim: Queue = field(default_factory=Queue) # user -> simulator. User can send many messages theoretically.
13    sim2tcp: Queue = field(default_factory=lambda: Queue(maxsize=1)) # sim -> user. One message at a time.
14    lock: threading.Lock = field(default_factory=threading.Lock)
15    client: socket.socket | None = None
16
17    def clear(self):
18        """clears the queues of this channel"""
19        with self.lock:
20            self.tcp2sim.queue.clear()
21            self.sim2tcp.queue.clear()

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