#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# pin_gr_stream.py - pin the CUDA client's GR push-buffer pages from a vsock
# stream produced by guest_pin_stream.py (guest truth via /proc/<pid>/pagemap).
#
# The guest scanner runs in the guest and streams only deltas, so the host
# learns a page the moment it is mapped - no guest-agent round trips on the
# critical path (the agent is only used once to install/start the scanner).
#
# Usage: tools/guest-diag/pin_gr_stream.py
# Env: PINWIN_LO/PINWIN_HI/PINWIN_PA_MAX as in pin_gr_window.py,
#      PINWIN_PORT vsock port (default 45678)
import os
import re
import socket
import subprocess
import sys
import threading
import time

HERE = os.path.dirname(os.path.abspath(__file__))
VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))
PARAM = "/sys/module/vgsp/parameters"
GUEST_SCRIPT = os.path.join(HERE, "guest_pin_stream.py")

PORT = int(os.environ.get("PINWIN_PORT", "45678"))
# 0x200000000 covers the GR GPFIFO rings/UVM channel mappings below the
# old 0x200400000 start; without them the host learns bogus vidmem pairs
# for the rings and never decodes GR pushes.
LO = os.environ.get("PINWIN_LO", "0x200000000")
HI = os.environ.get("PINWIN_HI", "0x206000000")
PAMAX = os.environ.get("PINWIN_PA_MAX", "0x180000000")
SEMSTREAM = os.environ.get("SEM_STREAM_FILE", "/tmp/opencode/sem_stream.txt")
# Liveness file checked by pin_supervisor.sh (mtime must stay < 75 s; the
# supervisor grants a 120 s startup grace on top).
HEARTBEAT_FILE = os.environ.get("PIN_GR_HEARTBEAT", "/tmp/opencode/pin_gr_heartbeat")
# Inter-tool re-assert request (written by semfix_uvmt.py's stall watchdog):
# when its mtime moves, re-assert every pair this tool has ever pinned.
REASSERT_REQ = os.environ.get("PIN_REASSERT_FILE",
                              "/tmp/opencode/pin_reassert_request")
# Stall watchdog: a known client with no pin/data activity for this long gets
# an automatic recovery (full re-assert + demand of the fixed counter/pool
# VAs). Pre-capture (no pin yet for this client) uses the tighter limit.
STALL_PRE_S = float(os.environ.get("PIN_STALL_PRE", "60"))
STALL_POST_S = float(os.environ.get("PIN_STALL_POST", "120"))
STALL_RETRY_S = float(os.environ.get("PIN_STALL_RETRY", "30"))
STALL_ESCALATE_S = float(os.environ.get("PIN_STALL_ESCALATE", "180"))


def log(msg):
    print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)


def _heartbeat():
    try:
        with open(HEARTBEAT_FILE, "w") as f:
            f.write("%.3f\n" % time.time())
    except OSError:
        pass


def _hb_loop():
    """Refresh the heartbeat on a fixed cadence, independent of data-path
    activity: recv can be legally idle for minutes between client phases and
    the supervisor used to read that silence as a hung tool (r198 restarts at
    34 s staleness while injection bursts were running)."""
    while True:
        _heartbeat()
        time.sleep(5)


def _mirror(line):
    """Append guest control lines (#pid/#sem) to the shared host file so the
    single-QMP watcher can consume them without guest-agent round trips."""
    try:
        with open(SEMSTREAM, "a") as f:
            f.write(line + "\n")
    except OSError:
        pass


def vexec(cmd, timeout=40):
    try:
        r = subprocess.run([VMEXEC, "exec", "--timeout", str(timeout - 5), cmd],
                           capture_output=True, text=True, timeout=timeout)
        return r.stdout
    except Exception:
        return ""


def _scanner_alive():
    """True when the in-guest scanner process still exists.

    install_and_start() pkill's and restarts the scanner, which resets every
    in-guest latch; probing first avoids killing a healthy scanner that is
    merely silent because no client exists yet (r198 pre-client churn)."""
    out = vexec("sudo sh -c 'kill -0 $(cat /tmp/guest_pin_stream.pid) "
                "2>/dev/null && echo ALIVE'")
    return "ALIVE" in (out or "")


def _demand_send(vas):
    """Append VA hints to the guest scanner's demand file (best-effort).

    Same channel pin_gr_stream._demand_loop uses for host-dmesg stalls; the
    stall watchdog uses it to ask the scanner to re-resolve the fixed
    pre-capture VAs (counter page / UVM pool) by content/pagemap."""
    lines = "".join("0x%x\n" % va for va in vas)
    try:
        vexec("printf '%%s' '%s' >> /tmp/guest_want_vas.txt" % lines, 30)
    except Exception:
        pass


