path: root/tools/guest-diag/pin_gr_stream.py

#!/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")


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


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 _start_cmd():
    with open(GUEST_SCRIPT) as f:
        b64 = subprocess.run(["base64", "-w0"], stdin=f,
                             capture_output=True, text=True).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):
        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):
        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                           capture_output=True, text=True)
        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):
        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                           capture_output=True, text=True)
        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+: (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))


def main():
    if not install_and_start():
        log("WARN: could not start guest scanner (will retry via vsock wait)")
    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.listen(4)
    log("listening on vsock port %d (window %s..%s)" % (PORT, LO, HI))

    seen = {}
    pin_total = 0
    buf = b""
    while True:
        try:
            conn, _ = srv.accept()
        except OSError as e:
            log("accept failed: %s" % e)
            time.sleep(1)
            continue
        conn.settimeout(5)
        log("guest connected")
        buf = b""
        runs = []
        while True:
            try:
                chunk = conn.recv(1 << 20)
            except socket.timeout:
                if runs:
                    _flush(runs, seen)
                    runs = []
                continue
            except OSError:
                break
            if not chunk:
                break
            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.search(r"(\d+)", text)
                    seen = {}
                    log("guest reports cdt pid %s" % (m.group(1) if m else "?"))
                    _mirror(text)
                    continue
                if text.startswith("#sem"):
                    _mirror(text)
                    continue
                if text.startswith("#pb"):
                    _mirror(text)
                    continue
                if text.startswith("#put"):
                    log(text)
                    _mirror(text)
                    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:
                pin_total += _flush(runs, seen)
                runs = []
        log("guest disconnected (pins this session: %d)" % pin_total)
        try:
            conn.close()
        except Exception:
            pass


def _flush(runs, seen):
    """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.
    """
    todo = []
    for va, pa, n in runs:
        # The CE window (0x120000000..0x121000000) is 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).
        ce_win = 0x120000000 <= va < 0x121000000
        if not ce_win and seen.get(va) == pa and n == 1:
            continue
        todo.append((va, pa, n))
    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]
        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()