#!/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 _nv_vmap_areas(kc):
    """nv_vm_map_pages vmalloc areas: [(lo, hi, pages)] (5 s cache).

    Every CE push-buffer rm_mem is a vmap created by nv_vm_map_pages; the
    demand pinner walks these areas to find the page holding a stalled
    push VA by content, without needing any kprobe event or object chain.
    """
    now = time.time()
    if getattr(_nv_vmap_areas, "ts", 0) > now - 5.0:
        return _nv_vmap_areas.areas
    areas = []
    try:
        for line in open("/proc/vmallocinfo"):
            if "nv_vm_map_pages" not in line:
                continue
            p = line.split()
            if len(p) < 2:
                continue
            try:
                lo, hi = p[0].split("-")
                lo = int(lo, 16)
                hi = int(hi, 16)
            except ValueError:
                continue
            pages = (hi - lo) >> 12
            if pages > 0:
                areas.append((lo, hi, pages))
    except OSError:
        pass
    _nv_vmap_areas.areas = areas
    _nv_vmap_areas.ts = now
    return areas


def resolve_va_by_content(kc, va, out):
    """Demand pin: find the page carrying GPU VA 'va' by content match.

    The host reports the CE push-buffer VAs whose GPFIFO entries it cannot
    read (dmesg 'stalled VA #N: 0x...'). The guest walks every
    nv_vm_map_pages vmap area and tests the candidate leaf page at the VA's
    in-window offset (plus one neighbour on either side, for the per-mem
    header-page offset ambiguity) against the host's method-stream prefix
    check. On a match it emits the exact pair, so the host can translate
    and consume the entry - no kprobe event, no object-chain discovery, no
    resolution lottery. Only the requested VA is pinned, so a wrong
    candidate cannot poison other VAs.
    """
    if not (PB_CE_LO <= va < PB_CE_HI):
        return False
    off = (va - PB_CE_LO) & 0xFFFFF
    k = off >> 12
    for (lo, _hi, pages) in _nv_vmap_areas(kc):
        for adj in (0, 1, -1):
            idx = k + adj
            if idx < 0 or idx >= pages:
                continue
            cpu = lo + idx * 0x1000
            d = kc.read(cpu, 64)
            if not d or not any(d[:32]):
                continue
            if not _push_prefix_ok(d):
                continue
            pa = kc.va_to_pa(cpu)
            if pa is None:
                continue
            out.append("0x%x,0x%x,1\n" % (va & ~0xFFF, pa))
            dbg("demand-va: 0x%x -> pa=0x%x (area off=%d adj=%+d)"
                % (va & ~0xFFF, pa, idx, adj))
            return True
    dbg("demand-va: 0x%x not found by content (%d areas)"
        % (va & ~0xFFF, len(_nv_vmap_areas(kc))))
    return False


