229 lines
7.5 KiB
Python
229 lines
7.5 KiB
Python
import logging
|
|
import multiprocessing as mp
|
|
import selectors
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
|
|
from ..config import config
|
|
from .csi_frame import CSI
|
|
from .protocols import MergedCSI
|
|
|
|
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()
|
|
self.selectors: list[selectors.BaseSelector] = []
|
|
|
|
def check_connection(self) -> bool:
|
|
try:
|
|
feitcsi_status = subprocess.run(
|
|
f"ssh root@{self.host[0]} pgrep feitcsi",
|
|
check=False,
|
|
stdout=subprocess.DEVNULL,
|
|
shell=True,
|
|
timeout=1,
|
|
)
|
|
return feitcsi_status.returncode == 0
|
|
except subprocess.TimeoutExpired:
|
|
return False
|
|
|
|
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())
|
|
|
|
for selector in self.selectors:
|
|
selector.register(self.server, selectors.EVENT_READ, data=self)
|
|
|
|
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:
|
|
self.logger.debug(f"Checking connection to {self.host[0]}")
|
|
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.collection_sample_rate}"
|
|
)
|
|
super().__init__(config.transmit_host, command)
|
|
|
|
|
|
class FeitReceiver(FeitHost):
|
|
"""
|
|
A class used to connect to the host running FeitCSI and receive CSI data over a UDP
|
|
socket.
|
|
"""
|
|
|
|
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 recv(self) -> CSI:
|
|
"""
|
|
Block until a CSI frame is received and return it.
|
|
"""
|
|
|
|
# Receive the data from the socket
|
|
data = self.server.recv(65535)
|
|
# Decode the data using the CSI frame format
|
|
csi = CSI(data)
|
|
return csi
|
|
|
|
def register(self, selector: selectors.BaseSelector) -> None:
|
|
"""
|
|
Register the receiver as a selector. This will allow us to multiplex the I/O
|
|
operations on a single thread.
|
|
|
|
This allows us to block until any of the receivers receive data
|
|
"""
|
|
self.selectors.append(selector)
|
|
|
|
if hasattr(self, "server"):
|
|
selector.register(self.server, selectors.EVENT_READ, data=self)
|
|
|
|
|
|
class CSIAntennaArray:
|
|
"""
|
|
Represents a set of receiver hosts that are used to receive CSI data, assumed to be
|
|
part of a linear antenna array.
|
|
|
|
This class is used to buffer a set of CSI data from multiple receivers, and once a
|
|
set of readings is ready for each receiver, it will be merged into a single sample
|
|
using the antenna order in the configuration.
|
|
|
|
When the data is ready, the following callbacks are called:
|
|
- `pre_merge_callback`: Called with the data before merging. This is useful for
|
|
per-antenna analysis, e.g. finding the correct antenna order
|
|
- `sample_callback`: Called with the merged data
|
|
"""
|
|
|
|
def __init__(self, receivers: list[FeitReceiver]) -> 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.receivers = receivers
|
|
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) -> None | MergedCSI:
|
|
if not self.is_ready():
|
|
self.logger.debug("Not all data is ready")
|
|
return None
|
|
self.last_processed = datetime.now()
|
|
|
|
frames = {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)
|
|
return MergedCSI(frames=frames, matrix=all_data)
|
|
|
|
def process_forever(self) -> Iterator[MergedCSI]:
|
|
# Register the receivers with the selector
|
|
self.logger.info("Starting to process data from receivers")
|
|
selector = selectors.DefaultSelector()
|
|
|
|
for receiver in self.receivers:
|
|
receiver.register(selector)
|
|
|
|
while self.active:
|
|
events = selector.select(timeout=0.1)
|
|
if not events:
|
|
self.logger.warning("No events received in the last 0.1 seconds")
|
|
continue
|
|
|
|
for key, _ in events:
|
|
receiver = key.data
|
|
self.add_data(receiver.host, receiver.recv())
|
|
|
|
sample = self.process_data()
|
|
if sample is not None:
|
|
yield sample
|
|
|
|
|
|
class RealtimeCSIProducer:
|
|
def __init__(self) -> None:
|
|
self.receivers = [FeitReceiver(ip) for ip in config.receive_hosts]
|
|
self.transmitter = FeitTransmitter()
|
|
|
|
# Start injecting CSI frames
|
|
self.buffer = CSIAntennaArray(self.receivers)
|
|
self.is_live = True
|
|
|
|
def __call__(self) -> Iterator[MergedCSI]:
|
|
return self.buffer.process_forever()
|
|
|
|
def stop(self) -> None:
|
|
for r in self.receivers:
|
|
r.active = False
|
|
self.transmitter.active = False
|
|
self.buffer.active = False
|