def _start_cmd():
    with open(GUEST_SCRIPT) as f:
        b64 = subprocess.run(["base64", "-w0"], stdin=f,
                             capture_output=True, text=True,
                             timeout=10).stdout
    prefixes = os.environ.get("CLIENT_PREFIXES", "cdt,va,victory")
    vma_every = os.environ.get("STREAM_VMA_EVERY", "1")
    # Kill EVERY existing scanner, not just the pid-file one: a retry after
    # a partially-failed first attempt (pid file not written yet) used to
    # leave the first scanner alive, and multiple scanners each re-assert
    # the whole CE pair set - r148 ran three and drove the guest to load 19
    # (clients wedged in userspace spins). The pattern is anchored
    # (^python3) and assembled at runtime via printf: an unanchored pattern
    # or a literal one would also match this command's own shell (its
    # cmdline contains the script path), and pkill would kill the shell
    # before it could start the new scanner - every start attempt then
    # failed silently (r149).
    return ("P=$(printf '^python3 /tmp/guest_pin_%%s.py' stream); "
            "sudo pkill -f \"$P\" 2>/dev/null; sleep 1; "
            "echo '%s' | base64 -d > /tmp/guest_pin_stream.py; "
            "sudo rm -f /tmp/guest_pin_stream.pid; "
            "sudo setsid nohup env CLIENT_PREFIXES='%s' "
            "STREAM_VMA_EVERY='%s' "
            "python3 /tmp/guest_pin_stream.py %d %d %s %s %s "
            ">/tmp/guest_pin_stream.log 2>&1 & sleep 2; "
            "sudo cat /tmp/guest_pin_stream.pid 2>/dev/null"
            % (b64, prefixes, vma_every, 2, PORT, LO, HI, PAMAX))


def install_and_start():
    # The guest-agent exec can fail transiently (agent busy just after boot,
    # or a stale qga client); retry a few times instead of giving up, because
    # a boot without the scanner streams no pins and the client then hangs.
    for attempt in range(5):
        _heartbeat()
        out = vexec(_start_cmd(), 60)
        if any(ch.isdigit() for ch in out):
            return True
        log("WARN: scanner start attempt %d failed; retrying" % (attempt + 1))
        time.sleep(5)
    return False


def inject(va, pa, ln):
    cmd = "printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add" % (va, pa, ln, PARAM)
    for _ in range(60):
        try:
            r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                               capture_output=True, text=True, timeout=10)
        except subprocess.TimeoutExpired:
            time.sleep(0.05)
            continue
        if r.returncode == 0:
            return True
        time.sleep(0.05)
    return False


def scrub_register(fin_pa, hsem_pa):
    """Register the PMA scrubber buffer words (guest-resolved PAs) with the
    service so it can author the scrub progress (third-client scrub freeze)."""
    cmd = ("printf '0x%x,0x%x' > %s/sema_scrubbuf_add"
           % (fin_pa, hsem_pa, PARAM))
    for _ in range(60):
        try:
            r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                               capture_output=True, text=True, timeout=10)
        except subprocess.TimeoutExpired:
            time.sleep(0.05)
            continue
        if r.returncode == 0:
            log("scrub registered fin=0x%x hsem=0x%x" % (fin_pa, hsem_pa))
            return True
        time.sleep(0.05)
    log("scrub register FAILED fin=0x%x hsem=0x%x" % (fin_pa, hsem_pa))
    return False


def _demand_loop():
    """Push the host's unreadable CE push VAs to the guest scanner.

    The service logs each distinct stalled VA once per boot ('stalled VA
    #N: 0x...'). Each sighting is tracked and re-sent every 20 s for five
    minutes: the guest resolves the VA by content against the
    nv_vm_map_pages areas (event-independent, no object chain), so a page
    that was not yet written when first asked is picked up on a retry.
    This closes the stall class where the one-shot #sem kprobe event was
    lost and the chain resolver never learns the client.
    """
    track = {}			# va -> [first_seen, last_sent]
    while True:
        time.sleep(3)
        try:
            r = subprocess.run(["sudo", "-n", "dmesg"],
                               capture_output=True, text=True, timeout=20)
        except Exception:
            continue
        now = time.time()
        for m in re.finditer(r"(?:stalled VA #\d+|pending release starved va)=(0x[0-9a-f]+)",
                             (r.stdout or "")[-200000:]):
            va = int(m.group(1), 16)
            track.setdefault(va, [now, 0.0])
        for va in list(track):
            first, last = track[va]
            if now - first > 300.0:
                del track[va]
        todo = [va for va, (f, l) in track.items() if now - l > 20.0]
        if not todo:
            continue
        lines = "".join("0x%x\n" % va for va in todo)
        try:
            vexec("printf '%%s' '%s' >> /tmp/guest_want_vas.txt" % lines, 30)
        except Exception:
            continue
        for va in todo:
            track[va][1] = now
        log("demand: sent %d stalled VA(s) to guest" % len(todo))