def resolve_mem_pins(kc, mem, out, default_pages=16, va_lo=0, va_hi=0,
                     page_off=0):
    """uvm_rm_mem kptr -> GPU-VA/GPA runs (direct map or page-table walk).

    va_lo/va_hi optionally restrict the emitted GPU VA range; the kcore
    chain resolver targets the CE push-buffer window only. page_off shifts
    the vmap page paired with GPU data page k: the CE push-buffer vmap is
    one page longer than its GPU data window on some allocations (a leading
    header page), so GPU data page k lives at vmap page k + page_off.
    page_off is derived per allocation (mem_page_off): a fixed offset is
    wrong (r116/r117 regressed client #1 with a hard +1).
    """
    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) and (not va_lo or v >= va_lo)
            and (not va_hi or v < va_hi)]
    if not cpu or not gpus:
        return 0
    size = _mem_size(kc, mem)
    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
    area = _vmalloc_pages_for(kc, cpu)
    npages = (max(1, min(size >> 12, 4098)) if size
              else area or default_pages)
    if area and page_off + npages > area:
        npages = max(1, area - page_off)
    runs = []
    cur = None
    for k in range(npages):
        pa = kc.va_to_pa(cpu + (k + page_off) * 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 _mem_size(kc, mem):
    """uvm_rm_mem size (bytes) from the 570.144 layout, with a fallback.

    The struct is type(4)+pad(4)+mask(8)+vas[33] (264 B)+proxy_vas+
    gpu_owner+size, so size lives at +0x128; +0x60 (inside vas[]) reads 0
    for unused processors and is kept only as a fallback.
    """
    for off in (0x128, 0x60):
        v = kc.read64(mem + off) or 0
        if 0 < v <= (256 << 20):
            return v
    return 0


def _push_prefix_ok(buf):
    """First-4KB method-stream plausibility, mirroring the host's check."""
    off = 0
    while off + 4 <= len(buf):
        hdr = struct.unpack_from('<I', buf, off)[0]
        if hdr:
            break
        off += 4
    if off + 4 > len(buf):
        return True		# all zeros: not refutable
    op = (hdr >> 29) & 0x7
    if op in (4, 7):
        return True
    if op not in (1, 3, 5):
        return False
    return ((hdr >> 16) & 0xfff) != 0


def _page_content(kc, va):
    pa = kc.va_to_pa(va & ~0xFFF)
    if pa is None:
        return None
    return kc.read(kc.page_offset_base + pa, 64)


def mem_page_off(kc, mem):
    """Per-allocation vmap offset for a uvm_rm_mem's GPU data pages.

    Derived, never fixed: the vmap area's page count (vmallocinfo) minus
    the rm_mem's data page count (size at +0x128). A 4097-page area with
    4096 data pages means one leading header page (offset 1); an area that
    exactly matches the data means offset 0. When the size is unreadable,
    the two candidate offsets are content-probed with the host's method-
    stream plausibility test and the decisively-passing one wins; an
    ambiguous result keeps 0 (the status-quo pairing).
    """
    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)
    if not cpu or _direct_map_pa(kc, cpu) is not None:
        return 0
    area = _vmalloc_pages_for(kc, cpu)
    if area <= 1:
        return 0
    size = _mem_size(kc, mem)
    if size:
        data = min(size >> 12, 4098)
        if 0 <= area - data <= 2:
            return area - data
        return 0
    g = []
    for off in (0, 1):
        if off >= area:
            g.append(0)
            continue
        buf = _page_content(kc, cpu + off * 0x1000)
        g.append(1 if (buf is not None and _push_prefix_ok(buf)) else 0)
    if g[0] != g[1]:
        return 0 if g[0] else 1
    return 0


def _ce_mem_pages(kc, mem):
    """CE mem walk parameters: (gpus, cpu, off, npages) or None.

    Walks the same page set resolve_mem_pins used: the whole vmalloc area
    (minus the header offset) when known, else size at +0x128, else the
    default. The area governs: the live pushes are written sequentially and
    may run past the size field, so a size-derived prefix stopped one page
    short of the write frontier (r144: last page 0x120030000 emitted, the
    live page 0x120031000 never in any pin set). all CE-window GPU VAs share
    the CPU mapping.
    """
    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 v and PB_CE_LO <= v < PB_CE_HI]
    if not cpu or not gpus:
        return None
    if _direct_map_pa(kc, cpu) is not None:
        return None		# never a vmap: not the push-buffer shape
    off = mem_page_off(kc, mem)
    area = _vmalloc_pages_for(kc, cpu)
    size = _mem_size(kc, mem)
    size_pages = max(1, min(size >> 12, 4098)) if size else 0
    area_pages = area - off if area and area > off else 0
    npages = min(max(size_pages, area_pages) or 64, 4098)
    return (gpus, cpu, off, npages)


