refactor(connector.py): rewrite Connector to use selectors
This commit is contained in:
@@ -2,6 +2,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging as lg
|
import logging as lg
|
||||||
|
import selectors
|
||||||
import socket
|
import socket
|
||||||
import time
|
import time
|
||||||
from typing import Callable
|
from typing import Callable
|
||||||
@@ -16,8 +17,6 @@ class Connector:
|
|||||||
host: str,
|
host: str,
|
||||||
port: int,
|
port: int,
|
||||||
*,
|
*,
|
||||||
connect_timeout: float = 5.0,
|
|
||||||
ack_timeout: float | None = None,
|
|
||||||
on_message: Callable[[Message], None],
|
on_message: Callable[[Message], None],
|
||||||
) -> None:
|
) -> None:
|
||||||
self.logger: lg.Logger = lg.getLogger(
|
self.logger: lg.Logger = lg.getLogger(
|
||||||
@@ -27,178 +26,125 @@ class Connector:
|
|||||||
|
|
||||||
self.host: str = host
|
self.host: str = host
|
||||||
self.port: int = port
|
self.port: int = port
|
||||||
self.socket_timeout: None = None
|
|
||||||
self.connect_timeout: float = connect_timeout
|
|
||||||
self.ack_timeout: float | None = ack_timeout
|
|
||||||
|
|
||||||
|
self.selector = selectors.DefaultSelector()
|
||||||
self.socket: socket.socket = socket.socket(
|
self.socket: socket.socket = socket.socket(
|
||||||
socket.AF_INET, socket.SOCK_STREAM
|
socket.AF_INET, socket.SOCK_STREAM
|
||||||
)
|
)
|
||||||
|
self.socket.setblocking(False)
|
||||||
|
|
||||||
|
self.selector.register(
|
||||||
|
self.socket,
|
||||||
|
selectors.EVENT_READ | selectors.EVENT_WRITE,
|
||||||
|
data=None,
|
||||||
|
)
|
||||||
|
|
||||||
self.mac_address: str = mac_address
|
self.mac_address: str = mac_address
|
||||||
|
|
||||||
|
self.inbound_buffer: bytes = b""
|
||||||
|
self.outbound_buffer: bytes = b""
|
||||||
|
|
||||||
self.on_message: Callable[[Message], None] = on_message
|
self.on_message: Callable[[Message], None] = on_message
|
||||||
|
|
||||||
def _send_ack(self) -> None:
|
def _send_outbound(self) -> None:
|
||||||
self.logger.debug("[>] Sending ACK...")
|
while self.outbound_buffer:
|
||||||
try:
|
|
||||||
self.socket.sendall(Message.ack().to_bytes())
|
|
||||||
self.logger.debug("[<] ACK sent")
|
|
||||||
except socket.error as e:
|
|
||||||
self.logger.error(f"[!] Failed to send ACK: {e}")
|
|
||||||
|
|
||||||
def _check_ack(self) -> bool:
|
|
||||||
self.logger.debug("[.] Waiting for ACK...")
|
|
||||||
try:
|
|
||||||
self.socket.settimeout(self.ack_timeout)
|
|
||||||
ack: bytes = self.socket.recv(1024)
|
|
||||||
self.socket.settimeout(self.socket_timeout)
|
|
||||||
|
|
||||||
if ack == Message.ack().to_bytes():
|
|
||||||
self.logger.debug("[<] ACK received")
|
|
||||||
return True
|
|
||||||
else:
|
|
||||||
self.logger.error(f"[!] Invalid ACK received: {ack}")
|
|
||||||
|
|
||||||
except TimeoutError as e:
|
|
||||||
self.logger.error(f"[!] ACK timeout: {e}")
|
|
||||||
|
|
||||||
except socket.error as e:
|
|
||||||
self.logger.error(f"[!] Failed to receive ACK: {e}")
|
|
||||||
|
|
||||||
return False
|
|
||||||
|
|
||||||
def connect(self) -> None:
|
|
||||||
retry_interval: int = 1
|
|
||||||
connected: bool = False
|
|
||||||
while not connected:
|
|
||||||
self.logger.debug(
|
|
||||||
f"[.] Connecting to {self.host}:{self.port} with timeout {self.connect_timeout}s..."
|
|
||||||
)
|
|
||||||
try:
|
try:
|
||||||
self.socket.settimeout(self.connect_timeout)
|
sent = self.socket.send(self.outbound_buffer)
|
||||||
self.socket.connect((self.host, self.port))
|
self.logger.debug(
|
||||||
self.socket.settimeout(self.socket_timeout)
|
f"[>] Sent {sent} bytes: {self.outbound_buffer[:sent]!r}"
|
||||||
self.logger.info(f"[+] Connected to {self.host}:{self.port}")
|
|
||||||
self.send_hello()
|
|
||||||
connected = True
|
|
||||||
except (
|
|
||||||
socket.timeout,
|
|
||||||
ConnectionRefusedError,
|
|
||||||
ConnectionAbortedError,
|
|
||||||
) as e:
|
|
||||||
self.logger.error(
|
|
||||||
f"[!] Connection to {self.host}:{self.port} failed: {e}"
|
|
||||||
)
|
)
|
||||||
self.logger.info(
|
self.outbound_buffer = self.outbound_buffer[sent:]
|
||||||
f"[.] Retrying connection in {retry_interval} s..."
|
except BlockingIOError:
|
||||||
)
|
# OS buffer full, wait for next EVENT_WRITE
|
||||||
time.sleep(retry_interval)
|
break
|
||||||
retry_interval = min(
|
except socket.error as e:
|
||||||
retry_interval * 2, 30
|
self.logger.error(f"[!] Socket error: {e}")
|
||||||
) # exponential backoff
|
self.reconnect()
|
||||||
|
break
|
||||||
|
|
||||||
def send(self, data: bytes, no_check_ack: bool = False) -> None:
|
def _receive_inbound(self) -> None:
|
||||||
self.logger.debug(f"[>] Sending data: {data}")
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
self.socket.sendall(data)
|
|
||||||
|
|
||||||
if no_check_ack:
|
|
||||||
self.logger.debug("[>] Data sent without ACK check")
|
|
||||||
break
|
|
||||||
else:
|
|
||||||
self.logger.info("[>] Data sent")
|
|
||||||
|
|
||||||
acknowledged: bool = self._check_ack()
|
|
||||||
if acknowledged:
|
|
||||||
self.logger.debug("[.] Data acknowledged")
|
|
||||||
break
|
|
||||||
else:
|
|
||||||
self.logger.warning(
|
|
||||||
"[!] Data not acknowledged, retrying..."
|
|
||||||
)
|
|
||||||
|
|
||||||
except BrokenPipeError as e:
|
|
||||||
self.logger.error(f"[!] Broken pipe: {e}")
|
|
||||||
self.logger.info("[.] Reconnecting...")
|
|
||||||
self.connect()
|
|
||||||
except (socket.error, ValueError) as e:
|
|
||||||
self.logger.error(f"[!] Failed to send data: {e}")
|
|
||||||
time.sleep(1)
|
|
||||||
|
|
||||||
def receive(self) -> bytes:
|
|
||||||
self.logger.debug("[.] Waiting to receive data...")
|
|
||||||
try:
|
try:
|
||||||
data: bytes = self.socket.recv(4096)
|
data: bytes = self.socket.recv(4096)
|
||||||
if not data:
|
if data:
|
||||||
self.logger.warning("[!] Received empty message")
|
self.logger.debug(f"[<] Received {len(data)} bytes: {data!r}")
|
||||||
return b""
|
self.inbound_buffer += data
|
||||||
self.logger.debug(f"[<] Received data: {data}")
|
else:
|
||||||
return data
|
self.logger.debug("[!] Connection closed by the server.")
|
||||||
|
self.reconnect()
|
||||||
except socket.error as e:
|
except socket.error as e:
|
||||||
self.logger.error(f"[!] Failed to receive data: {e}")
|
self.logger.error(f"[!] Socket error: {e}")
|
||||||
return b""
|
self.reconnect()
|
||||||
|
|
||||||
def close(self) -> None:
|
|
||||||
self.logger.debug("Closing connection...")
|
|
||||||
self.socket.close()
|
|
||||||
self.logger.info("Connection closed.")
|
|
||||||
|
|
||||||
def reconnect(self) -> None:
|
|
||||||
self.logger.debug("Reconnecting...")
|
|
||||||
self.close()
|
|
||||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
||||||
self.connect()
|
|
||||||
|
|
||||||
def send_hello(self) -> None:
|
def send_hello(self) -> None:
|
||||||
self.logger.debug("[.] Sending hello message...")
|
self.logger.debug("[*] Sending HELLO message...")
|
||||||
hello_message: Message = Message.hello(self.mac_address)
|
hello_message: bytes = Message.hello(self.mac_address).to_bytes()
|
||||||
acknowledged: bool = False
|
self.outbound_buffer += hello_message
|
||||||
while not acknowledged:
|
self._send_outbound()
|
||||||
self.send(hello_message.to_bytes(), no_check_ack=True)
|
|
||||||
self.logger.debug("[.] Hello message sent, waiting for ACK...")
|
|
||||||
acknowledged = self._check_ack()
|
|
||||||
if not acknowledged:
|
|
||||||
self.logger.warning(
|
|
||||||
"[!] Hello message not acknowledged, retrying..."
|
|
||||||
)
|
|
||||||
time.sleep(1)
|
|
||||||
|
|
||||||
def _loop(self) -> None:
|
def close(self) -> None:
|
||||||
self.logger.debug("Starting connector loop...")
|
self.logger.debug("[*] Closing connection...")
|
||||||
while True:
|
self.selector.unregister(self.socket)
|
||||||
time.sleep(0.1)
|
self.socket.close()
|
||||||
data: bytes = self.receive()
|
self.logger.debug("[.] Connection closed.")
|
||||||
if not data:
|
|
||||||
self.reconnect()
|
|
||||||
continue
|
|
||||||
for line in data.split(b"\n"):
|
|
||||||
line: bytes = line.strip()
|
|
||||||
|
|
||||||
if not line:
|
def reconnect(self) -> None:
|
||||||
continue
|
self.logger.debug("[*] Reconnecting...")
|
||||||
|
self.close()
|
||||||
|
|
||||||
self.logger.debug(f"[.] Raw message data: {line}")
|
# reinit socket
|
||||||
try:
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||||
message: Message = Message.from_bytes(line)
|
self.socket.setblocking(False)
|
||||||
except ValueError as e:
|
self.selector.register(
|
||||||
self.logger.error(f"[!] Failed to parse message: {e}")
|
self.socket,
|
||||||
continue
|
selectors.EVENT_READ | selectors.EVENT_WRITE,
|
||||||
self.logger.info(f"[*] Message received: {message}")
|
data=None,
|
||||||
self.on_message(message)
|
)
|
||||||
|
self.connect()
|
||||||
|
|
||||||
# if self._check_ack():
|
def connect(self) -> None:
|
||||||
# self.logger.debug("[.] ACK verified")
|
self.logger.debug(f"Connecting to {self.host}:{self.port}...")
|
||||||
# else:
|
connected: bool = False
|
||||||
# self.logger.error("[!] ACK verification failed")
|
delay: float = 1.0
|
||||||
|
while not connected:
|
||||||
|
try:
|
||||||
|
self.socket.connect((self.host, self.port))
|
||||||
|
connected = True
|
||||||
|
except BlockingIOError:
|
||||||
|
# Connection in progress
|
||||||
|
time.sleep(0.1)
|
||||||
|
except socket.error as e:
|
||||||
|
self.logger.error(f"[!] Connection error: {e}")
|
||||||
|
self.logger.debug(f"[.] Retrying in {delay} seconds...")
|
||||||
|
time.sleep(delay)
|
||||||
|
delay = min(delay * 2, 30) # exponential backoff
|
||||||
|
|
||||||
|
self.logger.debug("[*] Connected, sending HELLO...")
|
||||||
|
self.send_hello()
|
||||||
|
|
||||||
def run(self) -> None:
|
def run(self) -> None:
|
||||||
self.logger.debug("Running Connector...")
|
self.connect()
|
||||||
try:
|
try:
|
||||||
self.connect()
|
while True:
|
||||||
self._loop()
|
events = self.selector.select(timeout=1)
|
||||||
|
for key, mask in events:
|
||||||
|
if mask & selectors.EVENT_READ:
|
||||||
|
self._receive_inbound()
|
||||||
|
if mask & selectors.EVENT_WRITE:
|
||||||
|
self._send_outbound()
|
||||||
|
|
||||||
|
# Process inbound buffer for complete messages
|
||||||
|
while b"\n" in self.inbound_buffer:
|
||||||
|
message_bytes, self.inbound_buffer = (
|
||||||
|
self.inbound_buffer.split(b"\n", 1)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
message = Message.from_bytes(message_bytes)
|
||||||
|
self.on_message(message)
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.error(f"[!] Failed to parse message: {e}")
|
||||||
|
|
||||||
|
time.sleep(0.1)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
self.logger.info("Interrupted by user.")
|
self.logger.debug("[*] Interrupted by user.")
|
||||||
finally:
|
finally:
|
||||||
self.close()
|
self.close()
|
||||||
|
|||||||
Reference in New Issue
Block a user