poc4 wit rust
This commit is contained in:
Mateusz Gruszczyński
@@ -0,0 +1,287 @@
|
||||
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
|
||||
Reference in New Issue
Block a user