def _emit_ce_mem_delta(kc, st, out, fresh_map=False):
    """Emit the CE mem's live {va,pa} pairs as merged runs.

    Two emission modes share one merge step:

    * full re-assert (KC_FULLREM_S cadence, and every fresh_map pass): every
      page that translates is emitted, content-independent. This is the
      fix for the host's session flush: at a client switch the host wipes
      its whole sema map when the previous client's last channel closes
      ("session flush (all channels closed)"), and a pin batch that raced
      in for the new client is lost. The old content-gated delta never
      re-emitted those unchanged pairs, so the new client's CE pushes were
      permanently unreadable ("NO ENTRY" in sema-probe) and every GPFIFO
      entry was skipped after the 8 s grace (r144/r145 second-client stall,
      both mail and matmul). Re-asserting the full pair set every couple of
      seconds heals any number of flush wipes; merged contiguous runs keep
      the traffic to a few lines per pass. Zero-content pages are pinned
      deliberately: the pin carries the translation, the host reads the
      live content, and the client writes it later.

    * delta (between full passes): content-gated, emits pages whose PA is
      new or changed, so a remap or a freshly-faulted page is picked up
      without waiting for the next full pass. Never drops a previously
      emitted pair (the host's last-store-wins pin sees the truth, and a
      failed burst never clears the whole map).
    """
    gpus, cpu, off, npages = st["gpus"], st["cpu"], st["page_off"], st["npages"]
    last = st["last"]
    nowm = time.monotonic()
    full = nowm >= st.get("full_next", 0.0)
    if full:
        age = nowm - st.get("resolved_at", nowm)
        st["full_next"] = nowm + (KC_FULLREM_S if age < KC_FULLREM_FAST_S
                                  else KC_FULLREM_SLOW_S)
    if fresh_map:
        for k in range(npages):
            kc._pa_cache.pop(cpu + (k + off) * 0x1000, None)
    changed = []
    if full:
        # Re-walk pages that did not translate on an earlier pass: the
        # vmap's leaf sizes are huge (2 MB/1 GB), so most pages translate
        # immediately; a straggler (the vmap populated lazily) is retried
        # with its negative PA-cache entry popped, bounded per pass so the
        # sweep stays cheap.
        miss = st.setdefault("miss", set())
        retry = 0
        for k in range(npages):
            page = cpu + (k + off) * 0x1000
            pa = kc.va_to_pa(page)
            if pa is None:
                if k in miss and retry < 512:
                    kc._pa_cache.pop(page, None)
                    retry += 1
                else:
                    miss.add(k)
                continue
            miss.discard(k)
            for gv in gpus:
                changed.append((gv + k * 0x1000, pa))
    else:
        for k in range(npages):
            d = kc.read(cpu + (k + off) * 0x1000, 32)
            if not d or not any(d):
                continue		# zero/unreadable: nothing to pin yet
            pa = kc.va_to_pa(cpu + (k + off) * 0x1000)
            if pa is None:
                continue
            for gv in gpus:
                if last.get(gv + k * 0x1000) != pa:
                    changed.append((gv + k * 0x1000, pa))
    changed.sort()
    merged = []
    for va, pa in changed:
        if merged and va == merged[-1][0] + merged[-1][2] * 0x1000 and \
                pa == merged[-1][1] + merged[-1][2] * 0x1000:
            merged[-1][2] += 1
        else:
            merged.append([va, pa, 1])
    for va, pa, n in merged:
        out.append("0x%x,0x%x,%d\n" % (va, pa, n))
    for va, pa in changed:
        last[va] = pa
    st["emit_total"] = st.get("emit_total", 0) + len(merged)
    if full:
        st["full_total"] = st.get("full_total", 0) + 1
    return len(merged)


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,
                            page_off=mem_page_off(kc, mem))


# The CE push-buffer window is the one UVM region the client never maps into
# its own address space: the pagemap window scan and the pool-pin resolver
# can never see it. The uvbp kprobe was its only supplier, and that event is
# one-shot per push buffer while two trace readers race it - this scanner's
# pid-switch clear_trace() and semfix's trace-tail clear - so when the event
# loses the race the second client's ring entries stay untranslated and the
# service falls back to the ~250 ms/entry content rescue (the second-client
# stall, the same resolution-miss class as the ISSUE-0001 lottery). The
# channel->pool->manager->pushbuffer->rm_mem chain is walkable in kernel
# memory, so resolve the CE push buffer directly with kcore: no trace event,
# nothing to lose.
#
# The kptr-keyed one-shot latches below are NOT object identity: UVM frees
# and re-creates these rm_mems and channels during a session, and the slab
# reuses the same addresses. A resolution taken before such a move pins
# pages that are stale by the time the host reads them (observed r125: the
# pinned CE pages read all-zero while the channel's real pushes lived at
# new vmap pages of the same rm_mem kptr; every ring entry then stalled in
# the host's 8 s grace since no release could be decoded from zeros). So a
# resolved mem is periodically re-validated against the channel's own ring
# entries and re-resolved when it stops matching (KC_PB_STATE/_kc_pb_fresh),
# and when several mem candidates are reachable the one whose pages actually
# hold the ring-entry pushes is chosen (content anchor, _mem_content_ok).
PB_CE_LO = 0x120000000
PB_CE_HI = 0x121000000
KC_CHAIN_BACK = 0x2000
KC_CHAIN_SPAN = 0xA000
KC_CHAIN_MAX = 400
RESOLVED_KC_PB = set()
KC_PB_STATE = {}		# sem kptr -> resolved-mem state (fingerprint + pins)
KC_REFRESH_S = 2.0		# re-validate a resolved CE mem at most this often
KC_REEMIT_S = 2.0		# sweep the CE mem's live pages this often
KC_REMAP_S = 30.0		# full re-translation sweep (catches remaps)
# Full re-assert cadence: every pass at least this far apart re-emits the
# whole translatable pair set regardless of the delta, healing host-side
# session-flush wipes (see _emit_ce_mem_delta). Must stay well under the
# host's gsp_sema_ce_grace_ms (8 s) so a wiped-and-needed page is restored
# inside the hold window. The fast cadence applies to a mem's first
# KC_FULLREM_FAST_S seconds (the flush race lives at the client switch);
# after that the re-assert backs off to KC_FULLREM_SLOW_S - a fully
# fragmented pair set (descending PAs) emits one page per line, and 4 k
# lines every 2 s saturated a guest under multi-client load (r148: load 19,
# three scanners from an install race, clients wedged).
KC_FULLREM_S = float(os.environ.get("STREAM_FULLREM_S", "1.5"))
KC_FULLREM_FAST_S = float(os.environ.get("STREAM_FULLREM_FAST_S", "45"))
KC_FULLREM_SLOW_S = float(os.environ.get("STREAM_FULLREM_SLOW_S", "10.0"))


