cgse-common 0.16.14__py3-none-any.whl → 0.17.1__py3-none-any.whl

egse/socketdevice.py

@@ -2,29 +2,52 @@
 This module defines base classes and generic functions to work with sockets.
 """
 
+import asyncio
+import select
 import socket
+import time
+from typing import Optional
 
+from egse.device import AsyncDeviceInterface
+from egse.device import AsyncDeviceTransport
 from egse.device import DeviceConnectionError
 from egse.device import DeviceConnectionInterface
 from egse.device import DeviceTimeoutError
 from egse.device import DeviceTransport
 from egse.log import logger
+from egse.system import type_name
 
 
 class SocketDevice(DeviceConnectionInterface, DeviceTransport):
     """Base class that implements the socket interface."""
 
-    def __init__(self, hostname: str, port: int):
+    # We set a default connect timeout of 3.0 sec before connecting and reset
+    # to None (=blocking) after connecting. The reason for this is that when no
+    # device is available, e.g. during testing, the timeout will take about
+    # two minutes which is way too long. It needs to be evaluated if this
+    # approach is acceptable and not causing problems during production.
+
+    def __init__(
+        self,
+        hostname: str,
+        port: int,
+        connect_timeout: float = 3.0,
+        read_timeout: float | None = 1.0,
+        separator: str = b"\x03",
+    ):
         super().__init__()
         self.is_connection_open = False
         self.hostname = hostname
         self.port = port
+        self.connect_timeout = connect_timeout
+        self.read_timeout = read_timeout
+        self.separator = separator
         self.socket = None
 
     @property
     def device_name(self):
         """The name of the device that this interface connects to."""
-        raise NotImplementedError
+        return f"SocketDevice({self.hostname}:{self.port})"
 
     def connect(self):
         """
@@ -53,18 +76,12 @@ class SocketDevice(DeviceConnectionInterface, DeviceTransport):
 
         try:
             self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        except socket.error as e_socket:
-            raise ConnectionError(f"{self.device_name}: Failed to create socket.") from e_socket
-
-        # We set a timeout of 3 sec before connecting and reset to None
-        # (=blocking) after the connect. The reason for this is because when no
-        # device is available, e.g during testing, the timeout will take about
-        # two minutes which is way too long. It needs to be evaluated if this
-        # approach is acceptable and not causing problems during production.
+        except socket.error as exc:
+            raise ConnectionError(f"{self.device_name}: Failed to create socket.") from exc
 
         try:
             logger.debug(f'Connecting a socket to host "{self.hostname}" using port {self.port}')
-            self.socket.settimeout(3)
+            self.socket.settimeout(self.connect_timeout)
             self.socket.connect((self.hostname, self.port))
             self.socket.settimeout(None)
         except ConnectionRefusedError as exc:
@@ -120,32 +137,64 @@ class SocketDevice(DeviceConnectionInterface, DeviceTransport):
 
     def read(self) -> bytes:
         """
-
-        Returns:
-            A bytes object containing the received telemetry.
+        Read until ETX (b'\x03') or until `self.read_timeout` elapses.
+        Uses `select` to avoid blocking indefinitely when no data is available.
+        If `self.read_timeout` was set to None in the constructor, this will block anyway.
         """
-        idx, n_total = 0, 0
+        if not self.socket:
+            raise DeviceConnectionError(self.device_name, "Not connected")
+
         buf_size = 1024 * 4
-        response = bytes()
+        response = bytearray()
+
+        # If read_timeout is None we preserve blocking behaviour; otherwise enforce overall timeout.
+        if self.read_timeout is None:
+            end_time = None
+        else:
+            end_time = time.monotonic() + self.read_timeout
 
         try:
