diff --git a/src/collection/ingest.py b/src/collection/ingest.py index 315aa73..31c6030 100644 --- a/src/collection/ingest.py +++ b/src/collection/ingest.py @@ -2,6 +2,8 @@ import logging import time import socket import multiprocessing as mp +import threading +import subprocess import struct from typing import Callable from datetime import datetime @@ -17,45 +19,78 @@ Host = tuple[str, int] CSICallback = Callable[[npt.NDArray[np.complex128]], None] -class FeitTransmitter: +class FeitHost: + def __init__(self, host: Host, command: str): + self.command = command + self.host = host + self.logger = logging.getLogger( + f"{__name__}.{self.__class__.__name__}-{self.host[0]}" + ) + self.checker = threading.Thread(target=self.check_continuous) + self.checker.start() + + def check_connection(self): + 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): + 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): + """ + 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 True: + 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) + + +class FeitTransmitter(FeitHost): def __init__(self): - self.inject_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.inject_server.connect(config.INJECT_HOST) - inject_start_string = ( + command = ( f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} " 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}" ) - self.inject_server.send(b"stop\n") - self.inject_server.send(inject_start_string.encode()) - self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") - self.logger.info(f"Connected to {config.INJECT_HOST}") + super().__init__(config.INJECT_HOST, command) -class FeitReceiver: +class FeitReceiver(FeitHost): def __init__(self, host: Host): - self.host = host - self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.server.connect(host) - self.start_string = ( + command = ( f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} " f"--channel-width {config.CHANNEL_WIDTH} " f"--format {config.FRAME_FORMAT} " f"--mode measure" ) - self.server.send(b"stop\n") - self.server.send(self.start_string.encode()) - self.logger = logging.getLogger( - f"{__name__}.{self.__class__.__name__}-{self.host}" - ) - self.logger.info(f"Connected to {self.host}") + super().__init__(host, command) def listen(self, queue: "mp.Queue[CSI]"): prev_time = datetime.now() self.logger.info("Listening for CSI data") while True: + 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 @@ -157,7 +192,7 @@ def start_processing(csi_callback: CSICallback): processor = CSIProcessor({r.ip: r.queue for r in receivers}) receiver_processes = [ - mp.Process(target=r.receiver.listen, args=(r.queue,)) for r in receivers + threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers ] for proc in receiver_processes: