#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# 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 re
import socket
import struct
import subprocess
import sys
import threading
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.1"))
BATCH = 4096
REGION = 0x200000
# The GR GPFIFO rings sit at 0x20020c000 + 0x3000*rank (region offset 0xc000
# ..0x45000) and the per-channel push allocations follow; 512 KB per 2 MB
# region covers both without scanning holes that hold nothing.
SPAN = int(os.environ.get("STREAM_SPAN", "0x80000"), 16)
# CE push-buffer window: one 1 MB allocation per CE channel at
# 0x120000000 + 0x100000*rank; SPAN2 covers the 1 MB boundary inside a
# 2 MB region. Extend HI2 to 0x121400000: the UVM semaphore pool (the
# tracker sems the client's context-init wait polls, pool word at
# 0x121000000 + 4*idx) lives above the push buffers and must reach the host
# through this fast stream - otherwise only the slow QMP chain resolution
# pins it, and the pin lands after the client's RM thread-state deadline.
LO2 = int(os.environ.get("STREAM_LO2", "0x120000000"), 16)
HI2 = int(os.environ.get("STREAM_HI2", "0x121400000"), 16)
SPAN2 = int(os.environ.get("STREAM_SPAN2", "0x180000"), 16)
# UVM semaphore pool window: the tracker pool word lives at
# 0x121000000 + 4*idx, above the CE window head span (LO2+SPAN2 =
# 0x120180000), so scan_window never reaches it. Scan every 4 KB page of
# [POOL_LO, POOL_HI) directly - cheap (16 pages) and host-visible as pins.
POOL_LO = int(os.environ.get("STREAM_POOL_LO", "0x121000000"), 16)
POOL_HI = int(os.environ.get("STREAM_POOL_HI", "0x121010000"), 16)
# UVM system-memory AVMAs: the client's sysmem buffers (the DtoH staging
# buffer among them) are /dev/zero mappings at high GPU VAs (e.g.
# 0x79c21c800000). They are not inside the two fixed windows, and only the
# GMMU page tables can translate them, so stream their present pages as
# pins too. Restricted to the high VA range and to nvidia/zero mappings;
# the PAMAX filter rejects device (BAR) PFNs.
VMA_LO = int(os.environ.get("STREAM_VMA_LO", "0x40000000000"), 16)
VMA_EVERY = int(os.environ.get("STREAM_VMA_EVERY", "1"))
VMA_MAX_PAGES = int(os.environ.get("STREAM_VMA_MAX_PAGES", "16384"))
VMA_PATHS = tuple(p for p in os.environ.get(
    "STREAM_VMA_PATHS", "zero,nvidia").split(",") if p)
TRACE = ("/sys/kernel/tracing/trace"
         if os.path.exists("/sys/kernel/tracing/trace")
         else "/sys/kernel/debug/tracing/trace")
TRACE_PIPE = ("/sys/kernel/tracing/trace_pipe"
              if os.path.exists("/sys/kernel/tracing/trace_pipe")
              else "/sys/kernel/debug/tracing/trace_pipe")
SEM_EVERY = int(os.environ.get("STREAM_SEM_EVERY", "4"))


