diff --git a/config-test.yaml b/config-test.yaml new file mode 100644 index 0000000..9745265 --- /dev/null +++ b/config-test.yaml @@ -0,0 +1,26 @@ +# Test config - TCP push only, no HTTP trackers +output_dir: "./output" +poll_interval_sec: 0.5 +status_interval_sec: 10.0 +retina_tracker_path: "../retina-tracker" + +retry: + max_attempts: 5 + backoff_base_sec: 2 + extended_outage_sec: 60 + health_check_interval_sec: 30 + +api_forward: + enabled: false + +tcp_server: + enabled: true + host: "0.0.0.0" + port: 30050 + token: "shared-secret" + auto_provision: true + auto_provision_port_start: 31000 + auto_provision_port_end: 31099 + +# No pre-configured trackers - rely entirely on auto-provisioning +trackers: [] diff --git a/config.yaml b/config.yaml index 05a82bf..8c59e1d 100644 --- a/config.yaml +++ b/config.yaml @@ -25,9 +25,21 @@ api_forward: # headers: # Authorization: "Bearer your-token" +# TCP push server +tcp_server: + enabled: true + host: "0.0.0.0" + port: 30050 + token: "shared-secret" + auto_provision: true + auto_provision_port_start: 31000 + auto_provision_port_end: 31099 # max 100 auto-provisioned nodes + # Tracker instances - one per detection endpoint trackers: + # HTTP-polled node (legacy, explicit opt-in) - name: "radar3" + mode: "http" detection_url: "https://radar3.retnode.com/api/detection" config_url: "https://radar3.retnode.com/api/config" tcp_port: 30012 @@ -40,10 +52,14 @@ trackers: # enabled: true # url: "https://other-api.com/tracks" + # TCP-pushed node (explicit config with overrides) - name: "radar4" - detection_url: "https://radar4.30hours.dev/api/detection" - config_url: "https://radar4.30hours.dev/api/config" + mode: "tcp" + node_id: "ret7dd2cb0d" tcp_port: 30013 geolocator: enabled: true tcp_port: 31013 + # tracker_config: "./custom-radar4.yaml" + + # Auto-provisioned nodes need no entry at all diff --git a/tests/simulate_node.py b/tests/simulate_node.py new file mode 100644 index 0000000..e7ae7f8 --- /dev/null +++ b/tests/simulate_node.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 +import argparse +import json +import random +import socket +import sys +import threading +import time + + +def generate_detection_frame(node_id, token, timestamp): + num_detections = random.randint(1, 5) + return { + "node_id": node_id, + "token": token, + "timestamp": timestamp, + "delay": [round(random.uniform(10, 300), 2) for _ in range(num_detections)], + "doppler": [round(random.uniform(-150, 150), 2) for _ in range(num_detections)], + "snr": [round(random.uniform(5, 25), 2) for _ in range(num_detections)], + "adsb": [None] * num_detections, + } + + +def generate_heartbeat(node_id, token): + return {"node_id": node_id, "token": token, "type": "heartbeat"} + + +def check_connection(sock): + """Non-blocking check if server has closed the connection.""" + sock.setblocking(False) + try: + data = sock.recv(1) + if not data: + raise ConnectionResetError("Server closed connection") + except BlockingIOError: + pass # No data available = connection still open + finally: + sock.setblocking(True) + + +def send_message(sock, msg, verbose=False, label=""): + line = json.dumps(msg) + "\n" + try: + sock.sendall(line.encode("utf-8")) + except (BrokenPipeError, ConnectionResetError, OSError) as e: + if verbose: + print(f"[{label}] Send failed: {e}") + raise + if verbose: + print(f"[{label}] Sent: {line.rstrip()}") + + +def run_node(args, node_id): + label = node_id + token = "wrong-token" if args.invalid_token else args.token + frames_sent = 0 + timestamp = int(time.time() * 1000) + + while True: + sock = None + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + print(f"[{label}] Connecting to {args.host}:{args.port}...") + sock.connect((args.host, args.port)) + print(f"[{label}] Connected") + + if args.malformed: + garbage = b"\xff\xfe NOT JSON {{{\x00\n" + sock.sendall(garbage) + if args.verbose: + print(f"[{label}] Sent malformed data") + time.sleep(1) + return + + if args.buffer_burst and frames_sent > 0: + burst_count = args.buffer_burst + print(f"[{label}] Draining buffer: {burst_count} frames") + for _ in range(burst_count): + frame = generate_detection_frame(node_id, token, timestamp) + send_message(sock, frame, args.verbose, label) + timestamp += int(args.cpi * 1000) + frames_sent += 1 + if args.num_frames and frames_sent >= args.num_frames: + print(f"[{label}] Reached {args.num_frames} frames, stopping") + return + + last_send_time = time.time() + frames_this_connection = 0 + + while True: + if args.num_frames and frames_sent >= args.num_frames: + print(f"[{label}] Reached {args.num_frames} frames, stopping") + return + + if ( + args.disconnect_after + and frames_this_connection >= args.disconnect_after + ): + print( + f"[{label}] Disconnecting after {args.disconnect_after} frames" + ) + sock.close() + sock = None + delay = args.reconnect_delay + print(f"[{label}] Waiting {delay}s before reconnect...") + time.sleep(delay) + break + + now = time.time() + + if args.heartbeat_only: + if now - last_send_time >= 10: + hb = generate_heartbeat(node_id, token) + send_message(sock, hb, args.verbose, label) + last_send_time = now + frames_sent += 1 + frames_this_connection += 1 + time.sleep(0.5) + continue + + frame = generate_detection_frame(node_id, token, timestamp) + send_message(sock, frame, args.verbose, label) + timestamp += int(args.cpi * 1000) + frames_sent += 1 + frames_this_connection += 1 + last_send_time = now + + time.sleep(args.cpi) + check_connection(sock) + + except ConnectionRefusedError: + print( + f"[{label}] Connection refused, retrying in {args.reconnect_delay}s..." + ) + time.sleep(args.reconnect_delay) + except (ConnectionResetError, BrokenPipeError, OSError) as e: + print(f"[{label}] Connection lost: {e}") + if args.invalid_token: + print(f"[{label}] Server likely rejected invalid token") + return + print(f"[{label}] Reconnecting in {args.reconnect_delay}s...") + time.sleep(args.reconnect_delay) + finally: + if sock is not None: + try: + sock.close() + except OSError: + pass + + +def main(): + parser = argparse.ArgumentParser( + description="Simulate radar nodes pushing detection frames over TCP" + ) + parser.add_argument( + "--host", default="localhost", help="Server hostname (default: localhost)" + ) + parser.add_argument( + "--port", type=int, default=30050, help="Server port (default: 30050)" + ) + parser.add_argument( + "--node-id", default="radar-test", help="Node ID to send (default: radar-test)" + ) + parser.add_argument( + "--token", default="secret", help="Auth token (default: secret)" + ) + parser.add_argument( + "--cpi", + type=float, + default=0.5, + help="Interval between frames in seconds (default: 0.5)", + ) + parser.add_argument( + "--nodes", + type=int, + default=0, + help="Simulate N concurrent nodes (node IDs: node-1, node-2, etc.)", + ) + parser.add_argument( + "--disconnect-after", + type=int, + default=0, + help="Disconnect after N frames (default: 0 = never)", + ) + parser.add_argument( + "--reconnect-delay", + type=float, + default=5, + help="Seconds to wait before reconnecting (default: 5)", + ) + parser.add_argument( + "--buffer-burst", + type=int, + default=0, + help="Send N frames rapidly on reconnect (simulating buffer drain)", + ) + parser.add_argument( + "--invalid-token", + action="store_true", + help="Use wrong token (for testing rejection)", + ) + parser.add_argument( + "--malformed", action="store_true", help="Send garbage data instead of JSON" + ) + parser.add_argument( + "--heartbeat-only", + action="store_true", + help="Send only heartbeats, no detections", + ) + parser.add_argument( + "--num-frames", + type=int, + default=0, + help="Stop after N frames (default: 0 = unlimited)", + ) + parser.add_argument("--verbose", action="store_true", help="Print each frame sent") + + args = parser.parse_args() + + if args.nodes > 1: + threads = [] + for i in range(1, args.nodes + 1): + nid = f"node-{i}" + t = threading.Thread(target=run_node, args=(args, nid), daemon=True) + t.start() + threads.append(t) + print(f"Started thread for {nid}") + + try: + for t in threads: + t.join() + except KeyboardInterrupt: + print("\nInterrupted, shutting down...") + sys.exit(0) + else: + try: + run_node(args, args.node_id) + except KeyboardInterrupt: + print("\nInterrupted, shutting down...") + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/tests/test_auto_provision.py b/tests/test_auto_provision.py new file mode 100644 index 0000000..0b641c3 --- /dev/null +++ b/tests/test_auto_provision.py @@ -0,0 +1,348 @@ +"""Tests for TrackerManager auto-provisioning and TcpServer integration.""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +import pytest_asyncio + +from tracker_host.config import Config, RetryConfig, TcpServerConfig, TrackerConfig +from tracker_host.fetcher import DetectionFetcher +from tracker_host.manager import TrackerManager +from tracker_host.tcp_receiver import TcpReceiver + +TOKEN = "test-secret" + + +def _make_config( + trackers=None, + tcp_enabled=True, + auto_provision=True, + port_start=31000, + port_end=31099, + tcp_port=0, +): + return Config( + output_dir="./output", + poll_interval_sec=0.5, + status_interval_sec=30.0, + retry=RetryConfig(), + tcp_server=TcpServerConfig( + enabled=tcp_enabled, + host="127.0.0.1", + port=tcp_port, + token=TOKEN, + auto_provision=auto_provision, + auto_provision_port_start=port_start, + auto_provision_port_end=port_end, + ), + trackers=trackers or [], + ) + + +async def _get_free_port(): + server = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + server.close() + await server.wait_closed() + return port + + +async def _connect(port): + return await asyncio.open_connection("127.0.0.1", port) + + +async def _send_and_drain(writer, msg): + writer.write((json.dumps(msg) + "\n").encode()) + await writer.drain() + + +@pytest.fixture +def instance_patcher(): + with patch("tracker_host.manager.TrackerInstance") as mock_cls: + + def make_instance(**kwargs): + inst = MagicMock() + inst.start = AsyncMock() + inst.stop = AsyncMock() + inst.get_status = MagicMock(return_value="[running] ok") + inst.config = kwargs.get( + "tracker_config", TrackerConfig(name="mock", tcp_port=99999) + ) + return inst + + mock_cls.side_effect = lambda **kwargs: make_instance(**kwargs) + yield mock_cls + + +@pytest_asyncio.fixture +async def manager_with_tcp(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port) + manager = TrackerManager(config) + await manager.start() + yield manager, tcp_port, instance_patcher + await manager.stop() + + +@pytest.mark.asyncio +async def test_auto_provision_creates_instance(manager_with_tcp): + manager, tcp_port, mock_cls = manager_with_tcp + assert len(manager._instances) == 0 + + _, writer = await _connect(tcp_port) + frame = {"node_id": "new-node", "token": TOKEN, "timestamp": 100, "delay": [1.0]} + await _send_and_drain(writer, frame) + await asyncio.sleep(0.1) + + assert len(manager._instances) == 1 + call_kwargs = mock_cls.call_args + assert call_kwargs[1]["tracker_config"].name == "new-node" + assert call_kwargs[1]["tracker_config"].mode == "tcp" + assert isinstance(call_kwargs[1]["source"], TcpReceiver) + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_auto_provision_receiver_registered(manager_with_tcp): + manager, tcp_port, _ = manager_with_tcp + + _, writer = await _connect(tcp_port) + frame = {"node_id": "node-a", "token": TOKEN, "timestamp": 200, "delay": [2.0]} + await _send_and_drain(writer, frame) + await asyncio.sleep(0.1) + + assert "node-a" in manager._tcp_server._receivers + receiver = manager._tcp_server._receivers["node-a"] + result = await receiver.receive() + assert result is not None + assert result["timestamp"] == 200 + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_auto_provision_port_range(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port, port_start=32000, port_end=32002) + manager = TrackerManager(config) + await manager.start() + + try: + for i, name in enumerate(["a", "b", "c"]): + _, writer = await _connect(tcp_port) + await _send_and_drain( + writer, + {"node_id": name, "token": TOKEN, "timestamp": i, "delay": [1.0]}, + ) + await asyncio.sleep(0.1) + writer.close() + await writer.wait_closed() + + assert len(manager._instances) == 3 + ports = {inst.config.tcp_port for inst in manager._instances} + assert ports == {32000, 32001, 32002} + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_auto_provision_port_exhaustion(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port, port_start=33000, port_end=33000) + manager = TrackerManager(config) + await manager.start() + + try: + _, w1 = await _connect(tcp_port) + await _send_and_drain( + w1, {"node_id": "first", "token": TOKEN, "timestamp": 1, "delay": [1.0]} + ) + await asyncio.sleep(0.1) + assert len(manager._instances) == 1 + + _, w2 = await _connect(tcp_port) + await _send_and_drain( + w2, {"node_id": "second", "token": TOKEN, "timestamp": 2, "delay": [1.0]} + ) + await asyncio.sleep(0.1) + assert len(manager._instances) == 1 + + w1.close() + w2.close() + await w1.wait_closed() + await w2.wait_closed() + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_reconnecting_node_reattaches(manager_with_tcp): + manager, tcp_port, mock_cls = manager_with_tcp + + _, w1 = await _connect(tcp_port) + await _send_and_drain( + w1, {"node_id": "recon-node", "token": TOKEN, "timestamp": 1, "delay": [1.0]} + ) + await asyncio.sleep(0.1) + assert len(manager._instances) == 1 + creation_count = mock_cls.call_count + + receiver = manager._tcp_server._receivers["recon-node"] + r1 = await receiver.receive() + assert r1 is not None + assert r1["timestamp"] == 1 + + w1.close() + await w1.wait_closed() + await asyncio.sleep(0.1) + + _, w2 = await _connect(tcp_port) + await _send_and_drain( + w2, {"node_id": "recon-node", "token": TOKEN, "timestamp": 2, "delay": [2.0]} + ) + await asyncio.sleep(0.1) + + assert len(manager._instances) == 1 + assert mock_cls.call_count == creation_count + + r2 = await receiver.receive() + assert r2 is not None + assert r2["timestamp"] == 2 + + w2.close() + await w2.wait_closed() + + +@pytest.mark.asyncio +async def test_tcp_mode_tracker_gets_receiver(instance_patcher): + tcp_port = await _get_free_port() + trackers = [ + TrackerConfig(name="tcp-node", tcp_port=30099, mode="tcp", node_id="tcp-node"), + ] + config = _make_config(trackers=trackers, tcp_port=tcp_port) + manager = TrackerManager(config) + await manager.start() + + try: + assert len(manager._instances) == 1 + call_kwargs = instance_patcher.call_args + assert isinstance(call_kwargs[1]["source"], TcpReceiver) + assert "tcp-node" in manager._tcp_server._receivers + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_http_mode_tracker_gets_fetcher(instance_patcher): + tcp_port = await _get_free_port() + trackers = [ + TrackerConfig( + name="http-node", + tcp_port=30098, + mode="http", + detection_url="http://example.com/api/detection", + ), + ] + config = _make_config(trackers=trackers, tcp_port=tcp_port) + manager = TrackerManager(config) + await manager.start() + + try: + assert len(manager._instances) == 1 + call_kwargs = instance_patcher.call_args + assert isinstance(call_kwargs[1]["source"], DetectionFetcher) + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_tcp_server_started_when_enabled(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port, tcp_enabled=True) + manager = TrackerManager(config) + await manager.start() + + try: + assert manager._tcp_server is not None + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_tcp_server_not_started_when_disabled(instance_patcher): + config = _make_config(tcp_enabled=False) + manager = TrackerManager(config) + await manager.start() + + try: + assert manager._tcp_server is None + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_stop_clears_tcp_server(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port) + manager = TrackerManager(config) + await manager.start() + assert manager._tcp_server is not None + + await manager.stop() + assert manager._tcp_server is None + + +@pytest.mark.asyncio +async def test_auto_provision_disabled(instance_patcher): + tcp_port = await _get_free_port() + config = _make_config(tcp_port=tcp_port, auto_provision=False) + manager = TrackerManager(config) + await manager.start() + + try: + _, writer = await _connect(tcp_port) + await _send_and_drain( + writer, + {"node_id": "unknown", "token": TOKEN, "timestamp": 1, "delay": [1.0]}, + ) + await asyncio.sleep(0.1) + + assert len(manager._instances) == 0 + + writer.close() + await writer.wait_closed() + finally: + await manager.stop() + + +@pytest.mark.asyncio +async def test_mixed_http_and_tcp_trackers(instance_patcher): + tcp_port = await _get_free_port() + trackers = [ + TrackerConfig( + name="http-node", + tcp_port=30090, + mode="http", + detection_url="http://example.com/det", + ), + TrackerConfig(name="tcp-node", tcp_port=30091, mode="tcp", node_id="tcp-node"), + ] + config = _make_config(trackers=trackers, tcp_port=tcp_port) + manager = TrackerManager(config) + await manager.start() + + try: + assert len(manager._instances) == 2 + + first_call = instance_patcher.call_args_list[0] + assert isinstance(first_call[1]["source"], DetectionFetcher) + + second_call = instance_patcher.call_args_list[1] + assert isinstance(second_call[1]["source"], TcpReceiver) + finally: + await manager.stop() diff --git a/tests/test_tcp_receiver.py b/tests/test_tcp_receiver.py new file mode 100644 index 0000000..3d3eded --- /dev/null +++ b/tests/test_tcp_receiver.py @@ -0,0 +1,87 @@ +"""Tests for TcpReceiver.""" + +import pytest + +from tracker_host.tcp_receiver import MAX_QUEUE_SIZE, TcpReceiver + + +@pytest.fixture +def receiver(): + return TcpReceiver(node_id="test-node") + + +@pytest.mark.asyncio +async def test_receive_returns_frame_from_queue(receiver): + frame = {"timestamp": 1000, "delay": [1.0], "doppler": [2.0]} + receiver.put(frame) + result = await receiver.receive() + assert result == frame + + +@pytest.mark.asyncio +async def test_receive_returns_none_on_empty_queue(receiver): + result = await receiver.receive() + assert result is None + + +@pytest.mark.asyncio +async def test_timestamp_dedup_skips_duplicate(receiver): + frame1 = {"timestamp": 1000, "delay": [1.0]} + frame2 = {"timestamp": 1000, "delay": [2.0]} + frame3 = {"timestamp": 2000, "delay": [3.0]} + + receiver.put(frame1) + receiver.put(frame2) + receiver.put(frame3) + + result1 = await receiver.receive() + assert result1 == frame1 + + result2 = await receiver.receive() + assert result2 is None + + result3 = await receiver.receive() + assert result3 == frame3 + + +@pytest.mark.asyncio +async def test_put_drops_oldest_when_queue_full(receiver): + for i in range(MAX_QUEUE_SIZE): + receiver.put({"timestamp": i}) + + assert receiver._queue.full() + + receiver.put({"timestamp": 9999}) + + assert receiver._queue.qsize() == MAX_QUEUE_SIZE + + first = receiver._queue.get_nowait() + assert first["timestamp"] == 1 + + +@pytest.mark.asyncio +async def test_is_healthy_tracks_connected_state(receiver): + assert receiver.is_healthy is False + + receiver.set_connected(True) + assert receiver.is_healthy is True + + receiver.set_connected(False) + assert receiver.is_healthy is False + + +@pytest.mark.asyncio +async def test_close_clears_queue(receiver): + receiver.put({"timestamp": 1}) + receiver.put({"timestamp": 2}) + + await receiver.close() + assert receiver._queue.empty() + + +@pytest.mark.asyncio +async def test_frame_without_timestamp_always_returned(receiver): + frame = {"delay": [1.0]} + receiver.put(frame) + result = await receiver.receive() + assert result == frame diff --git a/tests/test_tcp_server.py b/tests/test_tcp_server.py new file mode 100644 index 0000000..1bb0c28 --- /dev/null +++ b/tests/test_tcp_server.py @@ -0,0 +1,258 @@ +"""Tests for TcpServer.""" + +import asyncio +import json + +import pytest +import pytest_asyncio + +from tracker_host.tcp_receiver import TcpReceiver +from tracker_host.tcp_server import TcpServer + +TOKEN = "test-secret" + + +async def _get_free_port(): + server = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + server.close() + await server.wait_closed() + return port + + +async def _connect(port): + reader, writer = await asyncio.open_connection("127.0.0.1", port) + return reader, writer + + +def _send_line(writer, msg): + writer.write((json.dumps(msg) + "\n").encode()) + + +async def _send_and_drain(writer, msg): + _send_line(writer, msg) + await writer.drain() + + +@pytest_asyncio.fixture +async def server_and_port(): + port = await _get_free_port() + server = TcpServer(token=TOKEN) + await server.start("127.0.0.1", port) + yield server, port + await server.stop() + + +@pytest.mark.asyncio +async def test_accepts_connections(server_and_port): + server, port = server_and_port + reader, writer = await _connect(port) + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "type": "heartbeat"} + ) + await asyncio.sleep(0.05) + assert not writer.is_closing() + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_validates_token(server_and_port): + server, port = server_and_port + receiver = TcpReceiver("n1") + server.register_receiver("n1", receiver) + + reader, writer = await _connect(port) + frame = {"node_id": "n1", "token": TOKEN, "timestamp": 100, "delay": [1.0]} + await _send_and_drain(writer, frame) + await asyncio.sleep(0.05) + + result = await receiver.receive() + assert result is not None + assert result["timestamp"] == 100 + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_rejects_invalid_token(server_and_port): + server, port = server_and_port + reader, writer = await _connect(port) + await _send_and_drain(writer, {"node_id": "n1", "token": "wrong"}) + await asyncio.sleep(0.1) + + data = await asyncio.wait_for(reader.read(1024), timeout=1.0) + assert data == b"" + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_routes_frames_to_correct_receiver(server_and_port): + server, port = server_and_port + r1 = TcpReceiver("n1") + r2 = TcpReceiver("n2") + server.register_receiver("n1", r1) + server.register_receiver("n2", r2) + + _, w1 = await _connect(port) + _, w2 = await _connect(port) + + await _send_and_drain( + w1, {"node_id": "n1", "token": TOKEN, "timestamp": 1, "delay": [1.0]} + ) + await _send_and_drain( + w2, {"node_id": "n2", "token": TOKEN, "timestamp": 2, "delay": [2.0]} + ) + await asyncio.sleep(0.05) + + result1 = await r1.receive() + result2 = await r2.receive() + + assert result1["timestamp"] == 1 + assert result2["timestamp"] == 2 + + w1.close() + w2.close() + await w1.wait_closed() + await w2.wait_closed() + + +@pytest.mark.asyncio +async def test_discards_heartbeat_messages(server_and_port): + server, port = server_and_port + receiver = TcpReceiver("n1") + server.register_receiver("n1", receiver) + + _, writer = await _connect(port) + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "type": "heartbeat"} + ) + await asyncio.sleep(0.05) + + result = await receiver.receive() + assert result is None + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_calls_provision_callback_for_unknown_node(server_and_port): + server, port = server_and_port + provisioned = [] + + def on_provision(node_id, frame): + receiver = TcpReceiver(node_id) + server.register_receiver(node_id, receiver) + provisioned.append((node_id, frame)) + + server.set_provision_callback(on_provision) + + _, writer = await _connect(port) + frame = {"node_id": "new-node", "token": TOKEN, "timestamp": 500, "delay": [3.0]} + await _send_and_drain(writer, frame) + await asyncio.sleep(0.05) + + assert len(provisioned) == 1 + assert provisioned[0][0] == "new-node" + assert provisioned[0][1]["timestamp"] == 500 + + result = await server._receivers["new-node"].receive() + assert result is not None + assert result["timestamp"] == 500 + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_disconnect_sets_receiver_unhealthy(server_and_port): + server, port = server_and_port + receiver = TcpReceiver("n1") + server.register_receiver("n1", receiver) + + _, writer = await _connect(port) + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "timestamp": 1, "delay": [1.0]} + ) + await asyncio.sleep(0.05) + + assert receiver.is_healthy is True + + writer.close() + await writer.wait_closed() + await asyncio.sleep(0.1) + + assert receiver.is_healthy is False + + +@pytest.mark.asyncio +async def test_enforces_max_message_size(server_and_port): + server, port = server_and_port + receiver = TcpReceiver("n1") + server.register_receiver("n1", receiver) + + _, writer = await _connect(port) + await _send_and_drain(writer, {"node_id": "n1", "token": TOKEN, "timestamp": 1}) + await asyncio.sleep(0.05) + + oversized = json.dumps( + {"node_id": "n1", "token": TOKEN, "timestamp": 999, "data": "x" * 1_100_000} + ) + chunk_size = 32768 + payload = (oversized + "\n").encode() + for i in range(0, len(payload), chunk_size): + writer.write(payload[i : i + chunk_size]) + await asyncio.sleep(0.005) + await asyncio.sleep(0.5) + + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "timestamp": 2, "delay": [1.0]} + ) + await asyncio.sleep(0.1) + + frames = [] + while True: + r = await receiver.receive() + if r is None: + break + frames.append(r) + + timestamps = [f["timestamp"] for f in frames] + assert 999 not in timestamps + assert 2 in timestamps + + writer.close() + await writer.wait_closed() + + +@pytest.mark.asyncio +async def test_handles_malformed_json(server_and_port): + server, port = server_and_port + receiver = TcpReceiver("n1") + server.register_receiver("n1", receiver) + + _, writer = await _connect(port) + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "timestamp": 1, "delay": [1.0]} + ) + await asyncio.sleep(0.05) + + writer.write(b"not valid json\n") + await writer.drain() + await asyncio.sleep(0.05) + + await _send_and_drain( + writer, {"node_id": "n1", "token": TOKEN, "timestamp": 2, "delay": [2.0]} + ) + await asyncio.sleep(0.05) + + r1 = await receiver.receive() + assert r1["timestamp"] == 1 + r2 = await receiver.receive() + assert r2["timestamp"] == 2 + + writer.close() + await writer.wait_closed() diff --git a/tracker_host/__main__.py b/tracker_host/__main__.py index 9580149..6ab81ea 100644 --- a/tracker_host/__main__.py +++ b/tracker_host/__main__.py @@ -41,7 +41,9 @@ def main() -> None: config_path = Path(args.config) if not config_path.exists(): print(f"Error: Config file not found: {config_path}", file=sys.stderr) - print("Create a config.yaml or specify a different path with -c", file=sys.stderr) + print( + "Create a config.yaml or specify a different path with -c", file=sys.stderr + ) sys.exit(1) # Run the manager diff --git a/tracker_host/config.py b/tracker_host/config.py index c9e59ff..629f150 100644 --- a/tracker_host/config.py +++ b/tracker_host/config.py @@ -40,14 +40,29 @@ class TrackerConfig: """Configuration for a single tracker instance.""" name: str - detection_url: str tcp_port: int + detection_url: Optional[str] = None + mode: Optional[str] = None + node_id: Optional[str] = None tracker_config: Optional[str] = None config_url: Optional[str] = None api_forward: Optional[ApiForwardConfig] = None geolocator: Optional[GeolocatorConfig] = None +@dataclass +class TcpServerConfig: + """TCP push server configuration.""" + + enabled: bool = False + host: str = "0.0.0.0" + port: int = 30050 + token: str = "" + auto_provision: bool = True + auto_provision_port_start: int = 31000 + auto_provision_port_end: int = 31099 + + @dataclass class Config: """Main configuration for tracker-host.""" @@ -62,6 +77,7 @@ class Config: config_poll_interval_sec: float = 60.0 retry: RetryConfig = field(default_factory=RetryConfig) api_forward: ApiForwardConfig = field(default_factory=ApiForwardConfig) + tcp_server: TcpServerConfig = field(default_factory=TcpServerConfig) trackers: list[TrackerConfig] = field(default_factory=list) retina_tracker_path: str = "../retina-tracker" retina_geolocator_path: str = "../retina-geolocator" @@ -108,6 +124,16 @@ def _parse_config(raw: dict) -> Config: headers=tapi.get("headers", {}), ) + detection_url = t.get("detection_url") + explicit_mode = t.get("mode") + + if explicit_mode is not None: + mode = explicit_mode + elif detection_url is not None: + mode = "http" + else: + mode = "tcp" + tracker_geo = None if "geolocator" in t: tgeo = t["geolocator"] @@ -120,8 +146,10 @@ def _parse_config(raw: dict) -> Config: trackers.append( TrackerConfig( name=t["name"], - detection_url=t["detection_url"], tcp_port=t["tcp_port"], + detection_url=detection_url, + mode=mode, + node_id=t.get("node_id"), tracker_config=t.get("tracker_config"), config_url=t.get("config_url"), api_forward=tracker_api, @@ -129,6 +157,19 @@ def _parse_config(raw: dict) -> Config: ) ) + tcp_server_raw = raw.get("tcp_server", {}) + tcp_server = TcpServerConfig( + enabled=tcp_server_raw.get("enabled", False), + host=tcp_server_raw.get("host", "0.0.0.0"), + port=tcp_server_raw.get("port", 30050), + token=tcp_server_raw.get("token", ""), + auto_provision=tcp_server_raw.get("auto_provision", True), + auto_provision_port_start=tcp_server_raw.get( + "auto_provision_port_start", 31000 + ), + auto_provision_port_end=tcp_server_raw.get("auto_provision_port_end", 31099), + ) + return Config( output_dir=raw.get("output_dir", "./output"), poll_interval_sec=raw.get("poll_interval_sec", 0.5), @@ -140,6 +181,7 @@ def _parse_config(raw: dict) -> Config: config_poll_interval_sec=raw.get("config_poll_interval_sec", 60.0), retry=retry, api_forward=api_forward, + tcp_server=tcp_server, trackers=trackers, retina_tracker_path=raw.get("retina_tracker_path", "../retina-tracker"), retina_geolocator_path=raw.get( diff --git a/tracker_host/detection_source.py b/tracker_host/detection_source.py new file mode 100644 index 0000000..11df839 --- /dev/null +++ b/tracker_host/detection_source.py @@ -0,0 +1,14 @@ +"""DetectionSource protocol for pluggable detection ingestion.""" + +from typing import Any, Optional, Protocol + + +class DetectionSource(Protocol): + async def receive(self) -> Optional[dict[str, Any]]: + """Return next detection frame, or None if no new data.""" + ... + + async def close(self) -> None: ... + + @property + def is_healthy(self) -> bool: ... diff --git a/tracker_host/fetcher.py b/tracker_host/fetcher.py index 37181e6..76bedf5 100644 --- a/tracker_host/fetcher.py +++ b/tracker_host/fetcher.py @@ -63,7 +63,11 @@ async def close(self) -> None: await self._session.close() self._session = None - async def fetch(self) -> Optional[dict[str, Any]]: + @property + def is_healthy(self) -> bool: + return self.state.is_healthy + + async def receive(self) -> Optional[dict[str, Any]]: """ Fetch detection data from the endpoint. @@ -119,7 +123,7 @@ async def _handle_failure(self, error: Exception) -> None: # Calculate backoff delay backoff = min( - self.retry_config.backoff_base_sec ** self.state.consecutive_failures, + self.retry_config.backoff_base_sec**self.state.consecutive_failures, 30.0, # Cap at 30 seconds ) diff --git a/tracker_host/instance.py b/tracker_host/instance.py index 2f0c408..43338ee 100644 --- a/tracker_host/instance.py +++ b/tracker_host/instance.py @@ -1,4 +1,4 @@ -"""TrackerInstance: manages one detection endpoint and one retina-tracker subprocess.""" +"""TrackerInstance: manages one detection source and one retina-tracker subprocess.""" import asyncio import json @@ -11,6 +11,7 @@ import aiohttp from .config import Config, TrackerConfig +from .detection_source import DetectionSource from .fetcher import DetectionFetcher, ExtendedOutageError from .geolocator import GeolocatorInstance from .output_handler import OutputHandler @@ -35,6 +36,7 @@ def __init__( self, tracker_config: TrackerConfig, global_config: Config, + source: DetectionSource, session: Optional[aiohttp.ClientSession] = None, ): self.config = tracker_config @@ -47,11 +49,7 @@ def __init__( api_config = tracker_config.api_forward or global_config.api_forward # Components - self.fetcher = DetectionFetcher( - url=tracker_config.detection_url, - retry_config=global_config.retry, - session=session, - ) + self.source = source self.output_handler = OutputHandler( name=self.name, output_dir=global_config.output_dir, @@ -145,7 +143,7 @@ async def stop(self) -> None: await self._stop_tracker_process() # Close handlers - await self.fetcher.close() + await self.source.close() await self.output_handler.close() async def _start_tracker_process(self) -> None: @@ -211,7 +209,7 @@ async def _fetch_loop(self) -> None: """Main loop: fetch detections and send to tracker.""" while self._running: try: - data = await self.fetcher.fetch() + data = await self.source.receive() if data is not None: await self._send_to_tracker(data) @@ -231,8 +229,9 @@ async def _fetch_loop(self) -> None: # Clear metrics since tracker is down self.output_handler.metrics.clear() - # Wait for recovery - await self.fetcher.wait_for_recovery() + # Wait for recovery (only DetectionFetcher raises ExtendedOutageError) + if isinstance(self.source, DetectionFetcher): + await self.source.wait_for_recovery() # Restart logger.info(f"Restarting tracker for {self.name}") diff --git a/tracker_host/manager.py b/tracker_host/manager.py index 6d08d59..ed9d541 100644 --- a/tracker_host/manager.py +++ b/tracker_host/manager.py @@ -4,12 +4,15 @@ import logging import signal from datetime import datetime -from typing import Optional +from typing import Any, Optional import aiohttp -from .config import Config, load_config +from .config import Config, TrackerConfig, load_config +from .fetcher import DetectionFetcher from .instance import TrackerInstance +from .tcp_receiver import TcpReceiver +from .tcp_server import TcpServer logger = logging.getLogger(__name__) @@ -23,6 +26,8 @@ def __init__(self, config: Config): self._session: Optional[aiohttp.ClientSession] = None self._running = False self._status_task: Optional[asyncio.Task] = None + self._tcp_server: Optional[TcpServer] = None + self._used_ports: set[int] = set() async def start(self) -> None: """Start the manager and all configured tracker instances.""" @@ -31,23 +36,69 @@ async def start(self) -> None: self._running = True - # Create shared session self._session = aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=10), headers={ "Accept": "application/json", - "Accept-Encoding": "gzip, deflate", # Avoid brotli issues + "Accept-Encoding": "gzip, deflate", }, ) - # Create instances + tcp_cfg = self.config.tcp_server + if tcp_cfg.enabled: + self._tcp_server = TcpServer(token=tcp_cfg.token) + await self._tcp_server.start(tcp_cfg.host, tcp_cfg.port) + if tcp_cfg.auto_provision: + self._tcp_server.set_provision_callback(self._provision_node) + for tracker_config in self.config.trackers: - instance = TrackerInstance( - tracker_config=tracker_config, - global_config=self.config, - session=self._session, - ) - self._instances.append(instance) + mode = tracker_config.mode or "http" + + if mode == "http": + if not tracker_config.detection_url: + logger.error( + f"Tracker '{tracker_config.name}' is http mode but has no detection_url, skipping" + ) + continue + + source = DetectionFetcher( + url=tracker_config.detection_url, + retry_config=self.config.retry, + session=self._session, + ) + instance = TrackerInstance( + tracker_config=tracker_config, + global_config=self.config, + source=source, + session=self._session, + ) + self._instances.append(instance) + self._used_ports.add(tracker_config.tcp_port) + + elif mode == "tcp": + if self._tcp_server is None: + logger.error( + f"Tracker '{tracker_config.name}' is tcp mode but tcp_server is not enabled, skipping" + ) + continue + + node_id = tracker_config.node_id or tracker_config.name + receiver = TcpReceiver(node_id) + self._tcp_server.register_receiver(node_id, receiver) + + instance = TrackerInstance( + tracker_config=tracker_config, + global_config=self.config, + source=receiver, + session=self._session, + ) + self._instances.append(instance) + self._used_ports.add(tracker_config.tcp_port) + + else: + logger.error( + f"Unknown mode '{mode}' for tracker '{tracker_config.name}', skipping" + ) logger.info(f"Starting {len(self._instances)} tracker instance(s)") @@ -62,10 +113,7 @@ async def start(self) -> None: f"Failed to start instance {instance.name}: {result}" ) - # Start status reporting - self._status_task = asyncio.create_task( - self._status_loop(), name="status-loop" - ) + self._status_task = asyncio.create_task(self._status_loop(), name="status-loop") logger.info("All tracker instances started") @@ -78,7 +126,10 @@ async def stop(self) -> None: logger.info("Stopping tracker manager...") - # Stop status task + if self._tcp_server is not None: + await self._tcp_server.stop() + self._tcp_server = None + if self._status_task is not None: self._status_task.cancel() try: @@ -86,21 +137,70 @@ async def stop(self) -> None: except asyncio.CancelledError: pass - # Stop all instances concurrently await asyncio.gather( *(instance.stop() for instance in self._instances), return_exceptions=True, ) self._instances.clear() + self._used_ports.clear() - # Close shared session if self._session is not None: await self._session.close() self._session = None logger.info("Tracker manager stopped") + def _provision_node(self, node_id: str, frame: dict[str, Any]) -> None: + """Auto-provision a new tracker instance for an unknown node_id.""" + if node_id in {(tc.node_id or tc.name) for tc in self.config.trackers}: + return + + for inst in self._instances: + existing_node_id = inst.config.node_id or inst.config.name + if existing_node_id == node_id: + return + + tcp_cfg = self.config.tcp_server + port = self._allocate_port( + tcp_cfg.auto_provision_port_start, tcp_cfg.auto_provision_port_end + ) + if port is None: + logger.error( + f"Cannot auto-provision node '{node_id}': port range exhausted" + ) + return + + tracker_config = TrackerConfig( + name=node_id, + tcp_port=port, + mode="tcp", + node_id=node_id, + ) + + receiver = TcpReceiver(node_id) + self._tcp_server.register_receiver(node_id, receiver) + + instance = TrackerInstance( + tracker_config=tracker_config, + global_config=self.config, + source=receiver, + session=self._session, + ) + self._instances.append(instance) + self._used_ports.add(port) + + asyncio.create_task(instance.start(), name=f"auto-provision-{node_id}") + + logger.info(f"Auto-provisioned tracker for node '{node_id}' on port {port}") + + def _allocate_port(self, start: int, end: int) -> Optional[int]: + """Find the next available port in the given range.""" + for port in range(start, end + 1): + if port not in self._used_ports: + return port + return None + async def _status_loop(self) -> None: """Periodically print status to console.""" while self._running: @@ -124,7 +224,6 @@ def _print_status(self) -> None: async def run_manager(config_path: str, verbose: bool = False) -> None: """Run the tracker manager with the given config.""" - # Set up logging level = logging.DEBUG if verbose else logging.INFO logging.basicConfig( level=level, @@ -132,17 +231,14 @@ async def run_manager(config_path: str, verbose: bool = False) -> None: datefmt="%Y-%m-%d %H:%M:%S", ) - # Load config config = load_config(config_path) - if not config.trackers: - logger.error("No trackers configured") + if not config.trackers and not config.tcp_server.enabled: + logger.error("No trackers configured and TCP server not enabled") return - # Create and start manager manager = TrackerManager(config) - # Set up signal handlers loop = asyncio.get_running_loop() stop_event = asyncio.Event() @@ -155,9 +251,6 @@ def handle_signal() -> None: try: await manager.start() - - # Wait for shutdown signal await stop_event.wait() - finally: await manager.stop() diff --git a/tracker_host/output_handler.py b/tracker_host/output_handler.py index cae9b0b..2f18bd5 100644 --- a/tracker_host/output_handler.py +++ b/tracker_host/output_handler.py @@ -168,7 +168,9 @@ def get_status(self) -> str: if count == 0: return f"{self.name}: no active tracks" - lengths_str = ", ".join(str(l) for l in sorted(lengths, reverse=True)[:10]) + lengths_str = ", ".join( + str(length) for length in sorted(lengths, reverse=True)[:10] + ) if count > 10: lengths_str += f", ... ({count - 10} more)" diff --git a/tracker_host/tcp_receiver.py b/tracker_host/tcp_receiver.py new file mode 100644 index 0000000..00b51ba --- /dev/null +++ b/tracker_host/tcp_receiver.py @@ -0,0 +1,58 @@ +"""TCP push-based detection source backed by an asyncio queue.""" + +import asyncio +import logging +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +MAX_QUEUE_SIZE = 2000 +RECEIVE_TIMEOUT = 0.5 + + +class TcpReceiver: + """Detection source that receives frames pushed by TcpServer.""" + + def __init__(self, node_id: str): + self.node_id = node_id + self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue( + maxsize=MAX_QUEUE_SIZE + ) + self._last_timestamp: Optional[int] = None + self._connected: bool = False + + async def receive(self) -> Optional[dict[str, Any]]: + try: + frame = await asyncio.wait_for(self._queue.get(), timeout=RECEIVE_TIMEOUT) + except (asyncio.TimeoutError, TimeoutError): + return None + + timestamp = frame.get("timestamp") + if timestamp is not None and timestamp == self._last_timestamp: + return None + + self._last_timestamp = timestamp + return frame + + async def close(self) -> None: + while not self._queue.empty(): + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + break + + @property + def is_healthy(self) -> bool: + return self._connected + + def set_connected(self, connected: bool) -> None: + self._connected = connected + + def put(self, frame: dict[str, Any]) -> None: + if self._queue.full(): + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + pass + logger.warning(f"Queue full for node {self.node_id}, dropped oldest frame") + self._queue.put_nowait(frame) diff --git a/tracker_host/tcp_server.py b/tracker_host/tcp_server.py new file mode 100644 index 0000000..a7ce014 --- /dev/null +++ b/tracker_host/tcp_server.py @@ -0,0 +1,150 @@ +"""Asyncio TCP server for receiving push-based detection frames from radar nodes.""" + +import asyncio +import json +import logging +import socket +from typing import Any, Callable, Optional + +from .tcp_receiver import TcpReceiver + +logger = logging.getLogger(__name__) + +MAX_LINE_LENGTH = 1_048_576 + + +class TcpServer: + """Single TCP server that routes detection frames to registered TcpReceivers.""" + + def __init__(self, token: str): + self._token = token + self._receivers: dict[str, TcpReceiver] = {} + self._provision_callback: Optional[Callable[[str, dict[str, Any]], None]] = None + self._server: Optional[asyncio.AbstractServer] = None + self._connections: dict[asyncio.Task, asyncio.StreamWriter] = {} + self._connection_node_ids: dict[asyncio.Task, str] = {} + + def register_receiver(self, node_id: str, receiver: TcpReceiver) -> None: + self._receivers[node_id] = receiver + + def set_provision_callback( + self, callback: Callable[[str, dict[str, Any]], None] + ) -> None: + self._provision_callback = callback + + async def start(self, host: str, port: int) -> None: + self._server = await asyncio.start_server(self._handle_connection, host, port) + logger.info(f"TCP server listening on {host}:{port}") + + async def stop(self) -> None: + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + for task, writer in list(self._connections.items()): + if not writer.is_closing(): + writer.close() + task.cancel() + + self._connections.clear() + self._connection_node_ids.clear() + + async def _handle_connection( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + sock = writer.get_extra_info("socket") + if sock is not None: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + + peername = writer.get_extra_info("peername") + task = asyncio.current_task() + self._connections[task] = writer + + buffer = "" + authenticated = False + node_id: Optional[str] = None + + try: + while True: + data = await reader.read(65536) + if not data: + break + + buffer += data.decode("utf-8", errors="replace") + + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + + if len(line) > MAX_LINE_LENGTH: + logger.warning( + f"Message exceeds max length from {peername}, dropping" + ) + continue + + if not line.strip(): + continue + + try: + msg = json.loads(line) + except json.JSONDecodeError: + logger.warning(f"Invalid JSON from {peername}: {line[:100]}") + continue + + if not authenticated: + msg_token = msg.get("token") + if msg_token != self._token: + logger.warning( + f"Invalid token from {peername}, closing connection" + ) + transport = writer.transport + if transport is not None: + transport.abort() + return + authenticated = True + + msg_node_id = msg.get("node_id") + if msg_node_id is None: + logger.warning(f"Message without node_id from {peername}") + continue + + if node_id is None: + node_id = msg_node_id + self._connection_node_ids[task] = node_id + if node_id in self._receivers: + self._receivers[node_id].set_connected(True) + + if msg.get("type") == "heartbeat": + continue + + if msg_node_id in self._receivers: + self._receivers[msg_node_id].put(msg) + elif self._provision_callback is not None: + self._provision_callback(msg_node_id, msg) + if msg_node_id in self._receivers: + self._receivers[msg_node_id].set_connected(True) + self._receivers[msg_node_id].put(msg) + + if len(buffer) > MAX_LINE_LENGTH: + logger.warning( + f"Buffer exceeds max length from {peername}, clearing" + ) + buffer = "" + + except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError): + pass + except Exception: + logger.exception(f"Error handling connection from {peername}") + finally: + if node_id is not None and node_id in self._receivers: + self._receivers[node_id].set_connected(False) + + self._connections.pop(task, None) + self._connection_node_ids.pop(task, None) + + if not writer.is_closing(): + writer.close() + try: + await writer.wait_closed() + except Exception: + pass