guest-diag: vsock pin stream (guest pagemap -> host sema_map_add)

The guest-agent exec path is too slow under driver-init load (a 15k-page
window scan took 17-35 s to complete), so GR window pins landed after the
channels' entries had been consumed and after the client's cuCtxCreate
timeout. The guest now scans /proc/<pid>/pagemap locally and streams only
delta runs over vsock (no agent on the hot path); the host listener pins
them as they arrive. The scanner writes a pidfile so restarts do not kill
the new instance.

3 files changed, 287 insertions, 0 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
new file mode 100644
index 0000000..169d048
--- /dev/null
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+# guest_pin_stream.py - stream the CUDA client's push-buffer VA->PA pairs to
+# the host over vsock, continuously and without the guest agent on the hot
+# path. Runs in the guest as root (pagemap reads).
+#
+# Protocol (text lines to the host):
+#   #pid <n>            the client pid just changed (host resets its state)
+#   0x<va>,0x<pa>,<n>   n contiguous pages (va+i*0x1000 -> pa+i*0x1000)
+#
+# The host side is tools/guest-diag/pin_gr_stream.py.
+#
+# Usage: sudo python3 guest_pin_stream.py [host-cid] [port] [lo-hex] [hi-hex] [pa-max-hex]
+import os
+import socket
+import sys
+import time
+
+CID = int(sys.argv[1]) if len(sys.argv) > 1 else 2
+PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 45678
+LO = int(sys.argv[3], 16) if len(sys.argv) > 3 else 0x200400000
+HI = int(sys.argv[4], 16) if len(sys.argv) > 4 else 0x206000000
+PAMAX = int(sys.argv[5], 16) if len(sys.argv) > 5 else 0x180000000
+INTERVAL = float(os.environ.get("STREAM_INTERVAL", "0.25"))
+BATCH = 4096
+
+
+def cdt_pid():
+    for d in os.listdir('/proc'):
+        if not d.isdigit():
+            continue
+        try:
+            with open('/proc/%s/comm' % d) as f:
+                if f.read().strip() == 'cdt':
+                    return int(d)
+        except OSError:
+            pass
+    return None
+
+
+def scan(pid):
+    """Present pages of the VA range as {va: pa}."""
+    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
+    out = {}
+    try:
+        va = LO & ~0xFFF
+        while va < HI:
+            n = min(BATCH, (HI - va) // 0x1000)
+            if n <= 0:
+                break
+            os.lseek(fd, (va // 4096) * 8, 0)
+            data = os.read(fd, n * 8)
+            got = len(data) // 8
+            for i in range(got):
+                e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
+                if not ((e >> 63) & 1):
+                    continue
+                pfn = e & ((1 << 55) - 1)
+                pa = pfn << 12
+                if not pa or (PAMAX and pa >= PAMAX):
+                    continue
+                out[va + i * 0x1000] = pa
+            va += got * 0x1000
+    finally:
+        os.close(fd)
+    return out
+
+
+def main():
+    with open('/tmp/guest_pin_stream.pid', 'w') as f:
+        f.write('%d\n' % os.getpid())
+    sock = None
+    seen = {}
+    last_pid = None
+    while True:
+        pid = cdt_pid()
+        if pid is None:
+            last_pid = None
+            seen = {}
+            time.sleep(INTERVAL)
+            continue
+        if pid != last_pid:
+            last_pid = pid
+            seen = {}
+            if sock is not None:
+                try:
+                    sock.sendall(("#pid %d\n" % pid).encode())
+                except OSError:
+                    sock.close()
+                    sock = None
+        try:
+            cur = scan(pid)
+        except OSError:
+            time.sleep(INTERVAL)
+            continue
+        new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
+        seen = cur
+        if not new:
+            time.sleep(INTERVAL)
+            continue
+        new.sort()
+        lines = []
+        va0, pa0, n = new[0][0], new[0][1], 1
+        for va, pa in new[1:]:
+            if va == va0 + n * 0x1000 and pa == pa0 + n * 0x1000:
+                n += 1
+                continue
+            lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
+            va0, pa0, n = va, pa, 1
+        lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
+        payload = ''.join(lines).encode()
+        if sock is None:
+            try:
+                sock = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
+                sock.settimeout(5)
+                sock.connect((CID, PORT))
+                sock.settimeout(None)
+                if last_pid is not None:
+                    sock.sendall(("#pid %d\n" % last_pid).encode())
+            except OSError:
+                if sock:
+                    sock.close()
+                sock = None
+                time.sleep(INTERVAL)
+                continue
+        try:
+            sock.sendall(payload)
+        except OSError:
+            sock.close()
+            sock = None
+        time.sleep(INTERVAL)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/tools/guest-diag/pin_gr_stream.py b/tools/guest-diag/pin_gr_stream.py
new file mode 100644
index 0000000..22c202b
--- /dev/null
+++ b/tools/guest-diag/pin_gr_stream.py
@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+# 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 time
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))
+PARAM = "/sys/module/open_iov_vfio/parameters"
+GUEST_SCRIPT = os.path.join(HERE, "guest_pin_stream.py")
+
+PORT = int(os.environ.get("PINWIN_PORT", "45678"))
+LO = os.environ.get("PINWIN_LO", "0x200400000")
+HI = os.environ.get("PINWIN_HI", "0x206000000")
+PAMAX = os.environ.get("PINWIN_PA_MAX", "0x180000000")
+
+
+def log(msg):
+    print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)
+
+
+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 install_and_start():
+    with open(GUEST_SCRIPT) as f:
+        b64 = subprocess.run(["base64", "-w0"], stdin=f,
+                             capture_output=True, text=True).stdout
+    cmd = ("echo '%s' | base64 -d > /tmp/guest_pin_stream.py; "
+           "[ -f /tmp/guest_pin_stream.pid ] && sudo kill "
+           "$(cat /tmp/guest_pin_stream.pid) 2>/dev/null; "
+           "sudo rm -f /tmp/guest_pin_stream.pid; "
+           "sudo setsid nohup 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, 2, PORT, LO, HI, PAMAX))
+    out = vexec(cmd, 60)
+    return any(ch.isdigit() for ch in out)
+
+
+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 main():
+    if not install_and_start():
+        log("WARN: could not start guest scanner (will retry via vsock wait)")
+    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 "?"))
+                    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, return number of pages pinned."""
+    pages = 0
+    for va, pa, n in runs:
+        if seen.get(va) == pa and n == 1:
+            continue
+        for i in range(n):
+            seen[va + i * 0x1000] = pa + i * 0x1000
+        if not inject(va, pa, n * 0x1000):
+            for i in range(n):
+                if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
+                    pages += 1
+            continue
+        pages += n
+    if pages:
+        log("pinned %d page(s) in %d run(s)" % (pages, len(runs)))
+        for va, pa, n in runs[:3]:
+            log("  va=0x%x -> pa=0x%x (%d)" % (va, pa, n))
+    return pages
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/guest-diag/push_guest_tools.sh b/tools/guest-diag/push_guest_tools.sh
index 34866d2..5e331e5 100755
--- a/tools/guest-diag/push_guest_tools.sh
+++ b/tools/guest-diag/push_guest_tools.sh
@@ -34,6 +34,7 @@ push_b64 "$HERE/sem_slots.py"   /tmp/sem_slots.py
 push_b64 "$HERE/read_va.py"     /tmp/read_va.py
 push_b64 "$HERE/guest_pmap.py"  /tmp/guest_pmap.py
 push_b64 "$HERE/guest_pmap_win.py" /tmp/guest_pmap_win.py
+push_b64 "$HERE/guest_pin_stream.py" /tmp/guest_pin_stream.py
 push_b64 "$HERE/track_waits.sh" /tmp/track_waits.sh
 push_b64 "$HERE/guest_kprobes.sh" /tmp/kprobes.sh