class Kcore:
    """Kernel-truth reader: /proc/kcore for contents + manual page-table walk.

    The host's QMP-based struct walk intermittently reads garbage while the
    client spins in user mode (PTI hides the kernel half from the sampled
    vCPU), which left the CE push-buffer / semaphore-pool pages unpinned on
    ~half of all boots (the deterministic cudaMalloc stall). The guest can
    read its own kernel memory reliably: /proc/kcore has an ELF segment for
    every kernel mapping, and the direct-map segment lets us walk the page
    tables (init_top_pgt -> ... -> PTE) to turn the vmalloc CPU alias of a
    UVM rm_mem into the guest physical addresses the host needs.
    """

    def __init__(self):
        self.ok = False
        self.fd = None
        self.segs = []
        self.phys_base = 0
        self.page_offset_base = 0
        self.pgd_pa = 0
        self._pa_cache = {}

    def load(self):
        try:
            self.fd = open("/proc/kcore", "rb")
            hdr = self.fd.read(64)
            if hdr[:4] != b"\x7fELF":
                return False
            e_phoff = struct.unpack_from("<Q", hdr, 32)[0]
            e_phentsize, e_phnum = struct.unpack_from("<HH", hdr, 54)
            segs = []
            for i in range(e_phnum):
                self.fd.seek(e_phoff + i * e_phentsize)
                ph = self.fd.read(56)
                p_type, _fl, p_off, p_va, _pa, p_sz = struct.unpack_from(
                    "<IIQQQQ", ph, 0)
                if p_type == 1:
                    segs.append((p_va, p_va + p_sz, p_off))
            self.segs = segs
            syms = {}
            for line in open("/proc/kallsyms"):
                parts = line.split()
                if len(parts) == 3:
                    syms.setdefault(parts[2], int(parts[0], 16))
            need = ("phys_base", "page_offset_base", "init_top_pgt")
            if any(s not in syms or not syms[s] for s in need):
                return False
            self.phys_base = self.read64(syms["phys_base"])
            self.page_offset_base = self.read64(syms["page_offset_base"])
            if not self.phys_base or not self.page_offset_base:
                return False
            self.pgd_pa = (syms["init_top_pgt"] - 0xFFFFFFFF80000000 +
                           self.phys_base)
            self.ok = True
            return True
        except OSError:
            return False

    def read(self, va, n):
        for lo, hi, off in self.segs:
            if lo <= va and va + n <= hi:
                self.fd.seek(off + (va - lo))
                d = self.fd.read(n)
                return d if d and len(d) == n else None
        return None

    def read64(self, va):
        d = self.read(va, 8)
        return struct.unpack("<Q", d)[0] if d else None

    def phys64(self, pa):
        return self.read64(self.page_offset_base + (pa & 0xFFFFFFFFF))

    def va_to_pa(self, va):
        page = va & ~0xFFF
        cached = self._pa_cache.get(page)
        if cached is not None:
            return cached or None
        pa = self.pgd_pa
        for level, shift in enumerate((39, 30, 21, 12)):
            e = self.phys64(pa + ((va >> shift) & 0x1FF) * 8)
            if e is None or not (e & 1):
                self._pa_cache[page] = 0
                return None
            if level == 3:
                pa = e & 0x000FFFFFFFF000
                break
            if e & 0x80:
                if level == 1:
                    pa = (e & 0x000FFFFFC0000000) + (va & ((1 << 30) - 1))
                    break
                if level == 2:
                    pa = (e & 0x000FFFFFFFE00000) + (va & ((1 << 21) - 1))
                    break
                self._pa_cache[page] = 0
                return None
            pa = e & 0x000FFFFFFFF000
        if not _is_ram_pa(pa):
            self._pa_cache[page] = 0
            return None
        self._pa_cache[page] = pa
        return pa


def is_kern_ptr(v):
    return v and (v >> 48) == 0xFFFF


def is_gpu_va(v):
    return v and 0x120000000 <= v < 0x200000000


# Guest-RAM windows registered with the host (debugfs gsp-phys-guest-ram):
# low 0..2 GiB and high 4..6 GiB. A CPU mapping of a channel block resolves
# through the guest page tables to a phys address in one of these windows
# when the backing is guest sysmem, but to a PCI BAR address (e.g.
# 0x70000xxxxx, the mdev BAR1 aperture) when the backing is the emulated
# vidmem. The host pin stream must only carry guest-RAM PAs: a BAR phys fed
# to sema_map_add as a "sysmem" pin poisoned the ring/USERD lookups (r40:
# rank2's ring resolved to pa=0x70000c0000 ap=SYS and read 66/66 garbage).
RAM_WINDOWS = ((0x0, 0x80000000), (0x100000000, 0x180000000))


def _is_ram_pa(pa):
    if not pa:
        return False
    for lo, hi in RAM_WINDOWS:
        if lo <= pa < hi:
            return True
    return False


def _is_gpu_va_wide(v):
    """GPU VAs above 0x200000000 are the GR share window; accept them too."""
    return v and (is_gpu_va(v) or 0x200000000 <= v < 0x208000000)


def read_u32(kc, va):
    d = kc.read(va, 4) if is_kern_ptr(va) else None
    return struct.unpack("<I", d)[0] if d else 0xFFFFFFFF


DBG_BUDGET = int(os.environ.get("STREAM_DBG_BUDGET", "600"))
_dbg_count = [0]


def dbg(msg):
    """Bounded resolver diagnostics to the guest log.

    The single-QGA guest agent is often wedged exactly when the second
    client stalls, so the scanner's own log is the reliable diagnostic
    channel; it is captured from the guest after the run.
    """
    if _dbg_count[0] >= DBG_BUDGET:
        return
    _dbg_count[0] += 1
    try:
        with open('/tmp/guest_pin_stream.log', 'a') as f:
            f.write('%s DBG %s\n' % (time.strftime('%T'), msg))
    except OSError:
        pass


def _direct_map_pa(kc, va):
    """Direct-map aliases translate linearly: pa = va - page_offset_base."""
    if kc.page_offset_base <= va < kc.page_offset_base + (1 << 40):
        return va - kc.page_offset_base
    return None


def _vmalloc_pages_for(kc, va):
    """Page count of the vmap/vmalloc area containing va (0 if unknown)."""
    now = time.time()
    if not hasattr(_vmalloc_pages_for, "areas") or \
            now - _vmalloc_pages_for.ts > 5.0:
        areas = []
        try:
            for line in open("/proc/vmallocinfo"):
                p = line.split()
                if len(p) < 3:
                    continue
                try:
                    lo, hi = p[0].split("-")
                    areas.append((int(lo, 16), int(hi, 16)))
                except ValueError:
                    continue
        except OSError:
            pass
        _vmalloc_pages_for.areas = areas
        _vmalloc_pages_for.ts = now
    base = va & ~0xFFF
    for lo, hi in _vmalloc_pages_for.areas:
        if lo <= base < hi:
            return (hi - base) >> 12
    return 0


def resolve_mem_pins(kc, mem, out, default_pages=16):
    """uvm_rm_mem kptr -> GPU-VA/GPA runs (direct map or page-table walk)."""
    if not is_kern_ptr(mem):
        return 0
    vas = [kc.read64(mem + 0x10 + i * 8) for i in range(16)]
    cpu = next((v for v in vas if is_kern_ptr(v)), None)
    gpus = [v for v in vas if is_gpu_va(v)]
    if not cpu or not gpus:
        return 0
    size = kc.read64(mem + 0x60) or 0
    dm = _direct_map_pa(kc, cpu)
    emitted = 0
    if dm is not None:
        npages = max(1, min(size >> 12, 4098)) if size else default_pages
        for gv in gpus:
            out.append("0x%x,0x%x,%d\n" % (gv, dm, npages))
            emitted += 1
        return emitted
    npages = (max(1, min(size >> 12, 4098)) if size
              else _vmalloc_pages_for(kc, cpu))
    if not npages:
        npages = 1
    runs = []
    cur = None
    for k in range(npages):
        pa = kc.va_to_pa(cpu + k * 0x1000)
        if pa is None:
            if cur:
                runs.append(cur)
                cur = None
            continue
        if cur and pa == cur[1] + cur[2] * 0x1000:
            cur[2] += 1
        else:
            if cur:
                runs.append(cur)
            cur = [k, pa, 1]
    if cur:
        runs.append(cur)
    for gv in gpus:
        for k, pa, n in runs:
            out.append("0x%x,0x%x,%d\n" % (gv + k * 0x1000, pa, n))
            emitted += 1
    return emitted