def _chain_mems_at(kc, w):
    """rm_mem candidates when kernel pointer w is a pool, manager or pb."""
    mems = []
    mgr = kc.read64(w)
    if is_kern_ptr(mgr):
        pb = kc.read64(mgr + 8)
        if is_kern_ptr(pb):
            mem = kc.read64(pb + 8)
            if is_kern_ptr(mem):
                mems.append(mem)
    pb = kc.read64(w + 8)
    if is_kern_ptr(pb):
        mem = kc.read64(pb + 8)
        if is_kern_ptr(mem):
            mems.append(mem)
    return mems


def _mem_gpu_cpu(kc, mem):
    """The CE-window GPU VA and CPU alias of an rm_mem (None/None when absent)."""
    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)
    gpu = next((v for v in vas
                if v and PB_CE_LO <= v < PB_CE_HI), None)
    return gpu, cpu


def _mem_content_score(kc, mem, gpu):
    """Rough 'is this the live CE push buffer' score for a mem candidate.

    The channel's CE rings live in (emulated) video memory, so the guest
    cannot read the ring entries to anchor a push VA; but the CE push
    buffer's guest-RAM pages are CPU-written method streams during channel
    setup and client work. Count nonzero u32 pairs across the first 128
    pages of the candidate's window: a live mem scores high, a freed page
    range that a reused kptr still describes reads zeros and scores 0.
    """
    cpu = None
    for i in range(16):
        v = kc.read64(mem + 0x10 + i * 8)
        if is_kern_ptr(v):
            cpu = v
            break
    if cpu is None:
        return -1
    score = 0
    for k in range(128):
        buf = kc.read(cpu + k * 0x1000, 0x1000)
        if not buf:
            continue
        for off in range(0, 0x1000, 16):
            if buf[off:off + 8] != b"\x00" * 8:
                score += 1
    return score


def _kc_pb_fresh(kc, sem):
    """The pinned CE mem still matches the channel's live backing.

    Cheap check: the mem's own {cpu,gpu} fields and the guest PA of its
    first window page must still match the recorded resolution. A vmap area
    never moves while the mem object lives, so either check failing means
    the object was freed and (possibly) re-created at the same kptr - the
    caller then re-resolves and re-emits the current pages.
    """
    st = KC_PB_STATE.get(sem)
    if not st:
        return False
    mem = st.get("mem")
    if not mem or not is_kern_ptr(mem):
        return False
    if (kc.read64(mem + 0x30) or 0) != st.get("cpu") or \
            (kc.read64(mem + 0x38) or 0) != st.get("gpu"):
        return False
    if not st.get("page_va"):
        return True
    if not st.get("page_pa"):
        # Unknown PA at store time: the manual page-table walk could not
        # translate the vmap page (it may have been read before the kernel
        # page tables were walkable). Treat the record as fresh instead of
        # failing every check - a permanently-failing freshness test made
        # the scanner drop the latch and re-resolve in a tight loop, and
        # each failed re-resolution left the host with no pins at all
        # (observed r131: "kc-chain stale: re-resolving" every second for
        # five sems while smiley sat in cudaMalloc).
        return True
    pa = kc.va_to_pa(st["page_va"])
    if pa is None or pa != st.get("page_pa"):
        return False
    return True