_DEMAND_STARTED = False
_HB_STARTED = False
_REASSERT_REQ_SEEN = 0.0


def main():
    global _DEMAND_STARTED, _HB_STARTED
    _heartbeat()
    if not _HB_STARTED:
        _HB_STARTED = True
        try:
            threading.Thread(target=_hb_loop, daemon=True).start()
        except Exception as e:
            log("WARN: heartbeat thread failed to start: %s" % e)
    if not install_and_start():
        log("WARN: could not start guest scanner (will retry via vsock wait)")
    if not _DEMAND_STARTED:
        _DEMAND_STARTED = True
        try:
            threading.Thread(target=_demand_loop, daemon=True).start()
        except Exception as e:
            log("WARN: demand thread failed to start: %s" % e)
    srv = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((socket.VMADDR_CID_ANY, PORT))
    srv.settimeout(10)
    srv.listen(4)
    log("listening on vsock port %d (window %s..%s)" % (PORT, LO, HI))

    seen = {}
    pin_total = 0
    last_conn = time.time()
    last_reinstall = 0.0
    # Stall-watchdog state for the connected client. client_pid is the
    # scanner-reported pid (None until a client appears); last_activity marks
    # the last pin/data/control event; pins_since_client counts pages pinned
    # for the current client (0 = pre-capture).
    client_pid = None
    client_start = None
    last_activity = time.time()
    pins_since_client = 0
    stall_since = None
    recover_next = 0.0
    escalated = False
    no_conn_warn_next = 0.0

    def _activity():
        nonlocal last_activity, stall_since, escalated
        last_activity = time.time()
        stall_since = None
        escalated = False

    while True:
        try:
            conn, _ = srv.accept()
        except socket.timeout:
            now = time.time()
            # No guest scanner connection for 30 s: it died or never started
            # (a boot without it streams no pins and every client then hangs).
            # Probe the scanner first instead of reinstalling blindly:
            # install_and_start() pkill's a live scanner and resets all
            # in-guest latches, and pre-client silence is normal (r198 churn).
            if now - last_conn > 30.0 and now - last_reinstall >= 60.0:
                last_reinstall = now
                if _scanner_alive():
                    log("no guest connection for %.0fs but scanner alive; "
                        "not reinstalling" % (now - last_conn))
                    if now - last_conn > 180.0 and now >= no_conn_warn_next:
                        no_conn_warn_next = now + 60.0
                        log("STALL WATCHDOG: no guest scanner connection for "
                            "%.0fs (scanner process alive) - pins/demand may "
                            "be stranded" % (now - last_conn))
                    continue
                log("reinstall attempt: no guest connection for %.0fs"
                    % (now - last_conn))
                try:
                    install_and_start()
                except Exception as e:
                    log("reinstall failed: %s" % e)
            continue
        except OSError as e:
            log("accept failed: %s" % e)
            time.sleep(1)
            continue
        last_conn = time.time()
        _heartbeat()
        conn.settimeout(5)
        log("guest connected")
        buf = b""
        runs = []
        try:
            while True:
                try:
                    chunk = conn.recv(1 << 20)
                except socket.timeout:
                    if runs:
                        novel = _novel_pages(runs, seen)
                        pages = _flush(runs, seen)
                        pin_total += pages
                        pins_since_client += pages
                        if novel:
                            _activity()
                        runs = []
                    _maybe_reassert_request(seen)
                    now = time.time()
                    # Stall watchdog: a known client with no pin/data
                    # activity gets a bounded, idempotent recovery (re-assert
                    # everything + demand the fixed counter/pool VAs), then a
                    # loud escalation if the stall persists. This is the
                    # "first cudaMalloc / first GR capture must not stall
                    # silently" guard. recover_next rate-limits attempts (the
                    # loop wakes every 5 s on the recv timeout).
                    if client_pid is not None:
                        limit = STALL_PRE_S if pins_since_client == 0 \
                            else STALL_POST_S
                        idle = now - last_activity
                        if idle >= limit and now >= recover_next:
                            recover_next = now + STALL_RETRY_S
                            if stall_since is None:
                                stall_since = now
                            log("STALL WATCHDOG: client pid %d - no pin "
                                "activity for %.0fs (pre-capture=%s) - "
                                "re-asserting all known pairs + demanding "
                                "counter/pool VAs"
                                % (client_pid, idle, pins_since_client == 0))
                            if seen:
                                _flush([], seen, full=True)
                            _demand_send([0x20400f000, 0x121000000])
                        if not escalated and stall_since is not None \
                                and now - stall_since >= STALL_ESCALATE_S:
                            escalated = True
                            log("STALL WATCHDOG: recovery FAILED for pid %d "
                                "after %.0fs - stall persists, needs manual "
                                "attention" % (client_pid, now - stall_since))
                    continue
                except OSError:
                    break
                if not chunk:
                    break
                _heartbeat()
                buf += chunk
                while b"\n" in buf:
                    line, buf = buf.split(b"\n", 1)
                    text = line.decode(errors="replace").strip()
                    if not text:
                        continue
                    if text.startswith("#pid"):
                        m = re.match(r"#pid\s+(\d+)(?:\s+(\d+))?", text)
                        seen = {}
                        client_pid = int(m.group(1)) if m else None
                        client_start = (int(m.group(2)) if m and m.group(2)
                                        else None)
                        pins_since_client = 0
                        _activity()
                        log("guest reports cdt pid %s%s"
                            % (client_pid if client_pid is not None else "?",
                               (" start=%s" % client_start
                                if client_start is not None else "")))
                        _mirror(text)
                        continue
                    if text.startswith("#nopid"):
                        # In-band client death: the guest scanner tells us the
                        # client is gone (with its identity). Drop the per-pair
                        # dedupe map so the next client's pairs are injected
                        # unconditionally; semfix_uvmt consumes the mirrored
                        # line and purges its per-client state too.
                        client_pid = None
                        client_start = None
                        seen = {}
                        pins_since_client = 0
                        _activity()
                        log("guest reports client gone (%s)" % text)
                        _mirror(text)
                        continue
                    if text.startswith("#sem"):
                        _mirror(text)
                        _activity()
                        continue
                    if text.startswith("#pb"):
                        _mirror(text)
                        _activity()
                        continue
                    if text.startswith("#put"):
                        log(text)
                        _mirror(text)
                        _activity()
                        continue
                    if text.startswith("#scrub"):
                        m = re.search(r"0x([0-9a-fA-F]+),0x([0-9a-fA-F]+)",
                                      text)
                        _mirror(text)
                        if m:
                            scrub_register(int(m.group(1), 16),
                                           int(m.group(2), 16))
                        continue
                    m = re.match(r"^(0x[0-9a-f]+),(0x[0-9a-f]+)(?:,(\d+))?$",
                                 text)
                    if m:
                        runs.append((int(m.group(1), 16),
                                     int(m.group(2), 16),
                                     int(m.group(3)) if m.group(3) else 1))
                if runs:
                    novel = _novel_pages(runs, seen)
                    pages = _flush(runs, seen)
                    pin_total += pages
                    pins_since_client += pages
                    if novel:
                        _activity()
                    runs = []
                # Poll the semfix re-assert request on every recv cycle, not
                # only on recv-timeouts: sustained scanner traffic starves the
                # timeout branch, so the request fired post-mortem in r203
                # (119 s late, `seen` already cleared by #nopid).
                _maybe_reassert_request(seen)
        except Exception as e:
            # A transient error on one connection must not kill the process:
            # log it, drop the connection and go back to accept().
            log("connection handler error: %s (continuing)" % e)
        log("guest disconnected (pins this session: %d)" % pin_total)
        try:
            conn.close()
        except Exception:
            pass


