First commit
This commit is contained in:
228
simulator/modules/sbot_interface/socket_server.py
Normal file
228
simulator/modules/sbot_interface/socket_server.py
Normal file
@@ -0,0 +1,228 @@
|
||||
"""A server for multiple devices that can be connected to the simulator."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import select
|
||||
import signal
|
||||
import socket
|
||||
import sys
|
||||
from threading import Event
|
||||
from typing import Protocol
|
||||
|
||||
from sbot_interface.devices.util import get_globals
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
g = get_globals()
|
||||
|
||||
|
||||
class Board(Protocol):
|
||||
"""The interface for all board simulators that can be connected to the simulator."""
|
||||
|
||||
asset_tag: str
|
||||
software_version: str
|
||||
|
||||
def handle_command(self, command: str) -> str | bytes:
|
||||
"""
|
||||
Process a command string and return the response.
|
||||
|
||||
Bytes type are treated as tag-length-value (TLV) encoded data.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class DeviceServer:
|
||||
"""
|
||||
A server for a single device that can be connected to the simulator.
|
||||
|
||||
The process_data method is called when data is received from the socket.
|
||||
Line-delimited commands are processed and responses are sent back.
|
||||
"""
|
||||
|
||||
def __init__(self, board: Board) -> None:
|
||||
self.board = board
|
||||
# create TCP socket server
|
||||
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.server_socket.bind(('127.0.0.1', 0))
|
||||
self.server_socket.listen(1) # only allow one connection per device
|
||||
self.server_socket.setblocking(True)
|
||||
LOGGER.info(
|
||||
f'Started server for {self.board_type} ({self.board.asset_tag}) '
|
||||
f'on port {self.port}'
|
||||
)
|
||||
|
||||
self.device_socket: socket.socket | None = None
|
||||
self.buffer = b''
|
||||
|
||||
def process_data(self, data: bytes) -> bytes | None:
|
||||
"""Process incoming data if a line has been received and return the response."""
|
||||
self.buffer += data
|
||||
if b'\n' in self.buffer:
|
||||
# Sleep to simulate processing time
|
||||
g.sleep(g.timestep / 1000)
|
||||
data, self.buffer = self.buffer.split(b'\n', 1)
|
||||
return self.run_command(data.decode().strip())
|
||||
else:
|
||||
return None
|
||||
|
||||
def run_command(self, command: str) -> bytes:
|
||||
"""
|
||||
Process a command and return the response.
|
||||
|
||||
Wraps the board's handle_command method and deals with exceptions and data types.
|
||||
"""
|
||||
LOGGER.debug(f'> {command}')
|
||||
try:
|
||||
response = self.board.handle_command(command)
|
||||
if isinstance(response, bytes):
|
||||
LOGGER.debug(f'< {len(response)} bytes')
|
||||
return response
|
||||
else:
|
||||
LOGGER.debug(f'< {response}')
|
||||
return response.encode() + b'\n'
|
||||
except Exception as e:
|
||||
LOGGER.exception(f'Error processing command: {command}')
|
||||
return f'NACK:{e}\n'.encode()
|
||||
|
||||
def flush_buffer(self) -> None:
|
||||
"""Clear the internal buffer of received data."""
|
||||
self.buffer = b''
|
||||
|
||||
def socket(self) -> socket.socket:
|
||||
"""
|
||||
Return the socket to select on.
|
||||
|
||||
If the device is connected, return the device socket.
|
||||
Otherwise, return the server socket.
|
||||
"""
|
||||
if self.device_socket is not None:
|
||||
# ignore the server socket while we are connected
|
||||
return self.device_socket
|
||||
else:
|
||||
return self.server_socket
|
||||
|
||||
def accept(self) -> None:
|
||||
"""Accept a connection from a device and set the device socket to blocking."""
|
||||
if self.device_socket is not None:
|
||||
self.disconnect_device()
|
||||
self.device_socket, _ = self.server_socket.accept()
|
||||
self.device_socket.setblocking(True)
|
||||
LOGGER.info(f'Connected to {self.asset_tag} from {self.device_socket.getpeername()}')
|
||||
|
||||
def disconnect_device(self) -> None:
|
||||
"""Close the device socket, flushing the buffer first."""
|
||||
self.flush_buffer()
|
||||
if self.device_socket is not None:
|
||||
self.device_socket.close()
|
||||
self.device_socket = None
|
||||
LOGGER.info(f'Disconnected from {self.asset_tag}')
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the server and client sockets."""
|
||||
self.disconnect_device()
|
||||
self.server_socket.close()
|
||||
|
||||
def __del__(self) -> None:
|
||||
self.close()
|
||||
|
||||
@property
|
||||
def port(self) -> int:
|
||||
"""Return the port number of the server socket."""
|
||||
if self.server_socket is None:
|
||||
return -1
|
||||
return int(self.server_socket.getsockname()[1])
|
||||
|
||||
@property
|
||||
def asset_tag(self) -> str:
|
||||
"""Return the asset tag of the board."""
|
||||
return self.board.asset_tag
|
||||
|
||||
@property
|
||||
def board_type(self) -> str:
|
||||
"""Return the class name of the board object."""
|
||||
return self.board.__class__.__name__
|
||||
|
||||
|
||||
class SocketServer:
|
||||
"""
|
||||
A server for multiple devices that can be connected to the simulator.
|
||||
|
||||
The run method blocks until the stop_event is set.
|
||||
"""
|
||||
|
||||
def __init__(self, devices: list[DeviceServer]) -> None:
|
||||
self.devices = devices
|
||||
self.stop_event = Event()
|
||||
g.stop_event = self.stop_event
|
||||
# flag to indicate that we are exiting because the usercode has completed
|
||||
self.completed = False
|
||||
|
||||
def run(self) -> None:
|
||||
"""
|
||||
Run the server, accepting connections and processing data.
|
||||
|
||||
This method blocks until the stop_event is set.
|
||||
"""
|
||||
while not self.stop_event.is_set():
|
||||
# select on all server sockets and device sockets
|
||||
sockets = [device.socket() for device in self.devices]
|
||||
|
||||
readable, _, _ = select.select(sockets, [], [], 0.5)
|
||||
|
||||
for device in self.devices:
|
||||
try:
|
||||
if device.server_socket in readable:
|
||||
device.accept()
|
||||
|
||||
if device.device_socket in readable and device.device_socket is not None:
|
||||
try:
|
||||
if sys.platform == 'win32':
|
||||
data = device.device_socket.recv(4096)
|
||||
else:
|
||||
data = device.device_socket.recv(4096, socket.MSG_DONTWAIT)
|
||||
except ConnectionError:
|
||||
device.disconnect_device()
|
||||
continue
|
||||
|
||||
if not data:
|
||||
device.disconnect_device()
|
||||
else:
|
||||
response = device.process_data(data)
|
||||
if response is not None:
|
||||
try:
|
||||
device.device_socket.sendall(response)
|
||||
except ConnectionError:
|
||||
device.disconnect_device()
|
||||
continue
|
||||
except Exception as e:
|
||||
LOGGER.exception(f"Failure in simulated boards: {e}")
|
||||
|
||||
LOGGER.info('Stopping server')
|
||||
for device in self.devices:
|
||||
device.close()
|
||||
|
||||
if self.stop_event.is_set() and self.completed is False:
|
||||
# Stop the usercode
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
|
||||
def links(self) -> dict[str, dict[str, str]]:
|
||||
"""Return a mapping of asset tags to ports, grouped by board type."""
|
||||
return {
|
||||
device.asset_tag: {
|
||||
'board_type': device.board_type,
|
||||
'port': str(device.port),
|
||||
}
|
||||
for device in self.devices
|
||||
}
|
||||
|
||||
def links_formatted(self, address: str = '127.0.0.1') -> str:
|
||||
"""
|
||||
Return a formatted string of all the links to the devices.
|
||||
|
||||
The format is 'socket://address:port/board_type/asset_tag'.
|
||||
Each link is separated by a newline.
|
||||
"""
|
||||
return '\n'.join(
|
||||
f"socket://{address}:{data['port']}/{data['board_type']}/{asset_tag}"
|
||||
for asset_tag, data in self.links().items()
|
||||
)
|
||||
Reference in New Issue
Block a user