def _emit_ring_pins(kc, gpu_va, entries_cpu, entries, out):
    """Pin the GPFIFO ring pages (and the pages after it, where the GPPUT
    buffer usually sits) from the ring's CPU mapping."""
    emitted = 0
    if not is_kern_ptr(entries_cpu) or not (32 <= entries <= 0x100000):
        return 0
    npages = min(((entries * 8) + 0xFFF) // 0x1000 + 2, 8)
    runs = []
    cur = None
    for k in range(npages):
        pa = kc.va_to_pa(entries_cpu + k * 0x1000)
        if pa is None:
            if cur:
                runs.append(cur)
                cur = None
            continue
        if cur and pa == cur[1] + cur[2] * 0x1000:
            cur[2] += 1
        else:
            if cur:
                runs.append(cur)
            cur = [k, pa, 1]
    if cur:
        runs.append(cur)
    for k, pa, n in runs:
        out.append("0x%x,0x%x,%d\n" % (gpu_va + k * 0x1000, pa, n))
        emitted += 1
    return emitted


def _is_ring_va(v):
    """True for a page-aligned GPFIFO ring GPU VA in any channel window."""
    if not v or (v & 0xFFF):
        return False
    for lo, hi in RING_WINDOWS:
        if lo <= v < hi:
            return True
    return False


def _chan_info_at(kc, pos):
    """Validate the UvmGpuChannelInfo implied by a gpFifoGpuVa at pos.

    pos is the address of the +0x58 gpFifoGpuVa field; the surrounding
    fields must read as a real channel_info (sane entry count, kernel CPU
    pointers) before the anchor is trusted. This lets the candidate search
    accept both ring windows without latching onto an unrelated page-aligned
    GPU VA that happens to sit in the same allocation (the pool page VA,
    0x121000000, is exactly such a value).
    """
    entries_cpu = kc.read64(pos - 0x48)
    entries = kc.read64(pos - 0x40)
    if entries:
        entries &= 0xFFFFFFFF
    gp_put = kc.read64(pos - 0x50)
    if not entries or entries > 0x100000:
        return None
    if not is_kern_ptr(entries_cpu) or not is_kern_ptr(gp_put):
        return None
    return (entries_cpu, entries, gp_put, kc.read64(pos + 0x08))


def _find_chan_anchor(kc, lo, span):
    """First validated chan anchor in [lo, lo+span).

    Returns (pos, ring_va, entries_cpu, entries, gp_put, gp_put_va) or None.
    """
    data = kc.read(lo, span)
    if not data:
        return None
    for off in range(0, len(data) - 8, 8):
        v = struct.unpack_from("<Q", data, off)[0]
        if not _is_ring_va(v):
            continue
        pos = lo + off
        got = _chan_info_at(kc, pos)
        if got:
            return (pos, v) + got
    return None


def resolve_channel_from_sem(kc, sem, out):
    """Tracking sem (embedded in uvm_channel_t) -> ring/GPPUT pins.

    The channel's tracking semaphore is a uvm_channel_t member updated on
    every push, and channel_info is a later union member of the same
    allocation, so scan forward from the sem (plus a little back) and accept
    a validated gpFifoGpuVa from any ring window. The original GR-only
    window [0x200200000,0x200400000) rejected the CE channels' rings
    (0x121010000 + 0x30000*rank), which is why no #put note ever fired for
    the CE channels the second session stalls on.
    """
    if not is_kern_ptr(sem):
        return 0
    r = _find_chan_anchor(kc, sem - 0x2000, 0xA000)
    if r is None:
        dbg("chan-nofind sem=0x%x" % sem)
        return 0
    pos, v, entries_cpu, entries, gp_put, gp_put_va = r
    dbg("chan-find sem=0x%x ring=0x%x put=%u" %
        (sem, v, read_u32(kc, gp_put)))
    emitted = _emit_ring_pins(kc, v, entries_cpu, entries, out)
    chan_put_note(v, gp_put)
    if is_kern_ptr(gp_put) and gp_put_va and _is_gpu_va_wide(gp_put_va):
        pa = kc.va_to_pa(gp_put)
        if pa is not None:
            out.append("0x%x,0x%x,1\n" % (gp_put_va & ~0xFFF,
                                           pa & ~0xFFF))
            emitted += 1
    return emitted


def resolve_sem_pins(kc, sem, out):
    vals = [kc.read64(sem + i * 8) for i in range(12)]
    if not vals[0] or not is_kern_ptr(vals[0]):
        return 0
    mem = kc.read64(vals[0])
    if not mem or not is_kern_ptr(mem):
        return 0
    return resolve_mem_pins(kc, mem, out, default_pages=16)


def resolve_pb_pins(kc, pb, out):
    mem = kc.read64(pb + 8)
    if not mem:
        return 0
    return resolve_mem_pins(kc, mem, out, default_pages=64)


# UvmGpuChannelInfo field offsets (nv_uvm_types.h, x86_64):
#   +0x00 gpGet, +0x08 gpPut, +0x10 gpFifoEntries, +0x18 numGpFifoEntries,
#   +0x58 gpFifoGpuVa, +0x60 gpPutGpuVa, +0x68 gpGetGpuVa.
# The offset of channel_info inside uvm_channel_t is config-dependent, so the
# struct is located by searching the channel for its gpFifoGpuVa value (a GPU
# VA in a ring window, page aligned).
GR_RING_LO = 0x200200000
GR_RING_HI = 0x200400000
# ISSUE-0002: the CE channels register gpfifo=0x121010000 + 0x30000*rank in
# the service's `channel ... registered gpfifo=...` lines, i.e. the CE
# GPFIFO rings live just above the semaphore pool, NOT in the GR window.
# The GR-only window rejected them, so the #put note never fired for the
# channels the second session stalls on. Both windows are accepted; every
# candidate is validated against the channel_info layout first.
CE_GPFIFO_LO = 0x121000000
CE_GPFIFO_HI = 0x121400000
RING_WINDOWS = ((GR_RING_LO, GR_RING_HI), (CE_GPFIFO_LO, CE_GPFIFO_HI))
CHAN_INFO_GPFIFO_VA = 0x58

# ISSUE-0002: channel CPU-side GPUT pointers, keyed by ring GPU VA. The main
# loop emits a "#put <ring_va> <value>" control line when the value changes,
# so the host can see whether the guest's real put advances while the
# service's GP_PUT read stays at 1 (no guest agent needed at stall time).
CHAN_PUTS = {}
PUT_LAST = {}


def chan_put_note(gpu_va, gp_put):
    if gpu_va and is_kern_ptr(gp_put):
        CHAN_PUTS[gpu_va] = gp_put


def resolve_channel_pins(kc, ch, out):
    """uvm_channel_t kptr -> pin the GPFIFO ring + GPPUT buffers.

    The ring and GP_PUT buffers are CPU-mapped in the guest kernel (UVM
    writes pushes and the put counter directly), so their guest-physical
    backing is reachable with the page-table walk. Without these pins the
    host cannot read GR channel entries (the stall class: RM event loop
    waits for a completion whose ring entry was never processed).
    """
    if not is_kern_ptr(ch):
        return 0
    r = _find_chan_anchor(kc, ch, 0x4000)
    if r is None:
        dbg("ch-nofind ch=0x%x" % ch)
        return 0
    pos, gpu_va, entries_cpu, entries, gp_put, gp_put_va = r
    dbg("ch-find ch=0x%x ring=0x%x put=%u entries=%u" %
        (ch, gpu_va, read_u32(kc, gp_put), entries))
    base = pos - CHAN_INFO_GPFIFO_VA
    emitted = 0
    if is_kern_ptr(entries_cpu):
        nbytes = entries * 8
        npages = min((nbytes + 0xFFF) // 0x1000 + 2, 8)
        runs = []
        cur = None
        for k in range(npages):
            pa = kc.va_to_pa(entries_cpu + k * 0x1000)
            if pa is None:
                if cur:
                    runs.append(cur)
                    cur = None
                continue
            if cur and pa == cur[1] + cur[2] * 0x1000:
                cur[2] += 1
            else:
                if cur:
                    runs.append(cur)
                cur = [k, pa, 1]
        if cur:
            runs.append(cur)
        for k, pa, n in runs:
            out.append("0x%x,0x%x,%d\n" % (gpu_va + k * 0x1000, pa, n))
            emitted += 1
    chan_put_note(gpu_va, gp_put)
    if is_kern_ptr(gp_put) and gp_put_va and _is_gpu_va_wide(gp_put_va):
        pa = kc.va_to_pa(gp_put)
        if pa is not None:
            out.append("0x%x,0x%x,1\n" % (gp_put_va & ~0xFFF, pa & ~0xFFF))
            emitted += 1
    gp_get = kc.read64(base + 0x00)
    gp_get_va = kc.read64(base + 0x68)
    if is_kern_ptr(gp_get) and gp_get_va and _is_gpu_va_wide(gp_get_va):
        pa = kc.va_to_pa(gp_get)
        if pa is not None:
            out.append("0x%x,0x%x,1\n" % (gp_get_va & ~0xFFF, pa & ~0xFFF))
            emitted += 1
    return emitted



def clear_trace():
    """Drop stale ring events so the fallback tail reader cannot re-learn a
    previous client's freed sems after a client switch. The new client's
    events are delivered fresh through the consuming trace_pipe reader."""
    try:
        with open(TRACE, "w") as f:
            f.write("")
    except OSError:
        pass


class TraceReader:
    """Stream uvmt sems / uvbp pushbuffers from the guest trace via trace_pipe.

    trace_pipe is a consuming stream: every event is delivered exactly once,
    with no ring-window race. The old tail-of-trace reader polled the last
    256 KB every 0.3 s, so a one-shot sem= event could age out of the window
    when the client's spin flooded the ring before the next poll - that is
    the second-client stall class (the previous client's freed sem was
    re-learned from the stale ring while the new client's sem was missed).
    """

    def __init__(self):
        self.lines = []
        self.seen = set()
        self.lock = threading.Lock()
        self.stop = False
        self.fd = None
        self.t = threading.Thread(target=self._run, daemon=True)
        self.t.start()

    def _open(self):
        try:
            self.fd = os.open(TRACE_PIPE, os.O_RDONLY)
        except OSError:
            self.fd = None

    def _run(self):
        while not self.stop:
            if self.fd is None:
                self._open()
                if self.fd is None:
                    time.sleep(1)
                    continue
            try:
                data = os.read(self.fd, 65536)
            except OSError:
                try:
                    os.close(self.fd)
                except OSError:
                    pass
                self.fd = None
                time.sleep(0.5)
                continue
            if not data:
                continue
            with self.lock:
                for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("s", v) not in self.seen:
                        self.seen.add(("s", v))
                        self.lines.append("#sem 0x%x\n" % v)
                for m in re.finditer(rb"pb=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("p", v) not in self.seen:
                        self.seen.add(("p", v))
                        self.lines.append("#pb 0x%x\n" % v)
                for m in re.finditer(rb"ch=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("c", v) not in self.seen:
                        self.seen.add(("c", v))
                        self.lines.append("#ch 0x%x\n" % v)

    def drain(self):
        with self.lock:
            out, self.lines = self.lines, []
        return out


def new_sems(seen):
    """New tracking-sem kptrs (uvmt lines) in the trace tail.

    Read in the guest so the host does not need a guest-agent round trip for
    every semaphore discovery; the trace is only tailed while a client runs.
    """
    try:
        out = subprocess.run(["tail", "-c", "262144", TRACE],
                             capture_output=True, timeout=10).stdout
    except Exception:
        return []
    found = []
    for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", out):
        v = int(m.group(1), 16)
        if v not in seen:
            seen.add(v)
            found.append(v)
    return found


def nvidia_fds(pid):
    """CUDA device fds this pid holds: 'uvm', 'gpu', or None.

    /dev/nvidia-uvm marks a compute client (the CUDA runtime opens it at
    context init; nvidia-smi and persistenced do not), so a process holding
    it is a real client regardless of its binary name.
    """
    kind = None
    try:
        fddir = '/proc/%s/fd' % pid
        for fd in os.listdir(fddir):
            try:
                tgt = os.readlink(os.path.join(fddir, fd))
            except OSError:
                continue
            if tgt.startswith('/dev/nvidia-uvm'):
                return 'uvm'
            if tgt.startswith('/dev/nvidia'):
                kind = 'gpu'
    except OSError:
        pass
    return kind


def cdt_pid():
    prefixes = tuple(p for p in os.environ.get("CLIENT_PREFIXES",
                                               "cdt,va,victory").split(",") if p)
    # Processes that hold a /dev/nvidia* fd but are not CUDA clients; the
    # scanner must never commit to them (a transient nvidia-smi at boot was
    # picked once and its exited pagemap read back EOF forever).
    non_clients = ("nvidia-smi", "nvidia-persistenced", "Xorg", "X")
    best = None
    uvm = None
    gpu = None
    for d in os.listdir('/proc'):
        if not d.isdigit():
            continue
        try:
            with open('/proc/%s/comm' % d) as f:
                name = f.read().strip()
            # Skip zombies: a defunct previous client keeps its comm and
            # would otherwise shadow the live client (pagemap reads then fail
            # and no pins are ever streamed).
            with open('/proc/%s/stat' % d) as f:
                st = f.read()
            rp = st.rfind(')')
            if rp >= 0 and st[rp + 2:rp + 3] == 'Z':
                continue
            if any(name.startswith(p) for p in prefixes):
                best = int(d)
                continue
            if any(name.startswith(p) for p in non_clients):
                continue
            # Arbitrary CUDA binaries (any comm) are clients too: match by
            # the CUDA device they hold open so streaming never depends on
            # the binary's name. Prefer a UVM holder (compute client) over a
            # plain /dev/nvidia0 holder (nvidia-smi, persistenced); the last
            # match wins, which is the newest client when several overlap.
            kind = nvidia_fds(d)
            if kind == 'uvm':
                uvm = int(d)
            elif kind == 'gpu':
                gpu = int(d)
        except OSError:
            pass
    if best is not None:
        return best
    if uvm is not None:
        return uvm
    return gpu


def scan_window(fd, lo, hi, span, out):
    base = lo & ~(REGION - 1)
    while base < hi:
        va = max(base, lo) & ~0xFFF
        end = min(base + span, hi)
        while va < end:
            n = min(BATCH, (end - va) // 0x1000)
            if n <= 0:
                break
            os.lseek(fd, (va // 4096) * 8, 0)
            data = os.read(fd, n * 8)
            got = len(data) // 8
            # EOF: the process's address space (or the pagemap of an exited
            # pid) does not reach this VA. Without this the loop reads the
            # same offset forever (got == 0 -> va never advances).
            if got <= 0:
                break
            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
        base += REGION


def scan_range(fd, lo, hi, out):
    """Present pages of one explicit VA range, every 4 KB page.

    Same pagemap present/PFN/PAMAX filtering as scan_window, but with no
    REGION alignment or head-span limit: walks [lo, hi) directly. Used for
    the UVM semaphore pool window, which sits past the CE head span.
    """
    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
        if got <= 0:
            break
        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


def scan_vmas(fd, pid, out):
    """Present pages of the client's high-VA nvidia//dev/zero mappings.

    These are UVM system-memory mappings (sysmem buffers, DtoH staging);
    the fast windows do not cover them and the host cannot resolve them
    from guest RAM alone. `fd` is the pagemap fd of the client.
    """
    n = 0
    try:
        maps = open('/proc/%d/maps' % pid)
    except OSError:
        return
    try:
        for line in maps:
            parts = line.split()
            if len(parts) < 6:
                continue
            path = parts[5]
            if not any(p in path for p in VMA_PATHS):
                continue
            try:
                a, b = parts[0].split('-')
                lo, hi = int(a, 16), int(b, 16)
            except ValueError:
                continue
            if lo < VMA_LO:
                continue
            va = (lo + 0xFFF) & ~0xFFF
            while va < hi and n < VMA_MAX_PAGES:
                cnt = min(256, (hi - va) // 0x1000)
                if cnt <= 0:
                    break
                os.lseek(fd, (va // 4096) * 8, 0)
                data = os.read(fd, cnt * 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
                n += got
    finally:
        maps.close()


def scan(pid):
    """Present pages in the first SPAN bytes of every REGION, as {va: pa}.

    The pushes sit at the start of each per-channel allocation (2/4 MB
    aligned); scanning only the allocation heads keeps the per-loop page
    table work small, which is what makes the stream fast while the client
    is mapping and the guest is loaded.

    Two windows: the GR window (rings/pushbuffers/counter page) and the CE
    push-buffer window (0x120000000+, one 1 MB allocation per CE channel).
    The CE window uses a wider span because the per-channel allocations are
    1 MB apart inside a 2 MB region.
    """
    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
    out = {}
    scan.loop += 1
    try:
        scan_window(fd, LO, HI, SPAN, out)
        if HI2 > LO2:
            scan_window(fd, LO2, HI2, SPAN2, out)
        if POOL_HI > POOL_LO:
            scan_range(fd, POOL_LO, POOL_HI, out)
        if VMA_EVERY and scan.loop % VMA_EVERY == 0:
            scan_vmas(fd, pid, out)
    finally:
        os.close(fd)
    return out


scan.loop = 0


def main():
    with open('/tmp/guest_pin_stream.pid', 'w') as f:
        f.write('%d\n' % os.getpid())
    sock = None
    seen = {}
    seen_sems = set()
    seen_pbs = set()
    seen_chs = set()
    resolved_sems = set()
    resolved_sem_chan = set()
    resolved_pbs = set()
    resolved_chs = set()
    retry_at = {}
    tracer = TraceReader()
    kc = Kcore()
    kc_ok = kc.load()
    try:
        with open('/tmp/guest_pin_stream.log', 'a') as f:
            f.write('%s resolver: kcore=%s (sem->pool/pushbuffer pins)\n'
                    % (time.strftime('%T'), 'ON' if kc_ok else 'OFF'))
    except OSError:
        pass
    last_pid = None
    last_log = 0.0
    loop = 0
    while True:
        pid = cdt_pid()
        loop += 1
        if pid is None:
            last_pid = None
            seen = {}
            seen_sems = set()
            seen_pbs = set()
            seen_chs = set()
            time.sleep(INTERVAL)
            continue
        if pid != last_pid:
            last_pid = pid
            seen = {}
            seen_sems = set()
            seen_pbs = set()
            seen_chs = set()
            resolved_sems = set()
            resolved_sem_chan = set()
            resolved_pbs = set()
            resolved_chs = set()
            retry_at = {}
            # Client switch: the old client's channel/put pointers are freed
            # memory. Clearing the #put map stops stale reads from being
            # emitted as if they were the new client's counters.
            CHAN_PUTS.clear()
            PUT_LAST.clear()
            clear_trace()
            if sock is not None:
                try:
                    sock.sendall(("#pid %d\n" % pid).encode())
                except OSError:
                    sock.close()
                    sock = None
        t0 = time.time()
        try:
            cur = scan(pid)
        except OSError:
            time.sleep(INTERVAL)
            continue
        dt = time.time() - t0
        # Diagnostic: the GR completion-counter page is the one page the
        # client's context-init wait polls; log its pagemap entry every loop
        # so a scan blind spot is visible in the guest log.
        cpm = None
        try:
            fdp = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
            os.lseek(fdp, (0x20400fec0 // 4096) * 8, 0)
            cpm = int.from_bytes(os.read(fdp, 8), 'little')
            os.close(fdp)
        except OSError:
            cpm = None
        diag = ('cpm=%s' % ('present' if (cpm or 0) >> 63 & 1 else 'absent')
                if cpm is not None else 'cpm=ERR')
        # Guaranteed fast path for the GR completion-counter page: if the
        # window scan misses it (observed once), add it explicitly so the host
        # can register the slots before the client's first wait.
        if cpm is not None and (cpm >> 63) & 1:
            cpa = (cpm & ((1 << 55) - 1)) << 12
            if cpa and (not PAMAX or cpa < PAMAX):
                cur[0x20400f000] = cpa
        new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
        pending_cur = cur
        lines = []
        for ctl in tracer.drain():
            parts = ctl.split()
            if ctl.startswith("#sem") and len(parts) >= 2:
                seen_sems.add(int(parts[1], 16))
            elif ctl.startswith("#pb") and len(parts) >= 2:
                seen_pbs.add(int(parts[1], 16))
            elif ctl.startswith("#ch") and len(parts) >= 2:
                seen_chs.add(int(parts[1], 16))
            lines.append(ctl)
        if kc_ok:
            now_r = time.time()
            added = 0
            for sem in list(seen_sems):
                if sem not in resolved_sems and now_r >= retry_at.get(sem, 0):
                    got = []
                    try:
                        resolve_sem_pins(kc, sem, got)
                    except Exception:
                        got = []
                    if got:
                        lines.extend(got)
                        resolved_sems.add(sem)
                        added += len(got)
                    else:
                        retry_at[sem] = now_r + 1.0
                # ISSUE-0002: the pool pin (mem) and the channel ring are
                # resolved independently. The sem is marked resolved as soon
                # as the pool page pins, but the channel_info anchor may only
                # appear (or only become findable) later, and the #put probe
                # must not be skipped for an already-resolved sem.
                if sem in resolved_sem_chan or \
                        now_r < retry_at.get(("cy", sem), 0):
                    continue
                got = []
                try:
                    resolve_channel_from_sem(kc, sem, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_sem_chan.add(sem)
                    added += len(got)
                else:
                    retry_at[("cy", sem)] = now_r + 2.0
            for pb in list(seen_pbs):
                if pb in resolved_pbs or now_r < retry_at.get(pb, 0):
                    continue
                got = []
                try:
                    resolve_pb_pins(kc, pb, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_pbs.add(pb)
                    added += len(got)
                else:
                    retry_at[pb] = now_r + 1.0
            for ch in list(seen_chs):
                if ch in resolved_chs or now_r < retry_at.get(ch, 0):
                    continue
                got = []
                try:
                    resolve_channel_pins(kc, ch, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_chs.add(ch)
                    added += len(got)
                else:
                    retry_at[ch] = now_r + 1.0
            if added:
                try:
                    with open('/tmp/guest_pin_stream.log', 'a') as f:
                        f.write('%s resolver pins: +%d run(s) '
                                '(sems=%d/%d pbs=%d/%d chs=%d/%d)\n'
                                % (time.strftime('%T'), added,
                                   len(resolved_sems), len(seen_sems),
                                   len(resolved_pbs), len(seen_pbs),
                                   len(resolved_chs), len(seen_chs)))
                except OSError:
                    pass
        # ISSUE-0002: stream the channel CPU-side GPUT on change (see
        # chan_put_note). Emitted here so the value is sent with this
        # loop's payload even when no new pages appeared.
        for gv, pp in list(CHAN_PUTS.items()):
            d = kc.read(pp, 4)
            if not d:
                continue
            pv = struct.unpack("<I", d)[0]
            if PUT_LAST.get(gv) != pv:
                PUT_LAST[gv] = pv
                lines.append("#put 0x%x %u\n" % (gv, pv))
        now = time.time()
        if now - last_log > 1.0:
            last_log = now
            try:
                with open('/tmp/guest_pin_stream.log', 'a') as f:
                    f.write('%s pid=%s present=%d new=%d sem=%d scan=%.3fs %s\n'
                            % (time.strftime('%T'), pid, len(cur), len(new),
                               len(seen_sems), dt, diag))
            except OSError:
                pass
        if not new and not lines:
            seen = pending_cur
            time.sleep(INTERVAL)
            continue
        if new:
            new.sort()
            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)
            seen = pending_cur
        except OSError:
            sock.close()
            sock = None
        time.sleep(INTERVAL)


if __name__ == '__main__':
    main()