def resolve_pb_chain(kc, sem, out):
    """Tracking-sem kptr -> CE push-buffer rm_mem pins (kcore chain walk).

    Returns True when a CE push-buffer mem was found for this sem, so the
    caller marks it resolved even when the mem was already pinned; new pin
    lines are appended to out. Bounded: one window read plus a few reads per
    candidate, and each rm_mem is resolved at most once per session (until a
    freshness failure drops it again, see _kc_pb_fresh).

    When several mem candidates are reachable, the content score decides: a
    stale candidate (freed memory read through a reused kptr) reads zeros
    and is skipped in favour of the live push-buffer pages, so it cannot
    shadow them on the host's last-store-wins pin table. An all-zero
    candidate set emits nothing: pinning zeros is what made the host read
    zero pushes and stall every CE entry; the drift retries (retry_at) and
    the periodic refresh pick the pages up as soon as content appears.
    """
    if not is_kern_ptr(sem):
        return False
    data = kc.read(sem - KC_CHAIN_BACK, KC_CHAIN_SPAN)
    if not data:
        return False
    words = struct.unpack_from("<%dQ" % (len(data) // 8), data, 0)
    cands = list(dict.fromkeys(w for w in words if is_kern_ptr(w)))
    # Channel-anchored expansion (r134 live dump): the tracking sem is a
    # member of a uvm_channel_t table; the channel records sit in the sem's
    # own window and carry their CE ring GPU VA ([0x121010000, 0x121400000))
    # plus back-to-back 32-bit object handles. The pool/manager/pushbuffer
    # pointers hang off those records' targets, which the bare sem-window
    # scan never reaches - the second-client stall is exactly this: the only
    # live uvmt event is the spinning sem, and no mem candidate is found
    # around it. For each ring-VA word, scan the surrounding record for
    # kptrs, take those targets, and try the chain shape on both. Additive
    # and bounded; content scoring still decides, so a wrong candidate
    # cannot shadow the live pages.
    extra = []
    for i, w in enumerate(words):
        if not w or not (0x121010000 <= w < 0x121400000):
            continue
        lo = max(0, i - 0x60)
        hi = min(len(words), i + 0x60)
        for c in words[lo:hi]:
            if not is_kern_ptr(c):
                continue
            extra.append(c)
            tgt = kc.read(c, 0x100)
            if tgt:
                tw = struct.unpack_from("<%dQ" % (len(tgt) // 8),
                                        tgt, 0)
                for w2 in tw:
                    if is_kern_ptr(w2):
                        extra.append(w2)
        if len(extra) >= 512:
            break
    if extra:
        head = []
        for w in extra:
            if w not in head and w not in cands:
                head.append(w)
        cands = head + cands
    mems = []
    for w in cands[:KC_CHAIN_MAX]:
        for mem in _chain_mems_at(kc, w):
            if mem not in mems:
                mems.append(mem)
    found = False
    scored = []
    for mem in mems:
        gpu, cpu = _mem_gpu_cpu(kc, mem)
        if not gpu or not cpu:
            continue
        found = True
        score = _mem_content_score(kc, mem, gpu)
        scored.append((score, mem, gpu, cpu))
    if not scored:
        return found
    best = max(scored, key=lambda c: c[0])
    if best[0] <= 0:
        dbg("kc-chain all candidates zero (sem=0x%x, n=%d): deferring"
            % (sem, len(scored)))
        return found
    if len(scored) > 1:
        dbg("kc-chain scores sem=0x%x: %s" %
            (sem, ' '.join('%#x=%d' % (m, s) for s, m, _g, _c in scored)))
    _score, mem, gpu, cpu = best
    if mem in RESOLVED_KC_PB:
        return True
    RESOLVED_KC_PB.add(mem)
    params = _ce_mem_pages(kc, mem)
    if not params:
        return found
    gpus, cpu, off, npages = params
    st = {"mem": mem, "gpu": gpu, "cpu": cpu, "gpus": gpus,
          "page_va": cpu, "page_pa": kc.va_to_pa(cpu) or 0,
          "page_off": off, "npages": npages, "last": {},
          "resolved_at": time.monotonic()}
    # Fresh translation for the first sweep: the Kcore PA cache may hold
    # pairs from before this mem's vmap area was (re)populated.
    kc._pa_cache.clear()
    n = _emit_ce_mem_delta(kc, st, out)
    KC_PB_STATE[sem] = st
    try:
        with open('/tmp/guest_pin_stream.log', 'a') as f:
            f.write('%s kc-chain CE push-buffer mem=0x%x sem=0x%x '
                    'pin-runs=%d off=%d npages=%d size=%d score=%d\n'
                    % (time.strftime('%T'), mem, sem, n, off, npages,
                       _mem_size(kc, mem), _score))
    except OSError:
        pass
    return True


# 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 _read_trace_snapshot(max_bytes=1 << 20, budget=0.5):
    """Best-effort non-blocking read of the trace ring file.

    Spawning `tail` here left uninterruptible D-state processes when the
    seq_file read blocked under load (observed r127: dozens of stuck tails,
    load 26); a direct os.read with O_NONBLOCK returns whatever is ready and
    never accumulates processes.
    """
    try:
        fd = os.open(TRACE, os.O_RDONLY | os.O_NONBLOCK)
    except OSError:
        return b""
    chunks = []
    got = 0
    t0 = time.time()
    try:
        while got < (8 << 20) and time.time() - t0 < budget:
            try:
                d = os.read(fd, 262144)
            except BlockingIOError:
                break
            except OSError:
                break
            if not d:
                break
            chunks.append(d)
            got += len(d)
    finally:
        os.close(fd)
    return b"".join(chunks)[-max_bytes:]


def seed_events_from_trace(seen_sems, seen_pbs, seen_chs):
    """One-shot tail scan of the trace ring after a pid switch / start.

    The uvmt/uvbp/uvch kprobe events are one-shot per object, and the
    trace_pipe reader only sees events that arrive after it attaches: a
    client whose channels are registered in the instant around the scanner's
    pid switch would otherwise never appear in the sets (observed r126:
    vecadd_loop ran for minutes with sem=0 while its CE entries paid the
    8 s grace). The ring keeps recent events, so scan its tail once and add
    every distinct kptr. Stale previous-client pointers that are also
    re-learned are harmless: their chains no longer resolve to content
    (the mem candidates are scored and deferred).
    """
    try:
        out = _read_trace_snapshot()
    except Exception:
        return 0
    n = 0
    for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", out):
        v = int(m.group(1), 16)
        if v not in seen_sems:
            seen_sems.add(v)
            n += 1
    for m in re.finditer(rb"pb=0x([0-9a-fA-F]+)", out):
        v = int(m.group(1), 16)
        if v not in seen_pbs:
            seen_pbs.add(v)
            n += 1
    for m in re.finditer(rb"ch=0x([0-9a-fA-F]+)", out):
        v = int(m.group(1), 16)
        if v not in seen_chs:
            seen_chs.add(v)
            n += 1
    return n


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 reset(self):
        """Forget parsed-event dedup marks on a client switch.

        The kernel reuses freed objects, so the new client's uvmt/uvbp kptr
        can equal one already seen for the previous client; keeping the old
        marks suppressed the new client's #pb line entirely (second-client
        stall). Queued lines are kept: an event committed just before the
        switch may belong to the new client.
        """
        with self.lock:
            self.seen = set()


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


# --------------------------------------------------------------- scrubber
# PMA scrubber channel discovery (third-client scrub freeze). The guest's
# PMA scrubber is an internal CeUtils CE channel (4096 GPFIFO entries,
# pbGpuVA 0x120000000, ring at pbGpuVA+0x64000, finish payload at
# pbGpuVA+0x6c004). It lives in a kernel vmap area, not in any user pagemap,
# and its GPU VAs collide numerically with the client CE push-buffer window
# in the host's global translation map - so the host cannot resolve it and
# must be told the finish/host-semaphore word PAs directly. Resolve them
# here (vmallocinfo + manual kernel page-table walk) and emit one
#   #scrub <fin_pa>,<hsem_pa>
# line; pin_gr_stream.py registers them through sema_scrubbuf_add.
SCRUB_AREA_SIZES = (0x6e000, 0x6d000, 0x6c000, 0x6b000, 0x6f000,
                    0x70000, 0x71000, 0x72000)
_scrub_sent = {}


def _scrub_ring_ok(kc, ring_pa):
    """A 906F GPFIFO ring page: all-zero (fresh) or entries whose GET names a
    push inside the channel's own 0x120000000 window (encoded as
    push_va - 0x100000000, 4-byte aligned)."""
    buf = kc.read(kc.page_offset_base + ring_pa, 0x1000)
    if not buf:
        return False
    for off in range(0, 0x100, 8):
        e0, _e1 = struct.unpack_from('<II', buf, off)
        if e0 == 0:
            continue
        if not (0x20000000 <= e0 < 0x20080000) or (e0 & 3):
            return False
    return True


def scrub_scan(kc, lines):
    """Emit the scrubber buffer words' PAs once (re-emitted if they move)."""
    try:
        f = open('/proc/vmallocinfo')
    except OSError:
        return
    with f:
        for line in f:
            if 'nvidia' not in line or 'ioremap' in line or 'vmap' not in line:
                continue
            parts = line.split()
            try:
                a, b = parts[0].split('-')
                lo, hi = int(a, 16), int(b, 16)
            except (ValueError, IndexError):
                continue
            if (hi - lo) not in SCRUB_AREA_SIZES:
                continue
            fin_pa = kc.va_to_pa(lo + 0x6c000)
            ring_pa = kc.va_to_pa(lo + 0x64000)
            if fin_pa is None or ring_pa is None:
                continue
            if not _scrub_ring_ok(kc, ring_pa):
                dbg('scrub candidate 0x%x rejected (ring shape)' % lo)
                continue
            fin = fin_pa + 4
            hsem = fin_pa
            key = (fin, hsem)
            if _scrub_sent.get('key') == key:
                return
            _scrub_sent['key'] = key
            lines.append('#scrub 0x%x,0x%x\n' % (fin, hsem))
            dbg('scrub buffer emitted lo=0x%x fin=0x%x hsem=0x%x'
                % (lo, fin, hsem))
            return


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")
    cands = []
    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 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. The newest candidate (highest pid) is the
            # client being streamed: preferring a name-prefix match first
            # pinned the scanner to a previous client that was still tearing
            # down while the new client had already started pushing (its
            # first pushes were unreadable, the expectation drifted, and the
            # stall followed).
            kind = nvidia_fds(d)
            if kind is not None or any(name.startswith(p) for p in prefixes):
                cands.append(int(d))
        except OSError:
            pass
    if not cands:
        return None
    return max(cands)


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) < 5:
                continue
            # Anonymous mappings (no path field) are the client's UVM sysmem
            # buffers when the driver maps them without a /dev/zero backing:
            # the second client's cudaMalloc'd data buffers show up exactly
            # like that (e.g. 75d7b0000000-75d7b0021000 rw-p 00000000 00:00 0)
            # and were skipped by the old len<6 check, so their pages never
            # reached the host and the forwarded launch could not resolve
            # them (all-zero results for the second client). Accept them at
            # the high-VA range; the VMA_LO bound and the pagemap present/PFN
            # filters keep the set to real GPU mappings.
            path = parts[5] if len(parts) >= 6 else ''
            if path and 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()
    resolved_kc_pb = set()
    retry_at = {}
    kc_refresh_at = {}
    kc_reemit_at = {}
    kc_remap_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
    scrub_next = 0.0
    demand_next = 0.0
    _want_state = {"mtime": None, "done": set()}
    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()
            resolved_kc_pb = set()
            retry_at = {}
            kc_refresh_at = {}
            kc_reemit_at = {}
            kc_remap_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()
            tracer.reset()
            RESOLVED_KC_PB.clear()
            KC_PB_STATE.clear()
            # One-shot: recover the new client's one-shot channel events
            # from the trace ring (they can fall in the switch instant and
            # be invisible to the pipe reader; the old clear_trace() made
            # that permanent - part of the r126 stall class).
            _seeded = seed_events_from_trace(seen_sems, seen_pbs, seen_chs)
            if _seeded:
                dbg("switch seed: +%d event kptr(s)" % _seeded)
            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 = []
        # Demand pins (host -> this scanner): the host appends the CE
        # push-buffer VAs of GPFIFO entries it cannot read to
        # /tmp/guest_want_vas.txt (written via the guest agent by
        # pin_gr_stream); resolve each by content against the
        # nv_vm_map_pages areas and emit the pair. This is the
        # event-independent hole filler for the stall class where the
        # one-shot #sem kprobe event was lost and the object-chain
        # resolver never learns the client.
        if kc_ok and time.time() >= demand_next:
            demand_next = time.time() + 1.0
            try:
                st_d = os.stat("/tmp/guest_want_vas.txt")
                if st_d.st_mtime != _want_state.get("mtime"):
                    with open("/tmp/guest_want_vas.txt") as fw:
                        wvas = [int(t, 16) for t in fw.read().split()
                                if t.startswith("0x")]
                    _want_state["mtime"] = st_d.st_mtime
                    for wv in wvas[-64:]:
                        if wv in _want_state["done"]:
                            continue
                        got = []
                        try:
                            ok = resolve_va_by_content(kc, wv, got)
                        except Exception:
                            ok = False
                            got = []
                        if ok and got:
                            lines.extend(got)
                            _want_state["done"].add(wv)
            except OSError:
                pass
        # PMA scrubber buffer discovery (throttled): resolve the finish/host-
        # semaphore word PAs and stream them for host registration.
        if kc_ok and time.time() >= scrub_next:
            scrub_next = time.time() + 5.0
            try:
                scrub_scan(kc, lines)
            except Exception as e:
                dbg('scrub scan failed: %s' % e)
        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
                # CE push-buffer pin freshness: a resolved mem is re-probed
                # against the channel's live ring content. UVM frees and
                # re-creates these rm_mems/channels mid-session and the slab
                # reuses the same kptrs, so a one-shot resolution can pin
                # pages that are stale (observed reading all zeros) while
                # the live pushes sit at new pages; dropping the latches
                # makes the resolution below re-emit the current truth.
                if sem in resolved_kc_pb and \
                        now_r >= kc_refresh_at.get(sem, 0):
                    kc_refresh_at[sem] = now_r + KC_REFRESH_S
                    try:
                        _fresh = _kc_pb_fresh(kc, sem)
                    except Exception:
                        _fresh = True
                    if not _fresh:
                        # Re-resolve, but never lose the latch on a failed
                        # attempt: a dropped latch means the host keeps no
                        # pins for this channel at all, which is worse than
                        # a possibly-stale pin (the host content-checks the
                        # first page and now rejects all-FF continuation
                        # chunks, so stale pages are bounded). On failure
                        # keep the old state and defer the next attempt.
                        _st = KC_PB_STATE.get(sem)
                        if _st and _st.get("mem"):
                            RESOLVED_KC_PB.discard(_st["mem"])
                        got = []
                        try:
                            _found2 = resolve_pb_chain(kc, sem, got)
                        except Exception:
                            _found2 = False
                            got = []
                        if _found2 and got:
                            _new = KC_PB_STATE.get(sem)
                            _newmem = _new.get("mem") if _new else None
                            dbg("kc-chain stale: re-resolved sem=0x%x "
                                "new-mem=%s runs=%d"
                                % (sem, "yes" if _newmem !=
                                   (_st.get("mem") if _st else None)
                                   else "no", len(got)))
                            lines.extend(got)
                            added += len(got)
                            resolved_kc_pb.add(sem)
                        else:
                            if _st and _st.get("mem"):
                                RESOLVED_KC_PB.add(_st["mem"])
                                if sem not in KC_PB_STATE:
                                    KC_PB_STATE[sem] = _st
                            kc_refresh_at[sem] = now_r + 5.0
                            dbg("kc-chain stale: keep old latch sem=0x%x"
                                % sem)
                    else:
                        # CE mem delta sweep (second-client pins): the one-shot
                        # emission pinned the mem once; pages the client writes
                        # later (or whose vmap pairing moved) were never
                        # re-emitted, so the host kept zeros for the other CE
                        # ranks' pushes. Sweep every KC_REEMIT_S, emitting only
                        # new/changed nonzero pairs; burst the Kcore PA cache
                        # every KC_REMAP_S so a remap is not shadowed.
                        _st = KC_PB_STATE.get(sem)
                        if _st is not None and \
                                now_r >= kc_reemit_at.get(sem, 0):
                            _fresh_map = now_r >= kc_remap_at.get(sem, 0)
                            if _fresh_map:
                                kc_remap_at[sem] = now_r + KC_REMAP_S
                            kc_reemit_at[sem] = now_r + KC_REEMIT_S
                            got = []
                            try:
                                _n2 = _emit_ce_mem_delta(kc, _st, got,
                                                         fresh_map=_fresh_map)
                            except Exception:
                                _n2 = 0
                                got = []
                            if got:
                                lines.extend(got)
                                added += len(got)
                                dbg("ce-delta sem=0x%x runs=%d mem=0x%x"
                                    % (sem, _n2, _st.get("mem", 0)))
                # CE push-buffer chain (second-client stall fix): resolve
                # pool -> manager -> pushbuffer -> rm_mem directly, so the
                # push-buffer pages never depend on the racy one-shot uvbp
                # event. Independent of the ring/pool gates below (the
                # channel gate's continue must not starve this path).
                if sem not in resolved_kc_pb and \
                        now_r >= retry_at.get(("pk", sem), 0):
                    got = []
                    try:
                        _found = resolve_pb_chain(kc, sem, got)
                    except Exception:
                        _found = False
                        got = []
                    if _found:
                        resolved_kc_pb.add(sem)
                        if got:
                            lines.extend(got)
                            added += len(got)
                    else:
                        retry_at[("pk", sem)] = now_r + 2.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
                # A uvch channel kptr is another early seed for the CE
                # push-buffer chain (its window contains the pool pointer):
                # resolve it here so the pins never depend on the uvmt sem
                # events alone.
                if ch in resolved_kc_pb or \
                        now_r < retry_at.get(("pk", ch), 0):
                    continue
                got = []
                try:
                    _found = resolve_pb_chain(kc, ch, got)
                except Exception:
                    _found = False
                    got = []
                if _found:
                    resolved_kc_pb.add(ch)
                    if got:
                        lines.extend(got)
                        added += len(got)
                else:
                    retry_at[("pk", ch)] = now_r + 2.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()