288 lines
12 KiB
Python
288 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
from .state import RuntimeStats
|
|
|
|
|
|
class RustTZSPReceiver:
|
|
"""Supervise the Rust TZSP data-plane and ingest its 1 Hz telemetry.
|
|
|
|
Packet bytes never cross into Python. The Rust process owns UDP reception,
|
|
TZSP decoding and TAP injection. Python receives only compact telemetry over
|
|
a Unix datagram socket, so UI/Redis work cannot back-pressure packet capture.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
binary: str,
|
|
telemetry_socket: str,
|
|
stats: RuntimeStats,
|
|
stop_event: threading.Event,
|
|
throughput_sink: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> None:
|
|
self.binary = str(binary)
|
|
self.telemetry_socket = str(telemetry_socket)
|
|
self.stats = stats
|
|
self.stop_event = stop_event
|
|
self._throughput_sink = throughput_sink
|
|
self._process: subprocess.Popen | None = None
|
|
self._socket: socket.socket | None = None
|
|
self._thread = threading.Thread(target=self._run_telemetry, name="tzsp-rust-telemetry", daemon=True)
|
|
self._lock = threading.RLock()
|
|
self._ready = threading.Event()
|
|
self._last: dict[str, Any] = {}
|
|
self._samples = 0
|
|
self._telemetry_errors = 0
|
|
self._started_at = time.monotonic()
|
|
|
|
def start(self) -> None:
|
|
if self._process is not None:
|
|
return
|
|
binary = Path(self.binary)
|
|
if not binary.is_file():
|
|
raise RuntimeError(f"Rust TZSP receiver binary not found: {self.binary}")
|
|
|
|
path = Path(self.telemetry_socket)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
path.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
|
|
sock.bind(self.telemetry_socket)
|
|
sock.settimeout(0.5)
|
|
self._socket = sock
|
|
self._thread.start()
|
|
|
|
env = os.environ.copy()
|
|
env["TZSP_TELEMETRY_SOCKET"] = self.telemetry_socket
|
|
self._process = subprocess.Popen([self.binary], env=env)
|
|
print(f"[tzsp] Rust data-plane started, pid={self._process.pid}", flush=True)
|
|
|
|
def wait_ready(self, timeout: float = 8.0) -> bool:
|
|
deadline = time.monotonic() + max(0.1, float(timeout))
|
|
while time.monotonic() < deadline:
|
|
process = self._process
|
|
if process is not None and process.poll() is not None:
|
|
return False
|
|
if self._ready.wait(timeout=min(0.1, max(0.0, deadline - time.monotonic()))):
|
|
return True
|
|
return False
|
|
|
|
def set_throughput_sink(self, sink: Callable[[dict[str, Any]], None] | None) -> None:
|
|
with self._lock:
|
|
self._throughput_sink = sink
|
|
|
|
def is_alive(self) -> bool:
|
|
process = self._process
|
|
return bool(process is not None and process.poll() is None and self._thread.is_alive())
|
|
|
|
@property
|
|
def pid(self) -> int | None:
|
|
process = self._process
|
|
return process.pid if process is not None and process.poll() is None else None
|
|
|
|
def close(self, timeout: float = 3.0) -> None:
|
|
process = self._process
|
|
if process is not None and process.poll() is None:
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=timeout)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
try:
|
|
process.wait(timeout=1)
|
|
except subprocess.TimeoutExpired:
|
|
pass
|
|
sock = self._socket
|
|
self._socket = None
|
|
if sock is not None:
|
|
try:
|
|
sock.close()
|
|
except OSError:
|
|
pass
|
|
if self._thread.is_alive():
|
|
self._thread.join(timeout=1.0)
|
|
try:
|
|
Path(self.telemetry_socket).unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
self._process = None
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
data = dict(self._last)
|
|
now_ms = int(time.time() * 1000)
|
|
ts_ms = int(data.get("ts_ms") or 0)
|
|
data.update(
|
|
{
|
|
"engine": "rust",
|
|
"process_alive": self.is_alive(),
|
|
"pid": self.pid or data.get("pid"),
|
|
"ready": bool(self._ready.is_set() and self.is_alive()),
|
|
"telemetry_age_ms": max(0, now_ms - ts_ms) if ts_ms else None,
|
|
"throughput_samples": self._samples,
|
|
"telemetry_errors": self._telemetry_errors,
|
|
}
|
|
)
|
|
return data
|
|
|
|
def current_throughput(self, window_seconds: int | None = None) -> dict[str, Any]:
|
|
with self._lock:
|
|
sample = dict(self._last)
|
|
now_ms = int(time.time() * 1000)
|
|
ts_ms = int(sample.get("ts_ms") or 0)
|
|
interval = max(float(sample.get("interval_ms") or 1000) / 1000.0, 0.001)
|
|
age_ms = max(0, now_ms - ts_ms) if ts_ms else 10**9
|
|
fresh = bool(ts_ms and age_ms <= max(3000, int(interval * 3000)))
|
|
if fresh:
|
|
total = round(max(int(sample.get("bytes_total") or 0), 0) * 8 / interval)
|
|
inbound = round(max(int(sample.get("bytes_in") or 0), 0) * 8 / interval)
|
|
outbound = round(max(int(sample.get("bytes_out") or 0), 0) * 8 / interval)
|
|
pps = round(max(int(sample.get("packets_total") or 0), 0) / interval, 2)
|
|
ingress_bps = round(max(int(sample.get("rx_bytes_interval") or 0), 0) * 8 / interval)
|
|
ingress_pps = round(max(int(sample.get("rx_datagrams_interval") or 0), 0) / interval, 2)
|
|
else:
|
|
total = inbound = outbound = ingress_bps = 0
|
|
pps = ingress_pps = 0.0
|
|
queue_depth = max(int(sample.get("queue_depth_batches") or 0), 0)
|
|
queue_capacity = max(int(sample.get("queue_capacity_batches") or 0), 0)
|
|
queue_fill_pct = round((queue_depth / queue_capacity) * 100.0, 1) if queue_capacity else 0.0
|
|
inspection_ratio_pct = (
|
|
round(min(100.0, (total / ingress_bps) * 100.0), 1)
|
|
if ingress_bps > 0
|
|
else (100.0 if total == 0 else 0.0)
|
|
)
|
|
loss_per_second = round(
|
|
(
|
|
max(int(sample.get("kernel_udp_drops_interval") or 0), 0)
|
|
+ max(int(sample.get("queue_drops_interval") or 0), 0)
|
|
+ max(int(sample.get("truncated_interval") or 0), 0)
|
|
)
|
|
/ interval,
|
|
2,
|
|
) if fresh else 0.0
|
|
return {
|
|
"window_seconds": int(window_seconds or 0),
|
|
"current_bps": total,
|
|
"current_in_bps": inbound,
|
|
"current_out_bps": outbound,
|
|
"current_other_bps": max(0, total - inbound - outbound),
|
|
"current_pps": pps,
|
|
"current_ingress_bps": ingress_bps,
|
|
"current_ingress_pps": ingress_pps,
|
|
"inspection_ratio_pct": inspection_ratio_pct,
|
|
"capture_efficiency_pct": float(sample.get("capture_efficiency_pct") or 0.0),
|
|
"loss_pps": loss_per_second,
|
|
"current_sample_ts_ms": ts_ms,
|
|
"current_sample_age_ms": age_ms if ts_ms else None,
|
|
"current_sample_fresh": fresh,
|
|
"receiver_engine": "rust",
|
|
"receiver_pid": self.pid,
|
|
"kernel_udp_drops": int(sample.get("kernel_udp_drops") or 0),
|
|
"kernel_udp_drops_interval": int(sample.get("kernel_udp_drops_interval") or 0),
|
|
"queue_dropped_datagrams": int(sample.get("queue_dropped_datagrams") or 0),
|
|
"queue_drops_interval": int(sample.get("queue_drops_interval") or 0),
|
|
"truncated_datagrams": int(sample.get("truncated_datagrams") or 0),
|
|
"truncated_interval": int(sample.get("truncated_interval") or 0),
|
|
"queue_depth_batches": queue_depth,
|
|
"queue_capacity_batches": queue_capacity,
|
|
"queue_capacity_bytes": int(sample.get("queue_capacity_bytes") or 0),
|
|
"queue_high_water_batches": int(sample.get("queue_high_water_batches") or 0),
|
|
"queue_fill_pct": queue_fill_pct,
|
|
"rx_thread_alive": bool(sample.get("rx_thread_alive", False)),
|
|
"worker_thread_alive": bool(sample.get("worker_thread_alive", False)),
|
|
"rcvbuf_bytes": int(sample.get("rcvbuf_bytes") or 0),
|
|
"batch_size": int(sample.get("batch_size") or 0),
|
|
"datagram_bytes": int(sample.get("datagram_bytes") or 0),
|
|
}
|
|
|
|
def overlay_current(self, payload: dict[str, Any], window_seconds: int | None = None) -> dict[str, Any]:
|
|
result = dict(payload)
|
|
result.update(self.current_throughput(window_seconds or int(result.get("window_seconds") or 0)))
|
|
return result
|
|
|
|
def _run_telemetry(self) -> None:
|
|
while not self.stop_event.is_set():
|
|
sock = self._socket
|
|
if sock is None:
|
|
break
|
|
try:
|
|
raw = sock.recv(64 * 1024)
|
|
except socket.timeout:
|
|
continue
|
|
except OSError:
|
|
if self.stop_event.is_set() or self._socket is None:
|
|
break
|
|
self._telemetry_errors += 1
|
|
continue
|
|
try:
|
|
message = json.loads(raw.decode("utf-8"))
|
|
if not isinstance(message, dict) or message.get("type") != "tzsp_sample":
|
|
continue
|
|
self._ingest(message)
|
|
except (UnicodeDecodeError, json.JSONDecodeError, TypeError, ValueError):
|
|
self._telemetry_errors += 1
|
|
|
|
def _ingest(self, message: dict[str, Any]) -> None:
|
|
with self._lock:
|
|
self._last = dict(message)
|
|
sink = self._throughput_sink
|
|
if message.get("ready"):
|
|
self._ready.set()
|
|
|
|
last_packet_ms = int(message.get("last_packet_ms") or 0)
|
|
last_packet_at = None
|
|
if last_packet_ms:
|
|
last_packet_at = datetime.fromtimestamp(last_packet_ms / 1000.0, tz=timezone.utc).isoformat()
|
|
self.stats.update_tzsp_receiver(
|
|
{
|
|
"tzsp_datagrams": int(message.get("tzsp_datagrams") or 0),
|
|
"tzsp_decode_errors": int(message.get("tzsp_decode_errors") or 0),
|
|
"tzsp_unsupported": int(message.get("tzsp_unsupported") or 0),
|
|
"frames_injected": int(message.get("frames_injected") or 0),
|
|
"inject_errors": int(message.get("inject_errors") or 0),
|
|
"tzsp_kernel_udp_drops": int(message.get("kernel_udp_drops") or 0),
|
|
"tzsp_queue_drops": int(message.get("queue_dropped_datagrams") or 0),
|
|
"tzsp_truncated_datagrams": int(message.get("truncated_datagrams") or 0),
|
|
"last_packet_at": last_packet_at,
|
|
}
|
|
)
|
|
|
|
sample = {
|
|
key: int(message.get(key) or 0)
|
|
for key in (
|
|
"ts_ms",
|
|
"interval_ms",
|
|
"bytes_total",
|
|
"bytes_in",
|
|
"bytes_out",
|
|
"bytes_internal",
|
|
"bytes_external",
|
|
"packets_total",
|
|
"packets_in",
|
|
"packets_out",
|
|
"packets_internal",
|
|
"packets_external",
|
|
)
|
|
}
|
|
self._samples += 1
|
|
if sink is not None and sample["interval_ms"] > 0:
|
|
try:
|
|
sink(sample)
|
|
except Exception:
|
|
# Telemetry persistence is best-effort and is deliberately never
|
|
# allowed to affect the independent Rust packet data-plane.
|
|
self._telemetry_errors += 1
|