def _maybe_reassert_request(seen):
    """Honour a full-re-assert request written by semfix_uvmt's watchdog."""
    global _REASSERT_REQ_SEEN
    try:
        m = os.stat(REASSERT_REQ).st_mtime
    except OSError:
        return
    if m > _REASSERT_REQ_SEEN:
        _REASSERT_REQ_SEEN = m
        if seen:
            log("semfix requested a full pin re-assert - re-asserting %d "
                "pair(s)" % len(seen))
            _flush([], seen, full=True)


def _coalesce_seen(seen):
    """Coalesce the seen {va: pa} per-page map into (va, pa, npages) runs."""
    items = sorted(seen.items())
    runs = []
    i = 0
    while i < len(items):
        va, pa = items[i]
        n = 1
        while (i + n < len(items)
               and items[i + n][0] == va + n * 0x1000
               and items[i + n][1] == pa + n * 0x1000):
            n += 1
        runs.append((va, pa, n))
        i += n
    return runs


def _novel_pages(runs, seen):
    """Count pages in runs whose translation is new or changed.

    Progress accounting for the stall watchdog: the guest re-streams the CE
    window and UVM pool every couple of seconds even when nothing moves, so a
    pure repeat of an already-known pair must not count as pin activity
    (r203: those repeats masked the watchdog through every hang)."""
    return sum(1 for va, pa, n in runs for i in range(n)
               if seen.get(va + i * 0x1000) != pa + i * 0x1000)


