import logging import multiprocessing as mp import socket import struct import subprocess import threading import time from datetime import datetime from typing import Callable, NamedTuple import numpy as np from ..config import config from .csi_frame import CSI from .protocols import CSICallback Host = tuple[str, int] class FeitHost: def __init__(self, host: Host, command: str) -> None: self.command = command self.host = host self.logger = logging.getLogger( f"{__name__}.{self.__class__.__name__}-{self.host[0]}" ) self.active = True self.checker = threading.Thread(target=self.check_continuous) self.checker.start() def check_connection(self) -> bool: if not config.HOST_CHECKING: # This may not always be possible to check when the host is a reverse-proxy # or doesn't have SSH access. This is the case when running on cloud servers return True feitcsi_status = subprocess.run( f"ssh root@{self.host[0]} pgrep feitcsi", check=False, stdout=subprocess.DEVNULL, shell=True, ) return feitcsi_status.returncode == 0 def connect(self) -> None: self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) self.server.connect(self.host) self.server.send(b"stop\n") self.server.send(self.command.encode()) self.logger.info(f"Connected to {self.host}") def check_continuous(self) -> None: """ Repeatedly check if the FeitCSI service is running Because FeitCSI is using TCP, we will get no information if the service stops, or the computer is not reachable. In order to make debugging easier, this checks and logs continuously if the service is running. """ last_status = False while self.active: if not self.check_connection(): self.logger.error(f"FeitCSI is not running on {self.host[0]}") last_status = False else: if not last_status: self.connect() last_status = True time.sleep(1) self.logger.info("Stopping host checker") class FeitTransmitter(FeitHost): def __init__(self) -> None: command = ( f"feitcsi --frequency {config.central_freq} " f"--channel-width {config.channel_width} " f"--format {config.frame_format} " f"--mode inject -s 1 --verbose " f"--inject-delay {1_000_000 // config.sample_rate}" ) super().__init__(config.transmit_host, command) class FeitReceiver(FeitHost): def __init__(self, host: Host) -> None: command = ( f"feitcsi --frequency {config.central_freq} " f"--channel-width {config.channel_width} " f"--format {config.frame_format} " f"--mode measure" ) super().__init__(host, command) def listen(self, queue: "mp.Queue[CSI]") -> None: prev_time = datetime.now() self.logger.info("Listening for CSI data") while self.active: while not hasattr(self, "server"): time.sleep(0.1) # This is the max size of a UDP packet. The size of the actual CSI # packet will depend on the frame format and channel width, which # changes the number of subcarriers data = self.server.recv(65535) try: csidata = CSI(data) self.logger.debug( f"Received CSI data after {datetime.now() - prev_time}" ) prev_time = datetime.now() queue.put(csidata) except struct.error: self.logger.error("Failed to parse CSI data") self.logger.info("Stopping CSI receiver") class CSIProcessor: def __init__( self, receiver_connections: dict[Host, "mp.Queue[CSI]"], ) -> None: self.pending_data: dict[Host, tuple[datetime, CSI]] = {} self.pending_data_lock = mp.Lock() self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") self.last_processed = datetime.now() self.connections = receiver_connections self.active = True def add_data(self, host: Host, data: CSI) -> None: if ( host in self.pending_data and self.last_processed < self.pending_data[host][0] ): self.logger.warning( f"Skipping data from {host} at {self.pending_data[host][0]}" ) with self.pending_data_lock: self.pending_data[host] = (datetime.now(), data) def is_ready(self) -> bool: for host in config.receive_hosts: if ( host not in self.pending_data or self.pending_data[host][0] <= self.last_processed ): return False return True def process_data( self, callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> None: self.last_processed = datetime.now() if pre_merge_callback is not None: pre_merge_callback( {host: data[1] for host, data in self.pending_data.items()} ) antenna_data = [ np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2) for ip, antenna in config.antennas.order ] # We have data from all servers all_data = np.concat(antenna_data, axis=1) if callback: callback(all_data) def process_forever( self, callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> None: try: while self.active: for ip, queue in self.connections.items(): while not queue.empty(): self.add_data(ip, queue.get()) if self.is_ready(): self.process_data(callback, pre_merge_callback=pre_merge_callback) else: self.logger.debug("Not all data is ready") time.sleep(0.0005) except KeyboardInterrupt: self.logger.info("Exiting CSI processing") class Receiver(NamedTuple): ip: Host receiver: FeitReceiver queue: "mp.Queue[CSI]" def start_processing( csi_callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> None: receivers = [ Receiver(ip, FeitReceiver(ip), mp.Queue(config.sample_rate)) for ip in config.receive_hosts ] transmitter = FeitTransmitter() # Start injecting CSI frames processor = CSIProcessor({r.ip: r.queue for r in receivers}) receiver_processes = [ threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers ] for proc in receiver_processes: proc.start() processing_thread = mp.Process( target=processor.process_forever, args=(csi_callback, pre_merge_callback) ) processing_thread.start() try: while True: time.sleep(100) except KeyboardInterrupt: for r in receivers: r.receiver.active = False transmitter.active = False processor.active = False return