-            for idx in range(100):
-                # time.sleep(0.1)  # Give the device time to fill the buffer
-                data = self.socket.recv(buf_size)
-                n = len(data)
-                n_total += n
-                response += data
-                # if n < buf_size:
-                #     break  # there is not more data in the buffer
-                if b"\x03" in response:
+            while True:
+                # compute remaining timeout for select, this is needed because we read in different parts
+                # until ETX is received, and we want to receive the complete messages including ETX within
+                # the read timeout.
+                if end_time is None:
+                    timeout = None
+                else:
+                    remaining = end_time - time.monotonic()
+                    if remaining <= 0.0:
+                        raise DeviceTimeoutError(self.device_name, "Socket read timed out")
+                    timeout = remaining
+
+                ready, _, _ = select.select([self.socket], [], [], timeout)
+
+                if not ready:
+                    # no socket ready within timeout
+                    raise DeviceTimeoutError(self.device_name, "Socket read timed out")
+
+                try:
+                    data = self.socket.recv(buf_size)
+                except OSError as exc:
+                    raise DeviceConnectionError(self.device_name, f"Caught {type_name(exc)}: {exc}") from exc
+
+                if not data:
+                    # remote closed connection (EOF)
+                    raise DeviceConnectionError(self.device_name, "Connection closed by peer")
+
+                response.extend(data)
+
+                if self.separator in response:
                     break
-        except socket.timeout as e_timeout:
-            logger.warning(f"Socket timeout error from {e_timeout}")
-            raise DeviceTimeoutError(self.device_name, "Socket timeout error") from e_timeout
 
-        # logger.debug(f"Total number of bytes received is {n_total}, idx={idx}")
+        except DeviceTimeoutError:
+            raise
+        except DeviceConnectionError:
+            raise
+        except Exception as exc:
+            # unexpected errors - translate to DeviceConnectionError
+            raise DeviceConnectionError(self.device_name, "Socket read error") from exc
 
-        return response
+        return bytes(response)
 
     def write(self, command: str):
         """
@@ -163,11 +212,11 @@ class SocketDevice(DeviceConnectionInterface, DeviceTransport):
 
         try:
             self.socket.sendall(command.encode())
-        except socket.timeout as e_timeout:
-            raise DeviceTimeoutError(self.device_name, "Socket timeout error") from e_timeout
-        except socket.error as e_socket:
+        except socket.timeout as exc:
+            raise DeviceTimeoutError(self.device_name, "Socket timeout error") from exc
+        except socket.error as exc:
             # Interpret any socket-related error as an I/O error
-            raise DeviceConnectionError(self.device_name, "Socket communication error.") from e_socket
+            raise DeviceConnectionError(self.device_name, "Socket communication error.") from exc
 
     def trans(self, command: str) -> bytes:
         """
@@ -198,8 +247,133 @@ class SocketDevice(DeviceConnectionInterface, DeviceTransport):
 
             return return_string
 
-        except socket.timeout as e_timeout:
-            raise DeviceTimeoutError(self.device_name, "Socket timeout error") from e_timeout
-        except socket.error as e_socket:
+        except socket.timeout as exc:
+            raise DeviceTimeoutError(self.device_name, "Socket timeout error") from exc
+        except socket.error as exc:
             # Interpret any socket-related error as an I/O error