def _flush(runs, seen, full=False):
    """Inject run lines (one batched sudo pass), return number of pages pinned.

    The old implementation spawned one sudo process per run; with ~400 runs
    that cost ~4 s, which landed the GR pins after the client's context-init
    deadline. The batch writes all runs in a single shell, retrying only the
    ones the kernel queue rejected (ENOSPC) individually.

    full=True re-asserts every pair ever pinned (coalesced), not just this
    batch's runs: the kernel session flush wipes the pin table when a client's
    channels close, and the stall watchdog / semfix re-assert request use this
    to heal it without touching the guest. Writes are idempotent.
    """
    todo = {}
    for va, pa, n in runs:
        # The CE window (0x120000000..0x121000000) and the UVM semaphore
        # pool page range (0x121000000..0x122000000) are re-asserted by the
        # guest every couple of seconds on purpose: the host's session flush
        # wipes the whole pin table when the previous client's last channel
        # closes, and a pin batch that raced in for the new client must be
        # allowed to land again even when the pair is unchanged. Skipping
        # repeats there stranded the second client's CE pushes (the r144/
        # r145 stall: sema-probe "NO ENTRY" for the live push pages) and the
        # pool/tracker page (pending release starved, every client's first
        # seconds).
        ce_win = 0x120000000 <= va < 0x122000000
        if not full and not ce_win and seen.get(va) == pa and n == 1:
            continue
        todo[(va, pa, n)] = True
    if full:
        for va, pa, n in _coalesce_seen(seen):
            todo[(va, pa, n)] = True
    todo = list(todo)
    if not todo:
        return 0

    def mark(va, pa, n):
        for i in range(n):
            seen[va + i * 0x1000] = pa + i * 0x1000

    script = []
    for va, pa, n in todo:
        script.append(
            "printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add 2>/dev/null || "
            "echo 'FAIL 0x%x 0x%x 0x%x'"
            % (va, pa, n * 0x1000, PARAM, va, pa, n * 0x1000))
    failed = set()
    CHUNK = 200
    for start in range(0, len(script), CHUNK):
        batch = script[start:start + CHUNK]
        _heartbeat()
        try:
            r = subprocess.run(["sudo", "-n", "sh", "-c", "\n".join(batch)],
                               capture_output=True, text=True, timeout=60)
            for line in r.stdout.splitlines():
                m = re.match(r"FAIL (0x[0-9a-f]+) (0x[0-9a-f]+) (0x[0-9a-f]+)",
                             line.strip())
                if m:
                    failed.add((int(m.group(1), 16), int(m.group(2), 16),
                                int(m.group(3), 16) // 0x1000))
        except Exception as e:
            log("batch inject failed: %s" % e)
            for va, pa, n in todo[start:start + CHUNK]:
                failed.add((va, pa, n))

    pages = 0
    for va, pa, n in todo:
        if (va, pa, n) in failed:
            continue
        mark(va, pa, n)
        pages += n
    for va, pa, n in failed:
        for i in range(n):
            if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
                mark(va + i * 0x1000, pa + i * 0x1000, 1)
                pages += 1
    if pages:
        log("pinned %d page(s) in %d run(s)" % (pages, len(todo)))
        for va, pa, n in todo[:3]:
            log("  va=0x%x -> pa=0x%x (%d)" % (va, pa, n))
    return pages


if __name__ == "__main__":
    # main() loops forever; only an unexpected exception exits it. Restart
    # instead of dying silently - a dead listener means no pins and every
    # later CUDA client hangs at first cudaMalloc.
    while True:
        try:
            main()
        except Exception as e:
            log("FATAL: main() exited: %r - restarting in 2s" % (e,))
            time.sleep(2)