| # Copyright 2020 The Pigweed Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| # use this file except in compliance with the License. You may obtain a copy of |
| # the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| # License for the specific language governing permissions and limitations under |
| # the License. |
| """Utilities for using HDLC with pw_rpc.""" |
| |
| import collections |
| from concurrent.futures import ThreadPoolExecutor |
| import io |
| import logging |
| from queue import SimpleQueue |
| import random |
| import sys |
| import threading |
| import time |
| import socket |
| import subprocess |
| from typing import (Any, BinaryIO, Callable, Deque, Dict, Iterable, List, |
| NoReturn, Optional, Sequence, Tuple, Union) |
| |
| from pw_protobuf_compiler import python_protos |
| import pw_rpc |
| from pw_rpc import callback_client |
| |
| from pw_hdlc.decode import Frame, FrameDecoder |
| from pw_hdlc import encode |
| |
| _LOG = logging.getLogger(__name__) |
| |
| STDOUT_ADDRESS = 1 |
| DEFAULT_ADDRESS = ord('R') |
| _VERBOSE = logging.DEBUG - 1 |
| |
| |
| def channel_output(writer: Callable[[bytes], Any], |
| address: int = DEFAULT_ADDRESS, |
| delay_s: float = 0) -> Callable[[bytes], None]: |
| """Returns a function that can be used as a channel output for pw_rpc.""" |
| |
| if delay_s: |
| |
| def slow_write(data: bytes) -> None: |
| """Slows down writes in case unbuffered serial is in use.""" |
| for byte in data: |
| time.sleep(delay_s) |
| writer(bytes([byte])) |
| |
| return lambda data: slow_write(encode.ui_frame(address, data)) |
| |
| def write_hdlc(data: bytes): |
| frame = encode.ui_frame(address, data) |
| _LOG.log(_VERBOSE, 'Write %2d B: %s', len(frame), frame) |
| writer(frame) |
| |
| return write_hdlc |
| |
| |
| def _handle_error(frame: Frame) -> None: |
| _LOG.error('Failed to parse frame: %s', frame.status.value) |
| _LOG.debug('%s', frame.data) |
| |
| |
| FrameHandlers = Dict[int, Callable[[Frame], Any]] |
| |
| |
| def read_and_process_data(read: Callable[[], bytes], |
| on_read_error: Callable[[Exception], Any], |
| frame_handlers: FrameHandlers, |
| error_handler: Callable[[Frame], |
| Any] = _handle_error, |
| handler_threads: Optional[int] = 1) -> NoReturn: |
| """Continuously reads and handles HDLC frames. |
| |
| Passes frames to an executor that calls frame handler functions in other |
| threads. |
| """ |
| def handle_frame(frame: Frame): |
| try: |
| if not frame.ok(): |
| error_handler(frame) |
| return |
| |
| try: |
| frame_handlers[frame.address](frame) |
| except KeyError: |
| _LOG.warning('Unhandled frame for address %d: %s', |
| frame.address, frame) |
| except: # pylint: disable=bare-except |
| _LOG.exception('Exception in HDLC frame handler thread') |
| |
| decoder = FrameDecoder() |
| |
| # Execute callbacks in a ThreadPoolExecutor to decouple reading the input |
| # stream from handling the data. That way, if a handler function takes a |
| # long time or crashes, this reading thread is not interrupted. |
| with ThreadPoolExecutor(max_workers=handler_threads) as executor: |
| while True: |
| try: |
| data = read() |
| except Exception as exc: # pylint: disable=broad-except |
| on_read_error(exc) |
| continue |
| |
| if data: |
| _LOG.log(_VERBOSE, 'Read %2d B: %s', len(data), data) |
| |
| for frame in decoder.process_valid_frames(data): |
| executor.submit(handle_frame, frame) |
| |
| |
| def write_to_file(data: bytes, output: BinaryIO = sys.stdout.buffer): |
| output.write(data + b'\n') |
| output.flush() |
| |
| |
| def default_channels(write: Callable[[bytes], Any]) -> List[pw_rpc.Channel]: |
| return [pw_rpc.Channel(1, channel_output(write))] |
| |
| |
| PathsModulesOrProtoLibrary = Union[Iterable[python_protos.PathOrModule], |
| python_protos.Library] |
| |
| |
| class HdlcRpcClient: |
| """An RPC client configured to run over HDLC.""" |
| def __init__(self, |
| read: Callable[[], bytes], |
| paths_or_modules: PathsModulesOrProtoLibrary, |
| channels: Iterable[pw_rpc.Channel], |
| output: Callable[[bytes], Any] = write_to_file, |
| client_impl: pw_rpc.client.ClientImpl = None, |
| *, |
| _incoming_packet_filter_for_testing: '_PacketFilter' = None): |
| """Creates an RPC client configured to communicate using HDLC. |
| |
| Args: |
| read: Function that reads bytes; e.g serial_device.read. |
| paths_or_modules: paths to .proto files or proto modules |
| channel: RPC channels to use for output |
| output: where to write "stdout" output from the device |
| """ |
| if isinstance(paths_or_modules, python_protos.Library): |
| self.protos = paths_or_modules |
| else: |
| self.protos = python_protos.Library.from_paths(paths_or_modules) |
| |
| if client_impl is None: |
| client_impl = callback_client.Impl() |
| |
| self.client = pw_rpc.Client.from_modules(client_impl, channels, |
| self.protos.modules()) |
| |
| self._test_filter = _incoming_packet_filter_for_testing |
| |
| frame_handlers: FrameHandlers = { |
| DEFAULT_ADDRESS: self._handle_rpc_packet, |
| STDOUT_ADDRESS: lambda frame: output(frame.data), |
| } |
| |
| # Start background thread that reads and processes RPC packets. |
| threading.Thread(target=read_and_process_data, |
| daemon=True, |
| args=(read, lambda exc: None, |
| frame_handlers)).start() |
| |
| def rpcs(self, channel_id: int = None) -> Any: |
| """Returns object for accessing services on the specified channel. |
| |
| This skips some intermediate layers to make it simpler to invoke RPCs |
| from an HdlcRpcClient. If only one channel is in use, the channel ID is |
| not necessary. |
| """ |
| if channel_id is None: |
| return next(iter(self.client.channels())).rpcs |
| |
| return self.client.channel(channel_id).rpcs |
| |
| def _handle_rpc_packet(self, frame: Frame) -> None: |
| if self._test_filter and not self._test_filter.keep_packet(frame.data): |
| return |
| |
| if not self.client.process_packet(frame.data): |
| _LOG.error('Packet not handled by RPC client: %s', frame.data) |
| |
| |
| def _try_connect(sock: socket.socket, port: int, attempts: int = 10) -> None: |
| """Tries to connect to the specified port up to the given number of times. |
| |
| This is helpful when connecting to a process that was started by this |
| script. The process may need time to start listening for connections, and |
| that length of time can vary. This retries with a short delay rather than |
| having to wait for the worst case delay every time. |
| """ |
| while True: |
| attempts -= 1 |
| time.sleep(0.001) |
| |
| try: |
| sock.connect(('localhost', port)) |
| return |
| except ConnectionRefusedError: |
| if attempts <= 0: |
| raise |
| |
| |
| class SocketSubprocess: |
| """Executes a subprocess and connects to it with a socket.""" |
| def __init__(self, command: Sequence, port: int) -> None: |
| self._server_process = subprocess.Popen(command, stdin=subprocess.PIPE) |
| self.stdin = self._server_process.stdin |
| |
| sock = None |
| |
| try: |
| sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| _try_connect(sock, port) |
| except: |
| if sock: |
| sock.close() |
| |
| self._server_process.terminate() |
| self._server_process.communicate() |
| raise |
| |
| self.socket: socket.socket = sock # 🧦 |
| |
| def close(self) -> None: |
| try: |
| self.socket.close() |
| finally: |
| self._server_process.terminate() |
| self._server_process.communicate() |
| |
| def __enter__(self) -> 'SocketSubprocess': |
| return self |
| |
| def __exit__(self, exc_type, exc_value, traceback) -> None: |
| self.close() |
| |
| |
| class _PacketFilter: |
| """Determines if a packet should be kept or dropped for testing purposes.""" |
| _Action = Callable[[int], Tuple[bool, bool]] |
| _KEEP = lambda _: (True, False) |
| _DROP = lambda _: (False, False) |
| |
| def __init__(self, name: str) -> None: |
| self.name = name |
| self.packet_count = 0 |
| self._actions: Deque[_PacketFilter._Action] = collections.deque() |
| |
| def reset(self) -> None: |
| self.packet_count = 0 |
| self._actions.clear() |
| |
| def keep(self, count: int) -> None: |
| """Keeps the next count packets.""" |
| self._actions.extend(_PacketFilter._KEEP for _ in range(count)) |
| |
| def drop(self, count: int) -> None: |
| """Drops the next count packets.""" |
| self._actions.extend(_PacketFilter._DROP for _ in range(count)) |
| |
| def drop_every(self, every: int) -> None: |
| """Drops every Nth packet forever.""" |
| self._actions.append(lambda count: (count % every != 0, True)) |
| |
| def randomly_drop(self, one_in: int, gen: random.Random) -> None: |
| """Drops packets randomly forever.""" |
| self._actions.append(lambda _: (gen.randrange(one_in) != 0, True)) |
| |
| def keep_packet(self, packet: bytes) -> bool: |
| """Returns whether the provided packet should be kept or dropped.""" |
| self.packet_count += 1 |
| |
| if not self._actions: |
| return True |
| |
| keep, repeat = self._actions[0](self.packet_count) |
| |
| if not repeat: |
| self._actions.popleft() |
| |
| if not keep: |
| _LOG.debug('Dropping %s packet %d for testing: %s', self.name, |
| self.packet_count, packet) |
| return keep |
| |
| |
| class _TestChannelOutput: |
| def __init__(self, send: Callable[[bytes], Any]) -> None: |
| self._send = send |
| self.packets = _PacketFilter('outgoing RPC') |
| |
| def __call__(self, data: bytes) -> None: |
| if self.packets.keep_packet(data): |
| self._send(data) |
| |
| |
| class HdlcRpcLocalServerAndClient: |
| """Runs an RPC server in a subprocess and connects to it over a socket. |
| |
| This can be used to run a local RPC server in an integration test. |
| """ |
| def __init__(self, |
| server_command: Sequence, |
| port: int, |
| protos: PathsModulesOrProtoLibrary, |
| *, |
| for_testing: bool = False) -> None: |
| """Creates a new HdlcRpcLocalServerAndClient. |
| |
| If for_testing=True, the HdlcRpcLocalServerAndClient will have |
| outgoing_packets and incoming_packets _PacketFilter members that can be |
| used to program packet loss for testing purposes. |
| """ |
| |
| self.server = SocketSubprocess(server_command, port) |
| |
| self._bytes_queue: 'SimpleQueue[bytes]' = SimpleQueue() |
| self._read_thread = threading.Thread(target=self._read_from_socket) |
| self._read_thread.start() |
| |
| self.output = io.BytesIO() |
| |
| self.channel_output: Any = self.server.socket.sendall |
| if for_testing: |
| self.channel_output = _TestChannelOutput(self.channel_output) |
| self.outgoing_packets = self.channel_output.packets |
| self.incoming_packets = _PacketFilter('incoming RPC') |
| incoming_filter: Optional[_PacketFilter] = self.incoming_packets |
| else: |
| incoming_filter = None |
| |
| self.client = HdlcRpcClient( |
| self._bytes_queue.get, |
| protos, |
| default_channels(self.channel_output), |
| self.output.write, |
| _incoming_packet_filter_for_testing=incoming_filter).client |
| |
| def _read_from_socket(self): |
| while True: |
| data = self.server.socket.recv(4096) |
| self._bytes_queue.put(data) |
| if not data: |
| return |
| |
| def close(self): |
| self.server.close() |
| self.output.close() |
| self._read_thread.join() |
| |
| def __enter__(self) -> 'HdlcRpcLocalServerAndClient': |
| return self |
| |
| def __exit__(self, exc_type, exc_value, traceback) -> None: |
| self.close() |