-            raise DeviceConnectionError(self.device_name, "Socket communication error.") from e_socket
+            raise DeviceConnectionError(self.device_name, "Socket communication error.") from exc
+
+
+class AsyncSocketDevice(AsyncDeviceInterface, AsyncDeviceTransport):
+    """
+    Async socket-backed device using asyncio streams.
+
+    - async connect() / disconnect()
+    - async read() reads until ETX (b'\\x03') or timeout
+    - async write() and async trans()
+    """
+
+    def __init__(
+        self,
+        hostname: str,
+        port: int,
+        connect_timeout: float = 3.0,
+        read_timeout: float | None = 1.0,
+        separator: str = b"\x03",
+    ):
+        super().__init__()
+        self.hostname = hostname
+        self.port = port
+        self.connect_timeout = connect_timeout
+        self.read_timeout = read_timeout
+        self.separator = separator
+        self.reader: Optional[asyncio.StreamReader] = None
+        self.writer: Optional[asyncio.StreamWriter] = None
+        self.is_connection_open = False
+
+    @property
+    def device_name(self) -> str:
+        # Override this property for a decent name
+        return f"AsyncSocketDevice({self.hostname}:{self.port})"
+
+    async def connect(self) -> None:
+        if self.is_connection_open:
+            logger.debug(f"{self.device_name}: already connected")
+            return
+
+        if not self.hostname:
+            raise ValueError(f"{self.device_name}: hostname is not initialized.")
+        if not self.port:
+            raise ValueError(f"{self.device_name}: port is not initialized.")
+
+        try:
+            logger.debug(f"{self.device_name}: connect() called; is_connection_open={self.is_connection_open}")
+            coro = asyncio.open_connection(self.hostname, self.port)
+            self.reader, self.writer = await asyncio.wait_for(coro, timeout=self.connect_timeout)
+            self.is_connection_open = True
+            logger.debug(f"{self.device_name}: connected -> peer={self.writer.get_extra_info('peername')}")
+
+        except asyncio.TimeoutError as exc:
+            await self._cleanup()
+            logger.warning(f"{self.device_name}: connect timed out")
+            raise DeviceTimeoutError(self.device_name, f"Connection to {self.hostname}:{self.port} timed out.") from exc
+        except Exception as exc:
+            await self._cleanup()
+            logger.warning(f"{self.device_name}: connect failed: {type_name(exc)} – {exc}")
+            raise DeviceConnectionError(self.device_name, f"Failed to connect to {self.hostname}:{self.port}") from exc
+
+    async def disconnect(self) -> None:
+        logger.debug(f"{self.device_name}: disconnect() called; writer_exists={bool(self.writer)}")
+        peer = None
+        try:
+            if self.writer and not self.writer.is_closing():
+                peer = self.writer.get_extra_info("peername")
+                self.writer.close()
+                # wait for close, but don't hang forever
+                try:
+                    await asyncio.wait_for(self.writer.wait_closed(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    logger.debug(f"{self.device_name}: wait_closed() timed out for peer={peer}")
+
+        finally:
+            await self._cleanup()
+            logger.debug(f"{self.device_name}: disconnected ({peer=})")
+
+    def is_connected(self) -> bool:
+        return bool(self.is_connection_open and self.writer and not self.writer.is_closing())
+
+    async def _cleanup(self) -> None:
+        self.reader = None
+        self.writer = None
+        self.is_connection_open = False
+
+    async def read(self) -> bytes:
+        if not self.reader:
+            raise DeviceConnectionError(self.device_name, "Not connected")
+        try:
+            # readuntil includes the separator; we keep it for parity with existing code
+            data = await asyncio.wait_for(self.reader.readuntil(separator=self.separator), timeout=self.read_timeout)
+            return data
+        except asyncio.IncompleteReadError as exc:
+            # EOF before separator
+            await self._cleanup()
+            raise DeviceConnectionError(self.device_name, "Connection closed while reading") from exc
+        except asyncio.TimeoutError as exc:
+            raise DeviceTimeoutError(self.device_name, "Socket read timed out") from exc
+        except Exception as exc:
+            await self._cleanup()
+            raise DeviceConnectionError(self.device_name, "Socket read error") from exc
+
+    async def write(self, command: str) -> None:
+        if not self.writer:
+            raise DeviceConnectionError(self.device_name, "Not connected")
+        try:
+            self.writer.write(command.encode())
+            await asyncio.wait_for(self.writer.drain(), timeout=self.read_timeout)
+        except asyncio.TimeoutError as exc:
+            raise DeviceTimeoutError(self.device_name, "Socket write timed out") from exc
+        except Exception as exc:
+            await self._cleanup()
+            raise DeviceConnectionError(self.device_name, "Socket write error") from exc
+
+    async def trans(self, command: str) -> bytes:
+        if not self.writer or not self.reader:
+            raise DeviceConnectionError(self.device_name, "Not connected")
+        try:
+            await self.write(command)
+            return await self.read()
+        except (DeviceTimeoutError, DeviceConnectionError):
+            raise
+        except Exception as exc:
+            await self._cleanup()
+            raise DeviceConnectionError(self.device_name, "Socket trans error") from exc