# ----------------------------------------------------------------------------
# Description : Transport layer (abstract, IP, file, dummy)
# Git repository : https://gitlab.com/qblox/packages/software/qblox_instruments.git
# Copyright (C) Qblox BV (2020)
# ----------------------------------------------------------------------------
# -- include -----------------------------------------------------------------
import asyncio
import contextlib
import os
import socket
import sys
from qblox_instruments.ieee488_2.helpers import LoopError
from qblox_instruments.ieee488_2.transport import Transport
# -- helpers -----------------------------------------------------------------
def _set_keepalive(
sock: socket.socket, after_idle_sec: int = 60, interval_sec: int = 60, max_fails: int = 5
) -> None:
"""
Instructs the TCP socket to send a heart beat every n seconds to detect
dead connections. It's the TCP equivalent of the IRC ping-pong protocol
and allows for better cleanup / detection of dead TCP connections.
It activates after 60 second (after_idle_sec) of idleness, then sends
a keepalive ping once every 60 seconds (interval_sec), and closes the
connection after 5 failed ping (max_fails), or 300 seconds by default.
Parameters
----------
after_idle_sec : int
Activate keepalive after n seconds.
interval_sec : int
Packet interval in seconds.
max_fails : int
Maximum number of failed packets.
"""
if os.name == "nt": # Windows
after_idle_sec *= 1000
interval_sec *= 1000
# pylint: disable=no-member
sock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, after_idle_sec, interval_sec))
elif sys.platform == "darwin": # MacOS
TCP_KEEPALIVE = 0x10 # From /usr/include, not exported by Python's socket module
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
sock.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, interval_sec)
else: # Linux
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, after_idle_sec)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, interval_sec)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, max_fails)
# -- class -------------------------------------------------------------------
[docs]
class IpTransport(Transport):
"""
Class for data transport of IP socket.
"""
__slots__ = ["_host", "_port", "_reader", "_snd_buf_size", "_timeout", "_writer"]
# ------------------------------------------------------------------------
[docs]
def __init__(
self,
host: str,
port: int = 5025,
timeout: float = 60.0,
snd_buf_size: int = 512 * 1024,
loop_from: Transport | None = None,
use_running_loop: bool = False,
) -> None:
"""
Create IP socket transport class.
Parameters
----------
host : str
Instrument IP address.
port : int
Instrument port.
timeout : float
Instrument call timeout in seconds.
snd_buf_size : int
Instrument buffer size for transmissions to instrument.
loop_from: Transport | None = None
Transport to use event loop from: bypasses `use_running_loop` checks.
use_running_loop : bool = False
Whether to use an existing event loop if it is already running;
use this only if you are supplying your own event loop and only invoking
asynchronous functions, or if the event loop is not running when invoking
synchronous command functions. Default to `False`.
"""
super().__init__(loop_from=loop_from, use_running_loop=use_running_loop)
self._host = host
self._port = port
self._snd_buf_size = snd_buf_size
self._timeout = timeout
self._reader, self._writer = self._open_connection()
# ------------------------------------------------------------------------
def _open_connection(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
# 1. Setup timeout (before connecting)
# 2. Enlarge buffer
# 3. Send immediately
# 4. Setup keep alive pinging
# 5. Connect
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(3)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self._snd_buf_size)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
_set_keepalive(sock)
sock.connect((self._host, self._port))
# 6. Transfer to asyncio; the `limit` argument determines the buffer size,
# particularly relevant for `readline()` below.
return self._run_in_loop(asyncio.open_connection(sock=sock, limit=1 * 1024 * 1024))
# ------------------------------------------------------------------------
def __del__(self) -> None:
"""Delete IP socket transport class."""
# Catch case where an exception was thrown in the constructor,
# and the `_writer` field hasn't been set yet.
if getattr(self, "_writer", None):
# Suppress expected errors and warnings during garbage collection raised
# when the remote device is already closed or has abruptly dropped the connection.
with (
# ConnectionError and TimeoutError happen during reconnection attempts,
# while LoopError could occur when a Cluster object is garbage collected
# before it had a chance to clean itself up.
contextlib.suppress(ConnectionError, TimeoutError, LoopError),
):
self._run_in_loop(self.close())
# ------------------------------------------------------------------------
@property
def host(self) -> str:
"""
Get host this socket is connected to.
Returns
-------
str
Current host this socket is connected to.
"""
return self._host
# ------------------------------------------------------------------------
@property
def timeout(self) -> float:
"""
Get current socket timeout.
Returns
-------
float
Current socket timeout in seconds.
"""
return self._timeout
# ------------------------------------------------------------------------
@timeout.setter
def timeout(self, value: float) -> None:
"""
Set socket timeout.
Parameters
----------
value : float
Socket timeout in seconds.
"""
self._timeout = value
# ------------------------------------------------------------------------
[docs]
def reconnect(self) -> None:
"""
Re-open IP connection to instrument.
"""
if getattr(self, "_writer", None):
self._run_in_loop(self.close())
self._reader, self._writer = self._open_connection()
# ------------------------------------------------------------------------
[docs]
async def close(self) -> None:
"""
Close IP socket.
"""
# Catch case where the writer has already been garbage collected.
# This seems to happen occasionally at the end of the firmware update process.
if self._writer is not None:
self._writer.close()
await asyncio.wait_for(self._writer.wait_closed(), self._timeout)
self._reader = self._writer = None
# ------------------------------------------------------------------------
[docs]
async def write(self, cmd_str: str) -> None:
"""
Write command to instrument over IP socket.
Parameters
----------
cmd_str : str
Command
"""
out_str = cmd_str + "\n"
await self.write_binary(out_str.encode("ascii"))
# ------------------------------------------------------------------------
[docs]
async def write_binary(self, *data: bytes) -> None:
"""
Write binary data to instrument over IP socket.
Parameters
----------
*data : bytes
Binary data
"""
for item in data:
self._writer.write(item)
await asyncio.wait_for(self._writer.drain(), self._timeout)
# ------------------------------------------------------------------------
[docs]
async def read_binary(self, size: int) -> bytes:
"""
Read binary data from instrument over IP socket.
Parameters
----------
size : int
Number of bytes
Returns
-------
bytes
Binary data array of length "size".
"""
data = await asyncio.wait_for(self._reader.read(size), self._timeout)
act_len = len(data)
exp_len = size
while act_len != exp_len:
data += await asyncio.wait_for(self._reader.read(exp_len - act_len), self._timeout)
act_len = len(data)
return data
# ------------------------------------------------------------------------
[docs]
async def readline(self) -> str:
"""
Read data from instrument over IP socket.
Returns
-------
str
String with data.
"""
# Note: maximum line length is limited to the `limit` argument in the constructor.
return (await asyncio.wait_for(self._reader.readline(), self._timeout)).decode("utf-8")
