Switch rpc console to use CancellableReader instead of a pure read ca…

…llback (#31267)

* Use a cancellable reader for socket communication

* Remove the auto-reconnect, simpler code

* Restyle

* Ensure we can close the reader

* Restyle

* Remove unused import

---------

Co-authored-by: Andrei Litvin <andreilitvin@google.com>

examples/common/pigweed/rpc_console/py/chip_rpc/console.py

@@ -39,7 +39,6 @@
 import argparse
 import logging
 import re
-import socket
 import sys
 import threading
 from collections import namedtuple
@@ -50,10 +49,10 @@
 import pw_cli.log
 from chip_rpc.plugins.device_toolbar import DeviceToolbar
 from chip_rpc.plugins.helper_scripts import HelperScripts
-from pw_console import PwConsoleEmbed
+from pw_console import PwConsoleEmbed, socket_client
 from pw_console.__main__ import create_temp_log_file
 from pw_console.pyserial_wrapper import SerialWithLogging
-from pw_hdlc.rpc import HdlcRpcClient, default_channels
+from pw_hdlc.rpc import HdlcRpcClient, SelectableReader, SerialReader, default_channels
 from pw_rpc import callback_client
 from pw_rpc.console_tools.console import ClientInfo, flattened_rpc_completions
 from pw_tokenizer import tokens
@@ -219,28 +218,8 @@ def _start_ipython_hdlc_terminal(client: HdlcRpcClient) -> None:
     interactive_console.setup_python_logging()
     # Don't send device logs to the root logger.
     _DEVICE_LOG.propagate = False
-    interactive_console.embed()
-
-
-class SocketClientImpl:
-    def __init__(self, config: str):
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        socket_server = ''
-        socket_port = 0
-
-        if config == 'default':
-            socket_server = SOCKET_SERVER
-            socket_port = SOCKET_PORT
-        else:
-            socket_server, socket_port_str = config.split(':')
-            socket_port = int(socket_port_str)
-        self.socket.connect((socket_server, socket_port))
-
-    def write(self, data: bytes):
-        self.socket.sendall(data)
-
-    def read(self, num_bytes: int = PW_RPC_MAX_PACKET_SIZE):
-        return self.socket.recv(num_bytes)
+    with client:
+        interactive_console.embed()
 
 
 def write_to_output(data: bytes,
@@ -329,12 +308,12 @@ def console(device: str, baudrate: int,
 
     if socket_addr is None:
         serial_device = serial_impl(device, baudrate, timeout=0.1)
-        def read(): return serial_device.read(8192)
+        reader = SerialReader(serial_device, 8192)
         write = serial_device.write
     else:
         try:
-            socket_device = SocketClientImpl(socket_addr)
-            read = socket_device.read
+            socket_device = socket_client.SocketClient(socket_addr)
+            reader = SelectableReader(socket_device)
             write = socket_device.write
         except ValueError:
             _LOG.exception('Failed to initialize socket at %s', socket_addr)
@@ -351,14 +330,14 @@ def read(): return serial_device.read(8192)
     if raw_serial:
         threading.Thread(target=_read_raw_serial,
                          daemon=True,
-                         args=(read,
+                         args=(reader,
                                lambda data: write_to_output(
                                    data, output, detokenizer),
                                )).start()
         _start_ipython_raw_terminal()
     else:
         _start_ipython_hdlc_terminal(
-            HdlcRpcClient(read, PROTOS, default_channels(write),
+            HdlcRpcClient(reader, PROTOS, default_channels(write),
                           lambda data: write_to_output(
                               data, output, detokenizer),
                           client_impl=callback_client_impl)