"""sshj.scan — one-shot + periodic bandwidth scan (B14). B2 step 3e: moved verbatim from sshj_cli.py (single-file era) — no behavior change. Owns the SPEED_* constants (SpeedScan is their only consumer). Checker arrives duck-typed via the constructor, so the only package edge is tailnet.ssh_host. """ import os import subprocess import threading import time from concurrent.futures import ThreadPoolExecutor from .tailnet import ssh_host # Speed: SSHJ_SPEED_SCAN=0 disables; size/interval from SSHJ_SPEED_BYTES / SPEED_SIZE = 2_000_000 # bytes per transfer leg (~16 Mbit, quick on LAN) SPEED_TRANSFER_TIMEOUT = 2.5 # seconds per leg (down, then up) SPEED_INTERVAL = 25 # seconds between speed-scan passes (background) SPEED_FIRST_DELAY = 10 # let the reachability probes settle first class SpeedScan(threading.Thread): """One-shot bandwidth per UP host (B14): after reachability is known, measure a small transfer each way — DOWN = remote streams N bytes (`head -c N /dev/zero`), UP = local streams N bytes into the remote (`cat >/dev/null`). Bounded by SPEED_TRANSFER_TIMEOUT per leg so a slow network can never stall the UI; results are per-leg None on failure (the band simply shows what was measured). Throttled: SPEED_INTERVAL between passes, first pass after SPEED_FIRST_DELAY. SSHJ_SPEED_SCAN=0 disables; the dry-run hook (SSHJ_SSH_CMD) also disables it — an `echo` fake ssh would only measure the echo pipe.""" def __init__(self, checker): super().__init__(daemon=True) self.checker = checker self.done = False self.enabled = (os.environ.get("SSHJ_SPEED_SCAN", "1") != "0" and not os.environ.get("SSHJ_SSH_CMD")) self.size = self._clamp(os.environ.get("SSHJ_SPEED_BYTES", ""), SPEED_SIZE, 64_000, 64_000_000) self.interval = self._clamp(os.environ.get("SSHJ_SPEED_INTERVAL", ""), SPEED_INTERVAL, 5, 600) self.first_delay = SPEED_FIRST_DELAY self.speed = {} # alias -> {"up": bps|None, "down": bps|None} self.lock = threading.Lock() @staticmethod def _clamp(env, dflt, lo, hi): try: return max(lo, min(hi, int(env))) except (ValueError, TypeError): return dflt @staticmethod def speed_probe(h, to=SPEED_TRANSFER_TIMEOUT, stop_event=None): """Transfer `SPEED_SIZE` bytes each way over one ssh connection. Returns {"up": bps|None, "down": bps|None}; None = leg failed or the remote lacks the required coreutils (degrade to latency-only). `stop_event`: on quit, an in-flight leg is killed so the pool's worker finishes instantly (concurrent.futures' atexit otherwise waits for it).""" tgt = ssh_host(h) port = ["-p", h["port"]] if h.get("port") not in ("", "22") else [] base = ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no", "-o", "LogLevel=ERROR", "-o", "ServerAliveInterval=1"] + port n = SPEED_SIZE out = {} # DOWN (local receive): remote streams N zero bytes. p = None try: t0 = time.time() p = subprocess.Popen(base + ["--", tgt, f"head -c {n} /dev/zero"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) end = time.time() + to while p.poll() is None: if (stop_event is not None and stop_event.is_set() or time.time() >= end): break time.sleep(0.2) if p.poll() is None: # stopped or timed out mid-transfer p.kill() p.wait() if p.returncode == 0: dt = time.time() - t0 if dt > 0: out["down"] = n / dt except Exception: pass finally: if p is not None and p.stdout is not None: try: p.stdout.close() except Exception: pass # UP (remote receive): local streams N zero bytes. p = None try: t0 = time.time() p = subprocess.Popen(base + ["--", tgt, "cat >/dev/null"], stdin=subprocess.PIPE, stderr=subprocess.DEVNULL) end = time.time() + to # Feed from a thread: if the ssh child dies with the pipe still # full, a main-thread p.stdin.write() would block forever # (nothing drains the pipe, and the blocked write holds the GIL # so the poll loop below could never kill it). fed = [0] def _feed(pipe): while (not stop_event or not stop_event.is_set() and time.time() <= end + to): try: k = pipe.write(b"\0" * 65536) except Exception: # pipe broke return if not k: return fed[0] += k ft = threading.Thread(target=_feed, args=(p.stdin,), daemon=True) ft.start() while p.poll() is None: if (stop_event is not None and stop_event.is_set() or time.time() >= end or not ft.is_alive()): break time.sleep(0.2) if p.poll() is None: # stopped, timed out, or feed broke p.kill() p.wait() if p.stdin is not None: try: p.stdin.close() except Exception: pass ft.join(timeout=1.0) if p.returncode == 0: dt = time.time() - t0 if dt > 0 and fed[0] > 0: out["up"] = fed[0] / dt except Exception: pass finally: if p is not None: for s in (p.stdin, p.stdout, p.stderr): if s is not None: try: s.close() except Exception: pass return out def speed_scan_targets(self, checker=None): """Scan every UP host once, concurrently (bounded). Pure-ish: uses checker.status/by_alias/order; returns the number scanned. Calls the class-level speed_probe (tests replace that symbol).""" c = checker or self.checker with c.lock: todo = [a for a in c.order if c.status.get(a) == "up" and a in c.by_alias and a not in self.speed] hosts = [c.by_alias[a] for a in todo] for h in hosts: self._scan_one(h) return len(hosts) def _scan_one(self, h): alias = h["alias"] pool = getattr(self.checker, "pool", None) stop_ev = getattr(self.checker, "stop_ev", None) if pool is not None and not isinstance(pool, type(None)): pool.submit(self._store, alias, SpeedScan.speed_probe(h, stop_event=stop_ev)) else: # no pool (test stand-in) — run inline self._store(alias, SpeedScan.speed_probe(h, stop_event=stop_ev)) def _store(self, alias, res): with self.lock: self.speed[alias] = res def get(self, alias): with self.lock: return self.speed.get(alias) def stop(self): self.done = True def run(self): if not self.enabled: return time.sleep(self.first_delay) while not self.done: self.speed_scan_targets() time.sleep(self.interval)